Adds mock implementation for TelemetryPlugin

Signed-off-by: Gagan Juneja <gjjuneja@amazon.com>

server/src/main/java/org/opensearch/node/Node.java

@@ -1500,7 +1500,7 @@ public synchronized void close() throws IOException {
         toClose.add(injector.getInstance(NodeEnvironment.class));
         toClose.add(stopWatch::stop);
         if (FeatureFlags.isEnabled(TELEMETRY)) {
-            toClose.add(() -> injector.getInstance(TracerFactory.class));
+            toClose.add(injector.getInstance(TracerFactory.class));
         }
 
         if (logger.isTraceEnabled()) {

test/framework/src/main/java/org/opensearch/test/OpenSearchIntegTestCase.java

@@ -152,6 +152,7 @@
 import org.opensearch.test.disruption.ServiceDisruptionScheme;
 import org.opensearch.test.store.MockFSIndexStore;
 import org.opensearch.test.transport.MockTransportService;
+import org.opensearch.tracing.MockTelemetryPlugin;
 import org.opensearch.transport.TransportInterceptor;
 import org.opensearch.transport.TransportRequest;
 import org.opensearch.transport.TransportRequestHandler;
@@ -2101,6 +2102,7 @@ protected Collection<Class<? extends Plugin>> getMockPlugins() {
         if (addMockGeoShapeFieldMapper()) {
             mocks.add(TestGeoShapeFieldMapperPlugin.class);
         }
+        mocks.add(MockTelemetryPlugin.class);
 
         return Collections.unmodifiableList(mocks);
     }

test/framework/src/main/java/org/opensearch/test/OpenSearchSingleNodeTestCase.java

@@ -66,6 +66,8 @@
 import org.opensearch.plugins.Plugin;
 import org.opensearch.script.MockScriptService;
 import org.opensearch.search.internal.SearchContext;
+import org.opensearch.telemetry.TelemetrySettings;
+import org.opensearch.tracing.MockTelemetryPlugin;
 import org.opensearch.transport.TransportSettings;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
@@ -242,6 +244,7 @@ private Node newNode() {
             .put(HierarchyCircuitBreakerService.USE_REAL_MEMORY_USAGE_SETTING.getKey(), false)
             .putList(DISCOVERY_SEED_HOSTS_SETTING.getKey()) // empty list disables a port scan for other nodes
             .putList(INITIAL_CLUSTER_MANAGER_NODES_SETTING.getKey(), nodeName)
+            .put(TelemetrySettings.TRACER_ENABLED_SETTING.getKey(), true)
             .put(nodeSettings()) // allow test cases to provide their own settings or override these
             .build();
 
@@ -254,6 +257,7 @@ private Node newNode() {
             plugins.add(MockHttpTransport.TestPlugin.class);
         }
         plugins.add(MockScriptService.TestPlugin.class);
+        plugins.add(MockTelemetryPlugin.class);
         Node node = new MockNode(settings, plugins, forbidPrivateIndexSettings());
         try {
             node.start();

test/framework/src/main/java/org/opensearch/test/OpenSearchTestCase.java

@@ -290,6 +290,7 @@ private static void setTestSysProps() {
 
         // Enable Netty leak detection and monitor logger for logged leak errors
         System.setProperty("io.netty.leakDetection.level", "paranoid");
+        System.setProperty("opensearch.experimental.feature.telemetry.enabled", "true");
     }
 
     protected final Logger logger = LogManager.getLogger(getClass());

test/framework/src/main/java/org/opensearch/tracing/MockSpan.java

@@ -0,0 +1,149 @@
+/*
+ * 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.tracing;
+
+import java.util.HashMap;
+import java.util.Map;
+import org.opensearch.common.Randomness;
+import org.opensearch.telemetry.tracing.AbstractSpan;
+import org.opensearch.telemetry.tracing.Span;
+
+/**
+ * MockSpan for testing and strict check validations. Not to be used for production cases.
+ */
+public class MockSpan extends AbstractSpan {
+    private final SpanProcessor spanProcessor;
+    private final Map<String, Object> metadata;
+    private final String traceId;
+    private final String spanId;
+    private boolean hasEnded;
+    private final Long startTime;
+    private Long endTime;
+
+    private final Object lock = new Object();
+
+    /**
+     * Base Constructor.
+     * @param spanName span name
+     * @param parentSpan parent span
+     * @param spanProcessor span processor
+     */
+    public MockSpan(String spanName, Span parentSpan, SpanProcessor spanProcessor) {
+        this(spanName, parentSpan, IdGenerator.generateTraceId(), IdGenerator.generateSpanId(), spanProcessor);
+    }
+
+    /**
+     * Constructor with traceId and SpanIds
+     * @param spanName  Span Name
+     * @param parentSpan  Parent Span
+     * @param traceId  Trace ID
+     * @param spanId  Span ID
+     * @param spanProcessor  Span Processor
+     */
+    public MockSpan(String spanName, Span parentSpan, String traceId, String spanId, SpanProcessor spanProcessor) {
+        super(spanName, parentSpan);
+        this.spanProcessor = spanProcessor;
+        metadata = new HashMap<>();
+        this.traceId = traceId;
+        this.spanId = spanId;
+        startTime = System.nanoTime();
+    }
+
+    @Override
+    public void endSpan() {
+        synchronized (lock) {
+            if (hasEnded) {
+                return;
+            }
+            endTime = System.nanoTime();
+            hasEnded = true;
+        }
+        spanProcessor.onEnd(this);
+    }
+
+    @Override
+    public void addAttribute(String key, String value) {
+        putMetadata(key, value);
+    }
+
+    @Override
+    public void addAttribute(String key, Long value) {
+        putMetadata(key, value);
+    }
+
+    @Override
+    public void addAttribute(String key, Double value) {
+        putMetadata(key, value);
+    }
+
+    @Override
+    public void addAttribute(String key, Boolean value) {
+        putMetadata(key, value);
+    }
+
+    @Override
+    public void addEvent(String event) {
+        putMetadata(event, null);
+    }
+
+    private void putMetadata(String key, Object value) {
+        metadata.put(key, value);
+    }
+
+    @Override
+    public String getTraceId() {
+        return traceId;
+    }
+
+    @Override
+    public String getSpanId() {
+        return spanId;
+    }
+
+    /**
+     * Returns whether the span is ended or not.
+     * @return span end status.
+     */
+    public boolean hasEnded() {
+        synchronized (lock) {
+            return hasEnded;
+        }
+    }
+
+    /**
+     * Returns the start time of the span.
+     * @return start time of the span.
+     */
+    public Long getStartTime() {
+        return startTime;
+    }
+
+    /**
+     * Returns the start time of the span.
+     * @return end time of the span.
+     */
+    public Long getEndTime() {
+        return endTime;
+    }
+
+    private static class IdGenerator {
+        private static String generateSpanId() {
+            long id = Randomness.get().nextLong();
+            return Long.toHexString(id);
+        }
+
+        private static String generateTraceId() {
+            long idHi = Randomness.get().nextLong();
+            long idLo = Randomness.get().nextLong();
+            long result = idLo | (idHi << 32);
+            return Long.toHexString(result);
+        }
+
+    }
+}

test/framework/src/main/java/org/opensearch/tracing/MockTelemetry.java

@@ -0,0 +1,37 @@
+/*
+ * 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.tracing;
+
+import org.opensearch.telemetry.Telemetry;
+import org.opensearch.telemetry.TelemetrySettings;
+import org.opensearch.telemetry.metrics.MetricsTelemetry;
+import org.opensearch.telemetry.tracing.TracingTelemetry;
+
+/**
+ * Mock {@link Telemetry} implementation for testing.
+ */
+public class MockTelemetry implements Telemetry {
+
+    private final TelemetrySettings settings;
+
+    public MockTelemetry(TelemetrySettings settings) {
+        this.settings = settings;
+    }
+
+    @Override
+    public TracingTelemetry getTracingTelemetry() {
+        return new MockTracingTelemetry();
+    }
+
+    @Override
+    public MetricsTelemetry getMetricsTelemetry() {
+        return new MetricsTelemetry() {
+        };
+    }
+}

test/framework/src/main/java/org/opensearch/tracing/MockTelemetryPlugin.java

@@ -0,0 +1,32 @@
+/*
+ * 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.tracing;
+
+import java.util.Optional;
+import org.opensearch.plugins.Plugin;
+import org.opensearch.plugins.TelemetryPlugin;
+import org.opensearch.telemetry.Telemetry;
+import org.opensearch.telemetry.TelemetrySettings;
+
+/**
+ * Mock {@link TelemetryPlugin} implementation for testing.
+ */
+public class MockTelemetryPlugin extends Plugin implements TelemetryPlugin {
+    private static final String MOCK_TRACER_NAME = "mock";
+
+    @Override
+    public Optional<Telemetry> getTelemetry(TelemetrySettings settings) {
+        return Optional.of(new MockTelemetry(settings));
+    }
+
+    @Override
+    public String getName() {
+        return MOCK_TRACER_NAME;
+    }
+}

test/framework/src/main/java/org/opensearch/tracing/MockTracingContextPropagator.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.tracing;
+
+import java.util.Locale;
+import java.util.Map;
+import java.util.function.BiConsumer;
+import org.opensearch.telemetry.tracing.Span;
+import org.opensearch.telemetry.tracing.TracingContextPropagator;
+
+/**
+ * Mock {@link TracingContextPropagator} to persist the span for internode communication.
+ */
+public class MockTracingContextPropagator implements TracingContextPropagator {
+
+    private static final String TRACE_PARENT = "traceparent";
+    private static final String SEPARATOR = "~";
+    private final SpanProcessor spanProcessor;
+
+    public MockTracingContextPropagator(SpanProcessor spanProcessor) {
+        this.spanProcessor = spanProcessor;
+    }
+
+    @Override
+    public Span extract(Map<String, String> props) {
+        String value = props.get(TRACE_PARENT);
+        if (value != null) {
+            String[] values = value.split(SEPARATOR);
+            String traceId = values[0];
+            String spanId = values[1];
+            return new MockSpan(null, null, traceId, spanId, spanProcessor);
+        } else {
+            return null;
+        }
+    }
+
+    @Override
+    public void inject(Span currentSpan, BiConsumer<String, String> setter) {
+        if (currentSpan instanceof MockSpan) {
+            String traceId = currentSpan.getTraceId();
+            String spanId = currentSpan.getSpanId();
+            String traceParent = String.format(Locale.ROOT, "%s%s%s", traceId, TRACE_PARENT, spanId);
+            setter.accept(TRACE_PARENT, traceParent);
+        }
+    }
+}

test/framework/src/main/java/org/opensearch/tracing/MockTracingTelemetry.java

@@ -0,0 +1,39 @@
+/*
+ * 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.tracing;
+
+import org.opensearch.telemetry.tracing.Span;
+import org.opensearch.telemetry.tracing.TracingContextPropagator;
+import org.opensearch.telemetry.tracing.TracingTelemetry;
+
+/**
+ * Mock {@link TracingTelemetry} implementation for testing.
+ */
+public class MockTracingTelemetry implements TracingTelemetry {
+
+    private final SpanProcessor spanProcessor = new StrictCheckSpanProcessor();
+
+    @Override
+    public Span createSpan(String spanName, Span parentSpan) {
+        Span span = new MockSpan(spanName, parentSpan, spanProcessor);
+        spanProcessor.onStart(span);
+        return span;
+    }
+
+    @Override
+    public TracingContextPropagator getContextPropagator() {
+        return new MockTracingContextPropagator(spanProcessor);
+    }
+
+    @Override
+    public void close() {
+        ((StrictCheckSpanProcessor) spanProcessor).ensureAllSpansAreClosed();
+        ((StrictCheckSpanProcessor) spanProcessor).clear();
+    }
+}