core/compiler: register Span names for code-generated methods (take 2)

This is a more favorable approach than #3467. Doing the registration
in MethodDescriptor should allow us to deregister in case the
generated stub and its MethodDescriptors are garbage-collected
routinely, e.g., if they are loaded by a separate ClassLoader.

3 files changed, 36 insertions, 0 deletions

diff --git a/testing-proto/build.gradle b/testing-proto/build.gradle
index f5fc71ece..bda787f93 100644
--- a/testing-proto/build.gradle
+++ b/testing-proto/build.gradle
@@ -13,6 +13,9 @@ buildscript {
 dependencies {
     compile project(':grpc-protobuf'),
             project(':grpc-stub')
+    testCompile libraries.opencensus_api,
+                libraries.truth
+    testRuntime libraries.opencensus_impl
 }
 
 configureProtoCompilation()
diff --git a/testing-proto/src/generated/main/grpc/io/grpc/testing/protobuf/SimpleServiceGrpc.java b/testing-proto/src/generated/main/grpc/io/grpc/testing/protobuf/SimpleServiceGrpc.java
index 59d2646e2..7d8845574 100644
--- a/testing-proto/src/generated/main/grpc/io/grpc/testing/protobuf/SimpleServiceGrpc.java
+++ b/testing-proto/src/generated/main/grpc/io/grpc/testing/protobuf/SimpleServiceGrpc.java
@@ -37,6 +37,7 @@ public final class SimpleServiceGrpc {
           .setType(io.grpc.MethodDescriptor.MethodType.UNARY)
           .setFullMethodName(generateFullMethodName(
               "grpc.testing.SimpleService", "UnaryRpc"))
+          .setRegisterForTracing(true)
           .setRequestMarshaller(io.grpc.protobuf.ProtoUtils.marshaller(
               io.grpc.testing.protobuf.SimpleRequest.getDefaultInstance()))
           .setResponseMarshaller(io.grpc.protobuf.ProtoUtils.marshaller(
@@ -50,6 +51,7 @@ public final class SimpleServiceGrpc {
           .setType(io.grpc.MethodDescriptor.MethodType.CLIENT_STREAMING)
           .setFullMethodName(generateFullMethodName(
               "grpc.testing.SimpleService", "ClientStreamingRpc"))
+          .setRegisterForTracing(true)
           .setRequestMarshaller(io.grpc.protobuf.ProtoUtils.marshaller(
               io.grpc.testing.protobuf.SimpleRequest.getDefaultInstance()))
           .setResponseMarshaller(io.grpc.protobuf.ProtoUtils.marshaller(
@@ -63,6 +65,7 @@ public final class SimpleServiceGrpc {
           .setType(io.grpc.MethodDescriptor.MethodType.SERVER_STREAMING)
           .setFullMethodName(generateFullMethodName(
               "grpc.testing.SimpleService", "ServerStreamingRpc"))
+          .setRegisterForTracing(true)
           .setRequestMarshaller(io.grpc.protobuf.ProtoUtils.marshaller(
               io.grpc.testing.protobuf.SimpleRequest.getDefaultInstance()))
           .setResponseMarshaller(io.grpc.protobuf.ProtoUtils.marshaller(
@@ -76,6 +79,7 @@ public final class SimpleServiceGrpc {
           .setType(io.grpc.MethodDescriptor.MethodType.BIDI_STREAMING)
           .setFullMethodName(generateFullMethodName(
               "grpc.testing.SimpleService", "BidiStreamingRpc"))
+          .setRegisterForTracing(true)
           .setRequestMarshaller(io.grpc.protobuf.ProtoUtils.marshaller(
               io.grpc.testing.protobuf.SimpleRequest.getDefaultInstance()))
           .setResponseMarshaller(io.grpc.protobuf.ProtoUtils.marshaller(
diff --git a/testing-proto/src/test/java/io/grpc/testing/protobuf/SimpleServiceTest.java b/testing-proto/src/test/java/io/grpc/testing/protobuf/SimpleServiceTest.java
index 771b20f26..300c166a1 100644
--- a/testing-proto/src/test/java/io/grpc/testing/protobuf/SimpleServiceTest.java
+++ b/testing-proto/src/test/java/io/grpc/testing/protobuf/SimpleServiceTest.java
@@ -16,13 +16,20 @@
 
 package io.grpc.testing.protobuf;
 
+import static com.google.common.truth.Truth.assertThat;
 import static io.grpc.MethodDescriptor.MethodType.BIDI_STREAMING;
 import static io.grpc.MethodDescriptor.MethodType.CLIENT_STREAMING;
 import static io.grpc.MethodDescriptor.MethodType.SERVER_STREAMING;
 import static io.grpc.MethodDescriptor.MethodType.UNARY;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
 
+import io.grpc.InternalMethodDescriptor;
 import io.grpc.MethodDescriptor;
+import io.opencensus.trace.Tracing;
+import io.opencensus.trace.export.SampledSpanStore;
+import java.util.ArrayList;
+import java.util.Set;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -51,4 +58,26 @@ public class SimpleServiceTest {
     genericTypeShouldMatchWhenAssigned = SimpleServiceGrpc.METHOD_BIDI_STREAMING_RPC;
     assertEquals(BIDI_STREAMING, genericTypeShouldMatchWhenAssigned.getType());
   }
+
+  @Test
+  public void registerSampledMethodsForTracing() throws Exception {
+    // Make sure SimpleServiceGrpc class is loaded
+    assertNotNull(Class.forName(SimpleServiceGrpc.class.getName()));
+
+    String[] methodNames = new String[] {
+      "grpc.testing.SimpleService/UnaryRpc",
+      "grpc.testing.SimpleService/ClientStreamingRpc",
+      "grpc.testing.SimpleService/ServerStreamingRpc",
+      "grpc.testing.SimpleService/BidiStreamingRpc"};
+
+    ArrayList<String> expectedSpans = new ArrayList<String>();
+    for (String methodName : methodNames) {
+      expectedSpans.add(InternalMethodDescriptor.generateTraceSpanName(false, methodName));
+      expectedSpans.add(InternalMethodDescriptor.generateTraceSpanName(true, methodName));
+    }
+
+    SampledSpanStore sampledStore = Tracing.getExportComponent().getSampledSpanStore();
+    Set<String> registeredSpans = sampledStore.getRegisteredSpanNamesForCollection();
+    assertThat(registeredSpans).containsAllIn(expectedSpans);
+  }
 }