core: refactor flags in CensusStatsModule. (#5095)

There are currently three boolean flags, and there will be one more
soon.  Put them all in the top-level class instead of passing them as
arguments on lower levels.

core/src/main/java/io/grpc/internal/AbstractManagedChannelImplBuilder.java

@@ -432,12 +432,12 @@ public abstract class AbstractManagedChannelImplBuilder
       temporarilyDisableRetry = true;
       CensusStatsModule censusStats = this.censusStatsOverride;
       if (censusStats == null) {
-        censusStats = new CensusStatsModule(GrpcUtil.STOPWATCH_SUPPLIER, true);
+        censusStats = new CensusStatsModule(
+            GrpcUtil.STOPWATCH_SUPPLIER, true, recordStartedRpcs, recordFinishedRpcs);
       }
       // First interceptor runs last (see ClientInterceptors.intercept()), so that no
       // other interceptor can override the tracer factory we set in CallOptions.
-      effectiveInterceptors.add(
+      effectiveInterceptors.add(0, censusStats.getClientInterceptor());
-          0, censusStats.getClientInterceptor(recordStartedRpcs, recordFinishedRpcs));
     }
     if (tracingEnabled) {
       temporarilyDisableRetry = true;

core/src/main/java/io/grpc/internal/AbstractServerImplBuilder.java

@@ -265,10 +265,10 @@ public abstract class AbstractServerImplBuilder<T extends AbstractServerImplBuil
     if (statsEnabled) {
       CensusStatsModule censusStats = this.censusStatsOverride;
       if (censusStats == null) {
-        censusStats = new CensusStatsModule(GrpcUtil.STOPWATCH_SUPPLIER, true);
+        censusStats = new CensusStatsModule(
+            GrpcUtil.STOPWATCH_SUPPLIER, true, recordStartedRpcs, recordFinishedRpcs);
       }
-      tracerFactories.add(
+      tracerFactories.add(censusStats.getServerTracerFactory());
-          censusStats.getServerTracerFactory(recordStartedRpcs, recordFinishedRpcs));
     }
     if (tracingEnabled) {
       CensusTracingModule censusTracing =

core/src/main/java/io/grpc/internal/CensusStatsModule.java

@@ -75,17 +75,20 @@ public final class CensusStatsModule {
   @VisibleForTesting
   final Metadata.Key<TagContext> statsHeader;
   private final boolean propagateTags;
+  private final boolean recordStartedRpcs;
+  private final boolean recordFinishedRpcs;
 
   /**
    * Creates a {@link CensusStatsModule} with the default OpenCensus implementation.
    */
-  CensusStatsModule(Supplier<Stopwatch> stopwatchSupplier, boolean propagateTags) {
+  CensusStatsModule(Supplier<Stopwatch> stopwatchSupplier,
+      boolean propagateTags, boolean recordStartedRpcs, boolean recordFinishedRpcs) {
     this(
         Tags.getTagger(),
         Tags.getTagPropagationComponent().getBinarySerializer(),
         Stats.getStatsRecorder(),
         stopwatchSupplier,
-        propagateTags);
+        propagateTags, recordStartedRpcs, recordFinishedRpcs);
   }
 
   /**
@@ -95,12 +98,14 @@ public final class CensusStatsModule {
       final Tagger tagger,
       final TagContextBinarySerializer tagCtxSerializer,
       StatsRecorder statsRecorder, Supplier<Stopwatch> stopwatchSupplier,
-      boolean propagateTags) {
+      boolean propagateTags, boolean recordStartedRpcs, boolean recordFinishedRpcs) {
     this.tagger = checkNotNull(tagger, "tagger");
     this.statsRecorder = checkNotNull(statsRecorder, "statsRecorder");
     checkNotNull(tagCtxSerializer, "tagCtxSerializer");
     this.stopwatchSupplier = checkNotNull(stopwatchSupplier, "stopwatchSupplier");
     this.propagateTags = propagateTags;
+    this.recordStartedRpcs = recordStartedRpcs;
+    this.recordFinishedRpcs = recordFinishedRpcs;
     this.statsHeader =
         Metadata.Key.of("grpc-tags-bin", new Metadata.BinaryMarshaller<TagContext>() {
             @Override
@@ -131,25 +136,22 @@ public final class CensusStatsModule {
    */
   @VisibleForTesting
   ClientCallTracer newClientCallTracer(
-      TagContext parentCtx, String fullMethodName,
+      TagContext parentCtx, String fullMethodName) {
-      boolean recordStartedRpcs, boolean recordFinishedRpcs) {
+    return new ClientCallTracer(this, parentCtx, fullMethodName);
-    return new ClientCallTracer(
-        this, parentCtx, fullMethodName, recordStartedRpcs, recordFinishedRpcs);
   }
 
   /**
    * Returns the server tracer factory.
    */
-  ServerStreamTracer.Factory getServerTracerFactory(
+  ServerStreamTracer.Factory getServerTracerFactory() {
-      boolean recordStartedRpcs, boolean recordFinishedRpcs) {
+    return new ServerTracerFactory();
-    return new ServerTracerFactory(recordStartedRpcs, recordFinishedRpcs);
   }
 
   /**
    * Returns the client interceptor that facilitates Census-based stats reporting.
    */
-  ClientInterceptor getClientInterceptor(boolean recordStartedRpcs, boolean recordFinishedRpcs) {
+  ClientInterceptor getClientInterceptor() {
-    return new StatsClientInterceptor(recordStartedRpcs, recordFinishedRpcs);
+    return new StatsClientInterceptor();
   }
 
   private static final class ClientTracer extends ClientStreamTracer {
@@ -275,8 +277,6 @@ public final class CensusStatsModule {
     }
   }
 
-
-
   @VisibleForTesting
   static final class ClientCallTracer extends ClientStreamTracer.Factory {
     @Nullable
@@ -314,24 +314,16 @@ public final class CensusStatsModule {
     private volatile int callEnded;
     private final TagContext parentCtx;
     private final TagContext startCtx;
-    private final boolean recordFinishedRpcs;
 
-    ClientCallTracer(
+    ClientCallTracer(CensusStatsModule module, TagContext parentCtx, String fullMethodName) {
-        CensusStatsModule module,
+      this.module = checkNotNull(module);
-        TagContext parentCtx,
-        String fullMethodName,
-        boolean recordStartedRpcs,
-        boolean recordFinishedRpcs) {
-      this.module = module;
       this.parentCtx = checkNotNull(parentCtx);
       TagValue methodTag = TagValue.create(fullMethodName);
-      this.startCtx =
+      this.startCtx = module.tagger.toBuilder(parentCtx)
-          module.tagger.toBuilder(parentCtx)
           .put(DeprecatedCensusConstants.RPC_METHOD, methodTag)
           .build();
       this.stopwatch = module.stopwatchSupplier.get().start();
-      this.recordFinishedRpcs = recordFinishedRpcs;
+      if (module.recordStartedRpcs) {
-      if (recordStartedRpcs) {
         module.statsRecorder.newMeasureMap()
             .put(DeprecatedCensusConstants.RPC_CLIENT_STARTED_COUNT, 1)
             .record(startCtx);
@@ -379,7 +371,7 @@ public final class CensusStatsModule {
         }
         callEnded = 1;
       }
-      if (!recordFinishedRpcs) {
+      if (!module.recordFinishedRpcs) {
         return;
       }
       stopwatch.stop();
@@ -482,8 +474,6 @@ public final class CensusStatsModule {
     private final TagContext parentCtx;
     private volatile int streamClosed;
     private final Stopwatch stopwatch;
-    private final Tagger tagger;
-    private final boolean recordFinishedRpcs;
     private volatile long outboundMessageCount;
     private volatile long inboundMessageCount;
     private volatile long outboundWireSize;
@@ -493,17 +483,11 @@ public final class CensusStatsModule {
 
     ServerTracer(
         CensusStatsModule module,
-        TagContext parentCtx,
+        TagContext parentCtx) {
-        Supplier<Stopwatch> stopwatchSupplier,
+      this.module = checkNotNull(module, "module");
-        Tagger tagger,
-        boolean recordStartedRpcs,
-        boolean recordFinishedRpcs) {
-      this.module = module;
       this.parentCtx = checkNotNull(parentCtx, "parentCtx");
-      this.stopwatch = stopwatchSupplier.get().start();
+      this.stopwatch = module.stopwatchSupplier.get().start();
-      this.tagger = tagger;
+      if (module.recordStartedRpcs) {
-      this.recordFinishedRpcs = recordFinishedRpcs;
-      if (recordStartedRpcs) {
         module.statsRecorder.newMeasureMap()
             .put(DeprecatedCensusConstants.RPC_SERVER_STARTED_COUNT, 1)
             .record(parentCtx);
@@ -588,7 +572,7 @@ public final class CensusStatsModule {
         }
         streamClosed = 1;
       }
-      if (!recordFinishedRpcs) {
+      if (!module.recordFinishedRpcs) {
         return;
       }
       stopwatch.stop();
@@ -624,7 +608,7 @@ public final class CensusStatsModule {
 
     @Override
     public Context filterContext(Context context) {
-      if (!tagger.empty().equals(parentCtx)) {
+      if (!module.tagger.empty().equals(parentCtx)) {
         return context.withValue(TAG_CONTEXT_KEY, parentCtx);
       }
       return context;
@@ -633,14 +617,6 @@ public final class CensusStatsModule {
 
   @VisibleForTesting
   final class ServerTracerFactory extends ServerStreamTracer.Factory {
-    private final boolean recordStartedRpcs;
-    private final boolean recordFinishedRpcs;
-
-    ServerTracerFactory(boolean recordStartedRpcs, boolean recordFinishedRpcs) {
-      this.recordStartedRpcs = recordStartedRpcs;
-      this.recordFinishedRpcs = recordFinishedRpcs;
-    }
-
     @Override
     public ServerStreamTracer newServerStreamTracer(String fullMethodName, Metadata headers) {
       TagContext parentCtx = headers.get(statsHeader);
@@ -653,34 +629,19 @@ public final class CensusStatsModule {
               .toBuilder(parentCtx)
               .put(DeprecatedCensusConstants.RPC_METHOD, methodTag)
               .build();
-      return new ServerTracer(
+      return new ServerTracer(CensusStatsModule.this, parentCtx);
-          CensusStatsModule.this,
-          parentCtx,
-          stopwatchSupplier,
-          tagger,
-          recordStartedRpcs,
-          recordFinishedRpcs);
     }
   }
 
   @VisibleForTesting
   final class StatsClientInterceptor implements ClientInterceptor {
-    private final boolean recordStartedRpcs;
-    private final boolean recordFinishedRpcs;
-
-    StatsClientInterceptor(boolean recordStartedRpcs, boolean recordFinishedRpcs) {
-      this.recordStartedRpcs = recordStartedRpcs;
-      this.recordFinishedRpcs = recordFinishedRpcs;
-    }
-
     @Override
     public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
         MethodDescriptor<ReqT, RespT> method, CallOptions callOptions, Channel next) {
       // New RPCs on client-side inherit the tag context from the current Context.
       TagContext parentCtx = tagger.getCurrentTagContext();
       final ClientCallTracer tracerFactory =
-          newClientCallTracer(parentCtx, method.getFullMethodName(),
+          newClientCallTracer(parentCtx, method.getFullMethodName());
-              recordStartedRpcs, recordFinishedRpcs);
       ClientCall<ReqT, RespT> call =
           next.newCall(method, callOptions.withStreamTracerFactory(tracerFactory));
       return new SimpleForwardingClientCall<ReqT, RespT>(call) {

core/src/test/java/io/grpc/internal/AbstractManagedChannelImplBuilderTest.java

@@ -414,7 +414,7 @@ public class AbstractManagedChannelImplBuilderTest {
               new FakeTagContextBinarySerializer(),
               new FakeStatsRecorder(),
               GrpcUtil.STOPWATCH_SUPPLIER,
-              true));
+              true, true, true));
     }
 
     Builder(SocketAddress directServerAddress, String authority) {
@@ -425,7 +425,7 @@ public class AbstractManagedChannelImplBuilderTest {
               new FakeTagContextBinarySerializer(),
               new FakeStatsRecorder(),
               GrpcUtil.STOPWATCH_SUPPLIER,
-              true));
+              true, true, true));
     }
 
     @Override

core/src/test/java/io/grpc/internal/AbstractServerImplBuilderTest.java

@@ -91,7 +91,7 @@ public class AbstractServerImplBuilderTest {
               new FakeTagContextBinarySerializer(),
               new FakeStatsRecorder(),
               GrpcUtil.STOPWATCH_SUPPLIER,
-              true));
+              true, true, true));
     }
 
     @Override

core/src/test/java/io/grpc/internal/CensusModulesTest.java

@@ -191,7 +191,8 @@ public class CensusModulesTest {
         .thenReturn(fakeClientSpanContext);
     censusStats =
         new CensusStatsModule(
-            tagger, tagCtxSerializer, statsRecorder, fakeClock.getStopwatchSupplier(), true);
+            tagger, tagCtxSerializer, statsRecorder, fakeClock.getStopwatchSupplier(),
+            true, true, true);
     censusTracing = new CensusTracingModule(tracer, mockTracingPropagationHandler);
   }
 
@@ -240,7 +241,7 @@ public class CensusModulesTest {
     Channel interceptedChannel =
         ClientInterceptors.intercept(
             grpcServerRule.getChannel(), callOptionsCaptureInterceptor,
-            censusStats.getClientInterceptor(true, true), censusTracing.getClientInterceptor());
+            censusStats.getClientInterceptor(), censusTracing.getClientInterceptor());
     ClientCall<String, String> call;
     if (nonDefaultContext) {
       Context ctx =
@@ -353,9 +354,13 @@ public class CensusModulesTest {
   }
 
   private void subtestClientBasicStatsDefaultContext(boolean recordStarts, boolean recordFinishes) {
+    CensusStatsModule localCensusStats =
+        new CensusStatsModule(
+            tagger, tagCtxSerializer, statsRecorder, fakeClock.getStopwatchSupplier(),
+            true, recordStarts, recordFinishes);
     CensusStatsModule.ClientCallTracer callTracer =
-        censusStats.newClientCallTracer(
+        localCensusStats.newClientCallTracer(
-            tagger.empty(), method.getFullMethodName(), recordStarts, recordFinishes);
+            tagger.empty(), method.getFullMethodName());
     Metadata headers = new Metadata();
     ClientStreamTracer tracer = callTracer.newClientStreamTracer(CallOptions.DEFAULT, headers);
 
@@ -490,8 +495,7 @@ public class CensusModulesTest {
   @Test
   public void clientStreamNeverCreatedStillRecordStats() {
     CensusStatsModule.ClientCallTracer callTracer =
-        censusStats.newClientCallTracer(
+        censusStats.newClientCallTracer(tagger.empty(), method.getFullMethodName());
-            tagger.empty(), method.getFullMethodName(), true, true);
 
     fakeClock.forwardTime(3000, MILLISECONDS);
     callTracer.callEnded(Status.DEADLINE_EXCEEDED.withDescription("3 seconds"));
@@ -593,10 +597,10 @@ public class CensusModulesTest {
             tagCtxSerializer,
             statsRecorder,
             fakeClock.getStopwatchSupplier(),
-            propagate);
+            propagate, recordStats, recordStats);
     Metadata headers = new Metadata();
     CensusStatsModule.ClientCallTracer callTracer =
-        census.newClientCallTracer(clientCtx, method.getFullMethodName(), recordStats, recordStats);
+        census.newClientCallTracer(clientCtx, method.getFullMethodName());
     // This propagates clientCtx to headers if propagates==true
     callTracer.newClientStreamTracer(CallOptions.DEFAULT, headers);
     if (recordStats) {
@@ -619,8 +623,7 @@ public class CensusModulesTest {
     }
 
     ServerStreamTracer serverTracer =
-        census.getServerTracerFactory(recordStats, recordStats).newServerStreamTracer(
+        census.getServerTracerFactory().newServerStreamTracer(method.getFullMethodName(), headers);
-            method.getFullMethodName(), headers);
     // Server tracer deserializes clientCtx from the headers, so that it records stats with the
     // propagated tags.
     Context serverContext = serverTracer.filterContext(Context.ROOT);
@@ -686,10 +689,12 @@ public class CensusModulesTest {
   @Test
   public void statsHeadersNotPropagateDefaultContext() {
     CensusStatsModule.ClientCallTracer callTracer =
-        censusStats.newClientCallTracer(tagger.empty(), method.getFullMethodName(), false, false);
+        censusStats.newClientCallTracer(tagger.empty(), method.getFullMethodName());
     Metadata headers = new Metadata();
     callTracer.newClientStreamTracer(CallOptions.DEFAULT, headers);
     assertFalse(headers.containsKey(censusStats.statsHeader));
+    // Clear recorded stats to satisfy the assertions in wrapUp() 
+    statsRecorder.rolloverRecords();
   }
 
   @Test
@@ -828,8 +833,11 @@ public class CensusModulesTest {
   }
 
   private void subtestServerBasicStatsNoHeaders(boolean recordStarts, boolean recordFinishes) {
-    ServerStreamTracer.Factory tracerFactory =
+    CensusStatsModule localCensusStats =
-        censusStats.getServerTracerFactory(recordStarts, recordFinishes);
+        new CensusStatsModule(
+            tagger, tagCtxSerializer, statsRecorder, fakeClock.getStopwatchSupplier(),
+            true, recordStarts, recordFinishes);
+    ServerStreamTracer.Factory tracerFactory = localCensusStats.getServerTracerFactory();
     ServerStreamTracer tracer =
         tracerFactory.newServerStreamTracer(method.getFullMethodName(), new Metadata());
 

interop-testing/src/main/java/io/grpc/testing/integration/AbstractInteropTest.java

@@ -228,7 +228,7 @@ public abstract class AbstractInteropTest {
             tagContextBinarySerializer,
             serverStatsRecorder,
             GrpcUtil.STOPWATCH_SUPPLIER,
-            true));
+            true, true, true));
     try {
       server = builder.build().start();
     } catch (IOException ex) {
@@ -330,7 +330,8 @@ public abstract class AbstractInteropTest {
 
   protected final CensusStatsModule createClientCensusStatsModule() {
     return new CensusStatsModule(
-        tagger, tagContextBinarySerializer, clientStatsRecorder, GrpcUtil.STOPWATCH_SUPPLIER, true);
+        tagger, tagContextBinarySerializer, clientStatsRecorder, GrpcUtil.STOPWATCH_SUPPLIER,
+        true, true, true);
   }
 
   /**