diff --git a/CHANGELOG.md b/CHANGELOG.md index f44c993f6faa0..a68a413f1f76e 100644 --- a/CHANGELOG.md +++ b/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)) diff --git a/libs/core/src/main/java/org/opensearch/core/transport/TransportMessage.java b/libs/core/src/main/java/org/opensearch/core/transport/TransportMessage.java index 941babda40aa3..2203e3f445387 100644 --- a/libs/core/src/main/java/org/opensearch/core/transport/TransportMessage.java +++ b/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 header = Collections.emptyMap(); + public void remoteAddress(TransportAddress remoteAddress) { this.remoteAddress = remoteAddress; } + public void setResponseHeaders(Map header) { + this.header = header; + } + public TransportAddress remoteAddress() { return remoteAddress; } + public Map getResponseHeaders() { + return header; + } + /** * Constructs a new empty transport message */ diff --git a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultTracer.java b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultTracer.java index 8f1a26d99e725..db256184be969 100644 --- a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultTracer.java +++ b/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 @@ public Span startSpan(SpanCreationContext context) { } else { parentSpan = getCurrentSpanInternal(); } + Span span = createSpan(context, parentSpan); - addDefaultAttributes(span); + addDefaultAttributes(parentSpan, span); return span; } @@ -97,14 +99,32 @@ private Span createSpan(SpanCreationContext spanCreationContext, Span parentSpan * 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> headers) { Optional propagatedSpan = tracingTelemetry.getContextPropagator().extractFromHeaders(headers); + addRequestAttributeToContext(spanCreationContext, headers); return startSpan(spanCreationContext.parent(propagatedSpan.map(SpanContext::new).orElse(null))); } + private void addRequestAttributeToContext(SpanCreationContext spanCreationContext, Map> 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 inferredAttribute = Optional.ofNullable(parentSpan.getAttributeString(SamplingAttributes.SAMPLER.getValue())); + if (inferredAttribute.isPresent()) { + currentSpan.addAttribute(SamplingAttributes.SAMPLER.getValue(), SamplingAttributes.INFERRED_SAMPLER.getValue()); + } + } + } } diff --git a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/Span.java b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/Span.java index 00b64492c281e..d7e17ae10a743 100644 --- a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/Span.java +++ b/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); } diff --git a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/attributes/SamplingAttributes.java b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/attributes/SamplingAttributes.java new file mode 100644 index 0000000000000..08a712a408124 --- /dev/null +++ b/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); + } +} diff --git a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/noop/NoopSpan.java b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/noop/NoopSpan.java index f41e11017d155..33d8e7dbeaacd 100644 --- a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/noop/NoopSpan.java +++ b/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; + } } diff --git a/libs/telemetry/src/test/java/org/opensearch/telemetry/tracing/DefaultTracerTests.java b/libs/telemetry/src/test/java/org/opensearch/telemetry/tracing/DefaultTracerTests.java index 2182b3ea28ac8..f088cd6054e76 100644 --- a/libs/telemetry/src/test/java/org/opensearch/telemetry/tracing/DefaultTracerTests.java +++ b/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>) new HashMap>().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(); + } + } } diff --git a/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelResourceProvider.java b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelResourceProvider.java index 475fc09d04bff..51150662ff9c6 100644 --- a/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelResourceProvider.java +++ b/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)) diff --git a/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelSpan.java b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelSpan.java index fc917968579e1..10b95c309061f 100644 --- a/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelSpan.java +++ b/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 isSpanSampled = Optional.ofNullable(getAttributeBoolean(SamplingAttributes.SAMPLED.getValue())); + Optional 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; } - } diff --git a/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/processor/OTelSpanProcessor.java b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/processor/OTelSpanProcessor.java new file mode 100644 index 0000000000000..42d302458bf4b --- /dev/null +++ b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/processor/OTelSpanProcessor.java @@ -0,0 +1,118 @@ +/* + * 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.attributes.SamplingAttributes; + +import java.util.Objects; + +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. + * + *

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. + * + *

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() + && Objects.equals( + span.getAttribute(AttributeKey.stringKey(SamplingAttributes.SAMPLER.getValue())), + SamplingAttributes.INFERRED_SAMPLER.getValue() + )) { + if (span.getAttribute(AttributeKey.booleanKey(SamplingAttributes.SAMPLED.getValue())) != 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(); + } +} diff --git a/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/processor/package-info.java b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/processor/package-info.java new file mode 100644 index 0000000000000..7dce3e15922e8 --- /dev/null +++ b/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; diff --git a/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/sampler/InferredActionSampler.java b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/sampler/InferredActionSampler.java new file mode 100644 index 0000000000000..04baef85cb14d --- /dev/null +++ b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/sampler/InferredActionSampler.java @@ -0,0 +1,87 @@ +/* + * 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.sampler; + +import org.opensearch.common.settings.Settings; +import org.opensearch.telemetry.TelemetrySettings; +import org.opensearch.telemetry.tracing.attributes.SamplingAttributes; +import org.opensearch.telemetry.tracing.samplingResult.OTelSamplingResult; + +import java.util.List; +import java.util.Objects; + +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.context.Context; +import io.opentelemetry.sdk.trace.data.LinkData; +import io.opentelemetry.sdk.trace.samplers.Sampler; +import io.opentelemetry.sdk.trace.samplers.SamplingResult; + +/** + * InferredActionSampler implements a probability sampling strategy with sampling ratio as 1.0. + */ +public class InferredActionSampler implements Sampler { + + private final Sampler fallbackSampler; + private final TelemetrySettings telemetrySettings; + private final Settings settings; + + /** + * Constructor + * @param telemetrySettings the telemetry settings + * @param settings the settings + * @param fallbackSampler the fallback sampler + */ + private InferredActionSampler(TelemetrySettings telemetrySettings, Settings settings, Sampler fallbackSampler) { + this.telemetrySettings = Objects.requireNonNull(telemetrySettings); + this.settings = Objects.requireNonNull(settings); + this.fallbackSampler = fallbackSampler; + } + + /** + * Create Inferred sampler. + * + * @param telemetrySettings the telemetry settings + * @param settings the settings + * @param fallbackSampler the fallback sampler + * @return the inferred sampler + */ + public static Sampler create(TelemetrySettings telemetrySettings, Settings settings, Sampler fallbackSampler) { + return new InferredActionSampler(telemetrySettings, settings, fallbackSampler); + } + + @Override + public SamplingResult shouldSample( + Context parentContext, + String traceId, + String name, + SpanKind spanKind, + Attributes attributes, + List parentLinks + ) { + boolean inferredSamplingAllowListed = telemetrySettings.getInferredSamplingAllowListed(); + if (inferredSamplingAllowListed) { + Attributes customSampleAttributes = Attributes.builder() + .put(SamplingAttributes.SAMPLER.getValue(), SamplingAttributes.INFERRED_SAMPLER.getValue()) + .putAll(attributes) + .build(); + SamplingResult result = SamplingResult.recordAndSample(); + return new OTelSamplingResult(result.getDecision(), customSampleAttributes); + } else if (fallbackSampler != null) { + return fallbackSampler.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks); + } + + return SamplingResult.drop(); + } + + @Override + public String getDescription() { + return "Inferred Action Sampler"; + } +} diff --git a/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/samplingResult/OTelSamplingResult.java b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/samplingResult/OTelSamplingResult.java new file mode 100644 index 0000000000000..f115452491dad --- /dev/null +++ b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/samplingResult/OTelSamplingResult.java @@ -0,0 +1,55 @@ +/* + * 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.samplingResult; + +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.sdk.trace.samplers.SamplingDecision; +import io.opentelemetry.sdk.trace.samplers.SamplingResult; + +/** + * Custom Sampling Result Class* + */ +public class OTelSamplingResult implements SamplingResult { + + private final SamplingDecision samplingDecision; + + private final Attributes attributes; + + /** + * Constructor* + * @param samplingDecision decision that needs to be added + * @param attributes attribute list that needs to be added + */ + public OTelSamplingResult(SamplingDecision samplingDecision, Attributes attributes) { + this.samplingDecision = samplingDecision; + this.attributes = attributes; + } + + /** + * Return decision on whether a span should be recorded, recorded and sampled or not recorded. + * + * @return sampling result. + */ + @Override + public SamplingDecision getDecision() { + return samplingDecision; + } + + /** + * Return tags which will be attached to the span. + * + * @return attributes added to span. These attributes should be added to the span only when + * {@linkplain #getDecision() the sampling decision} is {@link SamplingDecision#RECORD_ONLY} + * or {@link SamplingDecision#RECORD_AND_SAMPLE}. + */ + @Override + public Attributes getAttributes() { + return attributes; + } +} diff --git a/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/samplingResult/package-info.java b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/samplingResult/package-info.java new file mode 100644 index 0000000000000..0623a60556711 --- /dev/null +++ b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/samplingResult/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 is the package for defining custom sampling result* + */ +package org.opensearch.telemetry.tracing.samplingResult; diff --git a/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/OTelTelemetryPluginTests.java b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/OTelTelemetryPluginTests.java index 4a1301588dad2..3bca5814ee824 100644 --- a/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/OTelTelemetryPluginTests.java +++ b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/OTelTelemetryPluginTests.java @@ -36,6 +36,7 @@ import static org.opensearch.telemetry.OTelTelemetrySettings.TRACER_EXPORTER_MAX_QUEUE_SIZE_SETTING; import static org.opensearch.telemetry.OTelTelemetrySettings.TRACER_SAMPLER_ACTION_PROBABILITY; import static org.opensearch.telemetry.TelemetrySettings.TRACER_ENABLED_SETTING; +import static org.opensearch.telemetry.TelemetrySettings.TRACER_INFERRED_SAMPLER_ALLOWLISTED; import static org.opensearch.telemetry.TelemetrySettings.TRACER_SAMPLER_PROBABILITY; public class OTelTelemetryPluginTests extends OpenSearchTestCase { @@ -50,10 +51,16 @@ public class OTelTelemetryPluginTests extends OpenSearchTestCase { public void setup() { // TRACER_EXPORTER_DELAY_SETTING should always be less than 10 seconds because // io.opentelemetry.sdk.OpenTelemetrySdk.close waits only for 10 seconds for shutdown to complete. - Settings settings = Settings.builder().put(TRACER_EXPORTER_DELAY_SETTING.getKey(), "1s").build(); + Settings settings = Settings.builder().put(TRACER_EXPORTER_DELAY_SETTING.getKey(), "2s").build(); oTelTelemetryPlugin = new OTelTelemetryPlugin(settings); telemetry = oTelTelemetryPlugin.getTelemetry( - new TelemetrySettings(Settings.EMPTY, new ClusterSettings(settings, Set.of(TRACER_ENABLED_SETTING, TRACER_SAMPLER_PROBABILITY))) + new TelemetrySettings( + Settings.EMPTY, + new ClusterSettings( + settings, + Set.of(TRACER_ENABLED_SETTING, TRACER_SAMPLER_PROBABILITY, TRACER_INFERRED_SAMPLER_ALLOWLISTED) + ) + ) ); tracingTelemetry = telemetry.get().getTracingTelemetry(); metricsTelemetry = telemetry.get().getMetricsTelemetry(); diff --git a/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/OTelSpanTests.java b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/OTelSpanTests.java index fc92ab36908e1..24ce823b44567 100644 --- a/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/OTelSpanTests.java +++ b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/OTelSpanTests.java @@ -8,12 +8,15 @@ package org.opensearch.telemetry.tracing; +import org.opensearch.telemetry.tracing.attributes.SamplingAttributes; import org.opensearch.test.OpenSearchTestCase; +import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.SpanContext; import io.opentelemetry.api.trace.TraceFlags; import io.opentelemetry.api.trace.TraceState; +import io.opentelemetry.sdk.trace.ReadWriteSpan; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; @@ -85,6 +88,57 @@ public void testGetSpanId() { assertEquals(SPAN_ID, oTelSpan.getSpanId()); } + public void testGetSpanBoolean() { + Span mockSpan = getMockSpan(); + OTelSpan oTelSpan = new OTelSpan("spanName", mockSpan, null); + assertNull(oTelSpan.getAttributeBoolean("key")); + } + + public void testGetSpanString() { + Span mockSpan = getMockSpan(); + OTelSpan oTelSpan = new OTelSpan("spanName", mockSpan, null); + assertNull(oTelSpan.getAttributeString("key")); + } + + public void testGetSpanLong() { + Span mockSpan = getMockSpan(); + OTelSpan oTelSpan = new OTelSpan("spanName", mockSpan, null); + assertNull(oTelSpan.getAttributeLong("key")); + } + + public void testGetSpanDouble() { + Span mockSpan = getMockSpan(); + OTelSpan oTelSpan = new OTelSpan("spanName", mockSpan, null); + assertNull(oTelSpan.getAttributeDouble("key")); + } + + public void testSpanOutlier() { + ReadWriteSpan mockReadWriteSpan = mock(ReadWriteSpan.class); + Span mockSpan = getMockSpan(); + OTelSpan mockParent = new OTelSpan("parentSpan", mockSpan, null); + OTelSpan oTelSpan = new OTelSpan("spanName", mockReadWriteSpan, mockParent); + when(mockReadWriteSpan.getAttribute(AttributeKey.booleanKey(SamplingAttributes.SAMPLED.getValue()))).thenReturn(true); + when(mockReadWriteSpan.getAttribute(AttributeKey.stringKey(SamplingAttributes.SAMPLER.getValue()))).thenReturn( + SamplingAttributes.INFERRED_SAMPLER.getValue() + ); + oTelSpan.endSpan(); + verify(mockSpan).setAttribute(SamplingAttributes.SAMPLED.getValue(), true); + } + + public void testSpanAttributes() { + ReadWriteSpan mockReadWriteSpan = mock(ReadWriteSpan.class); + OTelSpan oTelSpan = new OTelSpan("spanName", mockReadWriteSpan, null); + oTelSpan.addAttribute("key", 0.0); + when(mockReadWriteSpan.getAttribute(AttributeKey.doubleKey("key"))).thenReturn(0.0); + verify(mockReadWriteSpan).setAttribute("key", 0.0); + assertEquals(0.0, (Object) oTelSpan.getAttributeDouble("key")); + + oTelSpan.addAttribute("key1", 0L); + when(mockReadWriteSpan.getAttribute(AttributeKey.longKey("key1"))).thenReturn(0L); + verify(mockReadWriteSpan).setAttribute("key1", 0L); + assertEquals(0L, (Object) oTelSpan.getAttributeLong("key1")); + } + private Span getMockSpan() { Span mockSpan = mock(Span.class); when(mockSpan.getSpanContext()).thenReturn(SpanContext.create(TRACE_ID, SPAN_ID, TraceFlags.getDefault(), TraceState.getDefault())); diff --git a/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/processor/MockSpanProcessor.java b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/processor/MockSpanProcessor.java new file mode 100644 index 0000000000000..fc165cc9b5d01 --- /dev/null +++ b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/processor/MockSpanProcessor.java @@ -0,0 +1,47 @@ +/* + * 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 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; + +public class MockSpanProcessor implements SpanProcessor { + + @Override + public void onStart(Context parentContext, ReadWriteSpan span) { + {} + } + + @Override + public boolean isStartRequired() { + return false; + } + + @Override + public void onEnd(ReadableSpan span) {} + + @Override + public boolean isEndRequired() { + return false; + } + + @Override + public CompletableResultCode shutdown() { + return CompletableResultCode.ofSuccess(); + } + + @Override + public CompletableResultCode forceFlush() { + return CompletableResultCode.ofSuccess(); + } + +} diff --git a/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/processor/OTelSpanProcessorTests.java b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/processor/OTelSpanProcessorTests.java new file mode 100644 index 0000000000000..d63f88fafd33f --- /dev/null +++ b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/processor/OTelSpanProcessorTests.java @@ -0,0 +1,95 @@ +/* + * 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.attributes.SamplingAttributes; +import org.opensearch.test.OpenSearchTestCase; + +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.trace.SpanContext; +import io.opentelemetry.api.trace.TraceFlags; +import io.opentelemetry.api.trace.TraceState; +import io.opentelemetry.context.Context; +import io.opentelemetry.sdk.trace.ReadWriteSpan; +import io.opentelemetry.sdk.trace.ReadableSpan; +import io.opentelemetry.sdk.trace.SpanProcessor; +import org.mockito.Mockito; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class OTelSpanProcessorTests extends OpenSearchTestCase { + + private OTelSpanProcessor oTelSpanProcessor; + private MockSpanProcessor mockSpanProcessor; + private static final String TRACE_ID = "4aa59968f31dcbff7807741afa9d7d62"; + private static final String SPAN_ID = "bea205cd25756b5e"; + + private ReadableSpan createEndedSpan() { + return Mockito.mock(ReadableSpan.class); + } + + public void testOTelSpanProcessorDelegation() { + mockSpanProcessor = new MockSpanProcessor(); + oTelSpanProcessor = new OTelSpanProcessor(mockSpanProcessor); + assertEquals(mockSpanProcessor.forceFlush(), oTelSpanProcessor.forceFlush()); + assertEquals(mockSpanProcessor.shutdown(), oTelSpanProcessor.shutdown()); + assertEquals(mockSpanProcessor.isEndRequired(), oTelSpanProcessor.isEndRequired()); + assertEquals(mockSpanProcessor.isStartRequired(), oTelSpanProcessor.isStartRequired()); + } + + public void testOnEndFunction() { + SpanProcessor mockProcessor = mock(SpanProcessor.class); + oTelSpanProcessor = new OTelSpanProcessor(mockProcessor); + ReadableSpan readableSpan = this.createEndedSpan(); + when(readableSpan.getSpanContext()).thenReturn( + SpanContext.create(TRACE_ID, SPAN_ID, TraceFlags.getDefault(), TraceState.getDefault()) + ); + oTelSpanProcessor.onEnd(readableSpan); + Mockito.verify(mockProcessor, Mockito.times(1)).onEnd(readableSpan); + } + + public void testOnStartFunction() { + SpanProcessor mockProcessor = mock(SpanProcessor.class); + Context spanContext = mock(Context.class); + oTelSpanProcessor = new OTelSpanProcessor(mockProcessor); + ReadWriteSpan readWriteSpan = mock(ReadWriteSpan.class); + oTelSpanProcessor.onStart(spanContext, readWriteSpan); + Mockito.verify(mockProcessor, Mockito.times(1)).onStart(spanContext, readWriteSpan); + } + + public void testOnEndFunctionWithInferredAttribute() { + SpanProcessor mockProcessor = mock(SpanProcessor.class); + oTelSpanProcessor = new OTelSpanProcessor(mockProcessor); + ReadableSpan readableSpan = this.createEndedSpan(); + when(readableSpan.getSpanContext()).thenReturn( + SpanContext.create(TRACE_ID, SPAN_ID, TraceFlags.getSampled(), TraceState.getDefault()) + ); + when(readableSpan.getAttribute(AttributeKey.stringKey(SamplingAttributes.SAMPLER.getValue()))).thenReturn( + SamplingAttributes.INFERRED_SAMPLER.getValue() + ); + oTelSpanProcessor.onEnd(readableSpan); + Mockito.verify(mockProcessor, Mockito.times(0)).onEnd(readableSpan); + } + + public void testOnEndFunctionWithInferredAttributeAndSampled() { + SpanProcessor mockProcessor = mock(SpanProcessor.class); + oTelSpanProcessor = new OTelSpanProcessor(mockProcessor); + ReadableSpan readableSpan = this.createEndedSpan(); + when(readableSpan.getSpanContext()).thenReturn( + SpanContext.create(TRACE_ID, SPAN_ID, TraceFlags.getSampled(), TraceState.getDefault()) + ); + when(readableSpan.getAttribute(AttributeKey.stringKey(SamplingAttributes.SAMPLER.getValue()))).thenReturn( + SamplingAttributes.INFERRED_SAMPLER.getValue() + ); + when(readableSpan.getAttribute(AttributeKey.booleanKey(SamplingAttributes.SAMPLED.getValue()))).thenReturn(true); + oTelSpanProcessor.onEnd(readableSpan); + Mockito.verify(mockProcessor, Mockito.times(1)).onEnd(readableSpan); + } +} diff --git a/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/InferredSamplerTests.java b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/InferredSamplerTests.java new file mode 100644 index 0000000000000..a1679ffdb0ee9 --- /dev/null +++ b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/InferredSamplerTests.java @@ -0,0 +1,140 @@ +/* + * 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.sampler; + +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Settings; +import org.opensearch.telemetry.TelemetrySettings; +import org.opensearch.telemetry.tracing.attributes.SamplingAttributes; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Collections; +import java.util.Set; + +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.context.Context; +import io.opentelemetry.sdk.trace.samplers.Sampler; +import io.opentelemetry.sdk.trace.samplers.SamplingResult; + +import static org.opensearch.telemetry.TelemetrySettings.TRACER_ENABLED_SETTING; +import static org.opensearch.telemetry.TelemetrySettings.TRACER_INFERRED_SAMPLER_ALLOWLISTED; +import static org.opensearch.telemetry.TelemetrySettings.TRACER_SAMPLER_PROBABILITY; +import static org.opensearch.telemetry.tracing.AttributeNames.TRANSPORT_ACTION; +import static org.mockito.Mockito.mock; + +public class InferredSamplerTests extends OpenSearchTestCase { + + public void testGetSamplerWithSettingDisabled() { + ClusterSettings clusterSettings = new ClusterSettings( + Settings.EMPTY, + Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING, TRACER_INFERRED_SAMPLER_ALLOWLISTED) + ); + + TelemetrySettings telemetrySettings = new TelemetrySettings(Settings.EMPTY, clusterSettings); + + // InferredActionSampler + Sampler inferredActionSampler = InferredActionSampler.create(telemetrySettings, Settings.EMPTY, null); + + SamplingResult result = inferredActionSampler.shouldSample( + mock(Context.class), + "00000000000000000000000000000000", + "spanName", + SpanKind.INTERNAL, + Attributes.builder().put(TRANSPORT_ACTION, "dummy_action").build(), + Collections.emptyList() + ); + + assertEquals(SamplingResult.drop(), result); + } + + public void testGetSamplerWithSettingEnabled() { + ClusterSettings clusterSettings = new ClusterSettings( + Settings.EMPTY, + Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING, TRACER_INFERRED_SAMPLER_ALLOWLISTED) + ); + + TelemetrySettings telemetrySettings = new TelemetrySettings(Settings.EMPTY, clusterSettings); + telemetrySettings.setInferredSamplingAllowListed(true); + + // InferredActionSampler + Sampler inferredActionSampler = InferredActionSampler.create(telemetrySettings, Settings.EMPTY, null); + + SamplingResult result = inferredActionSampler.shouldSample( + mock(Context.class), + "00000000000000000000000000000000", + "spanName", + SpanKind.INTERNAL, + Attributes.builder().put(TRANSPORT_ACTION, "dummy_action").build(), + Collections.emptyList() + ); + + assertEquals(SamplingResult.recordAndSample().getDecision(), result.getDecision()); + } + + public void testGetSamplerWithAddedAttributes() { + ClusterSettings clusterSettings = new ClusterSettings( + Settings.EMPTY, + Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING, TRACER_INFERRED_SAMPLER_ALLOWLISTED) + ); + + TelemetrySettings telemetrySettings = new TelemetrySettings(Settings.EMPTY, clusterSettings); + telemetrySettings.setInferredSamplingAllowListed(true); + + // InferredActionSampler + Sampler inferredActionSampler = InferredActionSampler.create(telemetrySettings, Settings.EMPTY, null); + + SamplingResult result = inferredActionSampler.shouldSample( + mock(Context.class), + "00000000000000000000000000000000", + "spanName", + SpanKind.INTERNAL, + Attributes.builder().put(TRANSPORT_ACTION, "dummy_action").build(), + Collections.emptyList() + ); + + assertTrue(result.getAttributes().asMap().containsKey(AttributeKey.stringKey(SamplingAttributes.SAMPLER.getValue()))); + assertEquals( + result.getAttributes().get(AttributeKey.stringKey(SamplingAttributes.SAMPLER.getValue())), + SamplingAttributes.INFERRED_SAMPLER.getValue() + ); + assertEquals("Inferred Action Sampler", inferredActionSampler.getDescription()); + } + + public void testFallBackSampler() { + ClusterSettings clusterSettings = new ClusterSettings( + Settings.EMPTY, + Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING, TRACER_INFERRED_SAMPLER_ALLOWLISTED) + ); + + TelemetrySettings telemetrySettings = new TelemetrySettings(Settings.EMPTY, clusterSettings); + + // InferredActionSampler + Sampler probabilisticTransportActionSampler = ProbabilisticTransportActionSampler.create(telemetrySettings, Settings.EMPTY, null); + Sampler inferredActionSampler = InferredActionSampler.create( + telemetrySettings, + Settings.EMPTY, + probabilisticTransportActionSampler + ); + + SamplingResult result = inferredActionSampler.shouldSample( + mock(Context.class), + "00000000000000000000000000000000", + "spanName", + SpanKind.INTERNAL, + Attributes.builder().put(TRANSPORT_ACTION, "dummy_action").build(), + Collections.emptyList() + ); + + // ProbabilisticTransportActionSampler + assertEquals(SamplingResult.recordAndSample(), result); + assertEquals(0.001, ((ProbabilisticTransportActionSampler) probabilisticTransportActionSampler).getSamplingRatio(), 0.000d); + } +} diff --git a/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/OTelSamplerFactoryTests.java b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/OTelSamplerFactoryTests.java index 39ccf299dfdc4..0a25239485d07 100644 --- a/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/OTelSamplerFactoryTests.java +++ b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/OTelSamplerFactoryTests.java @@ -18,12 +18,16 @@ import io.opentelemetry.sdk.trace.samplers.Sampler; import static org.opensearch.telemetry.TelemetrySettings.TRACER_ENABLED_SETTING; +import static org.opensearch.telemetry.TelemetrySettings.TRACER_INFERRED_SAMPLER_ALLOWLISTED; import static org.opensearch.telemetry.TelemetrySettings.TRACER_SAMPLER_PROBABILITY; public class OTelSamplerFactoryTests extends OpenSearchTestCase { public void testDefaultCreate() { - ClusterSettings clusterSettings = new ClusterSettings(Settings.EMPTY, Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING)); + ClusterSettings clusterSettings = new ClusterSettings( + Settings.EMPTY, + Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING, TRACER_INFERRED_SAMPLER_ALLOWLISTED) + ); TelemetrySettings telemetrySettings = new TelemetrySettings(Settings.EMPTY, clusterSettings); Sampler sampler = OTelSamplerFactory.create(telemetrySettings, Settings.EMPTY); assertEquals(sampler.getClass(), ProbabilisticTransportActionSampler.class); @@ -34,7 +38,10 @@ public void testCreateWithSingleSampler() { .put(OTelTelemetrySettings.OTEL_TRACER_SPAN_SAMPLER_CLASS_SETTINGS.getKey(), ProbabilisticSampler.class.getName()) .build(); - ClusterSettings clusterSettings = new ClusterSettings(settings, Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING)); + ClusterSettings clusterSettings = new ClusterSettings( + settings, + Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING, TRACER_INFERRED_SAMPLER_ALLOWLISTED) + ); TelemetrySettings telemetrySettings = new TelemetrySettings(settings, clusterSettings); Sampler sampler = OTelSamplerFactory.create(telemetrySettings, settings); assertTrue(sampler instanceof ProbabilisticSampler); diff --git a/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/ProbabilisticSamplerTests.java b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/ProbabilisticSamplerTests.java index a094cd0119f5e..0c84d9a63c5a9 100644 --- a/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/ProbabilisticSamplerTests.java +++ b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/ProbabilisticSamplerTests.java @@ -21,6 +21,7 @@ import static org.opensearch.telemetry.OTelTelemetrySettings.TRACER_EXPORTER_DELAY_SETTING; import static org.opensearch.telemetry.TelemetrySettings.TRACER_ENABLED_SETTING; +import static org.opensearch.telemetry.TelemetrySettings.TRACER_INFERRED_SAMPLER_ALLOWLISTED; import static org.opensearch.telemetry.TelemetrySettings.TRACER_SAMPLER_PROBABILITY; import static org.mockito.Mockito.mock; @@ -36,7 +37,7 @@ public void testDefaultGetSampler() { Settings settings = Settings.builder().put(TRACER_EXPORTER_DELAY_SETTING.getKey(), "1s").build(); TelemetrySettings telemetrySettings = new TelemetrySettings( Settings.EMPTY, - new ClusterSettings(settings, Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING)) + new ClusterSettings(settings, Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING, TRACER_INFERRED_SAMPLER_ALLOWLISTED)) ); // Probabilistic Sampler @@ -49,7 +50,7 @@ public void testGetSamplerWithUpdatedSamplingRatio() { Settings settings = Settings.builder().put(TRACER_EXPORTER_DELAY_SETTING.getKey(), "1s").build(); TelemetrySettings telemetrySettings = new TelemetrySettings( Settings.EMPTY, - new ClusterSettings(settings, Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING)) + new ClusterSettings(settings, Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING, TRACER_INFERRED_SAMPLER_ALLOWLISTED)) ); // Probabilistic Sampler diff --git a/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/ProbabilisticTransportActionSamplerTests.java b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/ProbabilisticTransportActionSamplerTests.java index 261b0252fef60..71f295d1eac85 100644 --- a/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/ProbabilisticTransportActionSamplerTests.java +++ b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/ProbabilisticTransportActionSamplerTests.java @@ -23,6 +23,7 @@ import io.opentelemetry.sdk.trace.samplers.SamplingResult; import static org.opensearch.telemetry.TelemetrySettings.TRACER_ENABLED_SETTING; +import static org.opensearch.telemetry.TelemetrySettings.TRACER_INFERRED_SAMPLER_ALLOWLISTED; import static org.opensearch.telemetry.TelemetrySettings.TRACER_SAMPLER_PROBABILITY; import static org.opensearch.telemetry.tracing.AttributeNames.TRANSPORT_ACTION; import static org.mockito.Mockito.mock; @@ -30,7 +31,10 @@ public class ProbabilisticTransportActionSamplerTests extends OpenSearchTestCase { public void testGetSamplerWithActionSamplingRatio() { - ClusterSettings clusterSettings = new ClusterSettings(Settings.EMPTY, Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING)); + ClusterSettings clusterSettings = new ClusterSettings( + Settings.EMPTY, + Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING, TRACER_INFERRED_SAMPLER_ALLOWLISTED) + ); TelemetrySettings telemetrySettings = new TelemetrySettings(Settings.EMPTY, clusterSettings); diff --git a/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/RequestSamplerTests.java b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/RequestSamplerTests.java index da234ca13dc9d..3016f9f3833ef 100644 --- a/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/RequestSamplerTests.java +++ b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/RequestSamplerTests.java @@ -25,6 +25,7 @@ import io.opentelemetry.sdk.trace.samplers.SamplingResult; import static org.opensearch.telemetry.TelemetrySettings.TRACER_ENABLED_SETTING; +import static org.opensearch.telemetry.TelemetrySettings.TRACER_INFERRED_SAMPLER_ALLOWLISTED; import static org.opensearch.telemetry.TelemetrySettings.TRACER_SAMPLER_PROBABILITY; import static org.opensearch.telemetry.tracing.AttributeNames.TRANSPORT_ACTION; import static org.mockito.Mockito.mock; @@ -37,7 +38,10 @@ public class RequestSamplerTests extends OpenSearchTestCase { @Before public void init() { - clusterSettings = new ClusterSettings(Settings.EMPTY, Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING)); + clusterSettings = new ClusterSettings( + Settings.EMPTY, + Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING, TRACER_INFERRED_SAMPLER_ALLOWLISTED) + ); telemetrySettings = new TelemetrySettings(Settings.EMPTY, clusterSettings); Sampler fallbackSampler = OTelSamplerFactory.create(telemetrySettings, Settings.EMPTY); requestSampler = new RequestSampler(fallbackSampler); diff --git a/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/samplingResult/OTelSamplingResultTests.java b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/samplingResult/OTelSamplingResultTests.java new file mode 100644 index 0000000000000..b85357532cc55 --- /dev/null +++ b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/samplingResult/OTelSamplingResultTests.java @@ -0,0 +1,86 @@ +/* + * 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.samplingResult; + +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Settings; +import org.opensearch.telemetry.TelemetrySettings; +import org.opensearch.telemetry.tracing.attributes.SamplingAttributes; +import org.opensearch.telemetry.tracing.sampler.InferredActionSampler; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Collections; +import java.util.Set; + +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.context.Context; +import io.opentelemetry.sdk.trace.samplers.Sampler; +import io.opentelemetry.sdk.trace.samplers.SamplingResult; + +import static org.opensearch.telemetry.TelemetrySettings.TRACER_ENABLED_SETTING; +import static org.opensearch.telemetry.TelemetrySettings.TRACER_INFERRED_SAMPLER_ALLOWLISTED; +import static org.opensearch.telemetry.TelemetrySettings.TRACER_SAMPLER_PROBABILITY; +import static org.opensearch.telemetry.tracing.AttributeNames.TRANSPORT_ACTION; +import static org.mockito.Mockito.mock; + +public class OTelSamplingResultTests extends OpenSearchTestCase { + + public void testSamplingResult() { + ClusterSettings clusterSettings = new ClusterSettings( + Settings.EMPTY, + Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING, TRACER_INFERRED_SAMPLER_ALLOWLISTED) + ); + + TelemetrySettings telemetrySettings = new TelemetrySettings(Settings.EMPTY, clusterSettings); + telemetrySettings.setInferredSamplingAllowListed(true); + + // InferredActionSampler + Sampler inferredActionSampler = InferredActionSampler.create(telemetrySettings, Settings.EMPTY, null); + + SamplingResult result = inferredActionSampler.shouldSample( + mock(Context.class), + "00000000000000000000000000000000", + "spanName", + SpanKind.INTERNAL, + Attributes.builder().put(TRANSPORT_ACTION, "dummy_action").build(), + Collections.emptyList() + ); + + assertTrue(result.getAttributes().asMap().containsKey(AttributeKey.stringKey(SamplingAttributes.SAMPLER.getValue()))); + assertEquals(result.getClass(), OTelSamplingResult.class); + } + + public void testSamplingDecisionAndAttributes() { + ClusterSettings clusterSettings = new ClusterSettings( + Settings.EMPTY, + Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING, TRACER_INFERRED_SAMPLER_ALLOWLISTED) + ); + + TelemetrySettings telemetrySettings = new TelemetrySettings(Settings.EMPTY, clusterSettings); + telemetrySettings.setInferredSamplingAllowListed(true); + + // InferredActionSampler + Sampler inferredActionSampler = InferredActionSampler.create(telemetrySettings, Settings.EMPTY, null); + + SamplingResult result = inferredActionSampler.shouldSample( + mock(Context.class), + "00000000000000000000000000000000", + "spanName", + SpanKind.INTERNAL, + Attributes.builder().put(TRANSPORT_ACTION, "dummy_action").build(), + Collections.emptyList() + ); + + OTelSamplingResult oTelSamplingResult = new OTelSamplingResult(result.getDecision(), result.getAttributes()); + assertEquals(result.getDecision(), oTelSamplingResult.getDecision()); + assertEquals(result.getAttributes(), oTelSamplingResult.getAttributes()); + } +} diff --git a/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java b/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java index 8760ee3c94309..a6b0afec07481 100644 --- a/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java +++ b/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java @@ -744,7 +744,8 @@ public void apply(Settings value, Settings current, Settings previous) { TelemetrySettings.TRACER_SAMPLER_PROBABILITY, TelemetrySettings.METRICS_PUBLISH_INTERVAL_SETTING, TelemetrySettings.TRACER_FEATURE_ENABLED_SETTING, - TelemetrySettings.METRICS_FEATURE_ENABLED_SETTING + TelemetrySettings.METRICS_FEATURE_ENABLED_SETTING, + TelemetrySettings.TRACER_INFERRED_SAMPLER_ALLOWLISTED ), List.of(FeatureFlags.PLUGGABLE_CACHE), List.of(CacheSettings.getConcreteStoreNameSettingForCacheType(CacheType.INDICES_REQUEST_CACHE)) diff --git a/server/src/main/java/org/opensearch/telemetry/TelemetrySettings.java b/server/src/main/java/org/opensearch/telemetry/TelemetrySettings.java index 4b8897a318531..cfed8e628df3a 100644 --- a/server/src/main/java/org/opensearch/telemetry/TelemetrySettings.java +++ b/server/src/main/java/org/opensearch/telemetry/TelemetrySettings.java @@ -42,6 +42,13 @@ public class TelemetrySettings { Setting.Property.Final ); + public static final Setting TRACER_INFERRED_SAMPLER_ALLOWLISTED = Setting.boolSetting( + "telemetry.inferred.sampler.allowlisted", + false, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + /** * Probability of sampler */ @@ -68,15 +75,18 @@ public class TelemetrySettings { private volatile double samplingProbability; private final boolean tracingFeatureEnabled; private final boolean metricsFeatureEnabled; + private volatile boolean inferredSamplingAllowListed; public TelemetrySettings(Settings settings, ClusterSettings clusterSettings) { this.tracingEnabled = TRACER_ENABLED_SETTING.get(settings); this.samplingProbability = TRACER_SAMPLER_PROBABILITY.get(settings); this.tracingFeatureEnabled = TRACER_FEATURE_ENABLED_SETTING.get(settings); this.metricsFeatureEnabled = METRICS_FEATURE_ENABLED_SETTING.get(settings); + this.inferredSamplingAllowListed = TRACER_INFERRED_SAMPLER_ALLOWLISTED.get(settings); clusterSettings.addSettingsUpdateConsumer(TRACER_ENABLED_SETTING, this::setTracingEnabled); clusterSettings.addSettingsUpdateConsumer(TRACER_SAMPLER_PROBABILITY, this::setSamplingProbability); + clusterSettings.addSettingsUpdateConsumer(TRACER_INFERRED_SAMPLER_ALLOWLISTED, this::setInferredSamplingAllowListed); } public void setTracingEnabled(boolean tracingEnabled) { @@ -87,6 +97,22 @@ public boolean isTracingEnabled() { return tracingEnabled; } + /** + * Set sampling allowListing + * @param inferredSamplingAllowListed boolean + */ + public void setInferredSamplingAllowListed(boolean inferredSamplingAllowListed) { + this.inferredSamplingAllowListed = inferredSamplingAllowListed; + } + + /** + * Get sampling allowListing + * @return boolean + */ + public boolean getInferredSamplingAllowListed() { + return inferredSamplingAllowListed; + } + /** * Set sampling ratio * @param samplingProbability double diff --git a/server/src/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java b/server/src/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java index 908164d1935a7..800a3734425ab 100644 --- a/server/src/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java +++ b/server/src/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java @@ -11,11 +11,13 @@ import org.opensearch.common.annotation.InternalApi; import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.common.util.concurrent.ThreadContextStatePropagator; +import org.opensearch.telemetry.tracing.attributes.SamplingAttributes; import java.util.Collections; import java.util.HashMap; import java.util.Map; import java.util.Objects; +import java.util.Optional; /** * Core's ThreadContext based TracerContextStorage implementation @@ -80,7 +82,22 @@ public Map headers(Map source) { if (source.containsKey(CURRENT_SPAN)) { final SpanReference current = (SpanReference) source.get(CURRENT_SPAN); if (current != null && current.getSpan() != null) { - tracingTelemetry.getContextPropagator().inject(current.getSpan(), (key, value) -> headers.put(key, value)); + tracingTelemetry.getContextPropagator().inject(current.getSpan(), headers::put); + + // We will be sending one more header with the response if the request is marked for sampling + Optional isSpanSampled = Optional.ofNullable( + current.getSpan().getAttributeBoolean(SamplingAttributes.SAMPLED.getValue()) + ); + if (isSpanSampled.isPresent()) { + headers.put(SamplingAttributes.SAMPLED.getValue(), "true"); + } + Optional isSpanInferredSampled = Optional.ofNullable( + current.getSpan().getAttributeString(SamplingAttributes.SAMPLER.getValue()) + ); + if (isSpanInferredSampled.isPresent() + && isSpanInferredSampled.get().equals(SamplingAttributes.INFERRED_SAMPLER.getValue())) { + headers.put(SamplingAttributes.SAMPLER.getValue(), isSpanInferredSampled.get()); + } } } diff --git a/server/src/main/java/org/opensearch/telemetry/tracing/channels/TraceableRestChannel.java b/server/src/main/java/org/opensearch/telemetry/tracing/channels/TraceableRestChannel.java index 32769dd1d848d..d4a908729cfc6 100644 --- a/server/src/main/java/org/opensearch/telemetry/tracing/channels/TraceableRestChannel.java +++ b/server/src/main/java/org/opensearch/telemetry/tracing/channels/TraceableRestChannel.java @@ -97,9 +97,9 @@ public boolean detailedErrorsEnabled() { @Override public void sendResponse(RestResponse response) { try (SpanScope scope = tracer.withSpanInScope(span)) { - delegate.sendResponse(response); - } finally { span.endSpan(); + } finally { + delegate.sendResponse(response); } } } diff --git a/server/src/main/java/org/opensearch/telemetry/tracing/channels/TraceableTcpTransportChannel.java b/server/src/main/java/org/opensearch/telemetry/tracing/channels/TraceableTcpTransportChannel.java index 45268b4807cd9..289be55566f80 100644 --- a/server/src/main/java/org/opensearch/telemetry/tracing/channels/TraceableTcpTransportChannel.java +++ b/server/src/main/java/org/opensearch/telemetry/tracing/channels/TraceableTcpTransportChannel.java @@ -86,22 +86,17 @@ public String getChannelType() { @Override public void sendResponse(TransportResponse response) throws IOException { try (SpanScope scope = tracer.withSpanInScope(span)) { - delegate.sendResponse(response); - } catch (final IOException ex) { - span.setError(ex); - throw ex; - } finally { span.endSpan(); + delegate.sendResponse(response); } } @Override public void sendResponse(Exception exception) throws IOException { try (SpanScope scope = tracer.withSpanInScope(span)) { - delegate.sendResponse(exception); - } finally { span.setError(exception); span.endSpan(); + delegate.sendResponse(exception); } } diff --git a/server/src/main/java/org/opensearch/telemetry/tracing/handler/TraceableTransportResponseHandler.java b/server/src/main/java/org/opensearch/telemetry/tracing/handler/TraceableTransportResponseHandler.java index eb9d53d2df51b..d99f3ed7c5f49 100644 --- a/server/src/main/java/org/opensearch/telemetry/tracing/handler/TraceableTransportResponseHandler.java +++ b/server/src/main/java/org/opensearch/telemetry/tracing/handler/TraceableTransportResponseHandler.java @@ -13,6 +13,7 @@ import org.opensearch.telemetry.tracing.Span; import org.opensearch.telemetry.tracing.SpanScope; import org.opensearch.telemetry.tracing.Tracer; +import org.opensearch.telemetry.tracing.attributes.SamplingAttributes; import org.opensearch.transport.TransportException; import org.opensearch.transport.TransportResponseHandler; @@ -69,19 +70,23 @@ public T read(StreamInput in) throws IOException { @Override public void handleResponse(T response) { try (SpanScope scope = tracer.withSpanInScope(span)) { - delegate.handleResponse(response); - } finally { + String sampleInformation = response.getResponseHeaders().getOrDefault(SamplingAttributes.SAMPLED.getValue(), ""); + if (sampleInformation.equals("true")) { + span.addAttribute(SamplingAttributes.SAMPLED.getValue(), true); + } span.endSpan(); + } finally { + delegate.handleResponse(response); } } @Override public void handleException(TransportException exp) { try (SpanScope scope = tracer.withSpanInScope(span)) { - delegate.handleException(exp); - } finally { span.setError(exp); span.endSpan(); + } finally { + delegate.handleException(exp); } } @@ -98,10 +103,10 @@ public String toString() { @Override public void handleRejection(Exception exp) { try (SpanScope scope = tracer.withSpanInScope(span)) { - delegate.handleRejection(exp); - } finally { span.setError(exp); span.endSpan(); + } finally { + delegate.handleRejection(exp); } } } diff --git a/server/src/main/java/org/opensearch/telemetry/tracing/listener/TraceableActionListener.java b/server/src/main/java/org/opensearch/telemetry/tracing/listener/TraceableActionListener.java index 0cb4ce71d05f8..ed85fd9c73638 100644 --- a/server/src/main/java/org/opensearch/telemetry/tracing/listener/TraceableActionListener.java +++ b/server/src/main/java/org/opensearch/telemetry/tracing/listener/TraceableActionListener.java @@ -56,21 +56,19 @@ public static ActionListener create(ActionListener responseHeader; try (ThreadContext.StoredContext existing = threadContext.stashContext()) { // Place the context with the headers from the message threadContext.setHeaders(header.getHeaders()); threadContext.putTransient("_remote_address", remoteAddress); + responseHeader = message.getHeader().getHeaders().v1(); if (header.isRequest()) { handleRequest(channel, header, message); } else { @@ -169,11 +171,11 @@ private void messageReceived(TcpChannel channel, InboundMessage message, long st if (header.isError()) { handlerResponseError(requestId, streamInput, handler); } else { - handleResponse(requestId, remoteAddress, streamInput, handler); + handleResponse(requestId, remoteAddress, streamInput, handler, responseHeader); } } else { assert header.isError() == false; - handleResponse(requestId, remoteAddress, EMPTY_STREAM_INPUT, handler); + handleResponse(requestId, remoteAddress, EMPTY_STREAM_INPUT, handler, responseHeader); } } @@ -391,12 +393,14 @@ private void handleResponse( final long requestId, InetSocketAddress remoteAddress, final StreamInput stream, - final TransportResponseHandler handler + final TransportResponseHandler handler, + Map responseHeader ) { final T response; try { response = handler.read(stream); response.remoteAddress(new TransportAddress(remoteAddress)); + response.setResponseHeaders(responseHeader); checkStreamIsFullyConsumed(requestId, handler, stream, false); } catch (Exception e) { final Exception serializationException = new TransportSerializationException( diff --git a/server/src/test/java/org/opensearch/telemetry/TelemetrySettingsTests.java b/server/src/test/java/org/opensearch/telemetry/TelemetrySettingsTests.java index 4c96f79b30d55..badc8d92345ac 100644 --- a/server/src/test/java/org/opensearch/telemetry/TelemetrySettingsTests.java +++ b/server/src/test/java/org/opensearch/telemetry/TelemetrySettingsTests.java @@ -15,12 +15,16 @@ import java.util.Set; import static org.opensearch.telemetry.TelemetrySettings.TRACER_ENABLED_SETTING; +import static org.opensearch.telemetry.TelemetrySettings.TRACER_INFERRED_SAMPLER_ALLOWLISTED; import static org.opensearch.telemetry.TelemetrySettings.TRACER_SAMPLER_PROBABILITY; public class TelemetrySettingsTests extends OpenSearchTestCase { public void testSetTracingEnabledOrDisabled() { - ClusterSettings clusterSettings = new ClusterSettings(Settings.EMPTY, Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING)); + ClusterSettings clusterSettings = new ClusterSettings( + Settings.EMPTY, + Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING, TRACER_INFERRED_SAMPLER_ALLOWLISTED) + ); TelemetrySettings telemetrySettings = new TelemetrySettings(Settings.EMPTY, clusterSettings); // Validation for tracingEnabled as true @@ -33,7 +37,10 @@ public void testSetTracingEnabledOrDisabled() { } public void testSetSamplingProbability() { - ClusterSettings clusterSettings = new ClusterSettings(Settings.EMPTY, Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING)); + ClusterSettings clusterSettings = new ClusterSettings( + Settings.EMPTY, + Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING, TRACER_INFERRED_SAMPLER_ALLOWLISTED) + ); TelemetrySettings telemetrySettings = new TelemetrySettings(Settings.EMPTY, clusterSettings); // Validating default sample rate i.e 1% @@ -49,7 +56,10 @@ public void testSetSamplingProbability() { } public void testGetSamplingProbability() { - ClusterSettings clusterSettings = new ClusterSettings(Settings.EMPTY, Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING)); + ClusterSettings clusterSettings = new ClusterSettings( + Settings.EMPTY, + Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING, TRACER_INFERRED_SAMPLER_ALLOWLISTED) + ); TelemetrySettings telemetrySettings = new TelemetrySettings(Settings.EMPTY, clusterSettings); // Validating default value of Sampling is 1% @@ -61,4 +71,34 @@ public void testGetSamplingProbability() { assertEquals(0.02, telemetrySettings.getSamplingProbability(), 0.00d); } + public void testGetInferredSetting() { + ClusterSettings clusterSettings = new ClusterSettings( + Settings.EMPTY, + Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING, TRACER_INFERRED_SAMPLER_ALLOWLISTED) + ); + TelemetrySettings telemetrySettings = new TelemetrySettings(Settings.EMPTY, clusterSettings); + + assertFalse(telemetrySettings.getInferredSamplingAllowListed()); + + clusterSettings.applySettings(Settings.builder().put("telemetry.inferred.sampler.allowlisted", "true").build()); + + // Validate inferred allowlist setting + assertTrue(telemetrySettings.getInferredSamplingAllowListed()); + } + + public void testSetInferredSetting() { + ClusterSettings clusterSettings = new ClusterSettings( + Settings.EMPTY, + Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING, TRACER_INFERRED_SAMPLER_ALLOWLISTED) + ); + TelemetrySettings telemetrySettings = new TelemetrySettings(Settings.EMPTY, clusterSettings); + + assertFalse(telemetrySettings.getInferredSamplingAllowListed()); + + telemetrySettings.setInferredSamplingAllowListed(true); + + // Validate inferred allowlist setting + assertTrue(telemetrySettings.getInferredSamplingAllowListed()); + } + } diff --git a/server/src/test/java/org/opensearch/telemetry/tracing/SpanBuilderTests.java b/server/src/test/java/org/opensearch/telemetry/tracing/SpanBuilderTests.java index 4b763e4bd4454..96ab2c9399095 100644 --- a/server/src/test/java/org/opensearch/telemetry/tracing/SpanBuilderTests.java +++ b/server/src/test/java/org/opensearch/telemetry/tracing/SpanBuilderTests.java @@ -115,6 +115,20 @@ public void testParentSpan() { assertEquals(parentSpanContext, context.getParent()); } + public void testNoopSpanAttributes() { + String spanName = "test-name"; + SpanContext parentSpanContext = new SpanContext(NoopSpan.INSTANCE); + SpanCreationContext context = SpanBuilder.from(spanName, parentSpanContext); + Attributes attributes = context.getAttributes(); + assertNull(attributes); + assertEquals(spanName, context.getSpanName()); + assertEquals(parentSpanContext, context.getParent()); + assertEquals("", parentSpanContext.getSpan().getAttributeString("mockKey")); + assertEquals(false, parentSpanContext.getSpan().getAttributeBoolean("mockKey")); + assertEquals(0L, (Object) parentSpanContext.getSpan().getAttributeLong("mockKey")); + assertEquals(0.0, (Object) parentSpanContext.getSpan().getAttributeDouble("mockKey")); + } + private static Transport.Connection createTransportConnection() { return new Transport.Connection() { @Override diff --git a/server/src/test/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorageTests.java b/server/src/test/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorageTests.java index bf11bcaf39a96..406fc1c2efec5 100644 --- a/server/src/test/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorageTests.java +++ b/server/src/test/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorageTests.java @@ -31,6 +31,7 @@ import static org.opensearch.telemetry.TelemetrySettings.TRACER_ENABLED_SETTING; import static org.opensearch.telemetry.TelemetrySettings.TRACER_FEATURE_ENABLED_SETTING; +import static org.opensearch.telemetry.TelemetrySettings.TRACER_INFERRED_SAMPLER_ALLOWLISTED; import static org.opensearch.telemetry.TelemetrySettings.TRACER_SAMPLER_PROBABILITY; import static org.hamcrest.CoreMatchers.instanceOf; import static org.hamcrest.CoreMatchers.is; @@ -56,7 +57,10 @@ public void setUp() throws Exception { final TelemetrySettings telemetrySettings = new TelemetrySettings( settings, - new ClusterSettings(Settings.EMPTY, Set.of(TRACER_ENABLED_SETTING, TRACER_SAMPLER_PROBABILITY)) + new ClusterSettings( + Settings.EMPTY, + Set.of(TRACER_ENABLED_SETTING, TRACER_SAMPLER_PROBABILITY, TRACER_INFERRED_SAMPLER_ALLOWLISTED) + ) ); final TracingTelemetry tracingTelemetry = new MockTracingTelemetry(); diff --git a/test/telemetry/src/main/java/org/opensearch/test/telemetry/tracing/MockSpan.java b/test/telemetry/src/main/java/org/opensearch/test/telemetry/tracing/MockSpan.java index c5d179f6412a8..254bf1692f53d 100644 --- a/test/telemetry/src/main/java/org/opensearch/test/telemetry/tracing/MockSpan.java +++ b/test/telemetry/src/main/java/org/opensearch/test/telemetry/tracing/MockSpan.java @@ -139,6 +139,26 @@ public String getSpanId() { return spanId; } + @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; + } + /** * Returns whether the span is ended or not. * @return span end status.