Adding new inferred sampler

CHANGELOG.md

@@ -107,6 +107,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add a counter to node stat api to track shard going from idle to non-idle ([#12768](https://github.com/opensearch-project/OpenSearch/pull/12768))
 - 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))
+- [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

@@ -45,14 +45,24 @@
 
     private TransportAddress remoteAddress;
 
+    private boolean sampled;
+
     public void remoteAddress(TransportAddress remoteAddress) {
         this.remoteAddress = remoteAddress;
     }
 
+    public void markResponseAsSampled() {
+        this.sampled = true;
+    }
+
     public TransportAddress remoteAddress() {
         return remoteAddress;
     }
 
+    public boolean sampled() {
+        return sampled;
+    }
+
     /**
      * Constructs a new empty transport message
      */

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

@@ -52,8 +52,9 @@
         } else {
             parentSpan = getCurrentSpanInternal();
         }
+
         Span span = createSpan(context, parentSpan);
-        addDefaultAttributes(span);
+        addDefaultAttributes(parentSpan, span);
         return span;
     }
 
@@ -97,14 +98,33 @@
      * Adds default attributes in the span
      * @param span the current active span
      */
-    protected void addDefaultAttributes(Span span) {
+    protected void addDefaultAttributes(Span parentSpan, Span span) {
+        addCommonParentAttributes(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(TracerContextStorage.INFERRED_SAMPLER)) {
+            spanCreationContext.getAttributes().addAttribute(TracerContextStorage.INFERRED_SAMPLER, true);
+        }
+    }
+
+    private void addCommonParentAttributes(Span parentSpan, Span currentSpan) {
+        // This work as common attribute propagator from parent to child
+        if (parentSpan != null) {
+            Optional<Boolean> inferredAttribute = Optional.ofNullable(
+                parentSpan.getAttributeBoolean(TracerContextStorage.INFERRED_SAMPLER)
+            );
+            if (inferredAttribute.isPresent()) {
+                currentSpan.addAttribute(TracerContextStorage.INFERRED_SAMPLER, true);
+            }
+        }
+    }
 }

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/TracerContextStorage.java

@@ -24,6 +24,16 @@ public interface TracerContextStorage<K, V> {
      */
     String CURRENT_SPAN = "current_span";
 
+    /**
+     * Key for storing sample information
+     */
+    String SAMPLED = "sampled";
+
+    /**
+     * Key for storing inferred Sampling information
+     */
+    String INFERRED_SAMPLER = "Inferred_sampler";
+
     /**
      * Fetches value corresponding to key
      * @param key of the tracing context

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

@@ -82,4 +82,24 @@
     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;
+    }
 }

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,10 @@
 
 package org.opensearch.telemetry.tracing;
 
+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 +34,20 @@
 
     @Override
     public void endSpan() {
+        if (getAttributeBoolean(TracerContextStorage.SAMPLED) != null && getAttributeBoolean(TracerContextStorage.SAMPLED)) {
+            markParentForSampling();
+        }
         delegateSpan.end();
     }
 
+    private void markParentForSampling() {
+        org.opensearch.telemetry.tracing.Span current_parent = getParentSpan();
+        while (current_parent != null && current_parent.getAttributeBoolean(TracerContextStorage.SAMPLED) == null) {
+            current_parent.addAttribute(TracerContextStorage.SAMPLED, true);
+            current_parent = current_parent.getParentSpan();
+        }
+    }
+
     @Override
     public void addAttribute(String key, String value) {
         delegateSpan.setAttribute(key, value);
@@ -77,8 +90,43 @@
         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;
     }
-
 }

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

@@ -62,7 +62,9 @@ private Span createOtelSpan(SpanCreationContext spanCreationContext, Span parent
             OTelAttributesConverter.convert(spanCreationContext.getAttributes()),
             OTelSpanKindConverter.convert(spanCreationContext.getSpanKind())
         );
+
         Span newSpan = new OTelSpan(spanCreationContext.getSpanName(), otelSpan, parentSpan);
+
         return newSpan;
     }
 

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

@@ -0,0 +1,113 @@
+/*
+ * 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.processor;
+
+import org.opensearch.telemetry.tracing.TracerContextStorage;
+
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.sdk.common.CompletableResultCode;
+import io.opentelemetry.sdk.trace.ReadWriteSpan;
+import io.opentelemetry.sdk.trace.ReadableSpan;
+import io.opentelemetry.sdk.trace.SpanProcessor;
+
+/**
+ * Implementation of the SpanProcessor and delegates to the configured processor.
+ */
+public class OtelSpanProcessor implements SpanProcessor {
+
+    private final SpanProcessor delegateProcessor;
+
+    /**
+     * *
+     * @param delegateProcessor the span processor to which this processor will delegate
+     */
+    public OtelSpanProcessor(SpanProcessor delegateProcessor) {
+        this.delegateProcessor = delegateProcessor;
+    }
+
+    /**
+     * Called when a {@link Span} is started, if the {@link
+     * Span#isRecording()} returns true.
+     *
+     * <p>This method is called synchronously on the execution thread, should not throw or block the
+     * execution thread.
+     *
+     * @param parentContext the parent {@code Context} of the span that just started.
+     * @param span          the {@code Span} that just started.
+     */
+    @Override
+    public void onStart(Context parentContext, ReadWriteSpan span) {
+        this.delegateProcessor.onStart(parentContext, span);
+    }
+
+    /**
+     * Returns {@code true} if this {@link SpanProcessor} requires start events.
+     *
+     * @return {@code true} if this {@link SpanProcessor} requires start events.
+     */
+    @Override
+    public boolean isStartRequired() {
+        return this.delegateProcessor.isStartRequired();
+    }
+
+    /**
+     * Called when a {@link Span} is ended, if the {@link
+     * Span#isRecording()} returns true.
+     *
+     * <p>This method is called synchronously on the execution thread, should not throw or block the
+     * execution thread.
+     *
+     * @param span the {@code Span} that just ended.
+     */
+    @Override
+    public void onEnd(ReadableSpan span) {
+        if (span != null
+            && span.getSpanContext().isSampled()
+            && Boolean.TRUE.equals(span.getAttribute(AttributeKey.booleanKey(TracerContextStorage.INFERRED_SAMPLER)))) {
+            if (span.getAttribute(AttributeKey.booleanKey(TracerContextStorage.SAMPLED)) != null) {
+                this.delegateProcessor.onEnd(span);
+            }
+        } else {
+            this.delegateProcessor.onEnd(span);
+        }
+    }
+
+    /**
+     * Returns {@code true} if this {@link SpanProcessor} requires end events.
+     *
+     * @return {@code true} if this {@link SpanProcessor} requires end events.
+     */
+    @Override
+    public boolean isEndRequired() {
+        return this.delegateProcessor.isEndRequired();
+    }
+
+    /**
+     * Processes all span events that have not yet been processed and closes used resources.
+     *
+     * @return a {@link CompletableResultCode} which completes when shutdown is finished.
+     */
+    @Override
+    public CompletableResultCode shutdown() {
+        return this.delegateProcessor.shutdown();
+    }
+
+    /**
+     * Processes all span events that have not yet been processed.
+     *
+     * @return a {@link CompletableResultCode} which completes when currently queued spans are
+     *     finished processing.
+     */
+    @Override
+    public CompletableResultCode forceFlush() {
+        return this.delegateProcessor.forceFlush();
+    }
+}

plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/processor/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.
+ */
+
+/**
+ * This package contains classes needed for processor.
+ */
+package org.opensearch.telemetry.tracing.processor;