diff --git a/core/src/main/java/io/grpc/InternalMethodDescriptor.java b/core/src/main/java/io/grpc/InternalMethodDescriptor.java index ca16bca7c0..a8788f6ec1 100644 --- a/core/src/main/java/io/grpc/InternalMethodDescriptor.java +++ b/core/src/main/java/io/grpc/InternalMethodDescriptor.java @@ -38,7 +38,10 @@ public final class InternalMethodDescriptor { md.setRawMethodName(transport.ordinal(), o); } - public static String generateTraceSpanName(boolean isServer, String fullMethodName) { - return MethodDescriptor.generateTraceSpanName(isServer, fullMethodName); + public interface RegisterForTracingCallback extends + MethodDescriptor.Registrations.RegisterForTracingCallback {} + + public static void setRegisterCallback(RegisterForTracingCallback registerCallback) { + MethodDescriptor.Registrations.setRegisterForTracingCallback(registerCallback); } } diff --git a/core/src/main/java/io/grpc/MethodDescriptor.java b/core/src/main/java/io/grpc/MethodDescriptor.java index c14e59c5b0..45f4e58945 100644 --- a/core/src/main/java/io/grpc/MethodDescriptor.java +++ b/core/src/main/java/io/grpc/MethodDescriptor.java @@ -17,12 +17,15 @@ package io.grpc; import static com.google.common.base.Preconditions.checkNotNull; +import static com.google.common.base.Preconditions.checkState; import com.google.common.base.Preconditions; -import io.opencensus.trace.Tracing; -import io.opencensus.trace.export.SampledSpanStore; import java.io.InputStream; -import java.util.ArrayList; +import java.lang.ref.ReferenceQueue; +import java.lang.ref.WeakReference; +import java.util.Collection; +import java.util.Iterator; +import java.util.LinkedList; import java.util.concurrent.atomic.AtomicReferenceArray; import javax.annotation.CheckReturnValue; import javax.annotation.Nullable; @@ -48,11 +51,16 @@ public final class MethodDescriptor { private final @Nullable Object schemaDescriptor; private final boolean idempotent; private final boolean safe; + // This field is not exposed, since the result would be misleading. Setting this value to true, + // ensures that the method is traced, but false means that it might still be traced, due to + // another descriptor tracing the same name. + private final boolean registerForTracing; // Must be set to InternalKnownTransport.values().length // Not referenced to break the dependency. private final AtomicReferenceArray rawMethodNames = new AtomicReferenceArray(1); + /** * Gets the cached "raw" method name for this Method Descriptor. The raw name is transport * specific, and should be set using {@link #setRawMethodName} by the transport. @@ -72,18 +80,6 @@ public final class MethodDescriptor { final void setRawMethodName(int transportOrdinal, Object o) { rawMethodNames.lazySet(transportOrdinal, o); } - - /** - * Convert a full method name to a tracing span name. - * - * @param isServer {@code false} if the span is on the client-side, {@code true} if on the - * server-side - * @param fullMethodName the method name as returned by {@link #getFullMethodName}. - */ - static String generateTraceSpanName(boolean isServer, String fullMethodName) { - String prefix = isServer ? "Recv" : "Sent"; - return prefix + "." + fullMethodName.replace('/', '.'); - } /** * The call type of a method. @@ -224,7 +220,7 @@ public final class MethodDescriptor { Marshaller requestMarshaller, Marshaller responseMarshaller) { return new MethodDescriptor( - type, fullMethodName, requestMarshaller, responseMarshaller, null, false, false); + type, fullMethodName, requestMarshaller, responseMarshaller, null, false, false, false); } private MethodDescriptor( @@ -234,7 +230,8 @@ public final class MethodDescriptor { Marshaller responseMarshaller, Object schemaDescriptor, boolean idempotent, - boolean safe) { + boolean safe, + boolean registerForTracing) { this.type = Preconditions.checkNotNull(type, "type"); this.fullMethodName = Preconditions.checkNotNull(fullMethodName, "fullMethodName"); @@ -243,8 +240,12 @@ public final class MethodDescriptor { this.schemaDescriptor = schemaDescriptor; this.idempotent = idempotent; this.safe = safe; + this.registerForTracing = registerForTracing; Preconditions.checkArgument(!safe || type == MethodType.UNARY, "Only unary methods can be specified safe"); + if (registerForTracing) { + Registrations.registerForTracing(this); + } } /** @@ -441,7 +442,8 @@ public final class MethodDescriptor { .setType(type) .setFullMethodName(fullMethodName) .setIdempotent(idempotent) - .setSafe(safe); + .setSafe(safe) + .setRegisterForTracing(registerForTracing); } /** @@ -563,15 +565,6 @@ public final class MethodDescriptor { */ @CheckReturnValue public MethodDescriptor build() { - if (registerForTracing) { - SampledSpanStore sampledStore = Tracing.getExportComponent().getSampledSpanStore(); - if (sampledStore != null) { - ArrayList spanNames = new ArrayList(2); - spanNames.add(generateTraceSpanName(false, fullMethodName)); - spanNames.add(generateTraceSpanName(true, fullMethodName)); - sampledStore.registerSpanNamesForCollection(spanNames); - } - } return new MethodDescriptor( type, fullMethodName, @@ -579,7 +572,73 @@ public final class MethodDescriptor { responseMarshaller, schemaDescriptor, idempotent, - safe); + safe, + registerForTracing); + } + } + + static final class Registrations { + + private static final ReferenceQueue> droppedMethodDescriptors = + new ReferenceQueue>(); + private static final Collection>> pendingRegistrations = + new LinkedList>>(); + + private static volatile RegisterForTracingCallback registerCallback; + + interface RegisterForTracingCallback { + void onRegister(MethodDescriptor md); + } + + /** + * Sets a callback for method descriptor builds. Called for descriptors with + * {@link MethodDescriptor#registerForTracing} set. This should only be called by + * {@link io.grpc.internal.CensusTracingModule} since it will only work for one invocation. + * + * @param registerCallback the callback to handle descriptor registration. + */ + static synchronized void setRegisterForTracingCallback( + RegisterForTracingCallback registerCallback) { + checkState(Registrations.registerCallback == null, "callback already present"); + Registrations.registerCallback = checkNotNull(registerCallback, "registerCallback"); + for (WeakReference> mdRef : pendingRegistrations) { + MethodDescriptor md = mdRef.get(); + if (md != null) { + mdRef.clear(); + registerCallback.onRegister(md); + } + } + drainDroppedMethodDescriptors(); + } + + private static synchronized void drainDroppedMethodDescriptors() { + boolean found = false; + while (droppedMethodDescriptors.poll() != null) { + found = true; + } + if (found) { + Iterator>> it = pendingRegistrations.iterator(); + while (it.hasNext()) { + if (it.next().get() == null) { + it.remove(); + } + } + } + } + + private static void registerForTracing(MethodDescriptor md) { + RegisterForTracingCallback reg; + if ((reg = registerCallback) == null) { + synchronized (MethodDescriptor.class) { + if ((reg = registerCallback) == null) { + pendingRegistrations.add( + new WeakReference>(md, droppedMethodDescriptors)); + drainDroppedMethodDescriptors(); + return; + } + } + } + reg.onRegister(md); } } } diff --git a/core/src/main/java/io/grpc/internal/CensusTracingModule.java b/core/src/main/java/io/grpc/internal/CensusTracingModule.java index 0b7c5445e3..c691a63765 100644 --- a/core/src/main/java/io/grpc/internal/CensusTracingModule.java +++ b/core/src/main/java/io/grpc/internal/CensusTracingModule.java @@ -29,6 +29,7 @@ import io.grpc.Context; import io.grpc.ForwardingClientCall.SimpleForwardingClientCall; import io.grpc.ForwardingClientCallListener.SimpleForwardingClientCallListener; import io.grpc.InternalMethodDescriptor; +import io.grpc.InternalMethodDescriptor.RegisterForTracingCallback; import io.grpc.Metadata; import io.grpc.MethodDescriptor; import io.grpc.ServerStreamTracer; @@ -39,7 +40,11 @@ import io.opencensus.trace.Span; import io.opencensus.trace.SpanContext; import io.opencensus.trace.Status; import io.opencensus.trace.Tracer; +import io.opencensus.trace.Tracing; +import io.opencensus.trace.export.SampledSpanStore; import io.opencensus.trace.propagation.BinaryFormat; +import java.util.ArrayList; +import java.util.List; import java.util.concurrent.atomic.AtomicBoolean; import java.util.logging.Level; import java.util.logging.Logger; @@ -65,6 +70,24 @@ final class CensusTracingModule { private final TracingClientInterceptor clientInterceptor = new TracingClientInterceptor(); private final ServerTracerFactory serverTracerFactory = new ServerTracerFactory(); + private static final boolean isRegistered = registerCensusTracer(); + + private static boolean registerCensusTracer() { + InternalMethodDescriptor.setRegisterCallback(new RegisterForTracingCallback() { + @Override + public void onRegister(MethodDescriptor md) { + SampledSpanStore sampledStore = Tracing.getExportComponent().getSampledSpanStore(); + if (sampledStore != null) { + List spanNames = new ArrayList(2); + spanNames.add(generateTraceSpanName(false, md.getFullMethodName())); + spanNames.add(generateTraceSpanName(true, md.getFullMethodName())); + sampledStore.registerSpanNamesForCollection(spanNames); + } + } + }); + return true; + } + CensusTracingModule( Tracer censusTracer, final BinaryFormat censusPropagationBinaryFormat) { this.censusTracer = checkNotNull(censusTracer, "censusTracer"); @@ -202,7 +225,7 @@ final class CensusTracingModule { this.span = censusTracer .spanBuilderWithExplicitParent( - InternalMethodDescriptor.generateTraceSpanName(false, fullMethodName), + generateTraceSpanName(false, fullMethodName), parentSpan) .setRecordEvents(true) .startSpan(); @@ -260,7 +283,7 @@ final class CensusTracingModule { this.span = censusTracer .spanBuilderWithRemoteParent( - InternalMethodDescriptor.generateTraceSpanName(true, fullMethodName), + generateTraceSpanName(true, fullMethodName), remoteSpan) .setRecordEvents(true) .startSpan(); @@ -345,4 +368,19 @@ final class CensusTracingModule { }; } } + + /** + * Convert a full method name to a tracing span name. + * + * @param isServer {@code false} if the span is on the client-side, {@code true} if on the + * server-side + * @param fullMethodName the method name as returned by + * {@link MethodDescriptor#getFullMethodName}. + */ + @VisibleForTesting + static String generateTraceSpanName(boolean isServer, String fullMethodName) { + String prefix = isServer ? "Recv" : "Sent"; + return prefix + "." + fullMethodName.replace('/', '.'); + } + } diff --git a/core/src/test/java/io/grpc/MethodDescriptorTest.java b/core/src/test/java/io/grpc/MethodDescriptorTest.java index 8612d09b35..9647214e4f 100644 --- a/core/src/test/java/io/grpc/MethodDescriptorTest.java +++ b/core/src/test/java/io/grpc/MethodDescriptorTest.java @@ -94,12 +94,4 @@ public class MethodDescriptorTest { // Never reached assert discard == null; } - - @Test - public void generateTraceSpanName() { - assertEquals( - "Sent.io.grpc.Foo", MethodDescriptor.generateTraceSpanName(false, "io.grpc/Foo")); - assertEquals( - "Recv.io.grpc.Bar", MethodDescriptor.generateTraceSpanName(true, "io.grpc/Bar")); - } } diff --git a/core/src/test/java/io/grpc/internal/CensusModulesTest.java b/core/src/test/java/io/grpc/internal/CensusModulesTest.java index 6e6e6f852f..e54e3751c6 100644 --- a/core/src/test/java/io/grpc/internal/CensusModulesTest.java +++ b/core/src/test/java/io/grpc/internal/CensusModulesTest.java @@ -725,6 +725,15 @@ public class CensusModulesTest { } } + + @Test + public void generateTraceSpanName() { + assertEquals( + "Sent.io.grpc.Foo", CensusTracingModule.generateTraceSpanName(false, "io.grpc/Foo")); + assertEquals( + "Recv.io.grpc.Bar", CensusTracingModule.generateTraceSpanName(true, "io.grpc/Bar")); + } + private static void assertNoServerContent(StatsTestUtils.MetricsRecord record) { assertNull(record.getMetric(RpcConstants.RPC_SERVER_ERROR_COUNT)); assertNull(record.getMetric(RpcConstants.RPC_SERVER_REQUEST_COUNT)); 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 300c166a1b..2071f68711 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 @@ -24,7 +24,6 @@ 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; @@ -61,8 +60,9 @@ public class SimpleServiceTest { @Test public void registerSampledMethodsForTracing() throws Exception { - // Make sure SimpleServiceGrpc class is loaded + // Make sure SimpleServiceGrpc and CensusTracingModule classes are loaded. assertNotNull(Class.forName(SimpleServiceGrpc.class.getName())); + assertNotNull(Class.forName("io.grpc.internal.CensusTracingModule")); String[] methodNames = new String[] { "grpc.testing.SimpleService/UnaryRpc", @@ -72,12 +72,20 @@ public class SimpleServiceTest { ArrayList expectedSpans = new ArrayList(); for (String methodName : methodNames) { - expectedSpans.add(InternalMethodDescriptor.generateTraceSpanName(false, methodName)); - expectedSpans.add(InternalMethodDescriptor.generateTraceSpanName(true, methodName)); + expectedSpans.add(generateTraceSpanName(false, methodName)); + expectedSpans.add(generateTraceSpanName(true, methodName)); } SampledSpanStore sampledStore = Tracing.getExportComponent().getSampledSpanStore(); Set registeredSpans = sampledStore.getRegisteredSpanNamesForCollection(); assertThat(registeredSpans).containsAllIn(expectedSpans); } + + /** + * Copy of {@link io.grpc.internal.CensusTracingModule#generateTraceSpanName} to break dependency. + */ + private static String generateTraceSpanName(boolean isServer, String fullMethodName) { + String prefix = isServer ? "Recv" : "Sent"; + return prefix + "." + fullMethodName.replace('/', '.'); + } }