Add wrapped tracer implementation (#8565) (#8622)

* Add wrapped tracer implementation



* Add changelog entry



* Add @opensearch.internal annotation



* Fix test



* Fix changelog entry



---------

Signed-off-by: suranjay <[email protected]>
Signed-off-by: Andriy Redko <[email protected]>
Co-authored-by: Andriy Redko <[email protected]>

CHANGELOG.md

@@ -104,6 +104,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - [Refactor] OpenSearchException streamables to a registry ([#7646](https://github.com/opensearch-project/OpenSearch/pull/7646))
 - [Refactor] Serverless common classes to libraries ([#8065](https://github.com/opensearch-project/OpenSearch/pull/8065))
 - [Refactor] StreamIO and OpenSearchException foundation to core library ([#8035](https://github.com/opensearch-project/OpenSearch/pull/8035))
+- Move span actions to Scope ([#8411](https://github.com/opensearch-project/OpenSearch/pull/8411))
+- Add wrapper tracer implementation ([#8565](https://github.com/opensearch-project/OpenSearch/pull/8565))
 - Fix painless casting bug, which crashes the OpenSearch process ([#8315](https://github.com/opensearch-project/OpenSearch/pull/8315))
 
 ### Deprecated

libs/telemetry/src/main/java/org/opensearch/telemetry/Telemetry.java

@@ -13,6 +13,8 @@
 
 /**
  * Interface defining telemetry
+ *
+ * @opensearch.internal
  */
 public interface Telemetry {
 

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

@@ -10,6 +10,8 @@
 
 /**
  * Base span
+ *
+ * @opensearch.internal
  */
 public abstract class AbstractSpan implements Span {
 

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

@@ -12,8 +12,10 @@
 
 /**
  * Default implementation of Scope
+ *
+ * @opensearch.internal
  */
-public class DefaultSpanScope implements SpanScope {
+final class DefaultSpanScope implements SpanScope {
 
     private final Span span;
 

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

@@ -16,9 +16,9 @@
  * The default tracer implementation. It handles tracing context propagation between spans by maintaining
  * current active span in its storage
  *
- *
+ *  @opensearch.internal
  */
-public class DefaultTracer implements Tracer {
+class DefaultTracer implements Tracer {
     static final String THREAD_NAME = "th_name";
 
     private final TracingTelemetry tracingTelemetry;

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

@@ -11,7 +11,9 @@
 /**
  * An interface that represents a tracing span.
  * Spans are created by the Tracer.startSpan method.
- * Span must be ended by calling Tracer.endSpan which internally calls Span's endSpan.
+ * Span must be ended by calling SpanScope.close which internally calls Span's endSpan.
+ *
+ * @opensearch.internal
  */
 public interface Span {
 

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

@@ -10,8 +10,10 @@
 
 /**
  * Wrapper class to hold reference of Span
+ *
+ * @opensearch.internal
  */
-public class SpanReference {
+final class SpanReference {
 
     private Span span;
 

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

@@ -12,6 +12,8 @@
  * Storage interface used for storing tracing context
  * @param <K> key type
  * @param <V> value type
+ *
+ * @opensearch.internal
  */
 public interface TracerContextStorage<K, V> {
     /**

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

@@ -13,6 +13,8 @@
 
 /**
  * Interface defining the tracing related context propagation
+ *
+ * @opensearch.internal
  */
 public interface TracingContextPropagator {
 

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

@@ -12,6 +12,8 @@
 
 /**
  * Interface for tracing telemetry providers
+ *
+ * @opensearch.internal
  */
 public interface TracingTelemetry extends Closeable {
 

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

@@ -12,6 +12,8 @@
 
 /**
  * No-op implementation of SpanScope
+ *
+ * @opensearch.internal
  */
 public final class NoopSpanScope implements SpanScope {
 

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

@@ -13,6 +13,8 @@
 
 /**
  * No-op implementation of Tracer
+ *
+ * @opensearch.internal
  */
 public class NoopTracer implements Tracer {
 

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

@@ -14,6 +14,8 @@
 
 /**
  * No-op implementation of TracerFactory
+ *
+ * @opensearch.internal
  */
 public class NoopTracerFactory extends TracerFactory {
     public NoopTracerFactory() {

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

@@ -18,6 +18,8 @@
 
 /**
  * Core's ThreadContext based TracerContextStorage implementation
+ *
+ * @opensearch.internal
  */
 public class ThreadContextBasedTracerContextStorage implements TracerContextStorage<String, Span>, ThreadContextStatePropagator {
 

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

@@ -21,7 +21,7 @@
 
 /**
  * TracerManager represents a single global class that is used to access tracers.
- *
+ * <p>
  * The Tracer singleton object can be retrieved using tracerManager.getTracer(). The TracerManager object
  * is created during class initialization and cannot subsequently be changed.
  */
@@ -30,21 +30,20 @@ public class TracerFactory implements Closeable {
     private static final Logger logger = LogManager.getLogger(TracerFactory.class);
 
     private final TelemetrySettings telemetrySettings;
-    private final Tracer defaultTracer;
+    private final Tracer tracer;
 
     public TracerFactory(TelemetrySettings telemetrySettings, Optional<Telemetry> telemetry, ThreadContext threadContext) {
         this.telemetrySettings = telemetrySettings;
-        this.defaultTracer = telemetry.map(Telemetry::getTracingTelemetry)
-            .map(tracingTelemetry -> createDefaultTracer(tracingTelemetry, threadContext))
-            .orElse(NoopTracer.INSTANCE);
+        this.tracer = tracer(telemetry, threadContext);
     }
 
     /**
      * Returns the tracer instance
+     *
      * @return tracer instance
      */
     public Tracer getTracer() {
-        return telemetrySettings.isTracingEnabled() ? defaultTracer : NoopTracer.INSTANCE;
+        return tracer;
     }
 
     /**
@@ -53,12 +52,19 @@ public Tracer getTracer() {
     @Override
     public void close() {
         try {
-            defaultTracer.close();
+            tracer.close();
         } catch (IOException e) {
             logger.warn("Error closing tracer", e);
         }
     }
 
+    private Tracer tracer(Optional<Telemetry> telemetry, ThreadContext threadContext) {
+        return telemetry.map(Telemetry::getTracingTelemetry)
+            .map(tracingTelemetry -> createDefaultTracer(tracingTelemetry, threadContext))
+            .map(defaultTracer -> createWrappedTracer(defaultTracer))
+            .orElse(NoopTracer.INSTANCE);
+    }
+
     private Tracer createDefaultTracer(TracingTelemetry tracingTelemetry, ThreadContext threadContext) {
         TracerContextStorage<String, Span> tracerContextStorage = new ThreadContextBasedTracerContextStorage(
             threadContext,
@@ -67,4 +73,8 @@ private Tracer createDefaultTracer(TracingTelemetry tracingTelemetry, ThreadCont
         return new DefaultTracer(tracingTelemetry, tracerContextStorage);
     }
 
+    private Tracer createWrappedTracer(Tracer defaultTracer) {
+        return new WrappedTracer(telemetrySettings, defaultTracer);
+    }
+
 }

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

@@ -0,0 +1,52 @@
+/*
+ * 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 org.opensearch.telemetry.TelemetrySettings;
+import org.opensearch.telemetry.tracing.noop.NoopTracer;
+
+import java.io.IOException;
+
+/**
+ * Wrapper implementation of Tracer. This delegates call to right tracer based on the tracer settings
+ *
+ * @opensearch.internal
+ */
+final class WrappedTracer implements Tracer {
+
+    private final Tracer defaultTracer;
+    private final TelemetrySettings telemetrySettings;
+
+    /**
+     * Creates WrappedTracer instance
+     *
+     * @param telemetrySettings telemetry settings
+     * @param defaultTracer default tracer instance
+     */
+    public WrappedTracer(TelemetrySettings telemetrySettings, Tracer defaultTracer) {
+        this.defaultTracer = defaultTracer;
+        this.telemetrySettings = telemetrySettings;
+    }
+
+    @Override
+    public SpanScope startSpan(String spanName) {
+        Tracer delegateTracer = getDelegateTracer();
+        return delegateTracer.startSpan(spanName);
+    }
+
+    @Override
+    public void close() throws IOException {
+        defaultTracer.close();
+    }
+
+    // visible for testing
+    Tracer getDelegateTracer() {
+        return telemetrySettings.isTracingEnabled() ? defaultTracer : NoopTracer.INSTANCE;
+    }
+}

server/src/test/java/org/opensearch/telemetry/tracing/TracerFactoryTests.java

@@ -36,27 +36,28 @@ public void close() {
         tracerFactory.close();
     }
 
-    public void testGetTracerWithTracingDisabledReturnsNoopTracer() {
+    public void testGetTracerWithUnavailableTracingTelemetryReturnsNoopTracer() {
         Settings settings = Settings.builder().put(TelemetrySettings.TRACER_ENABLED_SETTING.getKey(), false).build();
         TelemetrySettings telemetrySettings = new TelemetrySettings(settings, new ClusterSettings(settings, getClusterSettings()));
         Telemetry mockTelemetry = mock(Telemetry.class);
         when(mockTelemetry.getTracingTelemetry()).thenReturn(mock(TracingTelemetry.class));
-        tracerFactory = new TracerFactory(telemetrySettings, Optional.of(mockTelemetry), new ThreadContext(Settings.EMPTY));
+        tracerFactory = new TracerFactory(telemetrySettings, Optional.empty(), new ThreadContext(Settings.EMPTY));
 
         Tracer tracer = tracerFactory.getTracer();
+
         assertTrue(tracer instanceof NoopTracer);
         assertTrue(tracer.startSpan("foo") == SpanScope.NO_OP);
     }
 
-    public void testGetTracerWithTracingEnabledReturnsDefaultTracer() {
+    public void testGetTracerWithAvailableTracingTelemetryReturnsWrappedTracer() {
         Settings settings = Settings.builder().put(TelemetrySettings.TRACER_ENABLED_SETTING.getKey(), true).build();
         TelemetrySettings telemetrySettings = new TelemetrySettings(settings, new ClusterSettings(settings, getClusterSettings()));
         Telemetry mockTelemetry = mock(Telemetry.class);
         when(mockTelemetry.getTracingTelemetry()).thenReturn(mock(TracingTelemetry.class));
         tracerFactory = new TracerFactory(telemetrySettings, Optional.of(mockTelemetry), new ThreadContext(Settings.EMPTY));
 
         Tracer tracer = tracerFactory.getTracer();
-        assertTrue(tracer instanceof DefaultTracer);
+        assertTrue(tracer instanceof WrappedTracer);
 
     }
 

server/src/test/java/org/opensearch/telemetry/tracing/WrappedTracerTests.java

@@ -0,0 +1,69 @@
+/*
+ * 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 org.opensearch.common.settings.ClusterSettings;
+import org.opensearch.common.settings.Setting;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.common.util.FeatureFlags;
+import org.opensearch.telemetry.TelemetrySettings;
+import org.opensearch.telemetry.tracing.noop.NoopTracer;
+import org.opensearch.test.OpenSearchTestCase;
+
+import java.io.IOException;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+
+public class WrappedTracerTests extends OpenSearchTestCase {
+
+    public void testStartSpanWithTracingDisabledInvokesNoopTracer() throws Exception {
+        Settings settings = Settings.builder().put(TelemetrySettings.TRACER_ENABLED_SETTING.getKey(), false).build();
+        TelemetrySettings telemetrySettings = new TelemetrySettings(settings, new ClusterSettings(settings, getClusterSettings()));
+        DefaultTracer mockDefaultTracer = mock(DefaultTracer.class);
+
+        try (WrappedTracer wrappedTracer = new WrappedTracer(telemetrySettings, mockDefaultTracer)) {
+            wrappedTracer.startSpan("foo");
+            assertTrue(wrappedTracer.getDelegateTracer() instanceof NoopTracer);
+            verify(mockDefaultTracer, never()).startSpan("foo");
+        }
+    }
+
+    public void testStartSpanWithTracingEnabledInvokesDefaultTracer() throws Exception {
+        Settings settings = Settings.builder().put(TelemetrySettings.TRACER_ENABLED_SETTING.getKey(), true).build();
+        TelemetrySettings telemetrySettings = new TelemetrySettings(settings, new ClusterSettings(settings, getClusterSettings()));
+        DefaultTracer mockDefaultTracer = mock(DefaultTracer.class);
+
+        try (WrappedTracer wrappedTracer = new WrappedTracer(telemetrySettings, mockDefaultTracer)) {
+            wrappedTracer.startSpan("foo");
+
+            assertTrue(wrappedTracer.getDelegateTracer() instanceof DefaultTracer);
+            verify(mockDefaultTracer).startSpan("foo");
+        }
+    }
+
+    public void testClose() throws IOException {
+        DefaultTracer mockDefaultTracer = mock(DefaultTracer.class);
+        WrappedTracer wrappedTracer = new WrappedTracer(null, mockDefaultTracer);
+
+        wrappedTracer.close();
+
+        verify(mockDefaultTracer).close();
+    }
+
+    private Set<Setting<?>> getClusterSettings() {
+        Set<Setting<?>> allTracerSettings = new HashSet<>();
+        ClusterSettings.FEATURE_FLAGGED_CLUSTER_SETTINGS.get(List.of(FeatureFlags.TELEMETRY)).stream().forEach((allTracerSettings::add));
+        return allTracerSettings;
+    }
+}