Add mdc refresh interceptor

libthrift/pom.xml

@@ -7,7 +7,7 @@
     <parent>
         <artifactId>woody</artifactId>
         <groupId>dev.vality.woody</groupId>
-        <version>2.0.14</version>
+        <version>2.0.15</version>
     </parent>
 
     <artifactId>libthrift</artifactId>

pom.xml

@@ -13,7 +13,7 @@
     <packaging>pom</packaging>
     <groupId>dev.vality.woody</groupId>
     <artifactId>woody</artifactId>
-    <version>2.0.14</version>
+    <version>2.0.15</version>
 
     <name>Woody Java</name>
     <description>Java implementation for Woody spec</description>

woody-api/pom.xml

@@ -7,7 +7,7 @@
     <parent>
         <artifactId>woody</artifactId>
         <groupId>dev.vality.woody</groupId>
-        <version>2.0.14</version>
+        <version>2.0.15</version>
     </parent>
 
     <artifactId>woody-api</artifactId>

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),

woody-api/src/main/java/dev/vality/woody/api/AbstractServiceBuilder.java

@@ -53,6 +53,7 @@ protected <T> T createProxyService(Class<T> iface, T handler) {
     protected MethodCallTracer createEventTracer() {
         return new CompositeTracer(
                 TargetCallTracer.forServer(),
+                new MdcRefreshTracer(),
                 DeadlineTracer.forService(),
                 new EventTracer(getOnCallStartEventListener(),
                         getOnCallEndEventListener(),

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<String> keys = EXTENDED_MDC_KEYS.get();
+        Iterator<String> 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();
+            }
+        }
+    }
 }

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);
+    }
+
+}

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);
+    }
+}

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,

woody-thrift/pom.xml

@@ -7,7 +7,7 @@
     <parent>
         <artifactId>woody</artifactId>
         <groupId>dev.vality.woody</groupId>
-        <version>2.0.14</version>
+        <version>2.0.15</version>
     </parent>
 
     <artifactId>woody-thrift</artifactId>

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));
     }
 

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(),