HDDS-13679. Add trace and span IDs to log MDC via OpenTelemetry ContextStorageProvider

hadoop-hdds/common/src/main/java/org/apache/hadoop/hdds/tracing/TracingUtil.java

@@ -22,11 +22,13 @@
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.api.trace.StatusCode;
 import io.opentelemetry.api.trace.Tracer;
 import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
 import io.opentelemetry.context.Context;
+import io.opentelemetry.context.ContextStorage;
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.context.propagation.ContextPropagators;
 import io.opentelemetry.context.propagation.TextMapGetter;
@@ -48,6 +50,7 @@
 import org.apache.ratis.util.function.CheckedSupplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.slf4j.MDC;
 
 /**
  * Utility class to collect all the tracing helper methods.
@@ -62,11 +65,29 @@ public final class TracingUtil {
   private static Tracer tracer = OpenTelemetry.noop().getTracer("noop");
   private static volatile SdkTracerProvider sdkTracerProvider;
   private static BatchSpanProcessor batchSpanProcessor;
+  private static boolean mdcContextStorageInstalled;
   public static final String GLOBAL_TRACER_NAME = "ozone";
 
   private TracingUtil() {
   }
 
+  /**
+   * Adds the current trace and span IDs to the logging MDC.
+   * Only for Ozone services and CLIs, not the client library, so applications keep their own MDC.
+   * Must be called before OpenTelemetry {@link Context} is first used in the JVM.
+   */
+  public static synchronized void installMdcContextStorage() {
+    if (mdcContextStorageInstalled) {
+      return;
+    }
+    mdcContextStorageInstalled = true;
+    ContextStorage.addWrapper(MdcContextStorage::new);
+    if (!(ContextStorage.get() instanceof MdcContextStorage)) {
+      LOG.warn("OpenTelemetry context is already initialized, {} and {} will not be added to the logging MDC",
+          MdcContextStorage.MDC_TRACE_ID, MdcContextStorage.MDC_SPAN_ID);
+    }
+  }
+
   /**
    * Initialize the tracing with the given service name.
    */
@@ -128,6 +149,7 @@ public static synchronized void flushTracing() {
 
   /**
    * This function initializes tracing, runs the command in a span, and exports spans before returning for CLI spans.
+   * Only for Ozone CLIs: it also installs the MDC context storage, see {@link #installMdcContextStorage()}.
    */
   public static <R, E extends Exception> R execute(
       String serviceName,
@@ -143,6 +165,7 @@ public static <R, E extends Exception> R execute(
       SpanKind spanKind,
       ConfigurationSource conf,
       CheckedSupplier<R, E> supplier) throws E {
+    installMdcContextStorage();
     initTracing(serviceName, conf);
     try {
       return executeInNewSpan(spanName, spanKind, supplier);
@@ -648,4 +671,56 @@ public static TraceCloseable createActivatedSpanFromW3cHttpHeaders(
       span.end();
     };
   }
+
+  /**
+   * {@link ContextStorage} that keeps the logging MDC in sync with the current span.
+   * Installed through {@link #installMdcContextStorage()} by Ozone services and CLIs only,
+   * so applications using the Ozone client keep their own MDC and context storage.
+   * Every {@code makeCurrent()} / {@code Context.wrap(...)} goes through {@link #attach(Context)},
+   * so trace and span IDs are available to log patterns without changes at the call sites.
+   * The MDC is always derived from the current context (never saved and restored),
+   * so values inherited by child threads through the MDC are overwritten on the next attach or close.
+   */
+  static final class MdcContextStorage implements ContextStorage {
+    static final String MDC_TRACE_ID = "trace_id";
+    static final String MDC_SPAN_ID = "span_id";
+
+    private final ContextStorage delegate;
+
+    MdcContextStorage(ContextStorage delegate) {
+      this.delegate = delegate;
+    }
+
+    @Override
+    public Scope attach(Context toAttach) {
+      Scope scope = delegate.attach(toAttach);
+      updateMdc(delegate.current());
+      return () -> {
+        scope.close();
+        updateMdc(delegate.current());
+      };
+    }
+
+    @Override
+    public Context current() {
+      return delegate.current();
+    }
+
+    @Override
+    public Context root() {
+      return delegate.root();
+    }
+
+    private static void updateMdc(Context context) {
+      SpanContext spanContext = context == null ? SpanContext.getInvalid()
+          : Span.fromContext(context).getSpanContext();
+      if (spanContext.isValid()) {
+        MDC.put(MDC_TRACE_ID, spanContext.getTraceId());
+        MDC.put(MDC_SPAN_ID, spanContext.getSpanId());
+      } else {
+        MDC.remove(MDC_TRACE_ID);
+        MDC.remove(MDC_SPAN_ID);
+      }
+    }
+  }
 }

hadoop-hdds/common/src/test/java/org/apache/hadoop/hdds/tracing/TestMdcContextStorageInstall.java

@@ -0,0 +1,161 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hadoop.hdds.tracing;
+
+import static org.apache.hadoop.hdds.tracing.TracingUtil.MdcContextStorage.MDC_SPAN_ID;
+import static org.apache.hadoop.hdds.tracing.TracingUtil.MdcContextStorage.MDC_TRACE_ID;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.ContextStorage;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.sdk.OpenTelemetrySdk;
+import io.opentelemetry.sdk.trace.SdkTracerProvider;
+import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor;
+import io.opentelemetry.sdk.trace.export.SpanExporter;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import org.apache.hadoop.hdds.conf.InMemoryConfigurationForTesting;
+import org.apache.hadoop.hdds.conf.MutableConfigurationSource;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.slf4j.MDC;
+
+/**
+ * Tests {@link TracingUtil#installMdcContextStorage()}.
+ * When running together with other test classes, use fork mode per class.
+ */
+public class TestMdcContextStorageInstall {
+
+  @BeforeAll
+  public static void install() {
+    TracingUtil.installMdcContextStorage();
+    assertInstanceOf(TracingUtil.MdcContextStorage.class, ContextStorage.get(), "use fork mode per class");
+  }
+
+  @BeforeEach
+  public void setUp() {
+    TracingUtil.shutdownTracing();
+    GlobalOpenTelemetry.resetForTest();
+    MDC.clear();
+  }
+
+  @AfterEach
+  public void tearDown() {
+    TracingUtil.shutdownTracing();
+    GlobalOpenTelemetry.resetForTest();
+    MDC.clear();
+  }
+
+  private static void initTracing(boolean enabled) {
+    SdkTracerProvider provider = SdkTracerProvider.builder()
+        .addSpanProcessor(SimpleSpanProcessor.create(SpanExporter.noop())).build();
+    OpenTelemetrySdk sdk = OpenTelemetrySdk.builder().setTracerProvider(provider).build();
+    GlobalOpenTelemetry.set(sdk);
+    MutableConfigurationSource conf = new InMemoryConfigurationForTesting();
+    conf.setBoolean("ozone.tracing.enabled", enabled);
+    TracingUtil.initTracing("mdc-install", conf);
+  }
+
+  private static void assertMdcMatchesCurrentSpan() {
+    SpanContext current = Span.current().getSpanContext();
+    assertTrue(current.isValid());
+    assertEquals(current.getTraceId(), MDC.get(MDC_TRACE_ID));
+    assertEquals(current.getSpanId(), MDC.get(MDC_SPAN_ID));
+  }
+
+  private static void assertMdcEmpty() {
+    assertNull(MDC.get(MDC_TRACE_ID));
+    assertNull(MDC.get(MDC_SPAN_ID));
+  }
+
+  @Test
+  public void testActivatedSpanSetsMdc() {
+    initTracing(true);
+    try (TracingUtil.TraceCloseable ignored = TracingUtil.createActivatedSpan("outer")) {
+      assertMdcMatchesCurrentSpan();
+      String traceId = MDC.get(MDC_TRACE_ID);
+      String outerSpanId = MDC.get(MDC_SPAN_ID);
+      try (TracingUtil.TraceCloseable ignoredInner = TracingUtil.createActivatedSpan("inner")) {
+        assertMdcMatchesCurrentSpan();
+        assertEquals(traceId, MDC.get(MDC_TRACE_ID));
+        assertNotEquals(outerSpanId, MDC.get(MDC_SPAN_ID));
+      }
+      assertEquals(traceId, MDC.get(MDC_TRACE_ID));
+      assertEquals(outerSpanId, MDC.get(MDC_SPAN_ID));
+    }
+    assertMdcEmpty();
+  }
+
+  @Test
+  public void testImportedParentSpanSetsMdc() {
+    initTracing(true);
+    String carrier;
+    String parentTraceId;
+    try (TracingUtil.TraceCloseable ignored = TracingUtil.createActivatedSpan("client")) {
+      carrier = TracingUtil.exportCurrentSpan();
+      parentTraceId = Span.current().getSpanContext().getTraceId();
+    }
+    assertMdcEmpty();
+
+    Span server = TracingUtil.importAndCreateSpan("server", carrier);
+    try (Scope ignored = server.makeCurrent()) {
+      assertMdcMatchesCurrentSpan();
+      assertEquals(parentTraceId, MDC.get(MDC_TRACE_ID));
+    } finally {
+      server.end();
+    }
+    assertMdcEmpty();
+  }
+
+  @Test
+  public void testContextWrapPropagatesMdcToOtherThread() throws Exception {
+    initTracing(true);
+    ExecutorService executor = Executors.newSingleThreadExecutor();
+    try {
+      try (TracingUtil.TraceCloseable ignored = TracingUtil.createActivatedSpan("span")) {
+        String[] inTask = executor.submit(Context.current().wrap(
+            () -> new String[]{MDC.get(MDC_TRACE_ID), MDC.get(MDC_SPAN_ID)})).get();
+        assertEquals(MDC.get(MDC_TRACE_ID), inTask[0]);
+        assertEquals(MDC.get(MDC_SPAN_ID), inTask[1]);
+      }
+      assertNull(executor.submit(() -> MDC.get(MDC_TRACE_ID)).get());
+      assertNull(executor.submit(() -> MDC.get(MDC_SPAN_ID)).get());
+    } finally {
+      executor.shutdownNow();
+    }
+  }
+
+  @Test
+  public void testTracingDisabledLeavesMdcEmpty() {
+    initTracing(false);
+    try (TracingUtil.TraceCloseable ignored = TracingUtil.createActivatedSpan("span")) {
+      assertMdcEmpty();
+    }
+    assertMdcEmpty();
+  }
+}