Add plumbing for tracing (#633)

* Add plumbing for tracing

This has been extracted from #613

This doesn't provide any extra functionality. Instead it just sets up the foundation for integrating tracing in to gax.
The general idea is to add a tracing abstraction layer that all parts of gax can call into. This is accomplished by introducing
- ApiTracer interface that has methods for all annotations that a span should contain
- ApiTracerFactory interface: since an ApiTracer is stateful and has a 1:1 mapping with operations, a factory is introduced to cleanly switch implementations of the tracer

By default a NoopApiTracerFactory is configured that will return NoopApiTracers. In the future a parallel OpenCensusApiTracerFactory will be added. A client implementation will be able to opt into tracing by setting the factory in StubSettings. The actual ApiTracer will be propagated by CallContext.

* rename tracer method to subject action

* address feedback

* address feedback

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.ApiTracer;
+import com.google.api.gax.tracing.NoopApiTracer;
 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<ApiTracer> 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();
     }
 
+    ApiTracer 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,24 @@ public GrpcCallContext withRequestParamsDynamicHeaderOption(String requestParams
     return withCallOptions(newCallOptions);
   }
 
+  /** {@inheritDoc} */
+  @Override
+  @Nonnull
+  public ApiTracer getTracer() {
+    ApiTracer tracer = callOptions.getOption(TRACER_KEY);
+    if (tracer == null) {
+      tracer = NoopApiTracer.getInstance();
+    }
+    return tracer;
+  }
+
+  /** {@inheritDoc} */
+  @Override
+  public GrpcCallContext withTracer(@Nonnull ApiTracer tracer) {
+    Preconditions.checkNotNull(tracer);
+    return withCallOptions(callOptions.withOption(TRACER_KEY, tracer));
+  }
+
   @Override
   public int hashCode() {
     return Objects.hash(

gax-grpc/src/test/java/com/google/api/gax/grpc/GrpcCallContextTest.java

@@ -35,6 +35,7 @@
 import com.google.api.gax.rpc.testing.FakeCallContext;
 import com.google.api.gax.rpc.testing.FakeChannel;
 import com.google.api.gax.rpc.testing.FakeTransportChannel;
+import com.google.api.gax.tracing.ApiTracer;
 import com.google.auth.Credentials;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.truth.Truth;
@@ -297,6 +298,28 @@ public void testMergeWithExtraHeaders() {
     Truth.assertThat(gotExtraHeaders).containsExactlyEntriesIn(expectedExtraHeaders);
   }
 
+  @Test
+  public void testMergeWithTracer() {
+    ApiTracer explicitTracer = Mockito.mock(ApiTracer.class);
+    GrpcCallContext ctxWithExplicitTracer =
+        GrpcCallContext.createDefault().withTracer(explicitTracer);
+
+    GrpcCallContext ctxWithDefaultTracer = GrpcCallContext.createDefault();
+    ApiTracer defaultTracer = ctxWithDefaultTracer.getTracer();
+
+    // Explicit tracer overrides the default tracer.
+    Truth.assertThat(ctxWithDefaultTracer.merge(ctxWithExplicitTracer).getTracer())
+        .isSameAs(explicitTracer);
+
+    // Default tracer does not override an explicit tracer.
+    Truth.assertThat(ctxWithExplicitTracer.merge(ctxWithDefaultTracer).getTracer())
+        .isSameAs(explicitTracer);
+
+    // Default tracer does not override another default tracer.
+    Truth.assertThat(ctxWithDefaultTracer.merge(GrpcCallContext.createDefault()).getTracer())
+        .isSameAs(defaultTracer);
+  }
+
   private static Map<String, List<String>> createTestExtraHeaders(String... keyValues) {
     Map<String, List<String>> extraHeaders = new HashMap<>();
     for (int i = 0; i < keyValues.length; i += 2) {

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.ApiTracer;
+import com.google.api.gax.tracing.NoopApiTracer;
 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 ApiTracer 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,
+      ApiTracer 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);
 
+    ApiTracer 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, this.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, this.tracer);
   }
 
   @BetaApi("The surface for extra headers is not stable yet and may change in the future.")
@@ -230,11 +242,31 @@ public Credentials getCredentials() {
   }
 
   public HttpJsonCallContext withChannel(HttpJsonChannel newChannel) {
-    return new HttpJsonCallContext(newChannel, timeout, deadline, credentials, extraHeaders);
+    return new HttpJsonCallContext(
+        newChannel, timeout, deadline, credentials, extraHeaders, this.tracer);
   }
 
   public HttpJsonCallContext withDeadline(Instant newDeadline) {
-    return new HttpJsonCallContext(channel, timeout, newDeadline, credentials, extraHeaders);
+    return new HttpJsonCallContext(
+        channel, timeout, newDeadline, credentials, extraHeaders, this.tracer);
+  }
+
+  @Nonnull
+  @Override
+  public ApiTracer getTracer() {
+    if (tracer == null) {
+      return NoopApiTracer.getInstance();
+    }
+    return tracer;
+  }
+
+  /** {@inheritDoc} */
+  @Override
+  public HttpJsonCallContext withTracer(@Nonnull ApiTracer newTracer) {
+    Preconditions.checkNotNull(newTracer);
+
+    return new HttpJsonCallContext(
+        channel, timeout, deadline, credentials, extraHeaders, newTracer);
   }
 
   @Override
@@ -250,11 +282,12 @@ public boolean equals(Object o) {
         && Objects.equals(timeout, that.timeout)
         && Objects.equals(deadline, that.deadline)
         && Objects.equals(credentials, that.credentials)
-        && Objects.equals(extraHeaders, that.extraHeaders);
+        && Objects.equals(extraHeaders, that.extraHeaders)
+        && Objects.equals(tracer, that.tracer);
   }
 
   @Override
   public int hashCode() {
-    return Objects.hash(channel, timeout, deadline, credentials, extraHeaders);
+    return Objects.hash(channel, timeout, deadline, credentials, extraHeaders, tracer);
   }
 }

gax-httpjson/src/test/java/com/google/api/gax/httpjson/HttpJsonCallContextTest.java

@@ -32,6 +32,7 @@
 import com.google.api.gax.rpc.testing.FakeCallContext;
 import com.google.api.gax.rpc.testing.FakeChannel;
 import com.google.api.gax.rpc.testing.FakeTransportChannel;
+import com.google.api.gax.tracing.ApiTracer;
 import com.google.auth.Credentials;
 import com.google.common.truth.Truth;
 import org.junit.Rule;
@@ -152,4 +153,26 @@ public void testMergeWithTimeout() {
 
     Truth.assertThat(ctx1.merge(ctx2).getTimeout()).isEqualTo(timeout);
   }
+
+  @Test
+  public void testMergeWithTracer() {
+    ApiTracer explicitTracer = Mockito.mock(ApiTracer.class);
+    HttpJsonCallContext ctxWithExplicitTracer =
+        HttpJsonCallContext.createDefault().withTracer(explicitTracer);
+
+    HttpJsonCallContext ctxWithDefaultTracer = HttpJsonCallContext.createDefault();
+    ApiTracer defaultTracer = ctxWithDefaultTracer.getTracer();
+
+    // Explicit tracer overrides the default tracer.
+    Truth.assertThat(ctxWithDefaultTracer.merge(ctxWithExplicitTracer).getTracer())
+        .isSameAs(explicitTracer);
+
+    // Default tracer does not override an explicit tracer.
+    Truth.assertThat(ctxWithExplicitTracer.merge(ctxWithDefaultTracer).getTracer())
+        .isSameAs(explicitTracer);
+
+    // Default tracer does not override another default tracer.
+    Truth.assertThat(ctxWithDefaultTracer.merge(HttpJsonCallContext.createDefault()).getTracer())
+        .isSameAs(defaultTracer);
+  }
 }

gax/src/main/java/com/google/api/gax/retrying/NoopRetryingContext.java

@@ -31,6 +31,11 @@
 
 // TODO(igorbernstein2): Remove this class once RetryingExecutor#createFuture(Callable) is
 // deprecated and removed.
+
+import com.google.api.gax.tracing.ApiTracer;
+import com.google.api.gax.tracing.NoopApiTracer;
+import javax.annotation.Nonnull;
+
 /**
  * Backwards compatibility class to aid in transition to adding operation state to {@link
  * RetryingFuture} implementations.
@@ -39,4 +44,11 @@ class NoopRetryingContext implements RetryingContext {
   public static RetryingContext create() {
     return new NoopRetryingContext();
   }
+
+  /** {@inheritDoc} */
+  @Nonnull
+  @Override
+  public ApiTracer getTracer() {
+    return NoopApiTracer.getInstance();
+  }
 }

gax/src/main/java/com/google/api/gax/retrying/RetryingContext.java

@@ -30,11 +30,17 @@
 package com.google.api.gax.retrying;
 
 import com.google.api.core.BetaApi;
+import com.google.api.gax.tracing.ApiTracer;
+import javax.annotation.Nonnull;
 
 /**
  * Context for a retryable operation.
  *
  * <p>It provides state to individual {@link RetryingFuture}s via the {@link RetryingExecutor}.
  */
 @BetaApi("The surface for passing per operation state is not yet stable")
-public interface RetryingContext {}
+public interface RetryingContext {
+  /** Returns the {@link ApiTracer} associated with the current operation. */
+  @Nonnull
+  ApiTracer getTracer();
+}

gax/src/main/java/com/google/api/gax/rpc/ApiCallContext.java

@@ -32,9 +32,11 @@
 import com.google.api.core.BetaApi;
 import com.google.api.core.InternalExtensionOnly;
 import com.google.api.gax.retrying.RetryingContext;
+import com.google.api.gax.tracing.ApiTracer;
 import com.google.auth.Credentials;
 import java.util.List;
 import java.util.Map;
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import org.threeten.bp.Duration;
 
@@ -130,6 +132,27 @@ public interface ApiCallContext extends RetryingContext {
   @Nullable
   Duration getStreamIdleTimeout();
 
+  /**
+   * The {@link ApiTracer} that was previously set for this context.
+   *
+   * <p>The {@link ApiTracer} will be used to trace the current operation and to annotate various
+   * events like retries.
+   */
+  @BetaApi("The surface for tracing is not stable yet and may change in the future")
+  @Nonnull
+  ApiTracer getTracer();
+
+  /**
+   * Returns a new {@link ApiCallContext} with the given {@link ApiTracer}.
+   *
+   * <p>The {@link ApiTracer} will be used to trace the current operation and to annotate various
+   * events like retries.
+   *
+   * @param tracer the {@link ApiTracer} to set.
+   */
+  @BetaApi("The surface for tracing is not stable yet and may change in the future")
+  ApiCallContext withTracer(@Nonnull ApiTracer tracer);
+
   /** If inputContext is not null, returns it; if it is null, returns the present instance. */
   ApiCallContext nullToSelf(ApiCallContext inputContext);