core: Change gRPC to use io.opencensus:opencensus-api:0.5.1. (#3204)

build.gradle

@@ -184,6 +184,7 @@ subprojects {
                 google_auth_credentials: 'com.google.auth:google-auth-library-credentials:0.4.0',
                 okhttp: 'com.squareup.okhttp:okhttp:2.5.0',
                 okio: 'com.squareup.okio:okio:1.6.0',
+                opencensus_api: 'io.opencensus:opencensus-api:0.5.1',
                 instrumentation_api: 'com.google.instrumentation:instrumentation-api:0.4.3',
                 protobuf: "com.google.protobuf:protobuf-java:${protobufVersion}",
                 protobuf_lite: "com.google.protobuf:protobuf-lite:3.0.1",

core/BUILD.bazel

@@ -43,6 +43,7 @@ java_library(
         "@com_google_errorprone_error_prone_annotations//jar",
         "@com_google_guava//jar",
         "@com_google_instrumentation_api//jar",
+        "@io_opencensus_api//jar",
     ],
 )
 

core/build.gradle

@@ -1,12 +1,15 @@
 description = 'gRPC: Core'
 
 dependencies {
-    compile libraries.guava,
+    compile project(':grpc-context'),
+            libraries.guava,
             libraries.errorprone,
             libraries.jsr305,
-            project(':grpc-context'),
+            libraries.instrumentation_api,
-            libraries.instrumentation_api
+            libraries.opencensus_api
+
     testCompile project(':grpc-testing')
+
     signature "org.codehaus.mojo.signature:java16:1.1@signature"
 }
 

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

@@ -23,7 +23,6 @@ import com.google.common.base.Preconditions;
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.instrumentation.stats.Stats;
 import com.google.instrumentation.stats.StatsContextFactory;
-import com.google.instrumentation.trace.Tracing;
 import io.grpc.Attributes;
 import io.grpc.ClientInterceptor;
 import io.grpc.ClientStreamTracer;
@@ -36,6 +35,7 @@ import io.grpc.ManagedChannelBuilder;
 import io.grpc.NameResolver;
 import io.grpc.NameResolverProvider;
 import io.grpc.PickFirstBalancerFactory;
+import io.opencensus.trace.Tracing;
 import java.net.SocketAddress;
 import java.net.URI;
 import java.net.URISyntaxException;
@@ -348,7 +348,8 @@ public abstract class AbstractManagedChannelImplBuilder
     }
     if (enableTracing) {
       CensusTracingModule censusTracing =
-          new CensusTracingModule(Tracing.getTracer(), Tracing.getBinaryPropagationHandler());
+          new CensusTracingModule(Tracing.getTracer(),
+              Tracing.getPropagationComponent().getBinaryFormat());
       effectiveInterceptors.add(0, censusTracing.getClientInterceptor());
     }
     return effectiveInterceptors;

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

@@ -22,7 +22,6 @@ import com.google.common.annotations.VisibleForTesting;
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.instrumentation.stats.Stats;
 import com.google.instrumentation.stats.StatsContextFactory;
-import com.google.instrumentation.trace.Tracing;
 import io.grpc.BindableService;
 import io.grpc.CompressorRegistry;
 import io.grpc.Context;
@@ -37,6 +36,7 @@ import io.grpc.ServerMethodDefinition;
 import io.grpc.ServerServiceDefinition;
 import io.grpc.ServerStreamTracer;
 import io.grpc.ServerTransportFilter;
+import io.opencensus.trace.Tracing;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -206,7 +206,8 @@ public abstract class AbstractServerImplBuilder<T extends AbstractServerImplBuil
       tracerFactories.add(censusStats.getServerTracerFactory());
     }
     CensusTracingModule censusTracing =
-        new CensusTracingModule(Tracing.getTracer(), Tracing.getBinaryPropagationHandler());
+        new CensusTracingModule(Tracing.getTracer(),
+            Tracing.getPropagationComponent().getBinaryFormat());
     tracerFactories.add(censusTracing.getServerTracerFactory());
     tracerFactories.addAll(streamTracerFactories);
     return tracerFactories;

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

@@ -17,15 +17,9 @@
 package io.grpc.internal;
 
 import static com.google.common.base.Preconditions.checkNotNull;
-import static com.google.instrumentation.trace.ContextUtils.CONTEXT_SPAN_KEY;
+import static io.opencensus.trace.unsafe.ContextUtils.CONTEXT_SPAN_KEY;
 
 import com.google.common.annotations.VisibleForTesting;
-import com.google.instrumentation.trace.BinaryPropagationHandler;
-import com.google.instrumentation.trace.EndSpanOptions;
-import com.google.instrumentation.trace.Span;
-import com.google.instrumentation.trace.SpanContext;
-import com.google.instrumentation.trace.Status;
-import com.google.instrumentation.trace.Tracer;
 import io.grpc.CallOptions;
 import io.grpc.Channel;
 import io.grpc.ClientCall;
@@ -38,6 +32,12 @@ import io.grpc.Metadata;
 import io.grpc.MethodDescriptor;
 import io.grpc.ServerStreamTracer;
 import io.grpc.StreamTracer;
+import io.opencensus.trace.EndSpanOptions;
+import io.opencensus.trace.Span;
+import io.opencensus.trace.SpanContext;
+import io.opencensus.trace.Status;
+import io.opencensus.trace.Tracer;
+import io.opencensus.trace.propagation.BinaryFormat;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -60,28 +60,28 @@ final class CensusTracingModule {
   private static final ClientStreamTracer noopClientTracer = new ClientStreamTracer() {};
 
   private final Tracer censusTracer;
-  private final BinaryPropagationHandler censusTracingPropagationHandler;
+  private final BinaryFormat censusPropagationBinaryFormat;
   @VisibleForTesting
   final Metadata.Key<SpanContext> tracingHeader;
   private final TracingClientInterceptor clientInterceptor = new TracingClientInterceptor();
   private final ServerTracerFactory serverTracerFactory = new ServerTracerFactory();
 
   CensusTracingModule(
-      Tracer censusTracer, final BinaryPropagationHandler censusTracingPropagationHandler) {
+      Tracer censusTracer, final BinaryFormat censusPropagationBinaryFormat) {
     this.censusTracer = checkNotNull(censusTracer, "censusTracer");
-    this.censusTracingPropagationHandler =
+    this.censusPropagationBinaryFormat =
-        checkNotNull(censusTracingPropagationHandler, "censusTracingPropagationHandler");
+        checkNotNull(censusPropagationBinaryFormat, "censusPropagationBinaryFormat");
     this.tracingHeader =
         Metadata.Key.of("grpc-trace-bin", new Metadata.BinaryMarshaller<SpanContext>() {
             @Override
             public byte[] toBytes(SpanContext context) {
-              return censusTracingPropagationHandler.toBinaryValue(context);
+              return censusPropagationBinaryFormat.toBinaryValue(context);
             }
 
             @Override
             public SpanContext parseBytes(byte[] serialized) {
               try {
-                return censusTracingPropagationHandler.fromBinaryValue(serialized);
+                return censusPropagationBinaryFormat.fromBinaryValue(serialized);
               } catch (Exception e) {
                 logger.log(Level.FINE, "Failed to parse tracing header", e);
                 return SpanContext.INVALID;
@@ -195,7 +195,7 @@ final class CensusTracingModule {
       this.fullMethodName = checkNotNull(fullMethodName, "fullMethodName");
       this.span =
           censusTracer
-              .spanBuilder(parentSpan, makeSpanName("Sent", fullMethodName))
+              .spanBuilderWithExplicitParent(makeSpanName("Sent", fullMethodName), parentSpan)
               .setRecordEvents(true)
               .startSpan();
     }
@@ -230,7 +230,7 @@ final class CensusTracingModule {
       this.fullMethodName = checkNotNull(fullMethodName, "fullMethodName");
       this.span =
           censusTracer
-              .spanBuilderWithRemoteParent(remoteSpan, makeSpanName("Recv", fullMethodName))
+              .spanBuilderWithRemoteParent(makeSpanName("Recv", fullMethodName), remoteSpan)
               .setRecordEvents(true)
               .startSpan();
     }
@@ -251,6 +251,9 @@ final class CensusTracingModule {
 
     @Override
     public <ReqT, RespT> Context filterContext(Context context) {
+      // Access directly the unsafe trace API to create the new Context. This is a safe usage
+      // because gRPC always creates a new Context for each of the server calls and does not
+      // inherit from the parent Context.
       return context.withValue(CONTEXT_SPAN_KEY, span);
     }
   }
@@ -272,9 +275,11 @@ final class CensusTracingModule {
     public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
         MethodDescriptor<ReqT, RespT> method, CallOptions callOptions, Channel next) {
       // New RPCs on client-side inherit the tracing context from the current Context.
-      Span parentSpan = CONTEXT_SPAN_KEY.get();
+      // Safe usage of the unsafe trace API because CONTEXT_SPAN_KEY.get() returns the same value
+      // as Tracer.getCurrentSpan() except when no value available when the return value is null
+      // for the direct access and BlankSpan when Tracer API is used.
       final ClientCallTracer tracerFactory =
-          newClientCallTracer(parentSpan, method.getFullMethodName());
+          newClientCallTracer(CONTEXT_SPAN_KEY.get(), method.getFullMethodName());
       ClientCall<ReqT, RespT> call =
           next.newCall(method, callOptions.withStreamTracerFactory(tracerFactory));
       return new SimpleForwardingClientCall<ReqT, RespT>(call) {

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

@@ -28,7 +28,6 @@ import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyString;
-import static org.mockito.Matchers.argThat;
 import static org.mockito.Matchers.eq;
 import static org.mockito.Matchers.isNull;
 import static org.mockito.Matchers.same;
@@ -42,21 +41,6 @@ import static org.mockito.Mockito.when;
 import com.google.instrumentation.stats.RpcConstants;
 import com.google.instrumentation.stats.StatsContext;
 import com.google.instrumentation.stats.TagValue;
-import com.google.instrumentation.trace.Annotation;
-import com.google.instrumentation.trace.AttributeValue;
-import com.google.instrumentation.trace.BinaryPropagationHandler;
-import com.google.instrumentation.trace.ContextUtils;
-import com.google.instrumentation.trace.EndSpanOptions;
-import com.google.instrumentation.trace.Link;
-import com.google.instrumentation.trace.NetworkEvent;
-import com.google.instrumentation.trace.Span;
-import com.google.instrumentation.trace.SpanContext;
-import com.google.instrumentation.trace.SpanFactory;
-import com.google.instrumentation.trace.SpanId;
-import com.google.instrumentation.trace.StartSpanOptions;
-import com.google.instrumentation.trace.TraceId;
-import com.google.instrumentation.trace.TraceOptions;
-import com.google.instrumentation.trace.Tracer;
 import io.grpc.CallOptions;
 import io.grpc.Channel;
 import io.grpc.ClientCall;
@@ -74,9 +58,26 @@ import io.grpc.Status;
 import io.grpc.internal.testing.StatsTestUtils;
 import io.grpc.internal.testing.StatsTestUtils.FakeStatsContextFactory;
 import io.grpc.testing.GrpcServerRule;
+import io.opencensus.trace.Annotation;
+import io.opencensus.trace.AttributeValue;
+import io.opencensus.trace.EndSpanOptions;
+import io.opencensus.trace.Link;
+import io.opencensus.trace.NetworkEvent;
+import io.opencensus.trace.Sampler;
+import io.opencensus.trace.Span;
+import io.opencensus.trace.SpanBuilder;
+import io.opencensus.trace.SpanContext;
+import io.opencensus.trace.SpanId;
+import io.opencensus.trace.TraceId;
+import io.opencensus.trace.TraceOptions;
+import io.opencensus.trace.Tracer;
+import io.opencensus.trace.propagation.BinaryFormat;
+import io.opencensus.trace.unsafe.ContextUtils;
 import java.io.ByteArrayInputStream;
 import java.io.InputStream;
 import java.text.ParseException;
+import java.util.EnumSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Random;
 import java.util.concurrent.atomic.AtomicReference;
@@ -88,7 +89,6 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 import org.mockito.ArgumentCaptor;
-import org.mockito.ArgumentMatcher;
 import org.mockito.Captor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
@@ -140,45 +140,22 @@ public class CensusModulesTest {
           .build();
   private final FakeClock fakeClock = new FakeClock();
   private final FakeStatsContextFactory statsCtxFactory = new FakeStatsContextFactory();
-  private final Random random = new Random(0);
+  private final Random random = new Random(1234);
-  private final SpanContext fakeClientSpanContext =
+  private final Span fakeClientParentSpan = MockableSpan.generateRandomSpan(random);
-      SpanContext.create(
+  private final Span spyClientSpan = spy(MockableSpan.generateRandomSpan(random));
-          TraceId.generateRandomId(random), SpanId.generateRandomId(random),
+  private final SpanContext fakeClientSpanContext = spyClientSpan.getContext();
-          TraceOptions.builder().build());
+  private final Span spyServerSpan = spy(MockableSpan.generateRandomSpan(random));
-  private final SpanContext fakeClientParentSpanContext =
-      SpanContext.create(
-          TraceId.generateRandomId(random), SpanId.generateRandomId(random),
-          TraceOptions.builder().build());
-  private final SpanContext fakeServerSpanContext =
-      SpanContext.create(
-          TraceId.generateRandomId(random), SpanId.generateRandomId(random),
-          TraceOptions.builder().build());
-  private final SpanContext fakeServerParentSpanContext =
-      SpanContext.create(
-          TraceId.generateRandomId(random), SpanId.generateRandomId(random),
-          TraceOptions.builder().build());
-  private final Span fakeClientSpan = new FakeSpan(fakeClientSpanContext);
-  private final Span fakeServerSpan = new FakeSpan(fakeServerSpanContext);
-  private final Span fakeClientParentSpan = new FakeSpan(fakeClientParentSpanContext);
-  private final Span fakeServerParentSpan = new FakeSpan(fakeServerParentSpanContext);
-  private final Span spyClientSpan = spy(fakeClientSpan);
-  private final Span spyServerSpan = spy(fakeServerSpan);
   private final byte[] binarySpanContext = new byte[]{3, 1, 5};
-  private final ArgumentMatcher<StartSpanOptions> startSpanOptionsMatcher =
+  private final SpanBuilder spyClientSpanBuilder = spy(new MockableSpan.Builder());
-      new ArgumentMatcher<StartSpanOptions>() {
+  private final SpanBuilder spyServerSpanBuilder = spy(new MockableSpan.Builder());
-        @Override
-        public boolean matches(Object argument) {
-          return Boolean.TRUE.equals(((StartSpanOptions) argument).getRecordEvents());
-        }
-      };
 
   @Rule
   public final GrpcServerRule grpcServerRule = new GrpcServerRule().directExecutor();
 
   @Mock
-  private AccessibleSpanFactory mockSpanFactory;
+  private Tracer tracer;
   @Mock
-  private BinaryPropagationHandler mockTracingPropagationHandler;
+  private BinaryFormat mockTracingPropagationHandler;
   @Mock
   private ClientCall.Listener<String> mockClientCallListener;
   @Mock
@@ -190,7 +167,6 @@ public class CensusModulesTest {
   @Captor
   private ArgumentCaptor<Status> statusCaptor;
 
-  private Tracer tracer;
   private CensusStatsModule censusStats;
   private CensusTracingModule censusTracing;
 
@@ -198,17 +174,16 @@ public class CensusModulesTest {
   @SuppressWarnings("unchecked")
   public void setUp() throws Exception {
     MockitoAnnotations.initMocks(this);
-    when(mockSpanFactory.startSpan(any(Span.class), anyString(), any(StartSpanOptions.class)))
+    when(spyClientSpanBuilder.startSpan()).thenReturn(spyClientSpan);
-        .thenReturn(spyClientSpan);
+    when(tracer.spanBuilderWithExplicitParent(anyString(), any(Span.class)))
-    when(
+        .thenReturn(spyClientSpanBuilder);
-        mockSpanFactory.startSpanWithRemoteParent(
+    when(spyServerSpanBuilder.startSpan()).thenReturn(spyServerSpan);
-            any(SpanContext.class), anyString(), any(StartSpanOptions.class)))
+    when(tracer.spanBuilderWithRemoteParent(anyString(), any(SpanContext.class)))
-        .thenReturn(spyServerSpan);
+        .thenReturn(spyServerSpanBuilder);
     when(mockTracingPropagationHandler.toBinaryValue(any(SpanContext.class)))
         .thenReturn(binarySpanContext);
     when(mockTracingPropagationHandler.fromBinaryValue(any(byte[].class)))
-        .thenReturn(fakeServerParentSpanContext);
+        .thenReturn(fakeClientSpanContext);
-    tracer = new Tracer(mockSpanFactory) {};
     censusStats = new CensusStatsModule(statsCtxFactory, fakeClock.getStopwatchSupplier(), true);
     censusTracing = new CensusTracingModule(tracer, mockTracingPropagationHandler);
   }
@@ -295,13 +270,13 @@ public class CensusModulesTest {
     call.start(mockClientCallListener, headers);
     assertNull(statsCtxFactory.pollRecord());
     if (nonDefaultContext) {
-      verify(mockSpanFactory).startSpan(
+      verify(tracer).spanBuilderWithExplicitParent(
-          same(fakeClientParentSpan), eq("Sent.package1.service2.method3"),
+          eq("Sent.package1.service2.method3"), same(fakeClientParentSpan));
-          argThat(startSpanOptionsMatcher));
+      verify(spyClientSpanBuilder).setRecordEvents(eq(true));
     } else {
-      verify(mockSpanFactory).startSpan(
+      verify(tracer).spanBuilderWithExplicitParent(
-          isNull(Span.class), eq("Sent.package1.service2.method3"),
+          eq("Sent.package1.service2.method3"), isNull(Span.class));
-          argThat(startSpanOptionsMatcher));
+      verify(spyClientSpanBuilder).setRecordEvents(eq(true));
     }
     verify(spyClientSpan, never()).end(any(EndSpanOptions.class));
 
@@ -330,7 +305,7 @@ public class CensusModulesTest {
     verify(spyClientSpan).end(
         EndSpanOptions.builder()
             .setStatus(
-                com.google.instrumentation.trace.Status.PERMISSION_DENIED
+                io.opencensus.trace.Status.PERMISSION_DENIED
                     .withDescription("No you don't"))
             .build());
     verify(spyClientSpan, never()).end();
@@ -385,17 +360,17 @@ public class CensusModulesTest {
     CensusTracingModule.ClientCallTracer callTracer =
         censusTracing.newClientCallTracer(null, method.getFullMethodName());
     Metadata headers = new Metadata();
-    ClientStreamTracer tracer = callTracer.newClientStreamTracer(headers);
+    ClientStreamTracer clientStreamTracer = callTracer.newClientStreamTracer(headers);
-    verify(mockSpanFactory).startSpan(
+    verify(tracer).spanBuilderWithExplicitParent(
-        isNull(Span.class), eq("Sent.package1.service2.method3"), argThat(startSpanOptionsMatcher));
+        eq("Sent.package1.service2.method3"), isNull(Span.class));
     verify(spyClientSpan, never()).end(any(EndSpanOptions.class));
 
-    tracer.streamClosed(Status.OK);
+    clientStreamTracer.streamClosed(Status.OK);
     callTracer.callEnded(Status.OK);
 
     verify(spyClientSpan).end(
-        EndSpanOptions.builder().setStatus(com.google.instrumentation.trace.Status.OK).build());
+        EndSpanOptions.builder().setStatus(io.opencensus.trace.Status.OK).build());
-    verifyNoMoreInteractions(mockSpanFactory);
+    verifyNoMoreInteractions(tracer);
   }
 
   @Test
@@ -429,15 +404,15 @@ public class CensusModulesTest {
   public void clientStreamNeverCreatedStillRecordTracing() {
     CensusTracingModule.ClientCallTracer callTracer =
         censusTracing.newClientCallTracer(fakeClientParentSpan, method.getFullMethodName());
-    verify(mockSpanFactory).startSpan(
+    verify(tracer).spanBuilderWithExplicitParent(
-        same(fakeClientParentSpan), eq("Sent.package1.service2.method3"),
+        eq("Sent.package1.service2.method3"), same(fakeClientParentSpan));
-        argThat(startSpanOptionsMatcher));
+    verify(spyClientSpanBuilder).setRecordEvents(eq(true));
 
     callTracer.callEnded(Status.DEADLINE_EXCEEDED.withDescription("3 seconds"));
     verify(spyClientSpan).end(
         EndSpanOptions.builder()
             .setStatus(
-                com.google.instrumentation.trace.Status.DEADLINE_EXCEEDED
+                io.opencensus.trace.Status.DEADLINE_EXCEEDED
                     .withDescription("3 seconds"))
             .build());
     verify(spyClientSpan, never()).end();
@@ -550,19 +525,19 @@ public class CensusModulesTest {
 
     verify(mockTracingPropagationHandler).toBinaryValue(same(fakeClientSpanContext));
     verifyNoMoreInteractions(mockTracingPropagationHandler);
-    verify(mockSpanFactory).startSpan(
+    verify(tracer).spanBuilderWithExplicitParent(
-        same(fakeClientParentSpan), eq("Sent.package1.service2.method3"),
+        eq("Sent.package1.service2.method3"), same(fakeClientParentSpan));
-        argThat(startSpanOptionsMatcher));
+    verify(spyClientSpanBuilder).setRecordEvents(eq(true));
-    verifyNoMoreInteractions(mockSpanFactory);
+    verifyNoMoreInteractions(tracer);
     assertTrue(headers.containsKey(censusTracing.tracingHeader));
 
     ServerStreamTracer serverTracer =
         censusTracing.getServerTracerFactory().newServerStreamTracer(
             method.getFullMethodName(), headers);
     verify(mockTracingPropagationHandler).fromBinaryValue(same(binarySpanContext));
-    verify(mockSpanFactory).startSpanWithRemoteParent(
+    verify(tracer).spanBuilderWithRemoteParent(
-        same(fakeServerParentSpanContext), eq("Recv.package1.service2.method3"),
+        eq("Recv.package1.service2.method3"), same(spyClientSpan.getContext()));
-        argThat(startSpanOptionsMatcher));
+    verify(spyServerSpanBuilder).setRecordEvents(eq(true));
 
     Context filteredContext = serverTracer.filterContext(Context.ROOT);
     assertSame(spyServerSpan, ContextUtils.CONTEXT_SPAN_KEY.get(filteredContext));
@@ -574,7 +549,7 @@ public class CensusModulesTest {
     Metadata headers = new Metadata();
     headers.put(censusTracing.tracingHeader, fakeClientSpanContext);
     // mockTracingPropagationHandler was stubbed to always return fakeServerParentSpanContext
-    assertSame(fakeServerParentSpanContext, headers.get(censusTracing.tracingHeader));
+    assertSame(spyClientSpan.getContext(), headers.get(censusTracing.tracingHeader));
 
     // Make BinaryPropagationHandler always throw when parsing the header
     when(mockTracingPropagationHandler.fromBinaryValue(any(byte[].class)))
@@ -584,14 +559,14 @@ public class CensusModulesTest {
     assertNull(headers.get(censusTracing.tracingHeader));
     headers.put(censusTracing.tracingHeader, fakeClientSpanContext);
     assertSame(SpanContext.INVALID, headers.get(censusTracing.tracingHeader));
-    assertNotSame(fakeServerParentSpanContext, SpanContext.INVALID);
+    assertNotSame(spyClientSpan.getContext(), SpanContext.INVALID);
 
     // A null Span is used as the parent in this case
     censusTracing.getServerTracerFactory().newServerStreamTracer(
         method.getFullMethodName(), headers);
-    verify(mockSpanFactory).startSpanWithRemoteParent(
+    verify(tracer).spanBuilderWithRemoteParent(
-        isNull(SpanContext.class), eq("Recv.package1.service2.method3"),
+        eq("Recv.package1.service2.method3"), isNull(SpanContext.class));
-        argThat(startSpanOptionsMatcher));
+    verify(spyServerSpanBuilder).setRecordEvents(eq(true));
   }
 
   @Test
@@ -641,22 +616,22 @@ public class CensusModulesTest {
   @Test
   public void serverBasicTracingNoHeaders() {
     ServerStreamTracer.Factory tracerFactory = censusTracing.getServerTracerFactory();
-    ServerStreamTracer tracer =
+    ServerStreamTracer serverStreamTracer =
         tracerFactory.newServerStreamTracer(method.getFullMethodName(), new Metadata());
     verifyZeroInteractions(mockTracingPropagationHandler);
-    verify(mockSpanFactory).startSpanWithRemoteParent(
+    verify(tracer).spanBuilderWithRemoteParent(
-        isNull(SpanContext.class), eq("Recv.package1.service2.method3"),
+        eq("Recv.package1.service2.method3"), isNull(SpanContext.class));
-        argThat(startSpanOptionsMatcher));
+    verify(spyServerSpanBuilder).setRecordEvents(eq(true));
 
-    Context filteredContext = tracer.filterContext(Context.ROOT);
+    Context filteredContext = serverStreamTracer.filterContext(Context.ROOT);
     assertSame(spyServerSpan, ContextUtils.CONTEXT_SPAN_KEY.get(filteredContext));
 
     verify(spyServerSpan, never()).end(any(EndSpanOptions.class));
-    tracer.streamClosed(Status.CANCELLED);
+    serverStreamTracer.streamClosed(Status.CANCELLED);
 
     verify(spyServerSpan).end(
         EndSpanOptions.builder()
-            .setStatus(com.google.instrumentation.trace.Status.CANCELLED).build());
+            .setStatus(io.opencensus.trace.Status.CANCELLED).build());
     verify(spyServerSpan, never()).end();
   }
 
@@ -665,7 +640,7 @@ public class CensusModulesTest {
     // Without description
     for (Status.Code grpcCode : Status.Code.values()) {
       Status grpcStatus = Status.fromCode(grpcCode);
-      com.google.instrumentation.trace.Status tracingStatus =
+      io.opencensus.trace.Status tracingStatus =
           CensusTracingModule.convertStatus(grpcStatus);
       assertEquals(grpcCode.toString(), tracingStatus.getCanonicalCode().toString());
       assertNull(tracingStatus.getDescription());
@@ -674,7 +649,7 @@ public class CensusModulesTest {
     // With description
     for (Status.Code grpcCode : Status.Code.values()) {
       Status grpcStatus = Status.fromCode(grpcCode).withDescription("This is my description");
-      com.google.instrumentation.trace.Status tracingStatus =
+      io.opencensus.trace.Status tracingStatus =
           CensusTracingModule.convertStatus(grpcStatus);
       assertEquals(grpcCode.toString(), tracingStatus.getCanonicalCode().toString());
       assertEquals(grpcStatus.getDescription(), tracingStatus.getDescription());
@@ -701,43 +676,66 @@ public class CensusModulesTest {
     assertNull(record.getMetric(RpcConstants.RPC_CLIENT_UNCOMPRESSED_RESPONSE_BYTES));
   }
 
-  // Promote the visibility of SpanFactory's methods to allow mocking
+  // TODO(bdrutu): Remove this class after OpenCensus releases support for this class.
-  private abstract static class AccessibleSpanFactory extends SpanFactory {
+  private static class MockableSpan extends Span {
-    @Override
+    private static MockableSpan generateRandomSpan(Random random) {
-    public abstract Span startSpan(@Nullable Span parent, String name, StartSpanOptions options);
+      return new MockableSpan(
-
+          SpanContext.create(
-    @Override
+              TraceId.generateRandomId(random),
-    public abstract Span startSpanWithRemoteParent(
+              SpanId.generateRandomId(random),
-        @Nullable SpanContext remoteParent, String name, StartSpanOptions options);
+              TraceOptions.DEFAULT),
-  }
+          null);
-
-  private static class FakeSpan extends Span {
-    FakeSpan(SpanContext ctx) {
-      super(ctx, null);
     }
 
     @Override
-    public void addAttributes(Map<String, AttributeValue> attributes) {
+    public void addAttributes(Map<String, AttributeValue> attributes) {}
-    }
 
     @Override
-    public void addAnnotation(String description, Map<String, AttributeValue> attributes) {
+    public void addAnnotation(String description, Map<String, AttributeValue> attributes) {}
-    }
 
     @Override
-    public void addAnnotation(Annotation annotation) {
+    public void addAnnotation(Annotation annotation) {}
-    }
 
     @Override
-    public void addNetworkEvent(NetworkEvent networkEvent) {
+    public void addNetworkEvent(NetworkEvent networkEvent) {}
-    }
 
     @Override
-    public void addLink(Link link) {
+    public void addLink(Link link) {}
-    }
 
     @Override
-    public void end(EndSpanOptions options) {
+    public void end(EndSpanOptions options) {}
+
+    private MockableSpan(SpanContext context, @Nullable EnumSet<Options> options) {
+      super(context, options);
+    }
+
+    /**
+     * Mockable implementation for the {@link SpanBuilder} class.
+     *
+     * <p>Not {@code final} to allow easy mocking.
+     *
+     */
+    public static class Builder extends SpanBuilder {
+
+      @Override
+      public SpanBuilder setSampler(Sampler sampler) {
+        return this;
+      }
+
+      @Override
+      public SpanBuilder setParentLinks(List<Span> parentLinks) {
+        return this;
+      }
+
+      @Override
+      public SpanBuilder setRecordEvents(boolean recordEvents) {
+        return this;
+      }
+
+      @Override
+      public Span startSpan() {
+        return null;
+      }
     }
   }
 }

repositories.bzl

@@ -22,7 +22,8 @@ def grpc_java_repositories(
     omit_io_netty_handler=False,
     omit_io_netty_handler_proxy=False,
     omit_io_netty_resolver=False,
-    omit_io_netty_tcnative_boringssl_static=False):
+    omit_io_netty_tcnative_boringssl_static=False,
+    omit_io_opencensus_api=False):
   """Imports dependencies for grpc-java."""
   if not omit_com_google_api_grpc_google_common_protos:
     com_google_api_grpc_google_common_protos()
@@ -68,6 +69,8 @@ def grpc_java_repositories(
     io_netty_resolver()
   if not omit_io_netty_tcnative_boringssl_static:
     io_netty_tcnative_boringssl_static()
+  if not omit_io_opencensus_api:
+    io_opencensus_api()
   native.bind(
     name = "guava",
     actual = "@com_google_guava//jar",
@@ -237,3 +240,10 @@ def io_netty_tcnative_boringssl_static():
       artifact = "io.netty:netty-tcnative-boringssl-static:2.0.5.Final",
       sha1 = "321c1239ceb3faec04531ffcdeb1bc8e85408b12",
   )
+
+def io_opencensus_api():
+  native.maven_jar(
+      name = "io_opencensus_api",
+      artifact = "io.opencensus:opencensus-api:0.5.1",
+      sha1 = "cbd0a716a7d85ac34b83d86b13f0a6655e45c2ba",
+  )