From dacd66a9f67d7d98ea64339e019a5a08c9bdba4c Mon Sep 17 00:00:00 2001 From: Anatoly Karlov Date: Fri, 17 Oct 2025 13:42:16 +0700 Subject: [PATCH] add MdcRefreshInterceptor --- libthrift/pom.xml | 2 +- pom.xml | 2 +- woody-api/pom.xml | 2 +- .../woody/api/AbstractClientBuilder.java | 1 + .../woody/api/AbstractServiceBuilder.java | 1 + .../java/dev/vality/woody/api/MDCUtils.java | 54 ++++++++++++- .../interceptor/MdcRefreshInterceptor.java | 40 +++++++++ .../api/proxy/tracer/MdcRefreshTracer.java | 32 ++++++++ .../api/proxy/tracer/TargetCallTracer.java | 6 -- .../proxy/tracer/TargetCallTracerMdcTest.java | 81 ------------------- woody-thrift/pom.xml | 2 +- .../thrift/impl/http/THClientBuilder.java | 3 +- .../thrift/impl/http/THServiceBuilder.java | 6 +- .../ext/TransportExtensionBundles.java | 24 +++++- .../impl/http/MetadataMdcPropagationTest.java | 58 ++++++++++++- 15 files changed, 211 insertions(+), 103 deletions(-) create mode 100644 woody-api/src/main/java/dev/vality/woody/api/interceptor/MdcRefreshInterceptor.java create mode 100644 woody-api/src/main/java/dev/vality/woody/api/proxy/tracer/MdcRefreshTracer.java delete mode 100644 woody-api/src/test/java/dev/vality/woody/api/proxy/tracer/TargetCallTracerMdcTest.java diff --git a/libthrift/pom.xml b/libthrift/pom.xml index 485f97d0..18bb4415 100644 --- a/libthrift/pom.xml +++ b/libthrift/pom.xml @@ -7,7 +7,7 @@ woody dev.vality.woody - 2.0.14 + 2.0.15 libthrift diff --git a/pom.xml b/pom.xml index 58b9f1bb..1b729d64 100644 --- a/pom.xml +++ b/pom.xml @@ -13,7 +13,7 @@ pom dev.vality.woody woody - 2.0.14 + 2.0.15 Woody Java Java implementation for Woody spec diff --git a/woody-api/pom.xml b/woody-api/pom.xml index e3790c99..ecb6e1bd 100644 --- a/woody-api/pom.xml +++ b/woody-api/pom.xml @@ -7,7 +7,7 @@ woody dev.vality.woody - 2.0.14 + 2.0.15 woody-api diff --git a/woody-api/src/main/java/dev/vality/woody/api/AbstractClientBuilder.java b/woody-api/src/main/java/dev/vality/woody/api/AbstractClientBuilder.java index 185a16bc..c10e6018 100644 --- a/woody-api/src/main/java/dev/vality/woody/api/AbstractClientBuilder.java +++ b/woody-api/src/main/java/dev/vality/woody/api/AbstractClientBuilder.java @@ -110,6 +110,7 @@ protected MethodCallTracer createCallTracer(Class iface, Runnable listenerStub) return new ContextTracer(traceContext, new CompositeTracer( TargetCallTracer.forClient(), + new MdcRefreshTracer(), new ErrorMappingTracer(errorMapProcessor, errDefConsumer), new EventTracer(listenerStub, getOnCallEndEventListener(), getErrorListener()), new ErrorGenTracer(errorMapProcessor), diff --git a/woody-api/src/main/java/dev/vality/woody/api/AbstractServiceBuilder.java b/woody-api/src/main/java/dev/vality/woody/api/AbstractServiceBuilder.java index 88e0ef25..5c57c1b9 100644 --- a/woody-api/src/main/java/dev/vality/woody/api/AbstractServiceBuilder.java +++ b/woody-api/src/main/java/dev/vality/woody/api/AbstractServiceBuilder.java @@ -53,6 +53,7 @@ protected T createProxyService(Class iface, T handler) { protected MethodCallTracer createEventTracer() { return new CompositeTracer( TargetCallTracer.forServer(), + new MdcRefreshTracer(), DeadlineTracer.forService(), new EventTracer(getOnCallStartEventListener(), getOnCallEndEventListener(), diff --git a/woody-api/src/main/java/dev/vality/woody/api/MDCUtils.java b/woody-api/src/main/java/dev/vality/woody/api/MDCUtils.java index 36ca5722..21fe75c7 100644 --- a/woody-api/src/main/java/dev/vality/woody/api/MDCUtils.java +++ b/woody-api/src/main/java/dev/vality/woody/api/MDCUtils.java @@ -7,6 +7,7 @@ import java.time.Instant; import java.util.HashSet; +import java.util.Iterator; import java.util.Locale; import java.util.Objects; import java.util.Set; @@ -42,9 +43,22 @@ public static void putTraceData(TraceData traceData, ContextSpan contextSpan) { populateSpanIdentifiers(contextSpan.getSpan(), otelSpan); populateOtelIdentifiers(spanContext, otelSpan); - clearExtendedEntries(false, otelSpan); + boolean updatingClientSpan = traceData.getClientSpan() == contextSpan; + boolean updatingServiceSpan = traceData.getServiceSpan() == contextSpan; + if (isExtendedFieldsEnabled()) { + if (updatingClientSpan) { + clearExtendedEntriesWithPrefix(TRACE_RPC_CLIENT_PREFIX, otelSpan); + } + if (updatingServiceSpan) { + clearExtendedEntriesWithPrefix(TRACE_RPC_SERVER_PREFIX, otelSpan); + } + if (!updatingClientSpan && !updatingServiceSpan) { + clearExtendedEntries(false, otelSpan); + } populateExtendedFields(traceData, otelSpan); + } else { + clearExtendedEntries(false, otelSpan); } updateDeadlineEntries(traceData, contextSpan, otelSpan); @@ -222,7 +236,7 @@ private static void putMdcValue(String key, String value) { MDC.put(key, value != null ? value : ""); } - private static void removeExtendedEntry(io.opentelemetry.api.trace.Span otelSpan, String key) { + public static void removeExtendedEntry(io.opentelemetry.api.trace.Span otelSpan, String key) { MDC.remove(key); EXTENDED_MDC_KEYS.get().remove(key); if (otelSpan != null) { @@ -242,14 +256,30 @@ private static void updateDeadlineEntries(TraceData traceData, ContextSpan conte } } - removeExtendedEntry(otelSpan, TRACE_RPC_CLIENT_PREFIX + "deadline"); - removeExtendedEntry(otelSpan, TRACE_RPC_SERVER_PREFIX + "deadline"); + boolean updatingClientSpan = traceData != null && traceData.getClientSpan() == contextSpan; + boolean updatingServiceSpan = traceData != null && traceData.getServiceSpan() == contextSpan; if (!isExtendedFieldsEnabled()) { + if (updatingClientSpan || (!updatingClientSpan && !updatingServiceSpan)) { + removeExtendedEntry(otelSpan, TRACE_RPC_CLIENT_PREFIX + "deadline"); + } + if (updatingServiceSpan || (!updatingClientSpan && !updatingServiceSpan)) { + removeExtendedEntry(otelSpan, TRACE_RPC_SERVER_PREFIX + "deadline"); + } return; } if (traceData != null) { + if (updatingClientSpan) { + removeExtendedEntry(otelSpan, TRACE_RPC_CLIENT_PREFIX + "deadline"); + } + if (updatingServiceSpan) { + removeExtendedEntry(otelSpan, TRACE_RPC_SERVER_PREFIX + "deadline"); + } + if (!updatingClientSpan && !updatingServiceSpan) { + removeExtendedEntry(otelSpan, TRACE_RPC_CLIENT_PREFIX + "deadline"); + removeExtendedEntry(otelSpan, TRACE_RPC_SERVER_PREFIX + "deadline"); + } addDeadlineEntry(traceData.getClientSpan(), TRACE_RPC_CLIENT_PREFIX, otelSpan); addDeadlineEntry(traceData.getServiceSpan(), TRACE_RPC_SERVER_PREFIX, otelSpan); } @@ -287,4 +317,20 @@ private static void clearExtendedEntries(boolean removeThreadLocal, io.opentelem keys.clear(); } } + + private static void clearExtendedEntriesWithPrefix(String prefix, + io.opentelemetry.api.trace.Span otelSpan) { + Set keys = EXTENDED_MDC_KEYS.get(); + Iterator iterator = keys.iterator(); + while (iterator.hasNext()) { + String key = iterator.next(); + if (key.startsWith(prefix)) { + MDC.remove(key); + if (otelSpan != null) { + otelSpan.setAttribute(key, null); + } + iterator.remove(); + } + } + } } diff --git a/woody-api/src/main/java/dev/vality/woody/api/interceptor/MdcRefreshInterceptor.java b/woody-api/src/main/java/dev/vality/woody/api/interceptor/MdcRefreshInterceptor.java new file mode 100644 index 00000000..23f06bb0 --- /dev/null +++ b/woody-api/src/main/java/dev/vality/woody/api/interceptor/MdcRefreshInterceptor.java @@ -0,0 +1,40 @@ +package dev.vality.woody.api.interceptor; + +import dev.vality.woody.api.MDCUtils; +import dev.vality.woody.api.trace.ContextSpan; +import dev.vality.woody.api.trace.TraceData; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class MdcRefreshInterceptor implements CommonInterceptor { + + private static final Logger LOG = LoggerFactory.getLogger(MdcRefreshInterceptor.class); + + @Override + public boolean interceptRequest(TraceData traceData, Object providerContext, Object... contextParams) { + LOG.trace("MDC refresh on request phase"); + refresh(traceData); + return true; + } + + @Override + public boolean interceptResponse(TraceData traceData, Object providerContext, Object... contextParams) { + LOG.trace("MDC refresh on response phase"); + refresh(traceData); + return true; + } + + public static void refresh(TraceData traceData) { + if (traceData == null) { + MDCUtils.removeTraceData(); + return; + } + ContextSpan activeSpan = traceData.getActiveSpan(); + if (activeSpan == null || !activeSpan.isFilled()) { + LOG.trace("Active span is not filled; skipping MDC refresh"); + return; + } + MDCUtils.putTraceData(traceData, activeSpan); + } + +} \ No newline at end of file diff --git a/woody-api/src/main/java/dev/vality/woody/api/proxy/tracer/MdcRefreshTracer.java b/woody-api/src/main/java/dev/vality/woody/api/proxy/tracer/MdcRefreshTracer.java new file mode 100644 index 00000000..3a1d7889 --- /dev/null +++ b/woody-api/src/main/java/dev/vality/woody/api/proxy/tracer/MdcRefreshTracer.java @@ -0,0 +1,32 @@ +package dev.vality.woody.api.proxy.tracer; + +import dev.vality.woody.api.interceptor.MdcRefreshInterceptor; +import dev.vality.woody.api.proxy.InstanceMethodCaller; +import dev.vality.woody.api.trace.TraceData; +import dev.vality.woody.api.trace.context.TraceContext; + +/** + * Ensures MDC stays in sync around method invocations once metadata is populated. + */ +public class MdcRefreshTracer implements MethodCallTracer { + + @Override + public void beforeCall(Object[] args, InstanceMethodCaller caller) { + refresh(); + } + + @Override + public void afterCall(Object[] args, InstanceMethodCaller caller, Object result) { + refresh(); + } + + @Override + public void callError(Object[] args, InstanceMethodCaller caller, Throwable error) { + refresh(); + } + + private void refresh() { + TraceData traceData = TraceContext.getCurrentTraceData(); + MdcRefreshInterceptor.refresh(traceData); + } +} diff --git a/woody-api/src/main/java/dev/vality/woody/api/proxy/tracer/TargetCallTracer.java b/woody-api/src/main/java/dev/vality/woody/api/proxy/tracer/TargetCallTracer.java index 87469afd..5994cf2d 100644 --- a/woody-api/src/main/java/dev/vality/woody/api/proxy/tracer/TargetCallTracer.java +++ b/woody-api/src/main/java/dev/vality/woody/api/proxy/tracer/TargetCallTracer.java @@ -1,6 +1,5 @@ package dev.vality.woody.api.proxy.tracer; -import dev.vality.woody.api.MDCUtils; import dev.vality.woody.api.event.ClientEventType; import dev.vality.woody.api.event.ServiceEventType; import dev.vality.woody.api.proxy.InstanceMethodCaller; @@ -64,11 +63,6 @@ private void setBeforeCall(Metadata metadata, Object[] args, InstanceMethodCalle metadata.putValue(MetadataProperties.INSTANCE_METHOD_CALLER, caller); metadata.putValue(MetadataProperties.EVENT_TYPE, isClient ? ClientEventType.CALL_SERVICE : ServiceEventType.CALL_HANDLER); - TraceData currentTraceData = TraceContext.getCurrentTraceData(); - ContextSpan activeSpan = currentTraceData != null ? currentTraceData.getActiveSpan() : null; - if (currentTraceData != null && activeSpan != null) { - MDCUtils.putTraceData(currentTraceData, activeSpan); - } } private void setAfterCall(Metadata metadata, Object[] args, InstanceMethodCaller caller, Object result, diff --git a/woody-api/src/test/java/dev/vality/woody/api/proxy/tracer/TargetCallTracerMdcTest.java b/woody-api/src/test/java/dev/vality/woody/api/proxy/tracer/TargetCallTracerMdcTest.java deleted file mode 100644 index 0369eb87..00000000 --- a/woody-api/src/test/java/dev/vality/woody/api/proxy/tracer/TargetCallTracerMdcTest.java +++ /dev/null @@ -1,81 +0,0 @@ -package dev.vality.woody.api.proxy.tracer; - -import dev.vality.woody.api.MDCUtils; -import dev.vality.woody.api.event.ServiceEventType; -import dev.vality.woody.api.proxy.InstanceMethodCaller; -import dev.vality.woody.api.trace.ContextSpan; -import dev.vality.woody.api.trace.MetadataProperties; -import dev.vality.woody.api.trace.TraceData; -import dev.vality.woody.api.trace.context.TraceContext; -import org.junit.After; -import org.junit.Before; -import org.junit.Test; -import org.slf4j.MDC; - -import java.lang.reflect.Method; - -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertNull; - -public class TargetCallTracerMdcTest { - - private TraceData originalTraceData; - private TraceContext traceContext; - - @Before - public void setUp() { - originalTraceData = TraceContext.getCurrentTraceData(); - TraceContext.setCurrentTraceData(new TraceData()); - MDC.clear(); - } - - @After - public void tearDown() { - try { - if (traceContext != null) { - traceContext.destroy(false); - } - } finally { - TraceContext.setCurrentTraceData(originalTraceData); - MDCUtils.removeTraceData(); - } - } - - @Test - public void shouldPopulateRpcFieldsAfterTargetCallTracer() throws Exception { - TraceData traceData = TraceContext.getCurrentTraceData(); - ContextSpan serviceSpan = traceData.getServiceSpan(); - serviceSpan.getSpan().setTraceId("trace-1"); - serviceSpan.getSpan().setParentId("parent-1"); - serviceSpan.getSpan().setId("span-1"); - serviceSpan.getMetadata().putValue(MetadataProperties.CALL_NAME, "ServerCall"); - - traceContext = TraceContext.forService(); - traceContext.init(); - - assertNull("RPC service field must not be populated before TargetCallTracer", - MDC.get(MDCUtils.TRACE_RPC_SERVER_PREFIX + "service")); - - InstanceMethodCaller caller = createCaller("sampleHandler"); - TargetCallTracer.forServer().beforeCall(new Object[0], caller); - - assertEquals("SampleService", MDC.get(MDCUtils.TRACE_RPC_SERVER_PREFIX + "service")); - assertEquals("ServerCall", MDC.get(MDCUtils.TRACE_RPC_SERVER_PREFIX + "function")); - assertEquals("call handler", MDC.get(MDCUtils.TRACE_RPC_SERVER_PREFIX + "event")); - } - - private InstanceMethodCaller createCaller(String methodName) throws Exception { - Method method = SampleService.class.getDeclaredMethod(methodName); - return new InstanceMethodCaller(method) { - @Override - public Object call(Object source, Object[] args) { - return null; - } - }; - } - - private static class SampleService { - public void sampleHandler() { - } - } -} diff --git a/woody-thrift/pom.xml b/woody-thrift/pom.xml index eb0983e9..ef8b4468 100644 --- a/woody-thrift/pom.xml +++ b/woody-thrift/pom.xml @@ -7,7 +7,7 @@ woody dev.vality.woody - 2.0.14 + 2.0.15 woody-thrift diff --git a/woody-thrift/src/main/java/dev/vality/woody/thrift/impl/http/THClientBuilder.java b/woody-thrift/src/main/java/dev/vality/woody/thrift/impl/http/THClientBuilder.java index 0ff5a3a5..2e87b4d7 100644 --- a/woody-thrift/src/main/java/dev/vality/woody/thrift/impl/http/THClientBuilder.java +++ b/woody-thrift/src/main/java/dev/vality/woody/thrift/impl/http/THClientBuilder.java @@ -11,6 +11,7 @@ import dev.vality.woody.api.interceptor.CommonInterceptor; import dev.vality.woody.api.interceptor.CompositeInterceptor; import dev.vality.woody.api.interceptor.ContainerCommonInterceptor; +import dev.vality.woody.api.interceptor.MdcRefreshInterceptor; import dev.vality.woody.api.interceptor.ext.ExtensionBundle; import dev.vality.woody.api.provider.ProviderEventInterceptor; import dev.vality.woody.api.proxy.InvocationTargetProvider; @@ -24,7 +25,6 @@ import dev.vality.woody.thrift.impl.http.interceptor.THMessageInterceptor; import dev.vality.woody.thrift.impl.http.interceptor.THTransportInterceptor; import dev.vality.woody.thrift.impl.http.interceptor.ext.MetadataExtensionBundle; -import io.opentelemetry.sdk.resources.Resource; import org.apache.hc.client5.http.classic.HttpClient; import org.apache.hc.client5.http.impl.classic.CloseableHttpClient; import org.apache.hc.client5.http.impl.classic.HttpClients; @@ -245,6 +245,7 @@ protected CommonInterceptor createTransportInterceptor() { return new CompositeInterceptor( new ContainerCommonInterceptor(new THTransportInterceptor(extensionBundles, true, true), new THTransportInterceptor(extensionBundles, true, false)), + new MdcRefreshInterceptor(), new TransportEventInterceptor(getOnSendEventListener(), getOnReceiveEventListener(), null)); } diff --git a/woody-thrift/src/main/java/dev/vality/woody/thrift/impl/http/THServiceBuilder.java b/woody-thrift/src/main/java/dev/vality/woody/thrift/impl/http/THServiceBuilder.java index b9482795..ac9909ee 100644 --- a/woody-thrift/src/main/java/dev/vality/woody/thrift/impl/http/THServiceBuilder.java +++ b/woody-thrift/src/main/java/dev/vality/woody/thrift/impl/http/THServiceBuilder.java @@ -1,7 +1,6 @@ package dev.vality.woody.thrift.impl.http; import dev.vality.woody.api.AbstractServiceBuilder; -import dev.vality.woody.api.ServiceBuilder; import dev.vality.woody.api.event.CompositeServiceEventListener; import dev.vality.woody.api.event.ServiceEventListener; import dev.vality.woody.api.flow.error.WErrorDefinition; @@ -19,7 +18,6 @@ import dev.vality.woody.thrift.impl.http.interceptor.THMessageInterceptor; import dev.vality.woody.thrift.impl.http.interceptor.THTransportInterceptor; import dev.vality.woody.thrift.impl.http.interceptor.ext.MetadataExtensionBundle; -import io.opentelemetry.sdk.resources.Resource; import org.apache.thrift.TProcessor; import org.apache.thrift.protocol.TBinaryProtocol; import org.apache.thrift.protocol.TProtocolFactory; @@ -136,6 +134,10 @@ protected CommonInterceptor createInterceptor(THErrorMapProcessor errorMapProces isTransportLevel ? new THTransportInterceptor(extensionBundles, false, true) : null, new THTransportInterceptor(extensionBundles, false, false))); + if (isTransportLevel) { + interceptors.add(new MdcRefreshInterceptor()); + } + if (isTransportLevel) { //interceptors.add(new ProviderEventInterceptor(getOnSendEventListener(), null)); interceptors.add(new ContextInterceptor(TraceContext.forService(), diff --git a/woody-thrift/src/main/java/dev/vality/woody/thrift/impl/http/interceptor/ext/TransportExtensionBundles.java b/woody-thrift/src/main/java/dev/vality/woody/thrift/impl/http/interceptor/ext/TransportExtensionBundles.java index 6124dcfd..dedf74e4 100644 --- a/woody-thrift/src/main/java/dev/vality/woody/thrift/impl/http/interceptor/ext/TransportExtensionBundles.java +++ b/woody-thrift/src/main/java/dev/vality/woody/thrift/impl/http/interceptor/ext/TransportExtensionBundles.java @@ -21,9 +21,11 @@ import io.opentelemetry.semconv.HttpAttributes; import jakarta.servlet.http.HttpServletRequest; import jakarta.servlet.http.HttpServletResponse; +import org.apache.hc.client5.http.classic.methods.HttpUriRequestBase; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.net.URISyntaxException; import java.net.URL; import java.time.Instant; import java.time.format.DateTimeParseException; @@ -84,9 +86,8 @@ public class TransportExtensionBundles { public static final ExtensionBundle CALL_ENDPOINT_BUNDLE = createExtBundle(createCtxBundle((InterceptorExtension) reqCCtx -> { ContextSpan contextSpan = reqCCtx.getTraceData().getClientSpan(); - URL url = reqCCtx.getRequestCallEndpoint(); - contextSpan.getMetadata().putValue(MetadataProperties.CALL_ENDPOINT, - new UrlStringEndpoint(url == null ? null : url.toString())); + String endpoint = resolveClientEndpoint(reqCCtx); + contextSpan.getMetadata().putValue(MetadataProperties.CALL_ENDPOINT, new UrlStringEndpoint(endpoint)); }, respCCtx -> { }), createCtxBundle((InterceptorExtension) reqSCtx -> { HttpServletRequest request = reqSCtx.getProviderRequest(); @@ -100,6 +101,23 @@ public class TransportExtensionBundles { }, reqSCtx -> { })); + private static String resolveClientEndpoint(THCExtensionContext reqCCtx) { + URL url = reqCCtx.getRequestCallEndpoint(); + if (url != null) { + return url.toString(); + } + Object providerContext = reqCCtx.getProviderContext(); + if (providerContext instanceof HttpUriRequestBase) { + HttpUriRequestBase request = (HttpUriRequestBase) providerContext; + try { + return request.getUri() != null ? request.getUri().toString() : null; + } catch (URISyntaxException e) { + throw new RuntimeException("Failed to resolve client endpoint URI", e); + } + } + return null; + } + public static final ExtensionBundle TRANSPORT_INJECTION_BUNDLE = createExtBundle(createCtxBundle((InterceptorExtension) reqCCtx -> { reqCCtx.getTraceData().getClientSpan().getMetadata() diff --git a/woody-thrift/src/test/java/dev/vality/woody/thrift/impl/http/MetadataMdcPropagationTest.java b/woody-thrift/src/test/java/dev/vality/woody/thrift/impl/http/MetadataMdcPropagationTest.java index 811f61f1..e83b05d1 100644 --- a/woody-thrift/src/test/java/dev/vality/woody/thrift/impl/http/MetadataMdcPropagationTest.java +++ b/woody-thrift/src/test/java/dev/vality/woody/thrift/impl/http/MetadataMdcPropagationTest.java @@ -1,5 +1,8 @@ package dev.vality.woody.thrift.impl.http; +import dev.vality.woody.api.MDCUtils; +import dev.vality.woody.api.event.ClientEventListener; +import dev.vality.woody.api.event.ClientEventType; import dev.vality.woody.api.event.ServiceEventListener; import dev.vality.woody.api.generator.TimestampIdGenerator; import dev.vality.woody.api.trace.ContextUtils; @@ -7,6 +10,7 @@ import dev.vality.woody.api.trace.context.metadata.MetadataExtensionKit; import dev.vality.woody.rpc.Owner; import dev.vality.woody.rpc.OwnerServiceSrv; +import dev.vality.woody.thrift.impl.http.event.THClientEvent; import jakarta.servlet.Servlet; import org.apache.hc.client5.http.impl.classic.CloseableHttpClient; import org.apache.hc.client5.http.impl.classic.HttpClients; @@ -41,6 +45,9 @@ public class MetadataMdcPropagationTest extends AbstractTest { private final AtomicReference downstreamMetadataDeadline = new AtomicReference<>(); private final AtomicReference downstreamMdcId = new AtomicReference<>(); private final AtomicReference downstreamMdcDeadline = new AtomicReference<>(); + private final AtomicReference downstreamRpcServerService = new AtomicReference<>(); + private final AtomicReference downstreamRpcServerFunction = new AtomicReference<>(); + private final AtomicReference downstreamRpcServerUrl = new AtomicReference<>(); private final AtomicReference upstreamOtelTraceId = new AtomicReference<>(); private final AtomicReference downstreamOtelTraceId = new AtomicReference<>(); private final AtomicReference upstreamTraceState = new AtomicReference<>(); @@ -49,6 +56,13 @@ public class MetadataMdcPropagationTest extends AbstractTest { private final AtomicReference downstreamTraceParent = new AtomicReference<>(); private final AtomicReference responseTraceParent = new AtomicReference<>(); private final AtomicReference responseTraceState = new AtomicReference<>(); + private final AtomicReference upstreamRpcServerService = new AtomicReference<>(); + private final AtomicReference upstreamRpcServerFunction = new AtomicReference<>(); + private final AtomicReference upstreamRpcServerUrl = new AtomicReference<>(); + private final AtomicReference upstreamRpcClientService = new AtomicReference<>(); + private final AtomicReference upstreamRpcClientFunction = new AtomicReference<>(); + private final AtomicReference upstreamRpcClientUrl = new AtomicReference<>(); + private final AtomicReference upstreamClientPrefixAfterCall = new AtomicReference<>(); private OwnerServiceSrv.Iface downstreamClient; @@ -80,6 +94,9 @@ public Owner getOwner(int id) throws TException { TraceContext.getCurrentTraceData().getOtelSpan().getSpanContext().getTraceId()); downstreamTraceState.set(TraceContext.getCurrentTraceData().getInboundTraceState()); downstreamTraceParent.set(TraceContext.getCurrentTraceData().getInboundTraceParent()); + downstreamRpcServerService.set(MDC.get("rpc.server.service")); + downstreamRpcServerFunction.set(MDC.get("rpc.server.function")); + downstreamRpcServerUrl.set(MDC.get("rpc.server.url")); return new Owner(id, "downstream"); } }; @@ -99,9 +116,13 @@ public Owner getOwner(int id) throws TException { TraceContext.getCurrentTraceData().getOtelSpan().getSpanContext().getTraceId()); upstreamTraceState.set(TraceContext.getCurrentTraceData().getInboundTraceState()); upstreamTraceParent.set(TraceContext.getCurrentTraceData().getInboundTraceParent()); + upstreamRpcServerService.set(MDC.get("rpc.server.service")); + upstreamRpcServerFunction.set(MDC.get("rpc.server.function")); + upstreamRpcServerUrl.set(MDC.get("rpc.server.url")); Owner result = downstreamClient.getOwner(id); + upstreamClientPrefixAfterCall.set(MDC.get("rpc.client.service")); assertNotNull("Active trace context must be available", TraceContext.getCurrentTraceData()); return result; } @@ -118,8 +139,19 @@ public Owner getOwner(int id) throws TException { ((org.eclipse.jetty.server.handler.HandlerCollection) server.getHandler()).addHandler(context); context.start(); - downstreamClient = createThriftRPCClient(OwnerServiceSrv.Iface.class, new TimestampIdGenerator(), null, - getUrlString("/downstream")); + ClientEventListener clientEventListener = new ClientEventListener() { + @Override + public void notifyEvent(THClientEvent event) { + if (ClientEventType.CLIENT_SEND.equals(event.getEventType())) { + upstreamRpcClientService.set(MDC.get(MDCUtils.TRACE_RPC_CLIENT_PREFIX + "service")); + upstreamRpcClientFunction.set(MDC.get(MDCUtils.TRACE_RPC_CLIENT_PREFIX + "function")); + upstreamRpcClientUrl.set(MDC.get(MDCUtils.TRACE_RPC_CLIENT_PREFIX + "url")); + } + } + }; + + downstreamClient = createThriftRPCClient(OwnerServiceSrv.Iface.class, new TimestampIdGenerator(), + clientEventListener, null, getUrlString("/downstream")); } @Test @@ -145,6 +177,10 @@ public void shouldPropagateMetadataHeadersAndPopulateMdc() throws Exception { assertEquals(X_REQUEST_DEADLINE, downstreamMetadataDeadline.get()); assertEquals(X_REQUEST_ID, downstreamMdcId.get()); assertEquals(X_REQUEST_DEADLINE, downstreamMdcDeadline.get()); + assertEquals("OwnerService", downstreamRpcServerService.get()); + assertEquals("getOwner", downstreamRpcServerFunction.get()); + assertTrue("Server URL should contain downstream path", + downstreamRpcServerUrl.get() != null && downstreamRpcServerUrl.get().contains("/downstream")); String upstreamTraceId = upstreamOtelTraceId.get(); String downstreamTraceId = downstreamOtelTraceId.get(); assertNotNull(upstreamTraceId); @@ -164,6 +200,14 @@ public void shouldPropagateMetadataHeadersAndPopulateMdc() throws Exception { assertTrue("Response traceparent should contain the original trace ID", responseTraceParent.get().contains(TRACE_ID)); assertEquals(TRACE_STATE, responseTraceState.get()); + assertEquals("OwnerService", upstreamRpcServerService.get()); + assertEquals("getOwner", upstreamRpcServerFunction.get()); + assertTrue("Server URL should contain upstream path", + upstreamRpcServerUrl.get() != null && upstreamRpcServerUrl.get().contains("/upstream")); + assertEquals("OwnerService", upstreamRpcClientService.get()); + assertEquals("getOwner", upstreamRpcClientFunction.get()); + assertEquals(getUrlString("/downstream"), upstreamRpcClientUrl.get()); + assertEquals("OwnerService", upstreamClientPrefixAfterCall.get()); } private void injectHeaders(HttpRequest request, EntityDetails entity, HttpContext context) @@ -211,6 +255,9 @@ private void clearCapturedValues() { downstreamMetadataDeadline.set(null); downstreamMdcId.set(null); downstreamMdcDeadline.set(null); + downstreamRpcServerService.set(null); + downstreamRpcServerFunction.set(null); + downstreamRpcServerUrl.set(null); upstreamOtelTraceId.set(null); downstreamOtelTraceId.set(null); upstreamTraceState.set(null); @@ -219,5 +266,12 @@ private void clearCapturedValues() { downstreamTraceParent.set(null); responseTraceParent.set(null); responseTraceState.set(null); + upstreamRpcServerService.set(null); + upstreamRpcServerFunction.set(null); + upstreamRpcServerUrl.set(null); + upstreamRpcClientService.set(null); + upstreamRpcClientFunction.set(null); + upstreamRpcClientUrl.set(null); + upstreamClientPrefixAfterCall.set(null); } }