HDDS-14765. Improve Trace Hierarchy for Ozone Shell Put Key Command

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

@@ -0,0 +1,31 @@
+/*
+ * 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 java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * Annotation to mark methods that should be excluded from tracing.
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target(ElementType.METHOD)
+public @interface SkipTracing {
+}

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

@@ -25,6 +25,7 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Map.Entry;
+import org.apache.commons.lang3.tuple.Pair;
 
 /**
  * A Java proxy invocation handler to trace all the methods of the delegate
@@ -37,8 +38,7 @@ public class TraceAllMethod<T> implements InvocationHandler {
   /**
    * Cache for all the method objects of the delegate class.
    */
-  private final Map<String, Map<Class<?>[], Method>> methods = new HashMap<>();
-
+  private final Map<String, Map<Class<?>[], Pair<Boolean, Method>>> methods = new HashMap<>();
   private final T delegate;
 
   private final String name;
@@ -50,18 +50,31 @@ public TraceAllMethod(T delegate, String name) {
       if (method.getDeclaringClass().equals(Object.class)) {
         continue;
       }
+      boolean shouldSkip = method.isAnnotationPresent(SkipTracing.class);
       methods.computeIfAbsent(method.getName(), any -> new HashMap<>())
-          .put(method.getParameterTypes(), method);
+          .put(method.getParameterTypes(), Pair.of(shouldSkip, method));
     }
   }
 
   @Override
   public Object invoke(Object proxy, Method method, Object[] args)
       throws Throwable {
-    Method delegateMethod = findDelegatedMethod(method);
-    if (delegateMethod == null) {
-      throw new NoSuchMethodException("Method not found: " +
-        method.getName());
+    Pair<Boolean, Method> methodInfo = findDelegatedMethod(method);
+    if (methodInfo == null) {
+      throw new NoSuchMethodException("Method not found: " + method.getName());
+    }
+    boolean shouldSkip = methodInfo.getLeft();
+    Method delegateMethod = methodInfo.getRight();
+    if (shouldSkip) {
+      try {
+        return delegateMethod.invoke(delegate, args);
+      } catch (Exception ex) {
+        if (ex.getCause() != null) {
+          throw ex.getCause();
+        } else {
+          throw ex;
+        }
+      }
     }
 
     try (TracingUtil.TraceCloseable ignored = TracingUtil.createActivatedSpan(name + "." + method.getName())) {
@@ -77,8 +90,8 @@ public Object invoke(Object proxy, Method method, Object[] args)
     }
   }
 
-  private Method findDelegatedMethod(Method method) {
-    for (Entry<Class<?>[], Method> entry : methods.getOrDefault(
+  private Pair<Boolean, Method> findDelegatedMethod(Method method) {
+    for (Entry<Class<?>[], Pair<Boolean, Method>> entry : methods.getOrDefault(
         method.getName(), emptyMap()).entrySet()) {
       if (Arrays.equals(entry.getKey(), method.getParameterTypes())) {
         return entry.getValue();

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

@@ -179,10 +179,11 @@ public static boolean isTracingEnabled(
 
   /**
    * Execute {@code runnable} inside an activated new span.
+   * If a parent span exists in the current context, this becomes a child span.
    */
   public static <E extends Exception> void executeInNewSpan(String spanName,
       CheckedRunnable<E> runnable) throws E {
-    Span span = tracer.spanBuilder(spanName).setNoParent().startSpan();
+    Span span = buildSpan(spanName);
     executeInSpan(span, runnable);
   }
 
@@ -191,7 +192,7 @@ public static <E extends Exception> void executeInNewSpan(String spanName,
    */
   public static <R, E extends Exception> R executeInNewSpan(String spanName,
       CheckedSupplier<R, E> supplier) throws E {
-    Span span = tracer.spanBuilder(spanName).setNoParent().startSpan();
+    Span span = buildSpan(spanName);
     return executeInSpan(span, supplier);
   }
 
@@ -244,7 +245,7 @@ public static <E extends Exception> void executeAsChildSpan(String spanName,
    * in case of Exceptions.
    */
   public static TraceCloseable createActivatedSpan(String spanName) {
-    Span span = tracer.spanBuilder(spanName).setNoParent().startSpan();
+    Span span = buildSpan(spanName);
     Scope scope = span.makeCurrent();
     return () -> {
       scope.close();
@@ -299,4 +300,19 @@ private void parse(String carrier) {
       }
     }
   }
+
+  /**
+   * Creates a new span, using the current context as a parent if valid;
+   * otherwise, creates a root span.
+   */
+  private static Span buildSpan(String spanName) {
+    Context currentContext = Context.current();
+    Span parentSpan = Span.fromContext(currentContext);
+
+    if (parentSpan.getSpanContext().isValid()) {
+      return tracer.spanBuilder(spanName).setParent(currentContext).startSpan();
+    } else {
+      return tracer.spanBuilder(spanName).setNoParent().startSpan();
+    }
+  }
 }

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

@@ -20,6 +20,8 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
+import io.opentelemetry.api.trace.Span;
+import java.io.IOException;
 import org.junit.jupiter.api.Test;
 
 /**
@@ -52,6 +54,12 @@ default String defaultMethod() {
       return otherMethod("default");
     }
 
+    String skippedMethod();
+
+    void throwingMethod() throws IOException;
+
+    String normalMethod();
+
     String otherMethod(String name);
   }
 
@@ -60,9 +68,35 @@ default String defaultMethod() {
    */
   public static class ServiceImpl implements Service {
 
+    private boolean spanActive = false;
+
     @Override
     public String otherMethod(String name) {
       return "Hello " + name;
     }
+
+    @Override
+    @SkipTracing
+    public String skippedMethod() {
+      this.spanActive = Span.current().getSpanContext().isValid();
+      return "skipped";
+    }
+
+    @Override
+    @SkipTracing
+    public void throwingMethod() throws IOException {
+      this.spanActive = Span.current().getSpanContext().isValid();
+      throw new IOException("Original Exception");
+    }
+
+    @Override
+    public String normalMethod() {
+      this.spanActive = Span.current().getSpanContext().isValid();
+      return "normal";
+    }
+
+    public boolean wasSpanActive() {
+      return spanActive;
+    }
   }
 }

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

@@ -20,8 +20,12 @@
 import static org.apache.hadoop.hdds.tracing.TracingUtil.createProxy;
 import static org.apache.hadoop.hdds.tracing.TracingUtil.exportCurrentSpan;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 
+import java.io.IOException;
 import org.apache.hadoop.hdds.conf.InMemoryConfigurationForTesting;
 import org.apache.hadoop.hdds.conf.MutableConfigurationSource;
 import org.apache.hadoop.hdds.scm.ScmConfigKeys;
@@ -58,4 +62,47 @@ private static MutableConfigurationSource tracingEnabled() {
     return config;
   }
 
+  /**
+   * Test for checking if span was not created when a regular method
+   * in Service implementation has @SkipTracing.
+   */
+  @Test
+  public void testSkipTracingNoSpan() {
+    TracingUtil.initTracing("TestService", tracingEnabled());
+    ServiceImpl impl = new ServiceImpl();
+    Service serviceProxy = createProxy(impl, Service.class, tracingEnabled());
+
+    serviceProxy.skippedMethod();
+    assertFalse(impl.wasSpanActive(), "Span should NOT be created for @SkipTracing methods.");
+  }
+
+  /**
+   * Test for checking if span was not created when a method throws exception
+   * in Service implementation and has @SkipTracing.
+   */
+  @Test
+  public void testSkipTracingExceptionUnwrapped() {
+    TracingUtil.initTracing("TestService", tracingEnabled());
+    ServiceImpl impl = new ServiceImpl();
+    Service serviceProxy = createProxy(impl, Service.class, tracingEnabled());
+
+    IOException ex = assertThrows(IOException.class,
+        () -> serviceProxy.throwingMethod());
+    assertEquals("Original Exception", ex.getMessage());
+    assertFalse(impl.wasSpanActive(), "Span should NOT have been created for a @SkipTracing throwing method.");
+  }
+
+  /**
+   * Test for checking if span is created when a method in Service implementation
+   * does not have @SkipTracing.
+   */
+  @Test
+  public void testProxyNormalVsSkipped() {
+    TracingUtil.initTracing("TestService", tracingEnabled());
+    ServiceImpl impl = new ServiceImpl();
+    Service serviceProxy = createProxy(impl, Service.class, tracingEnabled());
+
+    serviceProxy.normalMethod();
+    assertTrue(impl.wasSpanActive(), "Normal method should have an active span.");
+  }
 }

hadoop-ozone/client/src/main/java/org/apache/hadoop/ozone/client/ObjectStore.java

@@ -28,7 +28,6 @@
 import org.apache.hadoop.crypto.key.KeyProvider;
 import org.apache.hadoop.hdds.conf.ConfigurationSource;
 import org.apache.hadoop.hdds.scm.client.HddsClientUtils;
-import org.apache.hadoop.hdds.tracing.TracingUtil;
 import org.apache.hadoop.io.Text;
 import org.apache.hadoop.ozone.OmUtils;
 import org.apache.hadoop.ozone.OzoneAcl;
@@ -83,7 +82,7 @@ public class ObjectStore {
    * @param proxy ClientProtocol proxy.
    */
   public ObjectStore(ConfigurationSource conf, ClientProtocol proxy) {
-    this.proxy = TracingUtil.createProxy(proxy, ClientProtocol.class, conf);
+    this.proxy = proxy;
     this.listCacheSize = HddsClientUtils.getListCacheSize(conf);
     s3BucketLayout = OmUtils.validateBucketLayout(
         conf.getTrimmed(

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

@@ -48,6 +48,7 @@
 import org.apache.hadoop.hdds.protocol.proto.HddsProtos.TransferLeadershipRequestProto;
 import org.apache.hadoop.hdds.protocol.proto.HddsProtos.UpgradeFinalizationStatus;
 import org.apache.hadoop.hdds.scm.container.common.helpers.ExcludeList;
+import org.apache.hadoop.hdds.tracing.SkipTracing;
 import org.apache.hadoop.hdds.tracing.TracingUtil;
 import org.apache.hadoop.io.Text;
 import org.apache.hadoop.ipc_.CallerContext;
@@ -2088,6 +2089,7 @@ public ThreadLocal<S3Auth> getS3CredentialsProvider() {
   }
 
   @Override
+  @SkipTracing
   public S3Auth getThreadLocalS3Auth() {
     return this.threadLocalS3Auth.get();
   }