diff --git a/gax-java/gax/src/main/java/com/google/api/gax/tracing/GoldenSignalsMetricsTracerFactory.java b/gax-java/gax/src/main/java/com/google/api/gax/tracing/GoldenSignalsMetricsTracerFactory.java index 1ef4e1112f..53b8998a86 100644 --- a/gax-java/gax/src/main/java/com/google/api/gax/tracing/GoldenSignalsMetricsTracerFactory.java +++ b/gax-java/gax/src/main/java/com/google/api/gax/tracing/GoldenSignalsMetricsTracerFactory.java @@ -42,13 +42,13 @@ @InternalApi public class GoldenSignalsMetricsTracerFactory implements ApiTracerFactory { - private ApiTracerContext apiTracerContext; + private ApiTracerContext clientLevelTracerContext; private final OpenTelemetry openTelemetry; private GoldenSignalsMetricsRecorder metricsRecorder; public GoldenSignalsMetricsTracerFactory(OpenTelemetry openTelemetry) { this.openTelemetry = openTelemetry; - this.apiTracerContext = ApiTracerContext.empty(); + this.clientLevelTracerContext = ApiTracerContext.empty(); } @Override @@ -58,15 +58,26 @@ public ApiTracer newTracer(ApiTracer parent, SpanName spanName, OperationType op // regular requests. return new BaseApiTracer(); } - return new GoldenSignalsMetricsTracer(metricsRecorder, apiTracerContext); + return new GoldenSignalsMetricsTracer(metricsRecorder, clientLevelTracerContext); + } + + @Override + public ApiTracer newTracer(ApiTracer parent, ApiTracerContext methodLevelTracerContext) { + if (metricsRecorder == null) { + // This should never happen, in case it happens, create a no-op api tracer to not block + // regular requests. + return new BaseApiTracer(); + } + ApiTracerContext mergedTracerContext = clientLevelTracerContext.merge(methodLevelTracerContext); + return new GoldenSignalsMetricsTracer(metricsRecorder, mergedTracerContext); } @Override public ApiTracerFactory withContext(ApiTracerContext context) { - this.apiTracerContext = context; + this.clientLevelTracerContext = context; this.metricsRecorder = new GoldenSignalsMetricsRecorder( - openTelemetry, apiTracerContext.libraryMetadata().artifactName()); + openTelemetry, clientLevelTracerContext.libraryMetadata().artifactName()); return this; } } diff --git a/gax-java/gax/src/test/java/com/google/api/gax/tracing/GoldenSignalsMetricsTracerFactoryTest.java b/gax-java/gax/src/test/java/com/google/api/gax/tracing/GoldenSignalsMetricsTracerFactoryTest.java index 8a39ff27f1..9726101c07 100644 --- a/gax-java/gax/src/test/java/com/google/api/gax/tracing/GoldenSignalsMetricsTracerFactoryTest.java +++ b/gax-java/gax/src/test/java/com/google/api/gax/tracing/GoldenSignalsMetricsTracerFactoryTest.java @@ -30,7 +30,8 @@ package com.google.api.gax.tracing; import static com.google.common.truth.Truth.assertThat; -import static org.mockito.Mockito.mock; +import static org.mockito.Answers.RETURNS_DEEP_STUBS; +import static org.mockito.Mockito.*; import io.opentelemetry.api.OpenTelemetry; import org.junit.jupiter.api.BeforeEach; @@ -46,7 +47,7 @@ void setUp() { } @Test - void newTracer_createsTracer_successfully() { + void newTracerWithSpanName_shouldCreateTracer_ifMetricsRecorderIsNotNull() { tracerFactory.withContext(ApiTracerContext.empty()); ApiTracer actual = tracerFactory.newTracer( @@ -55,10 +56,32 @@ void newTracer_createsTracer_successfully() { } @Test - void newTracer_createsBaseTracer_ifMetricsRecorderIsNull() { + void newTracerWithSpanName_shouldCreateBaseTracer_ifMetricsRecorderIsNull() { ApiTracer actual = tracerFactory.newTracer( mock(ApiTracer.class), mock(SpanName.class), ApiTracerFactory.OperationType.Unary); assertThat(actual).isInstanceOf(BaseApiTracer.class); } + + @Test + void newTracerWithApiTracerContext_shouldMergeApiTracerContext() { + ApiTracerContext clientLevelTracerContext = mock(ApiTracerContext.class, RETURNS_DEEP_STUBS); + ApiTracerContext methodLevelTracerContext = mock(ApiTracerContext.class); + when(clientLevelTracerContext.libraryMetadata().artifactName()).thenReturn("does not matter"); + when(clientLevelTracerContext.merge(methodLevelTracerContext)) + .thenReturn(clientLevelTracerContext); + + tracerFactory.withContext(clientLevelTracerContext); + ApiTracer actual = tracerFactory.newTracer(mock(ApiTracer.class), methodLevelTracerContext); + + verify(clientLevelTracerContext).merge(methodLevelTracerContext); + assertThat(actual).isInstanceOf(GoldenSignalsMetricsTracer.class); + } + + @Test + void newTracerWithApiTracerContext_shouldCreateBaseTracer_ifMetricsRecorderIsNull() { + ApiTracer actual = tracerFactory.newTracer(mock(ApiTracer.class), mock(ApiTracerContext.class)); + + assertThat(actual).isInstanceOf(BaseApiTracer.class); + } }