HDDS-14903. Fix Trace Hierarchy for Freon s3bg Command

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

@@ -323,10 +323,44 @@ public static TraceCloseable createActivatedSpan(String spanName) {
     };
   }
 
+  /**
+   * Create an active span whose parent is decoded from an encoded W3C
+   * traceparent string (produced by {@link #exportCurrentSpan()}).
+   * If {@code encodedParent} is null or empty, a root span is created.
+   * Scope and span are both closed when the returned TraceCloseable is closed.
+   */
+  public static TraceCloseable createActivatedSpanWithParent(
+      String spanName, String encodedParent) {
+    Span span = importAndCreateSpan(spanName, encodedParent);
+    Scope scope = span.makeCurrent();
+    return () -> {
+      scope.close();
+      span.end();
+    };
+  }
+
   public static Span getActiveSpan() {
     return Span.current();
   }
 
+  /**
+   * Helper to build the tracing carrier string from W3C headers.
+   * @param traceparent raw traceparent header
+   * @param tracestate raw tracestate header
+   * @return formatted carrier string for TextExtractor
+   */
+  public static String buildTraceContextCarrier(String traceparent, String tracestate) {
+    if (traceparent == null || traceparent.isEmpty()) {
+      return "";
+    }
+    StringBuilder sb = new StringBuilder("traceparent=")
+        .append(traceparent.trim());
+    if (tracestate != null && !tracestate.isEmpty()) {
+      sb.append(";tracestate=").append(tracestate.trim());
+    }
+    return sb.toString();
+  }
+
   /**
    * AutoCloseable interface for tracing span but no exception is thrown in close.
    */

hadoop-ozone/common/src/main/java/org/apache/hadoop/ozone/om/protocolPB/OzoneManagerProtocolClientSideTranslatorPB.java

@@ -2073,12 +2073,14 @@ public void cancelDelegationToken(Token<OzoneTokenIdentifier> token)
   }
 
   @Override
+  @SkipTracing
   public void setThreadLocalS3Auth(
       S3Auth s3Auth) {
     this.threadLocalS3Auth.set(s3Auth);
   }
 
   @Override
+  @SkipTracing
   public void clearThreadLocalS3Auth() {
     this.threadLocalS3Auth.remove();
   }

hadoop-ozone/freon/src/main/java/org/apache/hadoop/ozone/freon/FreonS3TraceContextRequestHandler.java

@@ -0,0 +1,47 @@
+/*
+ * 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.ozone.freon;
+
+import com.amazonaws.Request;
+import com.amazonaws.handlers.RequestHandler2;
+import org.apache.hadoop.hdds.tracing.TracingUtil;
+
+/**
+ * Adds W3C trace context headers to each outgoing S3 request so the S3 Gateway
+ * can attach its spans to the Freon task span.
+ * BeforeRequest extracts current span info and converts it into a string traceContext.
+ * This context is used to add headers to provide context to S3Gateway.
+ */
+public final class FreonS3TraceContextRequestHandler extends RequestHandler2 {
+
+  @Override
+  public void beforeRequest(Request<?> request) {
+    String traceContext = TracingUtil.exportCurrentSpan();
+    if (traceContext.isEmpty()) {
+      return;
+    }
+    for (String part : traceContext.split(";")) {
+      int eq = part.indexOf('=');
+      if (eq > 0) {
+        String key = part.substring(0, eq).trim();
+        String value = part.substring(eq + 1).trim();
+        request.addHeader(key, value);
+      }
+    }
+  }
+}

hadoop-ozone/freon/src/main/java/org/apache/hadoop/ozone/freon/S3EntityGenerator.java

@@ -46,11 +46,12 @@ protected void s3ClientInit() {
       amazonS3ClientBuilder
           .withPathStyleAccessEnabled(true)
           .withEndpointConfiguration(
-                  new EndpointConfiguration(endpoint, "us-east-1"));
+              new EndpointConfiguration(endpoint, "us-east-1"));
     } else {
       amazonS3ClientBuilder.withRegion(Regions.DEFAULT_REGION);
     }
 
+    amazonS3ClientBuilder.withRequestHandlers(new FreonS3TraceContextRequestHandler());
     s3 = amazonS3ClientBuilder.build();
   }
 

hadoop-ozone/s3gateway/src/main/java/org/apache/hadoop/ozone/s3/TracingFilter.java

@@ -46,34 +46,57 @@ public class TracingFilter implements ContainerRequestFilter,
   public void filter(ContainerRequestContext requestContext) {
     finishAndCloseActiveSpan();
 
-    TracingUtil.TraceCloseable activatedSpan =
-        TracingUtil.createActivatedSpan(resourceInfo.getResourceClass().getSimpleName() + "." +
-            resourceInfo.getResourceMethod().getName());
-    requestContext.setProperty(TRACING_SPAN_CLOSABLE, activatedSpan);
+    String traceparent = requestContext.getHeaderString("traceparent");
+    String tracestate  = requestContext.getHeaderString("tracestate");
+    String encodedParent = TracingUtil.buildTraceContextCarrier(traceparent, tracestate);
+
+    String spanName = resourceInfo.getResourceClass().getSimpleName() + "." +
+        resourceInfo.getResourceMethod().getName();
+
+    TracingUtil.TraceCloseable traceCloseable =
+        TracingUtil.createActivatedSpanWithParent(spanName, encodedParent);
+
+    requestContext.setProperty(TRACING_SPAN_CLOSABLE, traceCloseable);
   }
 
   @Override
   public void filter(ContainerRequestContext requestContext,
       ContainerResponseContext responseContext) {
     final TracingUtil.TraceCloseable spanClosable
         = (TracingUtil.TraceCloseable) requestContext.getProperty(TRACING_SPAN_CLOSABLE);
+    if (spanClosable == null) {
+      return;
+    }
     // HDDS-7064: Operation performed while writing StreamingOutput response
     // should only be closed once the StreamingOutput callback has completely
     // written the data to the destination
-    OutputStream out = responseContext.getEntityStream();
-    if (out != null) {
-      responseContext.setEntityStream(new WrappedOutputStream(out) {
-        @Override
-        public void close() throws IOException {
-          super.close();
-          finishAndClose(spanClosable);
-        }
-      });
+    if (isStreamingGetObject(requestContext)) {
+      OutputStream out = responseContext.getEntityStream();
+      if (out != null) {
+        responseContext.setEntityStream(new WrappedOutputStream(out) {
+          @Override
+          public void close() throws IOException {
+            super.close();
+            finishAndClose(spanClosable);
+          }
+        });
+      } else {
+        finishAndClose(spanClosable);
+      }
     } else {
       finishAndClose(spanClosable);
     }
   }
 
+  private boolean isStreamingGetObject(ContainerRequestContext req) {
+    if (!"GET".equalsIgnoreCase(req.getMethod())) {
+      return false;
+    }
+    String cls = resourceInfo.getResourceClass().getSimpleName();
+    String method = resourceInfo.getResourceMethod().getName();
+    return "ObjectEndpoint".equals(cls) && "get".equals(method);
+  }
+
   private static void finishAndClose(TracingUtil.TraceCloseable spanClosable) {
     try {
       spanClosable.close();