Merge remote-tracking branch 'origin/main' into observability/tracing-attr/http-method-definition

diegomarquezp · diegomarquezp · commit ac77f0a1d144 · 2026-03-11T16:08:30.000-04:00
diff --git a/gax-java/gax/src/main/java/com/google/api/gax/tracing/ApiTracerFactory.java b/gax-java/gax/src/main/java/com/google/api/gax/tracing/ApiTracerFactory.java
@@ -62,27 +62,15 @@ enum OperationType {
    */
   ApiTracer newTracer(ApiTracer parent, SpanName spanName, OperationType operationType);
 
-  /**
-   * Create a new {@link ApiTracer} that will be a child of the current context.
-   *
-   * @param parent the parent of this tracer
-   * @param tracerContext the method-definition-specific tracer context
-   * @param operationType the type of operation that the tracer will trace
-   */
-  default ApiTracer newTracer(
-      ApiTracer parent, ApiTracerContext tracerContext, OperationType operationType) {
-    SpanName spanName = SpanName.of(tracerContext);
-    return newTracer(parent, spanName, operationType);
-  }
-
   /**
    * Create a new {@link ApiTracer} that will be a child of the current context.
    *
    * @param parent the parent of this tracer
    * @param tracerContext the method-definition-specific tracer context
    */
   default ApiTracer newTracer(ApiTracer parent, ApiTracerContext tracerContext) {
-    return newTracer(parent, tracerContext, tracerContext.operationType());
+    SpanName spanName = SpanName.of(tracerContext);
+    return newTracer(parent, spanName, tracerContext.operationType());
   }
 
   /**
diff --git a/gax-java/gax/src/main/java/com/google/api/gax/tracing/MetricsTracerFactory.java b/gax-java/gax/src/main/java/com/google/api/gax/tracing/MetricsTracerFactory.java
@@ -31,6 +31,7 @@
 
 import com.google.api.core.BetaApi;
 import com.google.api.core.InternalApi;
+import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
 import java.util.Map;
 
@@ -66,6 +67,7 @@ public MetricsTracerFactory(MetricsRecorder metricsRecorder, Map<String, String>
 
   @Override
   public ApiTracer newTracer(ApiTracer parent, SpanName spanName, OperationType operationType) {
+    Preconditions.checkNotNull(operationType, "operationType cannot be null.");
     MetricsTracer metricsTracer =
         new MetricsTracer(
             MethodName.of(spanName.getClientName(), spanName.getMethodName()), metricsRecorder);
diff --git a/gax-java/gax/src/main/java/com/google/api/gax/tracing/OpencensusTracerFactory.java b/gax-java/gax/src/main/java/com/google/api/gax/tracing/OpencensusTracerFactory.java
@@ -96,6 +96,7 @@ public OpencensusTracerFactory(Map<String, String> spanAttributes) {
   /** {@inheritDoc } */
   @Override
   public ApiTracer newTracer(ApiTracer parent, SpanName spanName, OperationType operationType) {
+    Preconditions.checkNotNull(operationType, "operationType cannot be null.");
     // Default to the current in context span. This is used for outermost tracers that inherit
     // the caller's parent span.
     Span parentSpan = internalTracer.getCurrentSpan();
diff --git a/gax-java/gax/src/main/java/com/google/api/gax/tracing/SpanTracerFactory.java b/gax-java/gax/src/main/java/com/google/api/gax/tracing/SpanTracerFactory.java
@@ -76,8 +76,7 @@ public ApiTracer newTracer(ApiTracer parent, SpanName spanName, OperationType op
   }
 
   @Override
-  public ApiTracer newTracer(
-      ApiTracer parent, ApiTracerContext apiTracerContext, OperationType operationType) {
+  public ApiTracer newTracer(ApiTracer parent, ApiTracerContext apiTracerContext) {
     ApiTracerContext mergedContext = this.apiTracerContext.merge(apiTracerContext);
 
     String attemptSpanName;
diff --git a/gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedBatchingCallable.java b/gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedBatchingCallable.java
@@ -73,7 +73,8 @@ public TracedBatchingCallable(
       ApiTracerContext apiTracerContext,
       BatchingDescriptor<RequestT, ResponseT> batchingDescriptor) {
     this.tracerFactory = tracerFactory;
-    this.apiTracerContext = apiTracerContext;
+    this.apiTracerContext =
+        apiTracerContext.toBuilder().setOperationType(OperationType.Batching).build();
     this.batchingDescriptor = batchingDescriptor;
     this.innerCallable = innerCallable;
     this.spanName = SpanName.of(apiTracerContext);
diff --git a/gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedBidiCallable.java b/gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedBidiCallable.java
@@ -74,7 +74,9 @@ public TracedBidiCallable(
       @Nonnull ApiTracerContext apiTracerContext) {
     this.tracerFactory = Preconditions.checkNotNull(tracerFactory, "tracerFactory can't be null");
     this.apiTracerContext =
-        Preconditions.checkNotNull(apiTracerContext, "apiTracerContext can't be null");
+        Preconditions.checkNotNull(apiTracerContext, "apiTracerContext can't be null").toBuilder()
+            .setOperationType(OperationType.BidiStreaming)
+            .build();
     this.innerCallable = Preconditions.checkNotNull(innerCallable, "innerCallable can't be null");
     this.spanName = SpanName.of(this.apiTracerContext);
   }
diff --git a/gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedClientStreamingCallable.java b/gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedClientStreamingCallable.java
@@ -74,7 +74,9 @@ public TracedClientStreamingCallable(
       @Nonnull ApiTracerContext apiTracerContext) {
     this.tracerFactory = Preconditions.checkNotNull(tracerFactory, "tracerFactory can't be null");
     this.apiTracerContext =
-        Preconditions.checkNotNull(apiTracerContext, "apiTracerContext can't be null");
+        Preconditions.checkNotNull(apiTracerContext, "apiTracerContext can't be null").toBuilder()
+            .setOperationType(OperationType.ClientStreaming)
+            .build();
     this.innerCallable = Preconditions.checkNotNull(innerCallable, "innerCallable can't be null");
     this.spanName = SpanName.of(this.apiTracerContext);
   }
diff --git a/gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedServerStreamingCallable.java b/gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedServerStreamingCallable.java
@@ -70,7 +70,9 @@ public TracedServerStreamingCallable(
       @Nonnull ApiTracerContext apiTracerContext) {
     this.tracerFactory = Preconditions.checkNotNull(tracerFactory, "tracerFactory can't be null");
     this.apiTracerContext =
-        Preconditions.checkNotNull(apiTracerContext, "apiTracerContext can't be null");
+        Preconditions.checkNotNull(apiTracerContext, "apiTracerContext can't be null").toBuilder()
+            .setOperationType(OperationType.ServerStreaming)
+            .build();
     this.innerCallable = Preconditions.checkNotNull(innerCallable, "innerCallable can't be null");
     this.spanName = SpanName.of(this.apiTracerContext);
   }
diff --git a/gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedUnaryCallable.java b/gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedUnaryCallable.java
@@ -68,7 +68,8 @@ public TracedUnaryCallable(
       ApiTracerContext apiTracerContext) {
     this.innerCallable = innerCallable;
     this.tracerFactory = tracerFactory;
-    this.apiTracerContext = apiTracerContext;
+    this.apiTracerContext =
+        apiTracerContext.toBuilder().setOperationType(OperationType.Unary).build();
     this.spanName = SpanName.of(apiTracerContext);
   }
 
diff --git a/gax-java/gax/src/test/java/com/google/api/gax/tracing/SpanTracerFactoryTest.java b/gax-java/gax/src/test/java/com/google/api/gax/tracing/SpanTracerFactoryTest.java
@@ -73,7 +73,7 @@ void testNewTracer_createsSpanTracer(boolean useContext) {
               .setTransport(Transport.GRPC)
               .setLibraryMetadata(LibraryMetadata.empty())
               .build();
-      tracer = factory.newTracer(null, context, OperationType.Unary);
+      tracer = factory.newTracer(null, context);
     } else {
       tracer = factory.newTracer(null, SpanName.of("service", "method"), OperationType.Unary);
     }
@@ -98,7 +98,7 @@ void testNewTracer_addsAttributes(boolean useContext) {
               .setTransport(Transport.GRPC)
               .setLibraryMetadata(LibraryMetadata.empty())
               .build();
-      tracer = factory.newTracer(null, context, OperationType.Unary);
+      tracer = factory.newTracer(null, context);
     } else {
       tracer = factory.newTracer(null, SpanName.of("service", "method"), OperationType.Unary);
     }
@@ -132,7 +132,7 @@ void testWithContext_addsInferredAttributes(boolean useContext) {
               .setTransport(Transport.GRPC)
               .setLibraryMetadata(LibraryMetadata.empty())
               .build();
-      tracer = factoryWithContext.newTracer(null, callContext, OperationType.Unary);
+      tracer = factoryWithContext.newTracer(null, callContext);
     } else {
       tracer =
           factoryWithContext.newTracer(null, SpanName.of("service", "method"), OperationType.Unary);
@@ -164,7 +164,7 @@ void testWithContext_noEndpointContext_doesNotAddServerAddressAttribute(boolean
               .setTransport(Transport.GRPC)
               .setLibraryMetadata(LibraryMetadata.empty())
               .build();
-      tracer = factoryWithContext.newTracer(null, callContext, OperationType.Unary);
+      tracer = factoryWithContext.newTracer(null, callContext);
     } else {
       tracer =
           factoryWithContext.newTracer(null, SpanName.of("service", "method"), OperationType.Unary);
@@ -190,7 +190,7 @@ void testNewTracer_withContext_grpc_usesFullMethodName() {
             .build();
 
     SpanTracerFactory factory = new SpanTracerFactory(traceManager);
-    ApiTracer tracer = factory.newTracer(null, context, OperationType.Unary);
+    ApiTracer tracer = factory.newTracer(null, context);
 
     tracer.attemptStarted(null, 1);
 
@@ -216,7 +216,7 @@ void testNewTracer_withContext_http_usesHttpMethodAndPathTemplate(
             .build();
 
     SpanTracerFactory factory = new SpanTracerFactory(traceManager);
-    ApiTracer tracer = factory.newTracer(null, context, OperationType.Unary);
+    ApiTracer tracer = factory.newTracer(null, context);
 
     tracer.attemptStarted(null, 1);
 
@@ -250,7 +250,7 @@ void testNewTracer_mergesFactoryContext() {
             .setLibraryMetadata(LibraryMetadata.empty())
             .build();
 
-    ApiTracer tracer = factory.newTracer(null, callContext, OperationType.Unary);
+    ApiTracer tracer = factory.newTracer(null, callContext);
     tracer.attemptStarted(null, 1);
 
     ArgumentCaptor<Map<String, Object>> attributesCaptor = ArgumentCaptor.forClass(Map.class);
diff --git a/gax-java/gax/src/test/java/com/google/api/gax/tracing/TracedBatchingCallableTest.java b/gax-java/gax/src/test/java/com/google/api/gax/tracing/TracedBatchingCallableTest.java
@@ -30,6 +30,7 @@
 
 package com.google.api.gax.tracing;
 
+import static com.google.common.truth.Truth.assertThat;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
@@ -46,9 +47,11 @@
 import com.google.api.gax.rpc.testing.FakeCallContext;
 import com.google.api.gax.tracing.ApiTracerContext.Transport;
 import com.google.api.gax.tracing.ApiTracerFactory.OperationType;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
 
@@ -109,6 +112,31 @@ void testRootTracerCreated(boolean useContext) {
     }
   }
 
+  @Test
+  void testOperationTypeIsSet() {
+    when(tracerFactory.newTracer(any(ApiTracer.class), any(ApiTracerContext.class)))
+        .thenReturn(tracer);
+    tracedBatchingCallable =
+        new TracedBatchingCallable<>(
+            innerCallable, tracerFactory, TRACER_CONTEXT, batchingDescriptor);
+    ApiTracerContext contextWithWrongType =
+        TRACER_CONTEXT.toBuilder().setOperationType(OperationType.Unary).build();
+
+    tracedBatchingCallable =
+        new TracedBatchingCallable<>(
+            innerCallable, tracerFactory, contextWithWrongType, batchingDescriptor);
+
+    innerResult = SettableApiFuture.create();
+    when(innerCallable.futureCall(anyString(), any(ApiCallContext.class))).thenReturn(innerResult);
+
+    tracedBatchingCallable.futureCall("test", FakeCallContext.createDefault());
+
+    ArgumentCaptor<ApiTracerContext> contextCaptor =
+        ArgumentCaptor.forClass(ApiTracerContext.class);
+    verify(tracerFactory).newTracer(any(ApiTracer.class), contextCaptor.capture());
+    assertThat(contextCaptor.getValue().operationType()).isEqualTo(OperationType.Batching);
+  }
+
   @ParameterizedTest
   @ValueSource(booleans = {false, true})
   void testBatchAttributesStamped(boolean useContext) {
diff --git a/gax-java/gax/src/test/java/com/google/api/gax/tracing/TracedBidiCallableTest.java b/gax-java/gax/src/test/java/com/google/api/gax/tracing/TracedBidiCallableTest.java
@@ -50,9 +50,11 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.CancellationException;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
 
@@ -114,6 +116,23 @@ void testTracerCreated(boolean useContext) {
     }
   }
 
+  @Test
+  void testOperationTypeIsSet() {
+    init(true);
+    ApiTracerContext contextWithWrongType =
+        TRACER_CONTEXT.toBuilder().setOperationType(OperationType.Unary).build();
+
+    innerCallable = new FakeBidiCallable();
+    tracedCallable = new TracedBidiCallable<>(innerCallable, tracerFactory, contextWithWrongType);
+
+    tracedCallable.call(new FakeBidiObserver(), FakeCallContext.createDefault());
+
+    ArgumentCaptor<ApiTracerContext> contextCaptor =
+        ArgumentCaptor.forClass(ApiTracerContext.class);
+    verify(tracerFactory).newTracer(any(ApiTracer.class), contextCaptor.capture());
+    assertThat(contextCaptor.getValue().operationType()).isEqualTo(OperationType.BidiStreaming);
+  }
+
   @ParameterizedTest
   @ValueSource(booleans = {false, true})
   void testOperationCancelled(boolean useContext) {
diff --git a/gax-java/gax/src/test/java/com/google/api/gax/tracing/TracedClientStreamingCallableTest.java b/gax-java/gax/src/test/java/com/google/api/gax/tracing/TracedClientStreamingCallableTest.java
@@ -50,9 +50,11 @@
 import com.google.common.collect.Lists;
 import java.util.List;
 import java.util.concurrent.CancellationException;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
 
@@ -106,6 +108,24 @@ void testTracerCreated(boolean useContext) {
     }
   }
 
+  @Test
+  void testOperationTypeIsSet() {
+    init(true);
+    ApiTracerContext contextWithWrongType =
+        TRACER_CONTEXT.toBuilder().setOperationType(OperationType.Unary).build();
+
+    innerCallable = new FakeClientCallable();
+    tracedCallable =
+        new TracedClientStreamingCallable<>(innerCallable, tracerFactory, contextWithWrongType);
+
+    tracedCallable.clientStreamingCall(new FakeStreamObserver(), FakeCallContext.createDefault());
+
+    ArgumentCaptor<ApiTracerContext> contextCaptor =
+        ArgumentCaptor.forClass(ApiTracerContext.class);
+    verify(tracerFactory).newTracer(any(ApiTracer.class), contextCaptor.capture());
+    assertThat(contextCaptor.getValue().operationType()).isEqualTo(OperationType.ClientStreaming);
+  }
+
   @ParameterizedTest
   @ValueSource(booleans = {false, true})
   void testCallContextPropagated(boolean useContext) {
diff --git a/gax-java/gax/src/test/java/com/google/api/gax/tracing/TracedServerStreamingCallableTest.java b/gax-java/gax/src/test/java/com/google/api/gax/tracing/TracedServerStreamingCallableTest.java
@@ -40,7 +40,6 @@
 
 import com.google.api.gax.rpc.ApiCallContext;
 import com.google.api.gax.rpc.CancelledException;
-import com.google.api.gax.rpc.LibraryMetadata;
 import com.google.api.gax.rpc.ResponseObserver;
 import com.google.api.gax.rpc.ServerStreamingCallable;
 import com.google.api.gax.rpc.StatusCode.Code;
@@ -49,8 +48,8 @@
 import com.google.api.gax.rpc.testing.MockStreamingApi.MockResponseObserver;
 import com.google.api.gax.rpc.testing.MockStreamingApi.MockServerStreamingCall;
 import com.google.api.gax.rpc.testing.MockStreamingApi.MockServerStreamingCallable;
-import com.google.api.gax.tracing.ApiTracerContext.Transport;
 import com.google.api.gax.tracing.ApiTracerFactory.OperationType;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
@@ -60,16 +59,11 @@
 @ExtendWith(MockitoExtension.class)
 class TracedServerStreamingCallableTest {
   private static final SpanName SPAN_NAME = SpanName.of("FakeClient", "FakeRpc");
-  private static final ApiTracerContext TRACER_CONTEXT =
-      ApiTracerContext.newBuilder()
-          .setFullMethodName("FakeClient/FakeRpc")
-          .setTransport(Transport.GRPC)
-          .setLibraryMetadata(LibraryMetadata.empty())
-          .setOperationType(OperationType.ServerStreaming)
-          .build();
 
   @Mock private ApiTracerFactory tracerFactory;
   @Mock private ApiTracer tracer;
+  @Mock private ApiTracerContext tracerContext;
+  @Mock private ApiTracerContext.Builder tracerContextBuilder;
 
   private MockServerStreamingCallable<String, String> innerCallable;
   private TracedServerStreamingCallable<String, String> tracedCallable;
@@ -81,10 +75,14 @@ void init(boolean useContext) {
     innerCallable = new MockServerStreamingCallable<>();
     // Wire the mock tracer factory
     if (useContext) {
+      when(tracerContext.toBuilder()).thenReturn(tracerContextBuilder);
+      when(tracerContextBuilder.setOperationType(any())).thenReturn(tracerContextBuilder);
+      when(tracerContextBuilder.build()).thenReturn(tracerContext);
       when(tracerFactory.newTracer(any(ApiTracer.class), any(ApiTracerContext.class)))
           .thenReturn(tracer);
+      when(tracerContext.fullMethodName()).thenReturn("FakeClient/FakeRpc");
       tracedCallable =
-          new TracedServerStreamingCallable<>(innerCallable, tracerFactory, TRACER_CONTEXT);
+          new TracedServerStreamingCallable<>(innerCallable, tracerFactory, tracerContext);
     } else {
       when(tracerFactory.newTracer(
               any(ApiTracer.class), any(SpanName.class), eq(OperationType.ServerStreaming)))
@@ -102,13 +100,21 @@ void testTracerCreated(boolean useContext) {
     init(useContext);
     tracedCallable.call("test", responseObserver, callContext);
     if (useContext) {
-      verify(tracerFactory, times(1)).newTracer(parentTracer, TRACER_CONTEXT);
+      verify(tracerFactory, times(1)).newTracer(parentTracer, tracerContext);
     } else {
       verify(tracerFactory, times(1))
           .newTracer(parentTracer, SPAN_NAME, OperationType.ServerStreaming);
     }
   }
 
+  @Test
+  void testOperationTypeIsSet() {
+    init(true);
+    tracedCallable.call("test", responseObserver, callContext);
+
+    verify(tracerContextBuilder).setOperationType(OperationType.ServerStreaming);
+  }
+
   @ParameterizedTest
   @ValueSource(booleans = {false, true})
   void testResponseNotify(boolean useContext) {
@@ -189,7 +195,7 @@ void testSyncError(boolean useContext) {
     // Recreate the tracedCallable using the new inner callable
     if (useContext) {
       tracedCallable =
-          new TracedServerStreamingCallable<>(innerCallable, tracerFactory, TRACER_CONTEXT);
+          new TracedServerStreamingCallable<>(innerCallable, tracerFactory, tracerContext);
     } else {
       tracedCallable = new TracedServerStreamingCallable<>(innerCallable, tracerFactory, SPAN_NAME);
     }
diff --git a/gax-java/gax/src/test/java/com/google/api/gax/tracing/TracedUnaryCallableTest.java b/gax-java/gax/src/test/java/com/google/api/gax/tracing/TracedUnaryCallableTest.java

Original file line number	Diff line number	Diff line change
`@@ -76,8 +76,7 @@ public ApiTracer newTracer(ApiTracer parent, SpanName spanName, OperationType op`
`76`	`76`	`}`
`77`	`77`
`78`	`78`	`@Override`
`79`		`- public ApiTracer newTracer(`
`80`		`- ApiTracer parent, ApiTracerContext apiTracerContext, OperationType operationType) {`
	`79`	`+ public ApiTracer newTracer(ApiTracer parent, ApiTracerContext apiTracerContext) {`
`81`	`80`	`ApiTracerContext mergedContext = this.apiTracerContext.merge(apiTracerContext);`
`82`	`81`
`83`	`82`	`String attemptSpanName;`
Original file line number	Diff line number	Diff line change
`@@ -68,7 +68,8 @@ public TracedUnaryCallable(`
`68`	`68`	`ApiTracerContext apiTracerContext) {`
`69`	`69`	`this.innerCallable = innerCallable;`
`70`	`70`	`this.tracerFactory = tracerFactory;`
`71`		`- this.apiTracerContext = apiTracerContext;`
	`71`	`+ this.apiTracerContext =`
	`72`	`+ apiTracerContext.toBuilder().setOperationType(OperationType.Unary).build();`
`72`	`73`	`this.spanName = SpanName.of(apiTracerContext);`
`73`	`74`	`}`
`74`	`75`