core: add internal API to disable stats and/or tracing. (#3464)

core/src/main/java/io/grpc/inprocess/InProcessChannelBuilder.java

@@ -67,6 +67,10 @@ public final class InProcessChannelBuilder extends
   private InProcessChannelBuilder(String name) {
     super(new InProcessSocketAddress(name), "localhost");
     this.name = Preconditions.checkNotNull(name, "name");
+    // TODO(zhangkun83): InProcessTransport by-passes framer and deframer, thus message sizses are
+    // not counted.  Therefore, we disable stats for now.
+    // (https://github.com/grpc/grpc-java/issues/2284)
+    setStatsEnabled(false);
   }
 
   @Override
@@ -89,15 +93,6 @@ public final class InProcessChannelBuilder extends
     return new InProcessClientTransportFactory(name);
   }
 
-  @Override
-  @Internal
-  protected boolean recordsStats() {
-    // TODO(zhangkun83): InProcessTransport by-passes framer and deframer, thus message sizses are
-    // not counted.  Therefore, we disable stats for now.
-    // (https://github.com/grpc/grpc-java/issues/2284)
-    return false;
-  }
-
   /**
    * Creates InProcess transports. Exposed for internal use, as it should be private.
    */

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

@@ -25,7 +25,6 @@ import com.google.instrumentation.stats.Stats;
 import com.google.instrumentation.stats.StatsContextFactory;
 import io.grpc.Attributes;
 import io.grpc.ClientInterceptor;
-import io.grpc.ClientStreamTracer;
 import io.grpc.CompressorRegistry;
 import io.grpc.DecompressorRegistry;
 import io.grpc.EquivalentAddressGroup;
@@ -144,6 +143,9 @@ public abstract class AbstractManagedChannelImplBuilder
     return maxInboundMessageSize;
   }
 
+  private boolean statsEnabled = true;
+  private boolean tracingEnabled = true;
+
   @Nullable
   private StatsContextFactory statsFactory;
 
@@ -280,15 +282,17 @@ public abstract class AbstractManagedChannelImplBuilder
   }
 
   /**
-   * Indicates whether this transport will record stats with {@link ClientStreamTracer}.
-   *
-   * <p>By default it returns {@code true}.  If the transport doesn't record stats, it may override
-   * this method to return {@code false} so that the builder won't install the Census interceptor.
-   *
-   * <p>If it returns true when it shouldn't be, Census will receive incomplete stats.
+   * Disable or enable stats features.  Enabled by default.
    */
-  protected boolean recordsStats() {
-    return true;
+  protected void setStatsEnabled(boolean value) {
+    statsEnabled = value;
+  }
+
+  /**
+   * Disable or enable tracing features.  Enabled by default.
+   */
+  protected void setTracingEnabled(boolean value) {
+    tracingEnabled = value;
   }
 
   @VisibleForTesting
@@ -317,10 +321,11 @@ public abstract class AbstractManagedChannelImplBuilder
         getEffectiveInterceptors());
   }
 
-  private List<ClientInterceptor> getEffectiveInterceptors() {
+  @VisibleForTesting
+  final List<ClientInterceptor> getEffectiveInterceptors() {
     List<ClientInterceptor> effectiveInterceptors =
         new ArrayList<ClientInterceptor>(this.interceptors);
-    if (recordsStats()) {
+    if (statsEnabled) {
       StatsContextFactory statsCtxFactory =
           this.statsFactory != null ? this.statsFactory : Stats.getStatsContextFactory();
       if (statsCtxFactory != null) {
@@ -331,10 +336,12 @@ public abstract class AbstractManagedChannelImplBuilder
         effectiveInterceptors.add(0, censusStats.getClientInterceptor());
       }
     }
-    CensusTracingModule censusTracing =
-        new CensusTracingModule(Tracing.getTracer(),
-            Tracing.getPropagationComponent().getBinaryFormat());
-    effectiveInterceptors.add(0, censusTracing.getClientInterceptor());
+    if (tracingEnabled) {
+      CensusTracingModule censusTracing =
+          new CensusTracingModule(Tracing.getTracer(),
+              Tracing.getPropagationComponent().getBinaryFormat());
+      effectiveInterceptors.add(0, censusTracing.getClientInterceptor());
+    }
     return effectiveInterceptors;
   }
 

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

@@ -99,6 +99,9 @@ public abstract class AbstractServerImplBuilder<T extends AbstractServerImplBuil
   @Nullable
   private StatsContextFactory statsFactory;
 
+  private boolean statsEnabled = true;
+  private boolean tracingEnabled = true;
+
   @Override
   public final T directExecutor() {
     return executor(MoreExecutors.directExecutor());
@@ -185,6 +188,20 @@ public abstract class AbstractServerImplBuilder<T extends AbstractServerImplBuil
     return thisT();
   }
 
+  /**
+   * Disable or enable stats features.  Enabled by default.
+   */
+  protected void setStatsEnabled(boolean value) {
+    statsEnabled = value;
+  }
+
+  /**
+   * Disable or enable tracing features.  Enabled by default.
+   */
+  protected void setTracingEnabled(boolean value) {
+    tracingEnabled = value;
+  }
+
   @Override
   public Server build() {
     ServerImpl server = new ServerImpl(
@@ -197,21 +214,25 @@ public abstract class AbstractServerImplBuilder<T extends AbstractServerImplBuil
     return server;
   }
 
-  private List<ServerStreamTracer.Factory> getTracerFactories() {
+  @VisibleForTesting
+  final List<ServerStreamTracer.Factory> getTracerFactories() {
     ArrayList<ServerStreamTracer.Factory> tracerFactories =
         new ArrayList<ServerStreamTracer.Factory>();
-    StatsContextFactory statsFactory =
-        this.statsFactory != null ? this.statsFactory : Stats.getStatsContextFactory();
-    if (statsFactory != null) {
-      CensusStatsModule censusStats =
-          new CensusStatsModule(
-              statsFactory, GrpcUtil.STOPWATCH_SUPPLIER, true /** only matters on client-side **/);
-      tracerFactories.add(censusStats.getServerTracerFactory());
+    if (statsEnabled) {
+      StatsContextFactory statsFactory =
+          this.statsFactory != null ? this.statsFactory : Stats.getStatsContextFactory();
+      if (statsFactory != null) {
+        CensusStatsModule censusStats =
+            new CensusStatsModule(statsFactory, GrpcUtil.STOPWATCH_SUPPLIER, true);
+        tracerFactories.add(censusStats.getServerTracerFactory());
+      }
+    }
+    if (tracingEnabled) {
+      CensusTracingModule censusTracing =
+          new CensusTracingModule(Tracing.getTracer(),
+              Tracing.getPropagationComponent().getBinaryFormat());
+      tracerFactories.add(censusTracing.getServerTracerFactory());
     }
-    CensusTracingModule censusTracing =
-        new CensusTracingModule(Tracing.getTracer(),
-            Tracing.getPropagationComponent().getBinaryFormat());
-    tracerFactories.add(censusTracing.getServerTracerFactory());
     tracerFactories.addAll(streamTracerFactories);
     return tracerFactories;
   }

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

@@ -335,7 +335,8 @@ final class CensusStatsModule {
     }
   }
 
-  private final class ServerTracerFactory extends ServerStreamTracer.Factory {
+  @VisibleForTesting
+  final class ServerTracerFactory extends ServerStreamTracer.Factory {
     @Override
     public ServerStreamTracer newServerStreamTracer(String fullMethodName, Metadata headers) {
       StatsContext parentCtx = headers.get(statsHeader);
@@ -347,7 +348,8 @@ final class CensusStatsModule {
     }
   }
 
-  private class StatsClientInterceptor implements ClientInterceptor {
+  @VisibleForTesting
+  final class StatsClientInterceptor implements ClientInterceptor {
     @Override
     public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
         MethodDescriptor<ReqT, RespT> method, CallOptions callOptions, Channel next) {

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

@@ -254,7 +254,8 @@ final class CensusTracingModule {
     }
   }
 
-  private final class ServerTracerFactory extends ServerStreamTracer.Factory {
+  @VisibleForTesting
+  final class ServerTracerFactory extends ServerStreamTracer.Factory {
     @SuppressWarnings("ReferenceEquality")
     @Override
     public ServerStreamTracer newServerStreamTracer(String fullMethodName, Metadata headers) {
@@ -266,7 +267,8 @@ final class CensusTracingModule {
     }
   }
 
-  private class TracingClientInterceptor implements ClientInterceptor {
+  @VisibleForTesting
+  final class TracingClientInterceptor implements ClientInterceptor {
     @Override
     public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
         MethodDescriptor<ReqT, RespT> method, CallOptions callOptions, Channel next) {

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

@@ -16,6 +16,7 @@
 
 package io.grpc.internal;
 
+import static com.google.common.truth.Truth.assertThat;
 import static junit.framework.TestCase.assertFalse;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotEquals;
@@ -26,13 +27,22 @@ import static org.junit.Assert.fail;
 import static org.mockito.Mockito.mock;
 
 import com.google.common.util.concurrent.MoreExecutors;
+import com.google.instrumentation.stats.StatsContext;
+import com.google.instrumentation.stats.StatsContextFactory;
+import io.grpc.CallOptions;
+import io.grpc.Channel;
+import io.grpc.ClientCall;
+import io.grpc.ClientInterceptor;
 import io.grpc.CompressorRegistry;
 import io.grpc.DecompressorRegistry;
 import io.grpc.LoadBalancer;
+import io.grpc.MethodDescriptor;
 import io.grpc.NameResolver;
+import java.io.InputStream;
 import java.net.InetSocketAddress;
 import java.net.SocketAddress;
 import java.net.URI;
+import java.util.List;
 import java.util.concurrent.Executor;
 import java.util.concurrent.TimeUnit;
 import org.junit.Test;
@@ -42,6 +52,28 @@ import org.junit.runners.JUnit4;
 /** Unit tests for {@link AbstractManagedChannelImplBuilder}. */
 @RunWith(JUnit4.class)
 public class AbstractManagedChannelImplBuilderTest {
+  private static final ClientInterceptor DUMMY_USER_INTERCEPTOR =
+      new ClientInterceptor() {
+        @Override
+        public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
+            MethodDescriptor<ReqT, RespT> method, CallOptions callOptions, Channel next) {
+          return next.newCall(method, callOptions);
+        }
+      };
+
+  private static final StatsContextFactory DUMMY_STATS_FACTORY =
+      new StatsContextFactory() {
+        @Override
+        public StatsContext deserialize(InputStream input) {
+          throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public StatsContext getDefault() {
+          throw new UnsupportedOperationException();
+        }
+      };
+
   private Builder builder = new Builder("fake");
   private Builder directAddressBuilder = new Builder(new SocketAddress(){}, "fake");
 
@@ -228,6 +260,49 @@ public class AbstractManagedChannelImplBuilderTest {
     assertEquals(target, uri.toString());
   }
 
+  @Test
+  public void getEffectiveInterceptors_default() {
+    builder.intercept(DUMMY_USER_INTERCEPTOR);
+    List<ClientInterceptor> effectiveInterceptors = builder.getEffectiveInterceptors();
+    assertEquals(3, effectiveInterceptors.size());
+    assertThat(effectiveInterceptors.get(0))
+        .isInstanceOf(CensusTracingModule.TracingClientInterceptor.class);
+    assertThat(effectiveInterceptors.get(1))
+        .isInstanceOf(CensusStatsModule.StatsClientInterceptor.class);
+    assertThat(effectiveInterceptors.get(2)).isSameAs(DUMMY_USER_INTERCEPTOR);
+  }
+
+  @Test
+  public void getEffectiveInterceptors_disableStats() {
+    builder.intercept(DUMMY_USER_INTERCEPTOR);
+    builder.setStatsEnabled(false);
+    List<ClientInterceptor> effectiveInterceptors = builder.getEffectiveInterceptors();
+    assertEquals(2, effectiveInterceptors.size());
+    assertThat(effectiveInterceptors.get(0))
+        .isInstanceOf(CensusTracingModule.TracingClientInterceptor.class);
+    assertThat(effectiveInterceptors.get(1)).isSameAs(DUMMY_USER_INTERCEPTOR);
+  }
+
+  @Test
+  public void getEffectiveInterceptors_disableTracing() {
+    builder.intercept(DUMMY_USER_INTERCEPTOR);
+    builder.setTracingEnabled(false);
+    List<ClientInterceptor> effectiveInterceptors = builder.getEffectiveInterceptors();
+    assertEquals(2, effectiveInterceptors.size());
+    assertThat(effectiveInterceptors.get(0))
+        .isInstanceOf(CensusStatsModule.StatsClientInterceptor.class);
+    assertThat(effectiveInterceptors.get(1)).isSameAs(DUMMY_USER_INTERCEPTOR);
+  }
+
+  @Test
+  public void getEffectiveInterceptors_disableBoth() {
+    builder.intercept(DUMMY_USER_INTERCEPTOR);
+    builder.setStatsEnabled(false);
+    builder.setTracingEnabled(false);
+    List<ClientInterceptor> effectiveInterceptors = builder.getEffectiveInterceptors();
+    assertThat(effectiveInterceptors).containsExactly(DUMMY_USER_INTERCEPTOR);
+  }
+
   @Test
   public void idleTimeout() {
     Builder builder = new Builder("target");
@@ -260,10 +335,12 @@ public class AbstractManagedChannelImplBuilderTest {
   static class Builder extends AbstractManagedChannelImplBuilder<Builder> {
     Builder(String target) {
       super(target);
+      statsContextFactory(DUMMY_STATS_FACTORY);
     }
 
     Builder(SocketAddress directServerAddress, String authority) {
       super(directServerAddress, authority);
+      statsContextFactory(DUMMY_STATS_FACTORY);
     }
 
     @Override

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

@@ -0,0 +1,115 @@
+/*
+ * Copyright 2017, gRPC Authors All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.grpc.internal;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertEquals;
+
+import com.google.instrumentation.stats.StatsContext;
+import com.google.instrumentation.stats.StatsContextFactory;
+import io.grpc.Metadata;
+import io.grpc.ServerStreamTracer;
+import java.io.File;
+import java.io.InputStream;
+import java.util.List;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Unit tests for {@link AbstractServerImplBuilder}. */
+@RunWith(JUnit4.class)
+public class AbstractServerImplBuilderTest {
+  private static final StatsContextFactory DUMMY_STATS_FACTORY =
+      new StatsContextFactory() {
+        @Override
+        public StatsContext deserialize(InputStream input) {
+          throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public StatsContext getDefault() {
+          throw new UnsupportedOperationException();
+        }
+      };
+
+  private static final ServerStreamTracer.Factory DUMMY_USER_TRACER =
+      new ServerStreamTracer.Factory() {
+        @Override
+        public ServerStreamTracer newServerStreamTracer(String fullMethodName, Metadata headers) {
+          throw new UnsupportedOperationException();
+        }
+      };
+
+  private Builder builder = new Builder();
+
+  @Test
+  public void getTracerFactories_default() {
+    builder.addStreamTracerFactory(DUMMY_USER_TRACER);
+    List<ServerStreamTracer.Factory> factories = builder.getTracerFactories();
+    assertEquals(3, factories.size());
+    assertThat(factories.get(0)).isInstanceOf(CensusStatsModule.ServerTracerFactory.class);
+    assertThat(factories.get(1)).isInstanceOf(CensusTracingModule.ServerTracerFactory.class);
+    assertThat(factories.get(2)).isSameAs(DUMMY_USER_TRACER);
+  }
+
+  @Test
+  public void getTracerFactories_disableStats() {
+    builder.addStreamTracerFactory(DUMMY_USER_TRACER);
+    builder.setStatsEnabled(false);
+    List<ServerStreamTracer.Factory> factories = builder.getTracerFactories();
+    assertEquals(2, factories.size());
+    assertThat(factories.get(0)).isInstanceOf(CensusTracingModule.ServerTracerFactory.class);
+    assertThat(factories.get(1)).isSameAs(DUMMY_USER_TRACER);
+  }
+
+  @Test
+  public void getTracerFactories_disableTracing() {
+    builder.addStreamTracerFactory(DUMMY_USER_TRACER);
+    builder.setTracingEnabled(false);
+    List<ServerStreamTracer.Factory> factories = builder.getTracerFactories();
+    assertEquals(2, factories.size());
+    assertThat(factories.get(0)).isInstanceOf(CensusStatsModule.ServerTracerFactory.class);
+    assertThat(factories.get(1)).isSameAs(DUMMY_USER_TRACER);
+  }
+
+  @Test
+  public void getTracerFactories_disableBoth() {
+    builder.addStreamTracerFactory(DUMMY_USER_TRACER);
+    builder.setTracingEnabled(false);
+    builder.setStatsEnabled(false);
+    List<ServerStreamTracer.Factory> factories = builder.getTracerFactories();
+    assertThat(factories).containsExactly(DUMMY_USER_TRACER);
+  }
+
+  static class Builder extends AbstractServerImplBuilder<Builder> {
+    Builder() {
+      statsContextFactory(DUMMY_STATS_FACTORY);
+    }
+
+    @Override
+    protected io.grpc.internal.InternalServer buildTransportServer(
+        List<ServerStreamTracer.Factory> streamTracerFactories) {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public Builder useTransportSecurity(File certChain, File privateKey) {
+      throw new UnsupportedOperationException();
+    }
+  }
+
+}

netty/src/main/java/io/grpc/netty/InternalNettyChannelBuilder.java

@@ -60,5 +60,13 @@ public final class InternalNettyChannelBuilder {
     builder.setDynamicParamsFactory(factory);
   }
 
+  public static void setStatsEnabled(NettyChannelBuilder builder, boolean value) {
+    builder.setStatsEnabled(value);
+  }
+
+  public static void setTracingEnabled(NettyChannelBuilder builder, boolean value) {
+    builder.setTracingEnabled(value);
+  }
+
   private InternalNettyChannelBuilder() {}
 }

netty/src/main/java/io/grpc/netty/InternalNettyServerBuilder.java

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2016, gRPC Authors All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.grpc.netty;
+
+import io.grpc.Internal;
+
+/**
+ * Internal {@link InternalNettyServerBuilder} accessor.  This is intended for usage internal to
+ * the gRPC team.  If you *really* think you need to use this, contact the gRPC team first.
+ */
+@Internal
+public final class InternalNettyServerBuilder {
+
+  public static void setStatsEnabled(NettyServerBuilder builder, boolean value) {
+    builder.setStatsEnabled(value);
+  }
+
+  public static void setTracingEnabled(NettyServerBuilder builder, boolean value) {
+    builder.setTracingEnabled(value);
+  }
+
+  private InternalNettyServerBuilder() {}
+}

netty/src/main/java/io/grpc/netty/NettyChannelBuilder.java

@@ -407,6 +407,16 @@ public final class NettyChannelBuilder
     this.dynamicParamsFactory = checkNotNull(factory, "factory");
   }
 
+  @Override
+  protected void setTracingEnabled(boolean value) {
+    super.setTracingEnabled(value);
+  }
+
+  @Override
+  protected void setStatsEnabled(boolean value) {
+    super.setStatsEnabled(value);
+  }
+
   interface TransportCreationParamsFilterFactory {
     @CheckReturnValue
     TransportCreationParamsFilter create(

netty/src/main/java/io/grpc/netty/NettyServerBuilder.java

@@ -194,6 +194,16 @@ public final class NettyServerBuilder extends AbstractServerImplBuilder<NettySer
     return this;
   }
 
+  @Override
+  protected void setTracingEnabled(boolean value) {
+    super.setTracingEnabled(value);
+  }
+
+  @Override
+  protected void setStatsEnabled(boolean value) {
+    super.setStatsEnabled(value);
+  }
+
   /**
    * The maximum number of concurrent calls permitted for each incoming connection. Defaults to no
    * limit.