core/tracing: set always record events option for the RPC spans. (grp…

…c#2998)
FrancoisPoinsot · May 23, 2017 · 530b714 · 530b714
1 parent fae55a6
commit 530b714
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 12 deletions.
diff --git a/core/src/main/java/io/grpc/internal/CensusTracingModule.java b/core/src/main/java/io/grpc/internal/CensusTracingModule.java
@@ -209,7 +209,10 @@ final class ClientCallTracer extends ClientStreamTracer.Factory {
     ClientCallTracer(@Nullable Span parentSpan, String fullMethodName) {
       this.fullMethodName = checkNotNull(fullMethodName, "fullMethodName");
       this.span =
-          censusTracer.spanBuilder(parentSpan, makeSpanName("Sent", fullMethodName)).startSpan();
+          censusTracer
+              .spanBuilder(parentSpan, makeSpanName("Sent", fullMethodName))
+              .setRecordEvents(true)
+              .startSpan();
     }
 
     @Override
@@ -241,9 +244,10 @@ private final class ServerTracer extends ServerStreamTracer {
     ServerTracer(String fullMethodName, @Nullable SpanContext remoteSpan) {
       this.fullMethodName = checkNotNull(fullMethodName, "fullMethodName");
       this.span =
-          censusTracer.spanBuilderWithRemoteParent(
-              remoteSpan, makeSpanName("Recv", fullMethodName))
-          .startSpan();
+          censusTracer
+              .spanBuilderWithRemoteParent(remoteSpan, makeSpanName("Recv", fullMethodName))
+              .setRecordEvents(true)
+              .startSpan();
     }
 
     /**

diff --git a/core/src/test/java/io/grpc/internal/CensusModulesTest.java b/core/src/test/java/io/grpc/internal/CensusModulesTest.java
@@ -42,6 +42,7 @@
 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;
@@ -101,6 +102,7 @@
 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;
@@ -172,6 +174,13 @@ public String parse(InputStream stream) {
   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 =
+      new ArgumentMatcher<StartSpanOptions>() {
+        @Override
+        public boolean matches(Object argument) {
+          return Boolean.TRUE.equals(((StartSpanOptions) argument).getRecordEvents());
+        }
+      };
 
   @Rule
   public final GrpcServerRule grpcServerRule = new GrpcServerRule().directExecutor();
@@ -298,10 +307,11 @@ public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
     if (nonDefaultContext) {
       verify(mockSpanFactory).startSpan(
           same(fakeClientParentSpan), eq("Sent.package1.service2.method3"),
-          any(StartSpanOptions.class));
+          argThat(startSpanOptionsMatcher));
     } else {
       verify(mockSpanFactory).startSpan(
-          isNull(Span.class), eq("Sent.package1.service2.method3"), any(StartSpanOptions.class));
+          isNull(Span.class), eq("Sent.package1.service2.method3"),
+          argThat(startSpanOptionsMatcher));
     }
     verify(spyClientSpan, never()).end(any(EndSpanOptions.class));
 
@@ -387,7 +397,7 @@ public void clientBasicTracingDefaultSpan() {
     Metadata headers = new Metadata();
     ClientStreamTracer tracer = callTracer.newClientStreamTracer(headers);
     verify(mockSpanFactory).startSpan(
-        isNull(Span.class), eq("Sent.package1.service2.method3"), any(StartSpanOptions.class));
+        isNull(Span.class), eq("Sent.package1.service2.method3"), argThat(startSpanOptionsMatcher));
     verify(spyClientSpan, never()).end(any(EndSpanOptions.class));
 
     tracer.streamClosed(Status.OK);
@@ -431,7 +441,7 @@ public void clientStreamNeverCreatedStillRecordTracing() {
         censusTracing.newClientCallTracer(fakeClientParentSpan, method.getFullMethodName());
     verify(mockSpanFactory).startSpan(
         same(fakeClientParentSpan), eq("Sent.package1.service2.method3"),
-        any(StartSpanOptions.class));
+        argThat(startSpanOptionsMatcher));
 
     callTracer.callEnded(Status.DEADLINE_EXCEEDED.withDescription("3 seconds"));
     verify(spyClientSpan).end(
@@ -552,7 +562,7 @@ public void traceHeadersPropagateSpanContext() throws Exception {
     verifyNoMoreInteractions(mockTracingPropagationHandler);
     verify(mockSpanFactory).startSpan(
         same(fakeClientParentSpan), eq("Sent.package1.service2.method3"),
-        any(StartSpanOptions.class));
+        argThat(startSpanOptionsMatcher));
     verifyNoMoreInteractions(mockSpanFactory);
     assertTrue(headers.containsKey(censusTracing.tracingHeader));
 
@@ -562,7 +572,7 @@ public void traceHeadersPropagateSpanContext() throws Exception {
     verify(mockTracingPropagationHandler).fromBinaryValue(same(binarySpanContext));
     verify(mockSpanFactory).startSpanWithRemoteParent(
         same(fakeServerParentSpanContext), eq("Recv.package1.service2.method3"),
-        any(StartSpanOptions.class));
+        argThat(startSpanOptionsMatcher));
 
     Context filteredContext = serverTracer.filterContext(Context.ROOT);
     assertSame(spyServerSpan, ContextUtils.CONTEXT_SPAN_KEY.get(filteredContext));
@@ -591,7 +601,7 @@ public void traceHeaderMalformed() throws Exception {
         method.getFullMethodName(), headers);
     verify(mockSpanFactory).startSpanWithRemoteParent(
         isNull(SpanContext.class), eq("Recv.package1.service2.method3"),
-        any(StartSpanOptions.class));
+        argThat(startSpanOptionsMatcher));
   }
 
   @Test
@@ -646,7 +656,7 @@ public void serverBasicTracingNoHeaders() {
     verifyZeroInteractions(mockTracingPropagationHandler);
     verify(mockSpanFactory).startSpanWithRemoteParent(
         isNull(SpanContext.class), eq("Recv.package1.service2.method3"),
-        any(StartSpanOptions.class));
+        argThat(startSpanOptionsMatcher));
 
     Context filteredContext = tracer.filterContext(Context.ROOT);
     assertSame(spyServerSpan, ContextUtils.CONTEXT_SPAN_KEY.get(filteredContext));