Adds support for tracing runnable scenarios (#9187)

Signed-off-by: Gagan Juneja <[email protected]>

Author: Gaganjuneja

CHANGELOG.md

@@ -54,6 +54,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - [Refactor] MediaTypeParser to MediaTypeParserRegistry ([#8636](https://github.com/opensearch-project/OpenSearch/pull/8636))
 - Make MultiBucketConsumerService thread safe to use across slices during search ([#9047](https://github.com/opensearch-project/OpenSearch/pull/9047))
 - Removed blocking wait in TransportGetSnapshotsAction which was exhausting generic threadpool ([#8377](https://github.com/opensearch-project/OpenSearch/pull/8377))
+- Adds support for tracing runnable scenarios ([#8831](https://github.com/opensearch-project/OpenSearch/pull/8831))
 
 ### Deprecated
 

libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultTracer.java

@@ -37,7 +37,17 @@ public DefaultTracer(TracingTelemetry tracingTelemetry, TracerContextStorage<Str
 
     @Override
     public SpanScope startSpan(String spanName) {
-        Span span = createSpan(spanName, getCurrentSpan());
+        return startSpan(spanName, null);
+    }
+
+    @Override
+    public SpanScope startSpan(String spanName, SpanContext parentSpan) {
+        Span span = null;
+        if (parentSpan != null) {
+            span = createSpan(spanName, parentSpan.getSpan());
+        } else {
+            span = createSpan(spanName, getCurrentSpanInternal());
+        }
         setCurrentSpanInContext(span);
         addDefaultAttributes(span);
         return new DefaultSpanScope(span, (scopeSpan) -> endSpan(scopeSpan));
@@ -48,11 +58,15 @@ public void close() throws IOException {
         ((Closeable) tracingTelemetry).close();
     }
 
-    // Visible for testing
-    Span getCurrentSpan() {
+    private Span getCurrentSpanInternal() {
         return tracerContextStorage.get(TracerContextStorage.CURRENT_SPAN);
     }
 
+    public SpanContext getCurrentSpan() {
+        final Span currentSpan = tracerContextStorage.get(TracerContextStorage.CURRENT_SPAN);
+        return (currentSpan == null) ? null : new SpanContext(currentSpan);
+    }
+
     private void endSpan(Span span) {
         if (span != null) {
             span.endSpan();

libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/SpanContext.java

@@ -0,0 +1,29 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.telemetry.tracing;
+
+/**
+ * Wrapped Span will be exposed to the code outside of tracing package for sharing the {@link Span} without having access to
+ * its properties.
+ */
+public final class SpanContext {
+    private final Span span;
+
+    /**
+     * Constructor.
+     * @param span span to be wrapped.
+     */
+    public SpanContext(Span span) {
+        this.span = span;
+    }
+
+    Span getSpan() {
+        return span;
+    }
+}

libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/Tracer.java

@@ -26,4 +26,17 @@ public interface Tracer extends Closeable {
      */
     SpanScope startSpan(String spanName);
 
+    /**
+     *  Started the {@link Span} with the given name and parent.
+     * @param spanName span name.
+     * @param parentSpan parent span.
+     * @return scope of the span, must be closed with explicit close or with try-with-resource
+     */
+    SpanScope startSpan(String spanName, SpanContext parentSpan);
+
+    /**
+     * Returns the current span.
+     * @return current wrapped span.
+     */
+    SpanContext getCurrentSpan();
 }

libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/noop/NoopTracer.java

@@ -10,6 +10,7 @@
 
 import org.opensearch.telemetry.tracing.SpanScope;
 import org.opensearch.telemetry.tracing.Tracer;
+import org.opensearch.telemetry.tracing.SpanContext;
 
 /**
  * No-op implementation of Tracer
@@ -30,6 +31,16 @@ public SpanScope startSpan(String spanName) {
         return SpanScope.NO_OP;
     }
 
+    @Override
+    public SpanContext getCurrentSpan() {
+        return null;
+    }
+
+    @Override
+    public SpanScope startSpan(String spanName, SpanContext parentSpan) {
+        return SpanScope.NO_OP;
+    }
+
     @Override
     public void close() {
 

libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/runnable/TraceableRunnable.java

@@ -0,0 +1,44 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.telemetry.tracing.runnable;
+
+import org.opensearch.telemetry.tracing.SpanContext;
+import org.opensearch.telemetry.tracing.SpanScope;
+import org.opensearch.telemetry.tracing.Tracer;
+
+/**
+ * Wraps the runnable and add instrumentation to trace the {@link Runnable}
+ */
+public class TraceableRunnable implements Runnable {
+    private final Runnable runnable;
+    private final SpanContext parent;
+    private final Tracer tracer;
+    private final String spanName;
+
+    /**
+     * Constructor.
+     * @param tracer tracer
+     * @param spanName spanName
+     * @param parent parent Span.
+     * @param runnable runnable.
+     */
+    public TraceableRunnable(Tracer tracer, String spanName, SpanContext parent, Runnable runnable) {
+        this.tracer = tracer;
+        this.spanName = spanName;
+        this.parent = parent;
+        this.runnable = runnable;
+    }
+
+    @Override
+    public void run() {
+        try (SpanScope spanScope = tracer.startSpan(spanName, parent)) {
+            runnable.run();
+        }
+    }
+}

libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/runnable/package-info.java

@@ -0,0 +1,12 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+/**
+ * Contains tracing related classes
+ */
+package org.opensearch.telemetry.tracing.runnable;

libs/telemetry/src/test/java/org/opensearch/telemetry/tracing/DefaultTracerTests.java

@@ -9,9 +9,12 @@
 package org.opensearch.telemetry.tracing;
 
 import org.junit.Assert;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.common.util.concurrent.ThreadContext;
 import org.opensearch.test.OpenSearchTestCase;
 
 import java.io.IOException;
+import org.opensearch.test.telemetry.tracing.MockTracingTelemetry;
 
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -40,7 +43,37 @@ public void testCreateSpan() {
 
         defaultTracer.startSpan("span_name");
 
-        Assert.assertEquals("span_name", defaultTracer.getCurrentSpan().getSpanName());
+        Assert.assertEquals("span_name", defaultTracer.getCurrentSpan().getSpan().getSpanName());
+    }
+
+    public void testCreateSpanWithParent() {
+        TracingTelemetry tracingTelemetry = new MockTracingTelemetry();
+        DefaultTracer defaultTracer = new DefaultTracer(
+            tracingTelemetry,
+            new ThreadContextBasedTracerContextStorage(new ThreadContext(Settings.EMPTY), tracingTelemetry)
+        );
+
+        defaultTracer.startSpan("span_name", null);
+
+        SpanContext parentSpan = defaultTracer.getCurrentSpan();
+
+        defaultTracer.startSpan("span_name_1", parentSpan);
+
+        Assert.assertEquals("span_name_1", defaultTracer.getCurrentSpan().getSpan().getSpanName());
+        Assert.assertEquals(parentSpan.getSpan(), defaultTracer.getCurrentSpan().getSpan().getParentSpan());
+    }
+
+    public void testCreateSpanWithNullParent() {
+        TracingTelemetry tracingTelemetry = new MockTracingTelemetry();
+        DefaultTracer defaultTracer = new DefaultTracer(
+            tracingTelemetry,
+            new ThreadContextBasedTracerContextStorage(new ThreadContext(Settings.EMPTY), tracingTelemetry)
+        );
+
+        defaultTracer.startSpan("span_name", null);
+
+        Assert.assertEquals("span_name", defaultTracer.getCurrentSpan().getSpan().getSpanName());
+        Assert.assertEquals(null, defaultTracer.getCurrentSpan().getSpan().getParentSpan());
     }
 
     public void testEndSpanByClosingScope() {

libs/telemetry/src/test/java/org/opensearch/telemetry/tracing/TraceableRunnableTests.java

@@ -0,0 +1,60 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.telemetry.tracing;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.common.util.concurrent.ThreadContext;
+import org.opensearch.telemetry.tracing.runnable.TraceableRunnable;
+import org.opensearch.test.OpenSearchTestCase;
+import org.opensearch.test.telemetry.tracing.MockTracingTelemetry;
+
+public class TraceableRunnableTests extends OpenSearchTestCase {
+
+    private final ThreadContextBasedTracerContextStorage contextStorage = new ThreadContextBasedTracerContextStorage(
+        new ThreadContext(Settings.EMPTY),
+        new MockTracingTelemetry()
+    );
+
+    public void testRunnableWithNullParent() throws Exception {
+        String spanName = "testRunnable";
+        DefaultTracer defaultTracer = new DefaultTracer(new MockTracingTelemetry(), contextStorage);
+        final AtomicBoolean isRunnableCompleted = new AtomicBoolean(false);
+        TraceableRunnable traceableRunnable = new TraceableRunnable(
+            defaultTracer,
+            spanName,
+            null,
+            () -> { isRunnableCompleted.set(true); }
+        );
+        traceableRunnable.run();
+        assertTrue(isRunnableCompleted.get());
+        assertEquals(spanName, defaultTracer.getCurrentSpan().getSpan().getSpanName());
+        assertEquals(null, defaultTracer.getCurrentSpan().getSpan().getParentSpan());
+    }
+
+    public void testRunnableWithParent() throws Exception {
+        String spanName = "testRunnable";
+        String parentSpanName = "parentSpan";
+        DefaultTracer defaultTracer = new DefaultTracer(new MockTracingTelemetry(), contextStorage);
+        defaultTracer.startSpan(parentSpanName);
+        SpanContext parentSpan = defaultTracer.getCurrentSpan();
+        AtomicReference<SpanContext> currrntSpan = new AtomicReference<>(new SpanContext(null));
+        final AtomicBoolean isRunnableCompleted = new AtomicBoolean(false);
+        TraceableRunnable traceableRunnable = new TraceableRunnable(defaultTracer, spanName, parentSpan, () -> {
+            isRunnableCompleted.set(true);
+            currrntSpan.set(defaultTracer.getCurrentSpan());
+        });
+        traceableRunnable.run();
+        assertTrue(isRunnableCompleted.get());
+        assertEquals(spanName, currrntSpan.get().getSpan().getSpanName());
+        assertEquals(parentSpan.getSpan(), currrntSpan.get().getSpan().getParentSpan());
+        assertEquals(parentSpan.getSpan(), defaultTracer.getCurrentSpan().getSpan());
+    }
+}

server/src/main/java/org/opensearch/telemetry/tracing/WrappedTracer.java

@@ -36,8 +36,19 @@ public WrappedTracer(TelemetrySettings telemetrySettings, Tracer defaultTracer)
 
     @Override
     public SpanScope startSpan(String spanName) {
+        return startSpan(spanName, null);
+    }
+
+    @Override
+    public SpanContext getCurrentSpan() {
+        Tracer delegateTracer = getDelegateTracer();
+        return delegateTracer.getCurrentSpan();
+    }
+
+    @Override
+    public SpanScope startSpan(String spanName, SpanContext parentSpan) {
         Tracer delegateTracer = getDelegateTracer();
-        return delegateTracer.startSpan(spanName);
+        return delegateTracer.startSpan(spanName, parentSpan);
     }
 
     @Override