WIP: Opencensus integration

build.gradle

@@ -29,6 +29,7 @@ ext {
   grpcVersion = '1.15.0'
   commonProtosVersion = '1.12.0'
   authVersion = '0.11.0'
+  opencensusVersion = '0.15.0'
   // Project names not used for release
   nonReleaseProjects = ['benchmark']
   // Project names not using the default publication configuration
@@ -109,6 +110,7 @@ subprojects {
   ext {
     grpcVersion = grpcVersion
     commonProtosVersion = commonProtosVersion
+    opencensusVersion = opencensusVersion
 
     // Shortcuts for libraries we are using
     libraries = [
@@ -126,6 +128,7 @@ subprojects {
         authCredentials: "com.google.auth:google-auth-library-credentials:${authVersion}",
         commonProtos: "com.google.api.grpc:proto-google-common-protos:${commonProtosVersion}",
         apiCommon: "com.google.api:api-common:1.7.0",
+        opencensusApi: "io.opencensus:opencensus-api:${opencensusVersion}",
 
         // Testing
         junit: 'junit:junit:4.12',

gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcCallContext.java

@@ -34,18 +34,22 @@
 import com.google.api.gax.rpc.ApiCallContext;
 import com.google.api.gax.rpc.TransportChannel;
 import com.google.api.gax.rpc.internal.Headers;
+import com.google.api.gax.tracing.NoopTracer;
+import com.google.api.gax.tracing.Tracer;
 import com.google.auth.Credentials;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
 import io.grpc.CallCredentials;
 import io.grpc.CallOptions;
+import io.grpc.CallOptions.Key;
 import io.grpc.Channel;
 import io.grpc.Deadline;
 import io.grpc.Metadata;
 import io.grpc.auth.MoreCallCredentials;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import org.threeten.bp.Duration;
 
@@ -60,6 +64,8 @@
 @BetaApi("Reference ApiCallContext instead - this class is likely to experience breaking changes")
 @InternalExtensionOnly
 public final class GrpcCallContext implements ApiCallContext {
+  private static final CallOptions.Key<Tracer> TRACER_KEY = Key.create("gax.tracer");
+
   private final Channel channel;
   private final CallOptions callOptions;
   @Nullable private final Duration timeout;
@@ -254,6 +260,11 @@ public ApiCallContext merge(ApiCallContext inputCallContext) {
       newCallCredentials = this.callOptions.getCredentials();
     }
 
+    Tracer newTracer = grpcCallContext.callOptions.getOption(TRACER_KEY);
+    if (newTracer == null) {
+      newTracer = this.callOptions.getOption(TRACER_KEY);
+    }
+
     Duration newTimeout = grpcCallContext.timeout;
     if (newTimeout == null) {
       newTimeout = this.timeout;
@@ -283,6 +294,10 @@ public ApiCallContext merge(ApiCallContext inputCallContext) {
             .withCallCredentials(newCallCredentials)
             .withDeadline(newDeadline);
 
+    if (newTracer != null) {
+      newCallOptions = newCallOptions.withOption(TRACER_KEY, newTracer);
+    }
+
     return new GrpcCallContext(
         newChannel,
         newCallOptions,
@@ -370,6 +385,22 @@ public GrpcCallContext withRequestParamsDynamicHeaderOption(String requestParams
     return withCallOptions(newCallOptions);
   }
 
+  @Nonnull
+  @Override
+  public Tracer getTracer() {
+    Tracer tracer = callOptions.getOption(TRACER_KEY);
+    if (tracer == null) {
+      tracer = NoopTracer.create();
+    }
+    return tracer;
+  }
+
+  @Override
+  public ApiCallContext withTracer(@Nonnull Tracer tracer) {
+    Preconditions.checkNotNull(tracer);
+    return withCallOptions(callOptions.withOption(TRACER_KEY, tracer));
+  }
+
   @Override
   public int hashCode() {
     return Objects.hash(

gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcCallableFactory.java

@@ -46,9 +46,17 @@
 import com.google.api.gax.rpc.StreamingCallSettings;
 import com.google.api.gax.rpc.UnaryCallSettings;
 import com.google.api.gax.rpc.UnaryCallable;
+import com.google.api.gax.tracing.SpanName;
+import com.google.api.gax.tracing.TracedBatchCallable;
+import com.google.api.gax.tracing.TracedBidiCallable;
+import com.google.api.gax.tracing.TracedClientStreamingCallable;
+import com.google.api.gax.tracing.TracedOperationCallable;
+import com.google.api.gax.tracing.TracedServerStreamingCallable;
+import com.google.api.gax.tracing.TracedUnaryCallable;
 import com.google.common.collect.ImmutableSet;
 import com.google.longrunning.Operation;
 import com.google.longrunning.stub.OperationsStub;
+import io.grpc.MethodDescriptor;
 
 /** Class with utility methods to create grpc-based direct callables. */
 @BetaApi("The surface for use by generated code is not stable yet and may change in the future.")
@@ -90,6 +98,11 @@ public static <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> createUna
       ClientContext clientContext) {
     UnaryCallable<RequestT, ResponseT> callable =
         createBaseUnaryCallable(grpcCallSettings, callSettings, clientContext);
+
+    callable =
+        new TracedUnaryCallable<>(
+            callable, clientContext.getTracerFactory(), getSpanName(grpcCallSettings));
+
     return callable.withDefaultCallContext(clientContext.getDefaultCallContext());
   }
 
@@ -109,8 +122,13 @@ UnaryCallable<RequestT, PagedListResponseT> createPagedCallable(
           ClientContext clientContext) {
     UnaryCallable<RequestT, ResponseT> innerCallable =
         createBaseUnaryCallable(grpcCallSettings, pagedCallSettings, clientContext);
+
+    UnaryCallable<RequestT, ResponseT> tracedCallable =
+        new TracedUnaryCallable<>(
+            innerCallable, clientContext.getTracerFactory(), getSpanName(grpcCallSettings));
+
     UnaryCallable<RequestT, PagedListResponseT> pagedCallable =
-        Callables.paged(innerCallable, pagedCallSettings);
+        Callables.paged(tracedCallable, pagedCallSettings);
     return pagedCallable.withDefaultCallContext(clientContext.getDefaultCallContext());
   }
 
@@ -131,6 +149,14 @@ public static <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> createBat
       ClientContext clientContext) {
     UnaryCallable<RequestT, ResponseT> callable =
         createBaseUnaryCallable(grpcCallSettings, batchingCallSettings, clientContext);
+
+    callable =
+        new TracedBatchCallable<>(
+            callable,
+            clientContext.getTracerFactory(),
+            getSpanName(grpcCallSettings),
+            batchingCallSettings.getBatchingDescriptor());
+
     callable = Callables.batching(callable, batchingCallSettings, clientContext);
     return callable.withDefaultCallContext(clientContext.getDefaultCallContext());
   }
@@ -159,11 +185,27 @@ OperationCallable<RequestT, ResponseT, MetadataT> createOperationCallable(
             grpcCallSettings, operationCallSettings.getInitialCallSettings(), clientContext);
     UnaryCallable<RequestT, OperationSnapshot> initialCallable =
         new GrpcOperationSnapshotCallable<>(initialGrpcCallable);
+
+    SpanName overallSpanName = getSpanName(grpcCallSettings);
+
+    // Wrap the start of the operation in a sub-span
+    SpanName startSpanName =
+        overallSpanName.withMethodName(overallSpanName.getMethodName() + ".Start");
+
+    UnaryCallable<RequestT, OperationSnapshot> tracedInitialCallable =
+        new TracedUnaryCallable<>(initialCallable, clientContext.getTracerFactory(), startSpanName);
+
     LongRunningClient longRunningClient = new GrpcLongRunningClient(operationsStub);
     OperationCallable<RequestT, ResponseT, MetadataT> operationCallable =
         Callables.longRunningOperation(
-            initialCallable, operationCallSettings, clientContext, longRunningClient);
-    return operationCallable.withDefaultCallContext(clientContext.getDefaultCallContext());
+            tracedInitialCallable, operationCallSettings, clientContext, longRunningClient);
+
+    // Outer span
+    TracedOperationCallable<RequestT, ResponseT, MetadataT> tracedCallable =
+        new TracedOperationCallable<>(
+            operationCallable, clientContext.getTracerFactory(), overallSpanName);
+
+    return tracedCallable.withDefaultCallContext(clientContext.getDefaultCallContext());
   }
 
   /**
@@ -188,6 +230,10 @@ BidiStreamingCallable<RequestT, ResponseT> createBidiStreamingCallable(
     callable =
         new GrpcExceptionBidiStreamingCallable<>(callable, ImmutableSet.<StatusCode.Code>of());
 
+    callable =
+        new TracedBidiCallable<>(
+            callable, clientContext.getTracerFactory(), getSpanName(grpcCallSettings));
+
     return callable.withDefaultCallContext(clientContext.getDefaultCallContext());
   }
 
@@ -249,6 +295,10 @@ ServerStreamingCallable<RequestT, ResponseT> createServerStreamingCallable(
 
     callable = Callables.retrying(callable, streamingCallSettings, clientContext);
 
+    callable =
+        new TracedServerStreamingCallable<>(
+            callable, clientContext.getTracerFactory(), getSpanName(grpcCallSettings));
+
     return callable.withDefaultCallContext(clientContext.getDefaultCallContext());
   }
 
@@ -274,6 +324,23 @@ ClientStreamingCallable<RequestT, ResponseT> createClientStreamingCallable(
     callable =
         new GrpcExceptionClientStreamingCallable<>(callable, ImmutableSet.<StatusCode.Code>of());
 
+    callable =
+        new TracedClientStreamingCallable<>(
+            callable, clientContext.getTracerFactory(), getSpanName(grpcCallSettings));
+
     return callable.withDefaultCallContext(clientContext.getDefaultCallContext());
   }
+
+  private static SpanName getSpanName(GrpcCallSettings<?, ?> grpcCallSettings) {
+    MethodDescriptor<?, ?> methodDescriptor = grpcCallSettings.getMethodDescriptor();
+
+    int index = methodDescriptor.getFullMethodName().lastIndexOf('/');
+    String fullServiceName = methodDescriptor.getFullMethodName().substring(0, index);
+    String methodName = methodDescriptor.getFullMethodName().substring(index + 1);
+
+    int serviceIndex = fullServiceName.lastIndexOf('.');
+    String clientName = fullServiceName.substring(serviceIndex + 1);
+
+    return SpanName.of(clientName, methodName);
+  }
 }

gax-grpc/src/test/java/com/google/api/gax/grpc/testing/FakeMethodDescriptor.java

@@ -39,7 +39,7 @@ public class FakeMethodDescriptor {
   private FakeMethodDescriptor() {}
 
   public static <I, O> MethodDescriptor<I, O> create() {
-    return create(MethodDescriptor.MethodType.UNARY, "(default name)");
+    return create(MethodDescriptor.MethodType.UNARY, "FakeClient/fake-method");
   }
 
   public static <I, O> MethodDescriptor<I, O> create(

gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonCallContext.java

@@ -34,6 +34,8 @@
 import com.google.api.gax.rpc.ApiCallContext;
 import com.google.api.gax.rpc.TransportChannel;
 import com.google.api.gax.rpc.internal.Headers;
+import com.google.api.gax.tracing.NoopTracer;
+import com.google.api.gax.tracing.Tracer;
 import com.google.auth.Credentials;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
@@ -61,23 +63,27 @@ public final class HttpJsonCallContext implements ApiCallContext {
   private final Instant deadline;
   private final Credentials credentials;
   private final ImmutableMap<String, List<String>> extraHeaders;
+  private final Tracer tracer;
 
   /** Returns an empty instance. */
   public static HttpJsonCallContext createDefault() {
-    return new HttpJsonCallContext(null, null, null, null, ImmutableMap.<String, List<String>>of());
+    return new HttpJsonCallContext(
+        null, null, null, null, ImmutableMap.<String, List<String>>of(), null);
   }
 
   private HttpJsonCallContext(
       HttpJsonChannel channel,
       Duration timeout,
       Instant deadline,
       Credentials credentials,
-      ImmutableMap<String, List<String>> extraHeaders) {
+      ImmutableMap<String, List<String>> extraHeaders,
+      Tracer tracer) {
     this.channel = channel;
     this.timeout = timeout;
     this.deadline = deadline;
     this.credentials = credentials;
     this.extraHeaders = extraHeaders;
+    this.tracer = tracer;
   }
 
   /**
@@ -137,14 +143,19 @@ public HttpJsonCallContext merge(ApiCallContext inputCallContext) {
     ImmutableMap<String, List<String>> newExtraHeaders =
         Headers.mergeHeaders(extraHeaders, httpJsonCallContext.extraHeaders);
 
+    Tracer newTracer = httpJsonCallContext.tracer;
+    if (newTracer == null) {
+      newTracer = this.tracer;
+    }
+
     return new HttpJsonCallContext(
-        newChannel, newTimeout, newDeadline, newCredentials, newExtraHeaders);
+        newChannel, newTimeout, newDeadline, newCredentials, newExtraHeaders, newTracer);
   }
 
   @Override
   public HttpJsonCallContext withCredentials(Credentials newCredentials) {
     return new HttpJsonCallContext(
-        this.channel, this.timeout, this.deadline, newCredentials, this.extraHeaders);
+        this.channel, this.timeout, this.deadline, newCredentials, this.extraHeaders, this.tracer);
   }
 
   @Override
@@ -171,7 +182,7 @@ public HttpJsonCallContext withTimeout(Duration timeout) {
     }
 
     return new HttpJsonCallContext(
-        this.channel, timeout, this.deadline, this.credentials, this.extraHeaders);
+        this.channel, timeout, this.deadline, this.credentials, this.extraHeaders, tracer);
   }
 
   @Nullable
@@ -208,7 +219,8 @@ public ApiCallContext withExtraHeaders(Map<String, List<String>> extraHeaders) {
     Preconditions.checkNotNull(extraHeaders);
     ImmutableMap<String, List<String>> newExtraHeaders =
         Headers.mergeHeaders(this.extraHeaders, extraHeaders);
-    return new HttpJsonCallContext(channel, timeout, deadline, credentials, newExtraHeaders);
+    return new HttpJsonCallContext(
+        channel, timeout, deadline, credentials, newExtraHeaders, tracer);
   }
 
   @BetaApi("The surface for extra headers is not stable yet and may change in the future.")
@@ -230,11 +242,30 @@ public Credentials getCredentials() {
   }
 
   public HttpJsonCallContext withChannel(HttpJsonChannel newChannel) {
-    return new HttpJsonCallContext(newChannel, timeout, deadline, credentials, extraHeaders);
+    return new HttpJsonCallContext(
+        newChannel, timeout, deadline, credentials, extraHeaders, tracer);
   }
 
   public HttpJsonCallContext withDeadline(Instant newDeadline) {
-    return new HttpJsonCallContext(channel, timeout, newDeadline, credentials, extraHeaders);
+    return new HttpJsonCallContext(
+        channel, timeout, newDeadline, credentials, extraHeaders, tracer);
+  }
+
+  @Nonnull
+  @Override
+  public Tracer getTracer() {
+    if (tracer == null) {
+      return NoopTracer.create();
+    }
+    return tracer;
+  }
+
+  @Override
+  public ApiCallContext withTracer(@Nonnull Tracer newTracer) {
+    Preconditions.checkNotNull(newTracer);
+
+    return new HttpJsonCallContext(
+        channel, timeout, deadline, credentials, extraHeaders, newTracer);
   }
 
   @Override

gax/build.gradle

@@ -19,7 +19,8 @@ dependencies {
     libraries.jsr305,
     libraries.threetenbp,
     libraries.auth,
-    libraries.apiCommon
+    libraries.apiCommon,
+    libraries.opencensusApi
 
   compileOnly libraries.autovalue