Adding new inferred sampler

CHANGELOG.md

@@ -103,8 +103,24 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ## [Unreleased 2.x]
 ### Added
+- [Tiered caching] Introducing cache plugins and exposing Ehcache as one of the pluggable disk cache option ([#11874](https://github.com/opensearch-project/OpenSearch/pull/11874))
+- Add support for dependencies in plugin descriptor properties with semver range ([#11441](https://github.com/opensearch-project/OpenSearch/pull/11441))
+- Add community_id ingest processor ([#12121](https://github.com/opensearch-project/OpenSearch/pull/12121))
+- Introduce query level setting `index.query.max_nested_depth` limiting nested queries ([#3268](https://github.com/opensearch-project/OpenSearch/issues/3268)
+- Add toString methods to MultiSearchRequest, MultiGetRequest and CreateIndexRequest ([#12163](https://github.com/opensearch-project/OpenSearch/pull/12163))
+- Support for returning scores in matched queries ([#11626](https://github.com/opensearch-project/OpenSearch/pull/11626))
+- Add shard id property to SearchLookup for use in field types provided by plugins ([#1063](https://github.com/opensearch-project/OpenSearch/pull/1063))
+- Force merge API supports performing on primary shards only ([#11269](https://github.com/opensearch-project/OpenSearch/pull/11269))
+- [Tiered caching] Make IndicesRequestCache implementation configurable [EXPERIMENTAL] ([#12533](https://github.com/opensearch-project/OpenSearch/pull/12533))
+- Add kuromoji_completion analyzer and filter ([#4835](https://github.com/opensearch-project/OpenSearch/issues/4835))
+- The org.opensearch.bootstrap.Security should support codebase for JAR files with classifiers ([#12586](https://github.com/opensearch-project/OpenSearch/issues/12586))
+- [Metrics Framework] Adds support for asynchronous gauge metric type. ([#12642](https://github.com/opensearch-project/OpenSearch/issues/12642))
+- Make search query counters dynamic to support all query types ([#12601](https://github.com/opensearch-project/OpenSearch/pull/12601))
+- [Tiered caching] Add policies controlling which values can enter pluggable caches [EXPERIMENTAL] ([#12542](https://github.com/opensearch-project/OpenSearch/pull/12542))
+- [Tiered caching] Add Stale keys Management and CacheCleaner to IndicesRequestCache ([#12625](https://github.com/opensearch-project/OpenSearch/pull/12625))
 - 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))
 - [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 markSampled() {
+        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/DefaultSpanScope.java

@@ -9,7 +9,10 @@
 package org.opensearch.telemetry.tracing;
 
 import org.opensearch.common.annotation.InternalApi;
+import org.opensearch.telemetry.tracing.attributes.AttributeType;
 
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Objects;
 
 /**
@@ -24,6 +27,7 @@
     private final Span beforeSpan;
     private static final ThreadLocal<SpanScope> spanScopeThreadLocal = new ThreadLocal<>();
     private final TracerContextStorage<String, Span> tracerContextStorage;
+    private final Map<String, AttributeType> commonAttributeMap = new HashMap<>();
 
     /**
      * Constructor
@@ -40,6 +44,7 @@
         this.beforeSpan = beforeSpan;
         this.previousSpanScope = previousSpanScope;
         this.tracerContextStorage = tracerContextStorage;
+        initializeCommonPropagationAttributes();
     }
 
     /**
@@ -55,6 +60,13 @@
         return newSpanScope;
     }
 
+    /**
+     * Common attributes need to be taken from parent and propagated to child span*
+     */
+    private void initializeCommonPropagationAttributes() {
+        commonAttributeMap.put(TracerContextStorage.INFERRED_SAMPLER, AttributeType.BOOLEAN);
+    }
+
     @Override
     public void close() {
         detach();
@@ -64,9 +76,38 @@
     public SpanScope attach() {
         spanScopeThreadLocal.set(this);
         tracerContextStorage.put(TracerContextStorage.CURRENT_SPAN, this.span);
+        addCommonParentAttributes();
         return this;
     }
 
+    private void addCommonParentAttributes() {
+        // This work as common attribute propagator from parent to child
+        for (String attribute : commonAttributeMap.keySet()) {
+            if (beforeSpan != null && beforeSpan.getAttributes().containsKey(attribute)) {
+                AttributeType attributeValue = commonAttributeMap.get(attribute);
+                this.storeAttributeValue(attribute, attributeValue);
+            }
+        }
+    }
+
+    private void storeAttributeValue(String attribute, AttributeType attributeType) {
+        switch (attributeType) {
+            case BOOLEAN:
+                span.addAttribute(attribute, (Boolean) beforeSpan.getAttribute(attribute));
+                break;
+            case LONG:
+                span.addAttribute(attribute, (Long) beforeSpan.getAttribute(attribute));
+                break;
+            case STRING:
+                span.addAttribute(attribute, (String) beforeSpan.getAttribute(attribute));
+                break;
+            case DOUBLE:
+                span.addAttribute(attribute, (Double) beforeSpan.getAttribute(attribute));
+                break;
+            // Add more cases for other types if needed
+        }
+    }
+
     private void detach() {
         spanScopeThreadLocal.set(previousSpanScope);
         if (beforeSpan != null) {

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

@@ -104,7 +104,13 @@
     @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);
+        }
+    }
 }

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

@@ -10,6 +10,8 @@
 
 import org.opensearch.common.annotation.ExperimentalApi;
 
+import java.util.Map;
+
 /**
  * An interface that represents a tracing span.
  * Spans are created by the Tracer.startSpan method.
@@ -93,4 +95,16 @@ public interface Span {
      */
     String getSpanId();
 
+    /**
+     * Returns attribute.
+     * @param key key
+     * @return value
+     */
+    Object getAttribute(String key);
+
+    /**
+     * Returns the attributes as map.
+     * @return returns the attributes map.
+     */
+    Map<String, Object> getAttributes();
 }

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/attributes/AttributeType.java

@@ -0,0 +1,38 @@
+/*
+ * 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.PublicApi;
+
+/**
+ * Type of Attribute.
+ */
+@PublicApi(since = "2.14.0")
+public enum AttributeType {
+
+    /**
+     * Attribute type represents the Boolean attribute.
+     */
+    BOOLEAN,
+
+    /**
+     * Attribute type represents the Long attribute.
+     */
+    LONG,
+
+    /**
+     * Attribute type represents the String attribute.
+     */
+    STRING,
+
+    /**
+     * Attribute type represents the Double attribute.
+     */
+    DOUBLE
+}

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

@@ -11,6 +11,8 @@
 import org.opensearch.common.annotation.InternalApi;
 import org.opensearch.telemetry.tracing.Span;
 
+import java.util.Map;
+
 /**
  * No-op implementation of {@link org.opensearch.telemetry.tracing.Span}
  *
@@ -82,4 +84,25 @@
     public String getSpanId() {
         return "noop-span-id";
     }
+
+    /**
+     * Returns attribute.
+     *
+     * @param key key
+     * @return value
+     */
+    @Override
+    public Object getAttribute(String key) {
+        return null;
+    }
+
+    /**
+     * Returns the attributes as map.
+     *
+     * @return returns the attributes map.
+     */
+    @Override
+    public Map<String, Object> getAttributes() {
+        return null;
+    }
 }

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,6 +8,9 @@
 
 package org.opensearch.telemetry.tracing;
 
+import java.util.HashMap;
+import java.util.Map;
+
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.StatusCode;
 
@@ -19,6 +22,8 @@
 
     private final Span delegateSpan;
 
+    private final Map<String, Object> metadata;
+
     /**
      * Constructor
      * @param spanName span name
@@ -28,31 +33,47 @@
     public OTelSpan(String spanName, Span span, org.opensearch.telemetry.tracing.Span parentSpan) {
         super(spanName, parentSpan);
         this.delegateSpan = span;
+        this.metadata = new HashMap<>();
     }
 
     @Override
     public void endSpan() {
+        if (getAttributes().containsKey(TracerContextStorage.SAMPLED)) {
+            markParentForSampling();
+        }
         delegateSpan.end();
     }
 
+    private void markParentForSampling() {
+        org.opensearch.telemetry.tracing.Span current_parent = getParentSpan();
+        while (current_parent != null && !current_parent.getAttributes().containsKey(TracerContextStorage.SAMPLED)) {
+            current_parent.addAttribute(TracerContextStorage.SAMPLED, true);
+            current_parent = current_parent.getParentSpan();
+        }
+    }
+
     @Override
     public void addAttribute(String key, String value) {
         delegateSpan.setAttribute(key, value);
+        metadata.put(key, value);
     }
 
     @Override
     public void addAttribute(String key, Long value) {
         delegateSpan.setAttribute(key, value);
+        metadata.put(key, value);
     }
 
     @Override
     public void addAttribute(String key, Double value) {
         delegateSpan.setAttribute(key, value);
+        metadata.put(key, value);
     }
 
     @Override
     public void addAttribute(String key, Boolean value) {
         delegateSpan.setAttribute(key, value);
+        metadata.put(key, value);
     }
 
     @Override
@@ -77,8 +98,17 @@
         return delegateSpan.getSpanContext().getSpanId();
     }
 
+    @Override
+    public Object getAttribute(String key) {
+        return metadata.get(key);
+    }
+
+    @Override
+    public Map<String, Object> getAttributes() {
+        return metadata;
+    }
+
     io.opentelemetry.api.trace.Span getDelegateSpan() {
         return delegateSpan;
     }
-
 }

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

@@ -15,6 +15,7 @@
 import java.io.Closeable;
 import java.io.IOException;
 
+import io.opentelemetry.api.baggage.Baggage;
 import io.opentelemetry.api.trace.TracerProvider;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.sdk.OpenTelemetrySdk;
@@ -63,9 +64,22 @@
             OTelSpanKindConverter.convert(spanCreationContext.getSpanKind())
         );
         Span newSpan = new OTelSpan(spanCreationContext.getSpanName(), otelSpan, parentSpan);
+
+        // This attribute is added for inferred sampling
+        addInferredAttribute(newSpan, spanCreationContext);
+
         return newSpan;
     }
 
+    private void addInferredAttribute(Span newSpan, SpanCreationContext spanCreationContext) {
+        // If the current context has this attribute we need to add the same to the span as well.
+        if (Baggage.current().getEntryValue(TracerContextStorage.INFERRED_SAMPLER) != null
+            || (spanCreationContext.getAttributes() != null
+                && spanCreationContext.getAttributes().getAttributesMap().containsKey(TracerContextStorage.INFERRED_SAMPLER))) {
+            newSpan.addAttribute(TracerContextStorage.INFERRED_SAMPLER, true);
+        }
+    }
+
     io.opentelemetry.api.trace.Span otelSpan(
         String spanName,
         Span parentOTelSpan,