Skip to content

Commit 04daa07

Browse files
AgraVatorAgraVator
committed
cover missing cases
1 parent 96dccf4 commit 04daa07

2 files changed

Lines changed: 256 additions & 3 deletions

File tree

‎opentelemetry/src/main/java/io/grpc/opentelemetry/OpenTelemetryMetricsModule.java‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -632,17 +632,23 @@ public void streamClosed(Status status) {
632632
}
633633
io.opentelemetry.api.common.Attributes attributes = builder.build();
634634

635+
Context ctxToRecord = otelContext;
636+
Baggage currentBaggage = BAGGAGE_KEY.get();
637+
if (currentBaggage != null && !currentBaggage.isEmpty()) {
638+
ctxToRecord = ctxToRecord.with(currentBaggage);
639+
}
640+
635641
if (module.resource.serverCallDurationCounter() != null) {
636642
module.resource.serverCallDurationCounter()
637-
.record(elapsedTimeNanos * SECONDS_PER_NANO, attributes, otelContext);
643+
.record(elapsedTimeNanos * SECONDS_PER_NANO, attributes, ctxToRecord);
638644
}
639645
if (module.resource.serverTotalSentCompressedMessageSizeCounter() != null) {
640646
module.resource.serverTotalSentCompressedMessageSizeCounter()
641-
.record(outboundWireSize, attributes, otelContext);
647+
.record(outboundWireSize, attributes, ctxToRecord);
642648
}
643649
if (module.resource.serverTotalReceivedCompressedMessageSizeCounter() != null) {
644650
module.resource.serverTotalReceivedCompressedMessageSizeCounter()
645-
.record(inboundWireSize, attributes, otelContext);
651+
.record(inboundWireSize, attributes, ctxToRecord);
646652
}
647653
}
648654
}

‎opentelemetry/src/test/java/io/grpc/opentelemetry/OpenTelemetryMetricsModuleTest.java‎

Lines changed: 247 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
package io.grpc.opentelemetry;
1818

1919
import static io.grpc.ClientStreamTracer.NAME_RESOLUTION_DELAYED;
20+
import static io.grpc.opentelemetry.internal.OpenTelemetryConstants.BAGGAGE_KEY;
2021
import static io.grpc.opentelemetry.internal.OpenTelemetryConstants.LOCALITY_KEY;
2122
import static io.grpc.opentelemetry.internal.OpenTelemetryConstants.METHOD_KEY;
2223
import static io.grpc.opentelemetry.internal.OpenTelemetryConstants.STATUS_KEY;
@@ -42,12 +43,14 @@
4243
import io.grpc.ClientInterceptor;
4344
import io.grpc.ClientInterceptors;
4445
import io.grpc.ClientStreamTracer;
46+
import io.grpc.Contexts;
4547
import io.grpc.ForwardingClientCall;
4648
import io.grpc.KnownLength;
4749
import io.grpc.Metadata;
4850
import io.grpc.MethodDescriptor;
4951
import io.grpc.ServerCall;
5052
import io.grpc.ServerCallHandler;
53+
import io.grpc.ServerInterceptor;
5154
import io.grpc.ServerServiceDefinition;
5255
import io.grpc.ServerStreamTracer;
5356
import io.grpc.ServerStreamTracer.ServerCallInfo;
@@ -59,7 +62,9 @@
5962
import io.grpc.opentelemetry.GrpcOpenTelemetry.TargetFilter;
6063
import io.grpc.opentelemetry.OpenTelemetryMetricsModule.CallAttemptsTracerFactory;
6164
import io.grpc.opentelemetry.internal.OpenTelemetryConstants;
65+
import io.grpc.stub.ClientCalls;
6266
import io.grpc.stub.MetadataUtils;
67+
import io.grpc.stub.ServerCalls;
6368
import io.grpc.stub.StreamObserver;
6469
import io.grpc.testing.GrpcServerRule;
6570
import io.grpc.testing.protobuf.SimpleRequest;
@@ -83,6 +88,7 @@
8388
import java.io.IOException;
8489
import java.io.InputStream;
8590
import java.util.Arrays;
91+
import java.util.Collections;
8692
import java.util.List;
8793
import java.util.Map;
8894
import java.util.Optional;
@@ -2091,4 +2097,245 @@ public void set(Metadata carrier, String key, String value) {
20912097
carrier.put(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER), value);
20922098
}
20932099
}
2100+
2101+
@Test
2102+
public void clientMetric_baggagePropagation_externalExecutor() throws Exception {
2103+
String target = "target:///";
2104+
ExecutorService executor = Executors.newSingleThreadExecutor();
2105+
try {
2106+
// Mock the instrument to verify Context
2107+
DoubleHistogram mockHistogram = mock(DoubleHistogram.class);
2108+
OpenTelemetryMetricsResource resource = OpenTelemetryMetricsResource.builder()
2109+
.clientAttemptDurationCounter(mockHistogram)
2110+
.build();
2111+
2112+
OpenTelemetryMetricsModule module = newOpenTelemetryMetricsModule(resource);
2113+
2114+
MethodDescriptor<String, String> methodDescriptor =
2115+
MethodDescriptor.<String, String>newBuilder()
2116+
.setType(MethodDescriptor.MethodType.UNARY)
2117+
.setFullMethodName("service/method")
2118+
.setRequestMarshaller(MARSHALLER)
2119+
.setResponseMarshaller(MARSHALLER)
2120+
.setSampledToLocalTracing(true)
2121+
.build();
2122+
2123+
ServerServiceDefinition serviceDef = ServerServiceDefinition.builder("service")
2124+
.addMethod(methodDescriptor, ServerCalls.asyncUnaryCall(
2125+
new ServerCalls.UnaryMethod<String, String>() {
2126+
@Override
2127+
public void invoke(String req, StreamObserver<String> responseObserver) {
2128+
responseObserver.onNext("Response");
2129+
responseObserver.onCompleted();
2130+
}
2131+
}))
2132+
.build();
2133+
2134+
server = InProcessServerBuilder.forName("client-baggage-test")
2135+
.directExecutor()
2136+
.addService(serviceDef)
2137+
.build().start();
2138+
2139+
InProcessChannelBuilder channelBuilder =
2140+
InProcessChannelBuilder.forName("client-baggage-test")
2141+
.executor(executor);
2142+
2143+
// Use the module's interceptor
2144+
ClientInterceptor interceptor = module.getClientInterceptor(target);
2145+
channel = channelBuilder.intercept(interceptor).build();
2146+
2147+
Baggage baggage = Baggage.builder().put("client_key", "client_value").build();
2148+
try (io.opentelemetry.context.Scope scope = Context.current().with(baggage).makeCurrent()) {
2149+
ClientCalls.blockingUnaryCall(channel, methodDescriptor, CallOptions.DEFAULT, "Request");
2150+
}
2151+
2152+
ArgumentCaptor<Context> contextCaptor = ArgumentCaptor.forClass(Context.class);
2153+
2154+
// Use atLeastOnce() and timeout to allow for async execution
2155+
verify(mockHistogram, timeout(1000).atLeastOnce())
2156+
.record(anyDouble(), any(io.opentelemetry.api.common.Attributes.class),
2157+
contextCaptor.capture());
2158+
2159+
boolean found = false;
2160+
for (Context ctx : contextCaptor.getAllValues()) {
2161+
Baggage b = Baggage.fromContext(ctx);
2162+
if ("client_value".equals(b.getEntryValue("client_key"))) {
2163+
found = true;
2164+
break;
2165+
}
2166+
}
2167+
assertTrue("Client baggage not found in metrics context", found);
2168+
} finally {
2169+
executor.shutdown();
2170+
}
2171+
}
2172+
2173+
@Test
2174+
public void serverMetric_baggagePropagation_externalExecutor() throws Exception {
2175+
ExecutorService executor = Executors.newSingleThreadExecutor();
2176+
try {
2177+
// Mock the instrument
2178+
DoubleHistogram mockHistogram = mock(DoubleHistogram.class);
2179+
OpenTelemetryMetricsResource resource = OpenTelemetryMetricsResource.builder()
2180+
.serverCallDurationCounter(mockHistogram)
2181+
.build();
2182+
2183+
// Configure module with propagation
2184+
ContextPropagators propagators =
2185+
ContextPropagators.create(W3CBaggagePropagator.getInstance());
2186+
2187+
OpenTelemetryMetricsModule module = new OpenTelemetryMetricsModule(
2188+
fakeClock.getStopwatchSupplier(), resource, emptyList(), emptyList(), propagators);
2189+
2190+
MethodDescriptor<String, String> methodDescriptor =
2191+
MethodDescriptor.<String, String>newBuilder()
2192+
.setType(MethodDescriptor.MethodType.UNARY)
2193+
.setFullMethodName("service/method")
2194+
.setRequestMarshaller(MARSHALLER)
2195+
.setResponseMarshaller(MARSHALLER)
2196+
.setSampledToLocalTracing(true)
2197+
.build();
2198+
2199+
ServerServiceDefinition serviceDef = ServerServiceDefinition.builder("service")
2200+
.addMethod(methodDescriptor, ServerCalls.asyncUnaryCall(
2201+
new ServerCalls.UnaryMethod<String, String>() {
2202+
@Override
2203+
public void invoke(String req, StreamObserver<String> responseObserver) {
2204+
responseObserver.onNext("Response");
2205+
responseObserver.onCompleted();
2206+
}
2207+
}))
2208+
.build();
2209+
2210+
// Use external executor by setting it on builder
2211+
InProcessServerBuilder serverBuilder = InProcessServerBuilder.forName("server-baggage-test")
2212+
.executor(executor)
2213+
.addService(serviceDef)
2214+
.addStreamTracerFactory(module.getServerTracerFactory());
2215+
2216+
server = serverBuilder.build().start();
2217+
2218+
channel = InProcessChannelBuilder.forName("server-baggage-test")
2219+
.directExecutor()
2220+
.build();
2221+
2222+
// We need to inject Baggage into the call.
2223+
ClientInterceptor baggageInjector = new ClientInterceptor() {
2224+
@Override
2225+
public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
2226+
MethodDescriptor<ReqT, RespT> method, CallOptions callOptions, Channel next) {
2227+
return new ForwardingClientCall.SimpleForwardingClientCall<ReqT, RespT>(
2228+
next.newCall(method, callOptions)) {
2229+
@Override
2230+
public void start(Listener<RespT> responseListener, Metadata headers) {
2231+
// Inject baggage manually into headers
2232+
Baggage baggage = Baggage.builder().put("server_key", "server_value").build();
2233+
propagators.getTextMapPropagator().inject(Context.current().with(baggage), headers,
2234+
new MetadataSetter());
2235+
super.start(responseListener, headers);
2236+
}
2237+
};
2238+
}
2239+
};
2240+
2241+
ClientCalls.blockingUnaryCall(
2242+
ClientInterceptors.intercept(channel, baggageInjector),
2243+
methodDescriptor, CallOptions.DEFAULT, "Request");
2244+
2245+
ArgumentCaptor<Context> contextCaptor = ArgumentCaptor.forClass(Context.class);
2246+
2247+
verify(mockHistogram, timeout(1000).atLeastOnce())
2248+
.record(anyDouble(), any(io.opentelemetry.api.common.Attributes.class),
2249+
contextCaptor.capture());
2250+
2251+
boolean found = false;
2252+
for (Context ctx : contextCaptor.getAllValues()) {
2253+
Baggage b = Baggage.fromContext(ctx);
2254+
if ("server_value".equals(b.getEntryValue("server_key"))) {
2255+
found = true;
2256+
break;
2257+
}
2258+
}
2259+
assertTrue("Server baggage not found in metrics context", found);
2260+
} finally {
2261+
executor.shutdown();
2262+
}
2263+
}
2264+
2265+
@Test
2266+
public void serverMetric_interceptedBaggage() throws Exception {
2267+
// This test verifies if baggage added by a ServerInterceptor is visible to
2268+
// OpenTelemetry metrics Context.
2269+
2270+
OpenTelemetryMetricsModule module = new OpenTelemetryMetricsModule(
2271+
() -> Stopwatch.createUnstarted(),
2272+
resource,
2273+
Collections.emptyList(), /* optionalLabels */
2274+
Collections.emptyList(), /* plugins */
2275+
ContextPropagators.noop());
2276+
2277+
ServerInterceptor baggageInterceptor = new ServerInterceptor() {
2278+
@Override
2279+
public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
2280+
ServerCall<ReqT, RespT> call, Metadata headers, ServerCallHandler<ReqT, RespT> next) {
2281+
// Add baggage to the context
2282+
Baggage baggage = Baggage.builder()
2283+
.put("interceptor_key", "interceptor_value")
2284+
.build();
2285+
io.grpc.Context ctx = io.grpc.Context.current().withValue(BAGGAGE_KEY, baggage);
2286+
return Contexts.interceptCall(ctx, call, headers, next);
2287+
}
2288+
};
2289+
2290+
MethodDescriptor<String, String> methodDescriptor =
2291+
MethodDescriptor.<String, String>newBuilder()
2292+
.setType(MethodDescriptor.MethodType.UNARY)
2293+
// Matching existing method name in test class
2294+
.setFullMethodName("package1.service2/method3")
2295+
.setRequestMarshaller(MARSHALLER) // Use existing MARSHALLER
2296+
.setResponseMarshaller(MARSHALLER)
2297+
.setSampledToLocalTracing(true)
2298+
.build();
2299+
2300+
ServerServiceDefinition serviceDef = ServerServiceDefinition.builder("package1.service2")
2301+
.addMethod(methodDescriptor, ServerCalls.asyncUnaryCall(
2302+
new ServerCalls.UnaryMethod<String, String>() {
2303+
@Override
2304+
public void invoke(String req, StreamObserver<String> responseObserver) {
2305+
responseObserver.onNext("Response");
2306+
responseObserver.onCompleted();
2307+
}
2308+
}))
2309+
.build();
2310+
2311+
InProcessServerBuilder serverBuilder = InProcessServerBuilder.forName("interceptor-test")
2312+
.addService(serviceDef)
2313+
.intercept(baggageInterceptor)
2314+
.addStreamTracerFactory(module.getServerTracerFactory());
2315+
2316+
server = serverBuilder.build().start();
2317+
2318+
// Use a real channel but we don't need metrics on client side for this test
2319+
InProcessChannelBuilder channelBuilder = InProcessChannelBuilder.forName("interceptor-test");
2320+
channel = channelBuilder.build();
2321+
2322+
ClientCalls.blockingUnaryCall(channel, methodDescriptor, CallOptions.DEFAULT, "Request");
2323+
2324+
// Verify that record was called with a Context containing the baggage
2325+
// Reusing contextCaptor from the class
2326+
verify(mockServerCallDurationHistogram, timeout(1000).atLeastOnce())
2327+
.record(anyDouble(), any(io.opentelemetry.api.common.Attributes.class),
2328+
contextCaptor.capture());
2329+
2330+
boolean found = false;
2331+
for (io.opentelemetry.context.Context ctx : contextCaptor.getAllValues()) {
2332+
// Baggage from OTEL Context
2333+
Baggage otelBaggage = Baggage.fromContext(ctx);
2334+
if ("interceptor_value".equals(otelBaggage.getEntryValue("interceptor_key"))) {
2335+
found = true;
2336+
break;
2337+
}
2338+
}
2339+
assertThat(found).isTrue();
2340+
}
20942341
}

0 commit comments

Comments
 (0)