HDDS-17581. Add TracingUtil span helpers and span-kind aware tracing proxies

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

@@ -23,11 +23,12 @@
 import java.lang.annotation.Target;
 
 /**
- * Annotation that marks a proxied client method as SpanKind.INTERNAL.
+ * Annotation that marks a method of a proxied delegate class, or every method declared in that class, as
+ * SpanKind.INTERNAL.
  * Used for methods that do not send an outbound RPC.
  */
 @Retention(RetentionPolicy.RUNTIME)
-@Target(ElementType.METHOD)
+@Target({ElementType.METHOD, ElementType.TYPE})
 public @interface InternalSpanKind {
 }
 

hadoop-hdds/common/src/main/java/org/apache/hadoop/hdds/tracing/SpanAttributeKeys.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;
+
+/**
+ * Span attribute keys shared by Ozone services and clients.
+ * Fields that are also in the audit log (volume, bucket, key, ...) use the audit key unchanged instead.
+ */
+public final class SpanAttributeKeys {
+
+  /** Status or result code of a failed request, for example KEY_NOT_FOUND. */
+  public static final String STATUS = "ozone.status";
+
+  private SpanAttributeKeys() {
+  }
+}

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

@@ -19,8 +19,11 @@
 
 import static java.util.Collections.emptyMap;
 
+import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.context.Scope;
 import java.lang.reflect.InvocationHandler;
+import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.util.Arrays;
 import java.util.HashMap;
@@ -40,10 +43,9 @@ public class TraceAllMethod<T> implements InvocationHandler {
    */
   private final Map<String, Map<Class<?>[], DelegatedMethodInfo>> methods = new HashMap<>();
   private final T delegate;
-
   private final String name;
 
-  public TraceAllMethod(T delegate, String name) {
+  public TraceAllMethod(T delegate, String name, SpanKind spanKind) {
     this.delegate = delegate;
     this.name = name;
     for (Method method : delegate.getClass().getMethods()) {
@@ -52,48 +54,43 @@ public TraceAllMethod(T delegate, String name) {
       }
       boolean shouldSkip = method.isAnnotationPresent(SkipTracing.class);
       methods.computeIfAbsent(method.getName(), any -> new HashMap<>())
-          .put(method.getParameterTypes(), new DelegatedMethodInfo(shouldSkip, method));
+          .put(method.getParameterTypes(), new DelegatedMethodInfo(shouldSkip, findSpanKind(method, spanKind)));
     }
   }
 
   @Override
   public Object invoke(Object proxy, Method method, Object[] args)
       throws Throwable {
     DelegatedMethodInfo methodInfo = findDelegatedMethod(method);
-    if (methodInfo == null) {
-      throw new NoSuchMethodException("Method not found: " + method.getName());
+    if (methodInfo == null || methodInfo.shouldSkip) {
+      return invokeDelegate(method, args);
     }
-    boolean shouldSkip = methodInfo.shouldSkip;
-    Method delegateMethod = methodInfo.delegateMethod;
-    if (shouldSkip) {
-      try {
-        return delegateMethod.invoke(delegate, args);
-      } catch (Exception ex) {
-        if (ex.getCause() != null) {
-          throw ex.getCause();
-        } else {
-          throw ex;
-        }
-      }
+
+    Span span = TracingUtil.createSpan(name + "." + method.getName(), methodInfo.spanKind);
+    try (Scope ignored = span.makeCurrent()) {
+      return invokeDelegate(method, args);
+    } catch (Throwable ex) {
+      TracingUtil.recordException(span, ex);
+      throw ex;
+    } finally {
+      span.end();
     }
+  }
 
-    // if a call is within a process, not an outbound RPC to other processes (OM/SCM/DN) , it is marked as INTERNAL.
-    SpanKind spanKind = delegateMethod.isAnnotationPresent(InternalSpanKind.class)
-        ? SpanKind.INTERNAL
-        : SpanKind.CLIENT;
+  private Object invokeDelegate(Method method, Object[] args) throws Throwable {
+    try {
+      return method.invoke(delegate, args);
+    } catch (InvocationTargetException e) {
+      throw e.getCause();
+    }
+  }
 
-    try (TracingUtil.TraceCloseable ignored = TracingUtil.createActivatedSpan(
-        name + "." + method.getName(), spanKind)) {
-      try {
-        return delegateMethod.invoke(delegate, args);
-      } catch (Exception ex) {
-        if (ex.getCause() != null) {
-          throw ex.getCause();
-        } else {
-          throw ex;
-        }
-      }
+  private static SpanKind findSpanKind(Method method, SpanKind defaultSpanKind) {
+    if (method.isAnnotationPresent(InternalSpanKind.class)
+        || method.getDeclaringClass().isAnnotationPresent(InternalSpanKind.class)) {
+      return SpanKind.INTERNAL;
     }
+    return defaultSpanKind;
   }
 
   private DelegatedMethodInfo findDelegatedMethod(Method method) {
@@ -111,11 +108,11 @@ private DelegatedMethodInfo findDelegatedMethod(Method method) {
    */
   private static final class DelegatedMethodInfo {
     private final boolean shouldSkip;
-    private final Method delegateMethod;
+    private final SpanKind spanKind;
 
-    private DelegatedMethodInfo(boolean shouldSkip, Method delegateMethod) {
+    private DelegatedMethodInfo(boolean shouldSkip, SpanKind spanKind) {
       this.shouldSkip = shouldSkip;
-      this.delegateMethod = delegateMethod;
+      this.spanKind = spanKind;
     }
   }
 }

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

@@ -394,13 +394,34 @@ private static boolean hasUsableTracer() {
    */
   public static <T> T createProxy(
       T delegate, Class<T> itf, ConfigurationSource conf) {
+    return createProxy(delegate, itf, conf, SpanKind.INTERNAL);
+  }
+
+  /** Same as {@link #createProxy(Object, Class, ConfigurationSource)}, with CLIENT spans for an RPC delegate. */
+  public static <T> T createClientProxy(T delegate, Class<T> itf, ConfigurationSource conf) {
+    return createProxy(delegate, itf, conf, SpanKind.CLIENT);
+  }
+
+  /** Creates a proxy with spans named {@code <delegate class>.<method>}, also usable at startup without a span. */
+  public static <T> T createProxy(T delegate, Class<T> itf, SpanKind spanKind) {
+    if (!tracingEnabled && !applicationAware) {
+      return delegate;
+    }
+    return newProxy(delegate, itf, delegate.getClass().getSimpleName(), spanKind);
+  }
+
+  private static <T> T createProxy(T delegate, Class<T> itf, ConfigurationSource conf, SpanKind spanKind) {
     if (!shouldCreateTracingProxy(conf)) {
       return delegate;
     }
+    return newProxy(delegate, itf, itf.getSimpleName(), spanKind);
+  }
+
+  private static <T> T newProxy(T delegate, Class<T> itf, String name, SpanKind spanKind) {
     Class<?> aClass = delegate.getClass();
     return itf.cast(Proxy.newProxyInstance(aClass.getClassLoader(),
         new Class<?>[] {itf},
-        new TraceAllMethod<>(delegate, itf.getSimpleName())));
+        new TraceAllMethod<>(delegate, name, spanKind)));
   }
 
   public static boolean isTracingEnabled(ConfigurationSource conf) {
@@ -499,7 +520,7 @@ private static <R, E extends Exception> R executeInSpan(Span span,
       return supplier.get();
     } catch (Exception ex) {
       span.addEvent("Failed with exception: " + ex.getMessage());
-      span.setStatus(StatusCode.ERROR);
+      recordException(span, ex);
       throw ex;
     } finally {
       span.end();
@@ -515,7 +536,7 @@ private static <E extends Exception> void executeInSpan(Span span,
       runnable.run();
     } catch (Exception ex) {
       span.addEvent("Failed with exception: " + ex.getMessage());
-      span.setStatus(StatusCode.ERROR);
+      recordException(span, ex);
       throw ex;
     } finally {
       span.end();
@@ -556,6 +577,63 @@ public static Span getActiveSpan() {
     return Span.current();
   }
 
+  public static Span createSpan(String spanName) {
+    return createSpan(spanName, SpanKind.INTERNAL);
+  }
+
+  public static Span createSpan(String spanName, SpanKind spanKind) {
+    return buildSpan(spanName, spanKind);
+  }
+
+  public static void recordExceptionInCurrentSpan(Throwable t) {
+    recordException(Span.current(), t);
+  }
+
+  public static void recordException(Span span, Throwable t) {
+    if (t == null || !span.isRecording()) {
+      return;
+    }
+    span.recordException(t);
+    if (t.getMessage() == null) {
+      span.setStatus(StatusCode.ERROR);
+    } else {
+      span.setStatus(StatusCode.ERROR, t.getMessage());
+    }
+  }
+
+  public static void addAttributesToCurrentSpan(String key, String value) {
+    Span span = Span.current();
+    if (value != null && span.isRecording()) {
+      span.setAttribute(key, value);
+    }
+  }
+
+  public static void addAttributesToCurrentSpan(String key, long value) {
+    Span span = Span.current();
+    if (span.isRecording()) {
+      span.setAttribute(key, value);
+    }
+  }
+
+  public static void addAttributesToCurrentSpan(String key, boolean value) {
+    Span span = Span.current();
+    if (span.isRecording()) {
+      span.setAttribute(key, value);
+    }
+  }
+
+  public static void addEventToCurrentSpan(String name, String key, String value) {
+    Span span = Span.current();
+    if (!span.isRecording()) {
+      return;
+    }
+    if (value == null) {
+      span.addEvent(name);
+    } else {
+      span.addEvent(name, Attributes.of(AttributeKey.stringKey(key), value));
+    }
+  }
+
   /**
    * AutoCloseable interface for tracing span but no exception is thrown in close.
    */