Adding new inferred sampler

CHANGELOG.md

@@ -108,6 +108,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Allow setting KEYSTORE_PASSWORD through env variable ([#12865](https://github.com/opensearch-project/OpenSearch/pull/12865))
 - [Concurrent Segment Search] Perform buildAggregation concurrently and support Composite Aggregations ([#12697](https://github.com/opensearch-project/OpenSearch/pull/12697))
 - [Concurrent Segment Search] Disable concurrent segment search for system indices and throttled requests ([#12954](https://github.com/opensearch-project/OpenSearch/pull/12954))
+- Tracing Framework] Adds support for inferred sampling ([#12315](https://github.com/opensearch-project/OpenSearch/issues/12315))
 
 ### Dependencies
 - Bump `org.apache.commons:commons-configuration2` from 2.10.0 to 2.10.1 ([#12896](https://github.com/opensearch-project/OpenSearch/pull/12896))

libs/core/src/main/java/org/opensearch/core/transport/TransportMessage.java

@@ -36,6 +36,9 @@
 import org.opensearch.core.common.io.stream.Writeable;
 import org.opensearch.core.common.transport.TransportAddress;
 
+import java.util.Collections;
+import java.util.Map;
+
 /**
  * Message over the transport interface
  *
@@ -45,14 +48,24 @@ public abstract class TransportMessage implements Writeable {
 
     private TransportAddress remoteAddress;
 
+    private Map<String, String> header = Collections.emptyMap();
+
     public void remoteAddress(TransportAddress remoteAddress) {
         this.remoteAddress = remoteAddress;
     }
 
+    public void setResponseHeaders(Map<String, String> header) {
+        this.header = header;
+    }
+
     public TransportAddress remoteAddress() {
         return remoteAddress;
     }
 
+    public Map<String, String> getResponseHeaders() {
+        return header;
+    }
+
     /**
      * Constructs a new empty transport message
      */

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

@@ -9,6 +9,7 @@
 package org.opensearch.telemetry.tracing;
 
 import org.opensearch.common.annotation.InternalApi;
+import org.opensearch.telemetry.tracing.attributes.SamplingAttributes;
 
 import java.io.Closeable;
 import java.io.IOException;
@@ -52,8 +53,9 @@
         } else {
             parentSpan = getCurrentSpanInternal();
         }
+
         Span span = createSpan(context, parentSpan);
-        addDefaultAttributes(span);
+        addDefaultAttributes(parentSpan, span);
         return span;
     }
 
@@ -97,14 +99,32 @@
      * Adds default attributes in the span
      * @param span the current active span
      */
-    protected void addDefaultAttributes(Span span) {
+    protected void addDefaultAttributes(Span parentSpan, Span span) {
+        copyInheritableParentAttributes(parentSpan, span);
         span.addAttribute(THREAD_NAME, Thread.currentThread().getName());
     }
 
     @Override
     public Span startSpan(SpanCreationContext spanCreationContext, Map<String, Collection<String>> headers) {
         Optional<Span> propagatedSpan = tracingTelemetry.getContextPropagator().extractFromHeaders(headers);
+        addRequestAttributeToContext(spanCreationContext, headers);
         return startSpan(spanCreationContext.parent(propagatedSpan.map(SpanContext::new).orElse(null)));
     }
 
+    private void addRequestAttributeToContext(SpanCreationContext spanCreationContext, Map<String, Collection<String>> headers) {
+        if (headers != null && headers.containsKey(SamplingAttributes.SAMPLER.getValue())) {
+            spanCreationContext.getAttributes()
+                .addAttribute(SamplingAttributes.SAMPLER.getValue(), SamplingAttributes.INFERRED_SAMPLER.getValue());
+        }
+    }
+
+    private void copyInheritableParentAttributes(Span parentSpan, Span currentSpan) {
+        // This work as common attribute propagator from parent to child
+        if (parentSpan != null) {
+            Optional<String> inferredAttribute = Optional.ofNullable(parentSpan.getAttributeString(SamplingAttributes.SAMPLER.getValue()));
+            if (inferredAttribute.isPresent()) {
+                currentSpan.addAttribute(SamplingAttributes.SAMPLER.getValue(), SamplingAttributes.INFERRED_SAMPLER.getValue());
+            }
+        }
+    }
 }

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

@@ -93,4 +93,31 @@ public interface Span {
      */
     String getSpanId();
 
+    /**
+     * *
+     * @param key for which we need to look for value
+     * @return string attribute value
+     */
+    String getAttributeString(String key);
+
+    /**
+     * *
+     * @param key for which we need to look for value
+     * @return Boolean attribute value
+     */
+    Boolean getAttributeBoolean(String key);
+
+    /**
+     * *
+     * @param key for which we need to look for value
+     * @return Long attribute value
+     */
+    Long getAttributeLong(String key);
+
+    /**
+     * *
+     * @param key for which we need to look for value
+     * @return Double attribute value
+     */
+    Double getAttributeDouble(String key);
 }

libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/attributes/SamplingAttributes.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.attributes;
+
+import org.opensearch.common.annotation.InternalApi;
+
+import java.util.Locale;
+
+/**
+ * Enum for Inferred Sampling*
+ * @opensearch.internal*
+ */
+@InternalApi
+public enum SamplingAttributes {
+
+    /**
+     * Attribute added if the span is sampled by inferred sampler*
+     */
+    INFERRED_SAMPLER,
+
+    /**
+     * Attribute Added if the span is an outlier*
+     */
+    SAMPLED,
+
+    /**
+     * Sampler Used in the framework*
+     */
+    SAMPLER;
+
+    /**
+     * returns lower case enum value*
+     * @return String
+     */
+    public String getValue() {
+        return name().toLowerCase(Locale.ROOT);
+    }
+}

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

@@ -82,4 +82,24 @@ public String getTraceId() {
     public String getSpanId() {
         return "noop-span-id";
     }
+
+    @Override
+    public String getAttributeString(String key) {
+        return "";
+    }
+
+    @Override
+    public Boolean getAttributeBoolean(String key) {
+        return false;
+    }
+
+    @Override
+    public Long getAttributeLong(String key) {
+        return 0L;
+    }
+
+    @Override
+    public Double getAttributeDouble(String key) {
+        return 0.0;
+    }
 }

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

@@ -13,12 +13,17 @@
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.node.Node;
 import org.opensearch.telemetry.tracing.attributes.Attributes;
+import org.opensearch.telemetry.tracing.attributes.SamplingAttributes;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.test.telemetry.tracing.MockSpan;
 import org.opensearch.test.telemetry.tracing.MockTracingTelemetry;
 import org.opensearch.threadpool.ThreadPool;
 
 import java.io.IOException;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 
@@ -424,4 +429,59 @@ private SpanCreationContext buildSpanCreationContext(String spanName, Attributes
         }
         return spanCreationContext;
     }
+
+    public void testCreateSpanWithInferredAttributes() {
+        TracingTelemetry tracingTelemetry = new MockTracingTelemetry();
+        ThreadContext threadContext = new ThreadContext(Settings.EMPTY);
+        DefaultTracer defaultTracer = new DefaultTracer(
+            tracingTelemetry,
+            new ThreadContextBasedTracerContextStorage(threadContext, tracingTelemetry)
+        );
+
+        SpanCreationContext spanCreationContext = buildSpanCreationContext(
+            "span_name",
+            Attributes.create().addAttribute(SamplingAttributes.SAMPLER.getValue(), SamplingAttributes.INFERRED_SAMPLER.getValue()),
+            null
+        );
+
+        Span span = defaultTracer.startSpan(spanCreationContext);
+
+        assertThat(defaultTracer.getCurrentSpan(), is(nullValue()));
+        assertEquals(SamplingAttributes.INFERRED_SAMPLER.getValue(), ((MockSpan) span).getAttribute(SamplingAttributes.SAMPLER.getValue()));
+        span.endSpan();
+    }
+
+    @SuppressWarnings("unchecked")
+    public void testCreateSpanWithInferredSampledParent() {
+        TracingTelemetry tracingTelemetry = new MockTracingTelemetry();
+        DefaultTracer defaultTracer = new DefaultTracer(
+            tracingTelemetry,
+            new ThreadContextBasedTracerContextStorage(new ThreadContext(Settings.EMPTY), tracingTelemetry)
+        );
+
+        SpanCreationContext spanCreationContext = buildSpanCreationContext("span_name", null, null);
+
+        Span span = defaultTracer.startSpan(
+            spanCreationContext,
+            (Map<String, Collection<String>>) new HashMap<String, Collection<String>>().put(
+                SamplingAttributes.SAMPLER.getValue(),
+                Collections.singleton(SamplingAttributes.INFERRED_SAMPLER.getValue())
+            )
+        );
+
+        try (final SpanScope scope = defaultTracer.withSpanInScope(span)) {
+            SpanContext parentSpan = defaultTracer.getCurrentSpan();
+            SpanCreationContext spanCreationContext1 = buildSpanCreationContext("span_name_1", Attributes.EMPTY, parentSpan.getSpan());
+            Span span1 = defaultTracer.startSpan(spanCreationContext1);
+            assertEquals("span_name_1", span1.getSpanName());
+            assertEquals(parentSpan.getSpan(), span1.getParentSpan());
+            assertEquals(
+                SamplingAttributes.INFERRED_SAMPLER.getValue(),
+                ((MockSpan) span1).getAttribute(SamplingAttributes.SAMPLER.getValue())
+            );
+            span1.endSpan();
+        } finally {
+            span.endSpan();
+        }
+    }
 }

plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelResourceProvider.java

@@ -12,6 +12,7 @@
 import org.opensearch.telemetry.TelemetrySettings;
 import org.opensearch.telemetry.metrics.exporter.OTelMetricsExporterFactory;
 import org.opensearch.telemetry.tracing.exporter.OTelSpanExporterFactory;
+import org.opensearch.telemetry.tracing.processor.OTelSpanProcessor;
 import org.opensearch.telemetry.tracing.sampler.OTelSamplerFactory;
 import org.opensearch.telemetry.tracing.sampler.RequestSampler;
 
@@ -117,7 +118,11 @@ private static SdkTracerProvider createSdkTracerProvider(
             .build();
     }
 
-    private static BatchSpanProcessor spanProcessor(Settings settings, SpanExporter spanExporter) {
+    private static OTelSpanProcessor spanProcessor(Settings settings, SpanExporter spanExporter) {
+        return new OTelSpanProcessor(batchSpanProcessor(settings, spanExporter));
+    }
+
+    private static BatchSpanProcessor batchSpanProcessor(Settings settings, SpanExporter spanExporter) {
         return BatchSpanProcessor.builder(spanExporter)
             .setScheduleDelay(TRACER_EXPORTER_DELAY_SETTING.get(settings).getSeconds(), TimeUnit.SECONDS)
             .setMaxExportBatchSize(TRACER_EXPORTER_BATCH_SIZE_SETTING.get(settings))

plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelSpan.java

@@ -8,8 +8,14 @@
 
 package org.opensearch.telemetry.tracing;
 
+import org.opensearch.telemetry.tracing.attributes.SamplingAttributes;
+
+import java.util.Optional;
+
+import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.StatusCode;
+import io.opentelemetry.sdk.trace.ReadableSpan;
 
 /**
  * Default implementation of {@link Span} using Otel span. It keeps a reference of OpenTelemetry Span and handles span
@@ -32,9 +38,33 @@ public OTelSpan(String spanName, Span span, org.opensearch.telemetry.tracing.Spa
 
     @Override
     public void endSpan() {
+        if (isSpanOutlier()) {
+            markParentForSampling();
+        }
         delegateSpan.end();
     }
 
+    /*
+     * This is added temporarily will remove this after the evaluation framework PR.
+     * This Framework will be used to evaluate a span if that is an outlier or not.
+     */
+    private boolean isSpanOutlier() {
+        Optional<Boolean> isSpanSampled = Optional.ofNullable(getAttributeBoolean(SamplingAttributes.SAMPLED.getValue()));
+        Optional<String> isSpanInferredSampled = Optional.ofNullable(getAttributeString(SamplingAttributes.SAMPLER.getValue()));
+
+        return isSpanSampled.isPresent()
+            && isSpanInferredSampled.isPresent()
+            && isSpanInferredSampled.get().equals(SamplingAttributes.INFERRED_SAMPLER.getValue());
+    }
+
+    private void markParentForSampling() {
+        org.opensearch.telemetry.tracing.Span currentParent = getParentSpan();
+        while (currentParent != null && currentParent.getAttributeBoolean(SamplingAttributes.SAMPLED.getValue()) == null) {
+            currentParent.addAttribute(SamplingAttributes.SAMPLED.getValue(), true);
+            currentParent = currentParent.getParentSpan();
+        }
+    }
+
     @Override
     public void addAttribute(String key, String value) {
         delegateSpan.setAttribute(key, value);
@@ -77,8 +107,43 @@ public String getSpanId() {
         return delegateSpan.getSpanContext().getSpanId();
     }
 
+    @Override
+    public String getAttributeString(String key) {
+        if (delegateSpan != null && delegateSpan instanceof ReadableSpan) return ((ReadableSpan) delegateSpan).getAttribute(
+            AttributeKey.stringKey(key)
+        );
+
+        return null;
+    }
+
+    @Override
+    public Boolean getAttributeBoolean(String key) {
+        if (delegateSpan != null && delegateSpan instanceof ReadableSpan) {
+            return ((ReadableSpan) delegateSpan).getAttribute(AttributeKey.booleanKey(key));
+        }
+
+        return null;
+    }
+
+    @Override
+    public Long getAttributeLong(String key) {
+        if (delegateSpan != null && delegateSpan instanceof ReadableSpan) return ((ReadableSpan) delegateSpan).getAttribute(
+            AttributeKey.longKey(key)
+        );
+
+        return null;
+    }
+
+    @Override
+    public Double getAttributeDouble(String key) {
+        if (delegateSpan != null && delegateSpan instanceof ReadableSpan) return ((ReadableSpan) delegateSpan).getAttribute(
+            AttributeKey.doubleKey(key)
+        );
+
+        return null;
+    }
+
     io.opentelemetry.api.trace.Span getDelegateSpan() {
         return delegateSpan;
     }
-
 }