Opencensus Tracing: Start tracing unary callable. (#634)

* Start tracing unary callable.

This depends on #633 and is extracted from  #613.

This introduces TracedUnaryCallables that will create a trace and complete it. It is inserted as one of the outermost links in the callable chain to allow all other callables to contribute their annotations.
google-cloud-java that extend the chains (like cloud bigtable) will also use this class to trace their extensions.

* trace http json calls as well

* fix test

* oops

* a bit more work

* improve method name parsing

* address feedback

* trace operation cancellation and use static mockito imports in tests

* address feedback

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

@@ -30,6 +30,7 @@
 package com.google.api.gax.grpc;
 
 import com.google.api.core.BetaApi;
+import com.google.api.core.InternalApi;
 import com.google.api.gax.longrunning.OperationSnapshot;
 import com.google.api.gax.rpc.BatchingCallSettings;
 import com.google.api.gax.rpc.BidiStreamingCallable;
@@ -46,13 +47,23 @@
 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.TracedUnaryCallable;
+import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableSet;
 import com.google.longrunning.Operation;
 import com.google.longrunning.stub.OperationsStub;
+import io.grpc.MethodDescriptor;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import javax.annotation.Nonnull;
 
 /** 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.")
 public class GrpcCallableFactory {
+  // Used to extract service and method name from a grpc MethodDescriptor.
+  private static final Pattern FULL_METHOD_NAME_REGEX = Pattern.compile("^.*?([^./]+)/([^./]+)$");
+
   private GrpcCallableFactory() {}
 
   /**
@@ -90,6 +101,13 @@ 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.getMethodDescriptor()));
+
     return callable.withDefaultCallContext(clientContext.getDefaultCallContext());
   }
 
@@ -276,4 +294,12 @@ ClientStreamingCallable<RequestT, ResponseT> createClientStreamingCallable(
 
     return callable.withDefaultCallContext(clientContext.getDefaultCallContext());
   }
+
+  @InternalApi("Visible for testing")
+  static SpanName getSpanName(@Nonnull MethodDescriptor<?, ?> methodDescriptor) {
+    Matcher matcher = FULL_METHOD_NAME_REGEX.matcher(methodDescriptor.getFullMethodName());
+
+    Preconditions.checkArgument(matcher.matches(), "Invalid fullMethodName");
+    return SpanName.of(matcher.group(1), matcher.group(2));
+  }
 }

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

@@ -29,6 +29,8 @@
  */
 package com.google.api.gax.grpc;
 
+import static com.google.common.truth.Truth.assertThat;
+
 import com.google.api.gax.grpc.testing.FakeServiceGrpc;
 import com.google.api.gax.grpc.testing.FakeServiceImpl;
 import com.google.api.gax.grpc.testing.InProcessServer;
@@ -38,17 +40,24 @@
 import com.google.api.gax.rpc.ServerStreamingCallSettings;
 import com.google.api.gax.rpc.ServerStreamingCallable;
 import com.google.api.gax.rpc.StatusCode.Code;
+import com.google.api.gax.tracing.SpanName;
+import com.google.common.collect.ImmutableList;
 import com.google.common.truth.Truth;
 import com.google.type.Color;
 import com.google.type.Money;
 import io.grpc.CallOptions;
 import io.grpc.ManagedChannel;
+import io.grpc.MethodDescriptor;
+import io.grpc.MethodDescriptor.Marshaller;
+import io.grpc.MethodDescriptor.MethodType;
 import io.grpc.inprocess.InProcessChannelBuilder;
+import java.util.List;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
 import org.threeten.bp.Duration;
 
 @RunWith(JUnit4.class)
@@ -105,8 +114,8 @@ public void createServerStreamingCallableRetryableExceptions() throws Exception
     } catch (Throwable e) {
       actualError = e;
     }
-    Truth.assertThat(actualError).isInstanceOf(InvalidArgumentException.class);
-    Truth.assertThat(((InvalidArgumentException) actualError).isRetryable()).isFalse();
+    assertThat(actualError).isInstanceOf(InvalidArgumentException.class);
+    assertThat(((InvalidArgumentException) actualError).isRetryable()).isFalse();
 
     // Actual test: with config, invalid argument errors are retryable.
     ServerStreamingCallSettings<Color, Money> retryableSettings =
@@ -131,7 +140,64 @@ public void createServerStreamingCallableRetryableExceptions() throws Exception
     } catch (Throwable e) {
       actualError2 = e;
     }
-    Truth.assertThat(actualError2).isInstanceOf(InvalidArgumentException.class);
-    Truth.assertThat(((InvalidArgumentException) actualError2).isRetryable()).isTrue();
+    assertThat(actualError2).isInstanceOf(InvalidArgumentException.class);
+    assertThat(((InvalidArgumentException) actualError2).isRetryable()).isTrue();
+  }
+
+  @Test
+  public void testGetSpanName() {
+    @SuppressWarnings("unchecked")
+    MethodDescriptor descriptor =
+        MethodDescriptor.newBuilder()
+            .setType(MethodType.SERVER_STREAMING)
+            .setFullMethodName("google.bigtable.v2.Bigtable/ReadRows")
+            .setRequestMarshaller(Mockito.mock(Marshaller.class))
+            .setResponseMarshaller(Mockito.mock(Marshaller.class))
+            .build();
+
+    SpanName actualSpanName = GrpcCallableFactory.getSpanName(descriptor);
+    assertThat(actualSpanName).isEqualTo(SpanName.of("Bigtable", "ReadRows"));
+  }
+
+  @Test
+  public void testGetSpanNameUnqualified() {
+    @SuppressWarnings("unchecked")
+    MethodDescriptor descriptor =
+        MethodDescriptor.newBuilder()
+            .setType(MethodType.SERVER_STREAMING)
+            .setFullMethodName("UnqualifiedService/ReadRows")
+            .setRequestMarshaller(Mockito.mock(Marshaller.class))
+            .setResponseMarshaller(Mockito.mock(Marshaller.class))
+            .build();
+
+    SpanName actualSpanName = GrpcCallableFactory.getSpanName(descriptor);
+    assertThat(actualSpanName).isEqualTo(SpanName.of("UnqualifiedService", "ReadRows"));
+  }
+
+  @Test
+  public void testGetSpanNameInvalid() {
+    List<String> invalidNames = ImmutableList.of("BareMethod", "/MethodWithoutService");
+
+    for (String invalidName : invalidNames) {
+      @SuppressWarnings("unchecked")
+      MethodDescriptor descriptor =
+          MethodDescriptor.newBuilder()
+              .setType(MethodType.SERVER_STREAMING)
+              .setFullMethodName(invalidName)
+              .setRequestMarshaller(Mockito.mock(Marshaller.class))
+              .setResponseMarshaller(Mockito.mock(Marshaller.class))
+              .build();
+
+      IllegalArgumentException actualError = null;
+      try {
+        SpanName spanName = GrpcCallableFactory.getSpanName(descriptor);
+        Truth.assertWithMessage("Invalid method descriptor should not have a valid span name")
+            .fail("%s should not generate the spanName: %s", invalidName, spanName);
+      } catch (IllegalArgumentException e) {
+        actualError = e;
+      }
+
+      assertThat(actualError).isNotNull();
+    }
   }
 }

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

@@ -30,16 +30,27 @@
 package com.google.api.gax.httpjson;
 
 import com.google.api.core.BetaApi;
+import com.google.api.core.InternalApi;
 import com.google.api.gax.rpc.BatchingCallSettings;
 import com.google.api.gax.rpc.Callables;
 import com.google.api.gax.rpc.ClientContext;
 import com.google.api.gax.rpc.PagedCallSettings;
 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.TracedUnaryCallable;
+import com.google.common.base.Preconditions;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import javax.annotation.Nonnull;
 
 /** Class with utility methods to create http/json-based direct callables. */
 @BetaApi
 public class HttpJsonCallableFactory {
+  // Used to extract service and method name from a grpc MethodDescriptor.
+  // fullMethodName has the format: service.resource.action
+  // For example: compute.instances.addAccessConfig
+  private static final Pattern FULL_METHOD_NAME_REGEX = Pattern.compile("^(.+)\\.(.+)$");
 
   private HttpJsonCallableFactory() {}
 
@@ -74,6 +85,13 @@ public static <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> createUna
       ClientContext clientContext) {
     UnaryCallable<RequestT, ResponseT> innerCallable =
         createDirectUnaryCallable(httpJsonCallSettings);
+
+    innerCallable =
+        new TracedUnaryCallable<>(
+            innerCallable,
+            clientContext.getTracerFactory(),
+            getSpanName(httpJsonCallSettings.getMethodDescriptor()));
+
     return createUnaryCallable(innerCallable, callSettings, clientContext);
   }
 
@@ -117,4 +135,12 @@ public static <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> createBat
     callable = Callables.batching(callable, batchingCallSettings, clientContext);
     return callable.withDefaultCallContext(clientContext.getDefaultCallContext());
   }
+
+  @InternalApi("Visible for testing")
+  static SpanName getSpanName(@Nonnull ApiMethodDescriptor<?, ?> methodDescriptor) {
+    Matcher matcher = FULL_METHOD_NAME_REGEX.matcher(methodDescriptor.getFullMethodName());
+
+    Preconditions.checkArgument(matcher.matches(), "Invalid fullMethodName");
+    return SpanName.of(matcher.group(1), matcher.group(2));
+  }
 }

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

@@ -0,0 +1,96 @@
+/*
+ * Copyright 2018 Google LLC
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ *
+ *     * Redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above
+ * copyright notice, this list of conditions and the following disclaimer
+ * in the documentation and/or other materials provided with the
+ * distribution.
+ *     * Neither the name of Google LLC nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+package com.google.api.gax.httpjson;
+
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.Truth.assertWithMessage;
+
+import com.google.api.client.http.HttpMethods;
+import com.google.api.gax.tracing.SpanName;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
+
+@RunWith(JUnit4.class)
+public class HttpJsonCallableFactoryTest {
+  @Test
+  public void testGetSpanName() {
+    Map<String, SpanName> validNames =
+        ImmutableMap.of(
+            "compute.projects.disableXpnHost", SpanName.of("compute.projects", "disableXpnHost"),
+            "client.method", SpanName.of("client", "method"));
+
+    for (Entry<String, SpanName> entry : validNames.entrySet()) {
+      @SuppressWarnings("unchecked")
+      ApiMethodDescriptor descriptor =
+          ApiMethodDescriptor.newBuilder()
+              .setFullMethodName(entry.getKey())
+              .setHttpMethod(HttpMethods.POST)
+              .setRequestFormatter(Mockito.mock(HttpRequestFormatter.class))
+              .setResponseParser(Mockito.mock(HttpResponseParser.class))
+              .build();
+
+      SpanName actualSpanName = HttpJsonCallableFactory.getSpanName(descriptor);
+      assertThat(actualSpanName).isEqualTo(entry.getValue());
+    }
+  }
+
+  @Test
+  public void testGetSpanNameInvalid() {
+    List<String> invalidNames = ImmutableList.of("no_split", ".no_client");
+
+    for (String invalidName : invalidNames) {
+      @SuppressWarnings("unchecked")
+      ApiMethodDescriptor descriptor =
+          ApiMethodDescriptor.newBuilder()
+              .setFullMethodName(invalidName)
+              .setHttpMethod(HttpMethods.POST)
+              .setRequestFormatter(Mockito.mock(HttpRequestFormatter.class))
+              .setResponseParser(Mockito.mock(HttpResponseParser.class))
+              .build();
+
+      IllegalArgumentException actualError = null;
+      try {
+        SpanName spanName = HttpJsonCallableFactory.getSpanName(descriptor);
+        assertWithMessage("Invalid method descriptor should not have a valid span name")
+            .fail("%s should not generate the spanName: %s", invalidName, spanName);
+      } catch (IllegalArgumentException e) {
+        actualError = e;
+      }
+      assertThat(actualError).isNotNull();
+    }
+  }
+}

gax/src/main/java/com/google/api/gax/tracing/ApiTracer.java

@@ -56,6 +56,12 @@ public interface ApiTracer {
    */
   void operationSucceeded();
 
+  /**
+   * Signals that the operation was cancelled by the user. The tracer is now considered closed and
+   * should no longer be used.
+   */
+  void operationCancelled();
+
   /**
    * Signals that the overall operation has failed and no further attempts will be made. The tracer
    * is now considered closed and should no longer be used.

gax/src/main/java/com/google/api/gax/tracing/NoopApiTracer.java

@@ -65,6 +65,11 @@ public void operationSucceeded() {
     // noop
   }
 
+  @Override
+  public void operationCancelled() {
+    // noop
+  }
+
   @Override
   public void operationFailed(Throwable error) {
     // noop