From e02fc95633f512ccca8526c0263e5d6458be9149 Mon Sep 17 00:00:00 2001 From: Nishchay Malhotra Date: Tue, 26 Mar 2024 22:22:35 +0530 Subject: [PATCH 01/11] Adding new inferred sampler Signed-off-by: Nishchay Malhotra --- .../core/transport/TransportMessage.java | 10 ++ .../telemetry/tracing/DefaultSpanScope.java | 41 ++++++++ .../telemetry/tracing/DefaultTracer.java | 6 ++ .../opensearch/telemetry/tracing/Span.java | 14 +++ .../tracing/TracerContextStorage.java | 10 ++ .../tracing/attributes/AttributeType.java | 38 ++++++++ .../telemetry/tracing/noop/NoopSpan.java | 23 +++++ .../tracing/OTelResourceProvider.java | 7 +- .../telemetry/tracing/OTelSpan.java | 32 ++++++- .../tracing/OTelTracingTelemetry.java | 13 +++ .../tracing/processor/OtelSpanProcessor.java | 91 ++++++++++++++++++ .../tracing/processor/package-info.java | 12 +++ .../sampler/InferredActionSampler.java | 94 +++++++++++++++++++ .../common/settings/ClusterSettings.java | 3 +- .../telemetry/TelemetrySettings.java | 26 +++++ ...hreadContextBasedTracerContextStorage.java | 21 ++++- .../channels/TraceableRestChannel.java | 4 +- .../TraceableTcpTransportChannel.java | 9 +- .../TraceableTransportResponseHandler.java | 18 ++-- .../listener/TraceableActionListener.java | 10 +- .../opensearch/transport/InboundHandler.java | 16 +++- .../test/telemetry/tracing/MockSpan.java | 2 + 22 files changed, 472 insertions(+), 28 deletions(-) create mode 100644 libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/attributes/AttributeType.java create mode 100644 plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/processor/OtelSpanProcessor.java create mode 100644 plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/processor/package-info.java create mode 100644 plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/sampler/InferredActionSampler.java 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..203908ff3dcb4 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 @@ -45,14 +45,24 @@ public abstract class TransportMessage implements Writeable { 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 */ diff --git a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultSpanScope.java b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultSpanScope.java index 93600da510977..63f95f7782dd7 100644 --- a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultSpanScope.java +++ b/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 @@ class DefaultSpanScope implements SpanScope { private final Span beforeSpan; private static final ThreadLocal spanScopeThreadLocal = new ThreadLocal<>(); private final TracerContextStorage tracerContextStorage; + private final Map commonAttributeMap = new HashMap<>(); /** * Constructor @@ -40,6 +44,7 @@ private DefaultSpanScope( this.beforeSpan = beforeSpan; this.previousSpanScope = previousSpanScope; this.tracerContextStorage = tracerContextStorage; + initializeCommonPropagationAttributes(); } /** @@ -55,6 +60,13 @@ public static SpanScope create(Span span, TracerContextStorage tra 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 void close() { 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) { 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..3c669d6946ad6 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 @@ -104,7 +104,13 @@ protected void addDefaultAttributes(Span span) { @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.containsKey(TracerContextStorage.INFERRED_SAMPLER)) { + spanCreationContext.getAttributes().addAttribute(TracerContextStorage.INFERRED_SAMPLER, true); + } + } } 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..a0dbd1d0ec524 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 @@ -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 getAttributes(); } diff --git a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/TracerContextStorage.java b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/TracerContextStorage.java index 958d054948483..b08baa4be4cd1 100644 --- a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/TracerContextStorage.java +++ b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/TracerContextStorage.java @@ -24,6 +24,16 @@ public interface TracerContextStorage { */ 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 diff --git a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/attributes/AttributeType.java b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/attributes/AttributeType.java new file mode 100644 index 0000000000000..134a8251e00c7 --- /dev/null +++ b/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 +} 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..30d59bb5d5d7d 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 @@ -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 getTraceId() { 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 getAttributes() { + return null; + } } 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..d461a24081184 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..367713d6d5e91 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,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 @@ class OTelSpan extends AbstractSpan { private final Span delegateSpan; + private final Map metadata; + /** * Constructor * @param spanName span name @@ -28,31 +33,47 @@ class OTelSpan extends AbstractSpan { 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 @@ public String getSpanId() { return delegateSpan.getSpanContext().getSpanId(); } + @Override + public Object getAttribute(String key) { + return metadata.get(key); + } + + @Override + public Map getAttributes() { + return metadata; + } + io.opentelemetry.api.trace.Span getDelegateSpan() { return delegateSpan; } - } diff --git a/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelTracingTelemetry.java b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelTracingTelemetry.java index af39617a8c744..91056cfd95fe3 100644 --- a/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelTracingTelemetry.java +++ b/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,21 @@ private Span createOtelSpan(SpanCreationContext spanCreationContext, Span parent 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().getAttributesMap().containsKey(TracerContextStorage.INFERRED_SAMPLER)) { + newSpan.addAttribute(TracerContextStorage.INFERRED_SAMPLER, true); + } + } + io.opentelemetry.api.trace.Span otelSpan( String spanName, Span parentOTelSpan, 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..f441b8b81a741 --- /dev/null +++ b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/processor/OtelSpanProcessor.java @@ -0,0 +1,91 @@ +/* + * 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.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() + && 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(); + } +} 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..c55594cd34a6e --- /dev/null +++ b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/sampler/InferredActionSampler.java @@ -0,0 +1,94 @@ +/* + * 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.TracerContextStorage; + +import java.util.List; +import java.util.Objects; + +import io.opentelemetry.api.baggage.Baggage; +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 Sampler actionSampler; + private final TelemetrySettings telemetrySettings; + private final Settings settings; + private final double samplingRatio; + + /** + * 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.samplingRatio = 1.0; + this.actionSampler = Sampler.traceIdRatioBased(samplingRatio); + 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) { + // Using baggage to store the common sampler attribute for context propagation + Baggage.fromContext(parentContext) + .toBuilder() + .put(TracerContextStorage.INFERRED_SAMPLER, "true") + .build() + .storeInContext(parentContext) + .makeCurrent(); + return actionSampler.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks); + } 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/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java b/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java index a424294371422..7f74cd735e9c0 100644 --- a/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java +++ b/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java @@ -730,7 +730,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..c20a21495b8f0 100644 --- a/server/src/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java +++ b/server/src/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java @@ -80,7 +80,26 @@ 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 + if (headers.containsKey("traceparent")) { + final String currentTrace = headers.get("traceparent"); + if (currentTrace.length() > 0) { + String[] traceParent = currentTrace.split("-"); + if (traceParent.length > 2) { + String traceID = traceParent[1]; + if (current.getSpan().getTraceId().equals(traceID)) { + if (current.getSpan().getAttributes().containsKey(SAMPLED)) { + headers.put(SAMPLED, "true"); + } + if (current.getSpan().getAttributes().containsKey(INFERRED_SAMPLER)) { + headers.put(INFERRED_SAMPLER, "true"); + } + } + } + } } } 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..3e5bdbcc4ee7b 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.TracerContextStorage; 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 { + if (response.sampled()) { + span.addAttribute(TracerContextStorage.SAMPLED, true); + } span.endSpan(); + } finally { + delegate.handleResponse(response); } } @Override public void handleException(TransportException exp) { try (SpanScope scope = tracer.withSpanInScope(span)) { - delegate.handleException(exp); - } finally { + span.addAttribute(TracerContextStorage.SAMPLED, true); span.setError(exp); span.endSpan(); + } finally { + delegate.handleException(exp); } } @@ -98,10 +103,11 @@ public String toString() { @Override public void handleRejection(Exception exp) { try (SpanScope scope = tracer.withSpanInScope(span)) { - delegate.handleRejection(exp); - } finally { + span.addAttribute(TracerContextStorage.SAMPLED, true); 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 void handleResponse( final long requestId, InetSocketAddress remoteAddress, final StreamInput stream, - final TransportResponseHandler handler + final TransportResponseHandler handler, + boolean responseSampled ) { final T response; try { response = handler.read(stream); response.remoteAddress(new TransportAddress(remoteAddress)); + if (responseSampled) { + response.markSampled(); + } checkStreamIsFullyConsumed(requestId, handler, stream, false); } catch (Exception e) { final Exception serializationException = new TransportSerializationException( 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..6d9721e1a3d1b 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 @@ -191,6 +191,7 @@ private static String generateTraceId() { * @param key key * @return value */ + @Override public Object getAttribute(String key) { return metadata.get(key); } @@ -199,6 +200,7 @@ public Object getAttribute(String key) { * Returns the attributes as map. * @return returns the attributes map. */ + @Override public Map getAttributes() { return metadata; } From 78dcc6bc2e27a4100e4b117155edd034bac4fa86 Mon Sep 17 00:00:00 2001 From: Nishchay Malhotra Date: Tue, 26 Mar 2024 22:46:26 +0530 Subject: [PATCH 02/11] Added header null check and changelog Signed-off-by: Nishchay Malhotra --- CHANGELOG.md | 1 + .../java/org/opensearch/telemetry/tracing/DefaultTracer.java | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a56ea61554bfc..9d8eb0032c0fd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -116,6 +116,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - 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)) +- [Tracing Framework] Adds support for inferred sampling ([#12315](https://github.com/opensearch-project/OpenSearch/issues/12315)) ### Dependencies - Bump `peter-evans/find-comment` from 2 to 3 ([#12288](https://github.com/opensearch-project/OpenSearch/pull/12288)) 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 3c669d6946ad6..ff50d3f450861 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 @@ -109,7 +109,7 @@ public Span startSpan(SpanCreationContext spanCreationContext, Map> headers) { - if (headers.containsKey(TracerContextStorage.INFERRED_SAMPLER)) { + if (headers !=null && headers.containsKey(TracerContextStorage.INFERRED_SAMPLER)) { spanCreationContext.getAttributes().addAttribute(TracerContextStorage.INFERRED_SAMPLER, true); } } From ad8d7f3530e4b3734e302c7e020fc81d2a0a55d6 Mon Sep 17 00:00:00 2001 From: Nishchay Malhotra Date: Wed, 27 Mar 2024 09:13:05 +0530 Subject: [PATCH 03/11] Adding spotless apply Signed-off-by: Nishchay Malhotra --- CHANGELOG.md | 2 +- .../java/org/opensearch/telemetry/tracing/DefaultTracer.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c2dd3ddcf83a1..291ecd71b975d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -120,7 +120,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - [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)) +- [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/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultTracer.java b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultTracer.java index ff50d3f450861..9a27eb27f4fcd 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 @@ -109,7 +109,7 @@ public Span startSpan(SpanCreationContext spanCreationContext, Map> headers) { - if (headers !=null && headers.containsKey(TracerContextStorage.INFERRED_SAMPLER)) { + if (headers != null && headers.containsKey(TracerContextStorage.INFERRED_SAMPLER)) { spanCreationContext.getAttributes().addAttribute(TracerContextStorage.INFERRED_SAMPLER, true); } } From 5c4128b8a24e6673de3a316d18f50bc50edfa377 Mon Sep 17 00:00:00 2001 From: Nishchay Malhotra Date: Wed, 27 Mar 2024 10:08:30 +0530 Subject: [PATCH 04/11] Fixing test cases Signed-off-by: Nishchay Malhotra --- .../telemetry/OTelTelemetryPluginTests.java | 9 ++++++++- .../tracing/sampler/OTelSamplerFactoryTests.java | 11 +++++++++-- .../sampler/ProbabilisticSamplerTests.java | 5 +++-- ...ProbabilisticTransportActionSamplerTests.java | 6 +++++- .../tracing/sampler/RequestSamplerTests.java | 6 +++++- .../telemetry/TelemetrySettingsTests.java | 16 +++++++++++++--- ...eadContextBasedTracerContextStorageTests.java | 6 +++++- 7 files changed, 48 insertions(+), 11 deletions(-) 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..9bd97c8a9cbc6 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 { @@ -53,7 +54,13 @@ public void setup() { Settings settings = Settings.builder().put(TRACER_EXPORTER_DELAY_SETTING.getKey(), "1s").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/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/server/src/test/java/org/opensearch/telemetry/TelemetrySettingsTests.java b/server/src/test/java/org/opensearch/telemetry/TelemetrySettingsTests.java index 4c96f79b30d55..98af6ec3ea6d9 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% 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(); From 47bef1422671969533659b164b606054a6582471 Mon Sep 17 00:00:00 2001 From: Nishchay Malhotra Date: Wed, 27 Mar 2024 11:16:55 +0530 Subject: [PATCH 05/11] Adding fix for span creation npe Signed-off-by: Nishchay Malhotra --- .../org/opensearch/telemetry/tracing/OTelTracingTelemetry.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelTracingTelemetry.java b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelTracingTelemetry.java index 91056cfd95fe3..90f471bbb8a13 100644 --- a/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelTracingTelemetry.java +++ b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelTracingTelemetry.java @@ -74,7 +74,8 @@ private Span createOtelSpan(SpanCreationContext spanCreationContext, Span parent 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().getAttributesMap().containsKey(TracerContextStorage.INFERRED_SAMPLER)) { + || (spanCreationContext.getAttributes() != null + && spanCreationContext.getAttributes().getAttributesMap().containsKey(TracerContextStorage.INFERRED_SAMPLER))) { newSpan.addAttribute(TracerContextStorage.INFERRED_SAMPLER, true); } } From 4f0586a6bb6f07c1cfb9ddf9befe34c1c32f50ca Mon Sep 17 00:00:00 2001 From: Nishchay Malhotra Date: Wed, 27 Mar 2024 12:08:49 +0530 Subject: [PATCH 06/11] Fix batch span processor worker leak Signed-off-by: Nishchay Malhotra --- .../tracing/processor/OtelSpanProcessor.java | 22 +++++++++++++++++++ .../telemetry/OTelTelemetryPluginTests.java | 2 +- 2 files changed, 23 insertions(+), 1 deletion(-) 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 index f441b8b81a741..c7bdc19afb5d3 100644 --- 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 @@ -13,6 +13,7 @@ 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; @@ -88,4 +89,25 @@ public void onEnd(ReadableSpan span) { 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/test/java/org/opensearch/telemetry/OTelTelemetryPluginTests.java b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/OTelTelemetryPluginTests.java index 9bd97c8a9cbc6..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 @@ -51,7 +51,7 @@ 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( From 6bd52e93e0549d5e950ffa89bf2d47053484d328 Mon Sep 17 00:00:00 2001 From: Nishchay Malhotra Date: Thu, 28 Mar 2024 17:16:13 +0530 Subject: [PATCH 07/11] Adding get attribute as readable span attribute Signed-off-by: Nishchay Malhotra --- .../core/transport/TransportMessage.java | 2 +- .../telemetry/tracing/DefaultSpanScope.java | 41 -------------- .../telemetry/tracing/DefaultTracer.java | 18 +++++- .../opensearch/telemetry/tracing/Span.java | 31 ++++++++--- .../tracing/attributes/AttributeType.java | 38 ------------- .../telemetry/tracing/noop/NoopSpan.java | 31 +++++------ .../telemetry/tracing/OTelSpan.java | 50 +++++++++++------ .../tracing/OTelTracingTelemetry.java | 14 +---- .../sampler/InferredActionSampler.java | 27 ++++----- .../samplingResult/OTelSamplingResult.java | 55 +++++++++++++++++++ .../tracing/samplingResult/package-info.java | 12 ++++ ...hreadContextBasedTracerContextStorage.java | 24 +++----- .../opensearch/transport/InboundHandler.java | 2 +- .../test/telemetry/tracing/MockSpan.java | 22 +++++++- 14 files changed, 193 insertions(+), 174 deletions(-) delete mode 100644 libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/attributes/AttributeType.java create mode 100644 plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/samplingResult/OTelSamplingResult.java create mode 100644 plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/samplingResult/package-info.java 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 203908ff3dcb4..616c723a3c769 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 @@ -51,7 +51,7 @@ public void remoteAddress(TransportAddress remoteAddress) { this.remoteAddress = remoteAddress; } - public void markSampled() { + public void markResponseAsSampled() { this.sampled = true; } diff --git a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultSpanScope.java b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultSpanScope.java index 63f95f7782dd7..93600da510977 100644 --- a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultSpanScope.java +++ b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultSpanScope.java @@ -9,10 +9,7 @@ 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; /** @@ -27,7 +24,6 @@ class DefaultSpanScope implements SpanScope { private final Span beforeSpan; private static final ThreadLocal spanScopeThreadLocal = new ThreadLocal<>(); private final TracerContextStorage tracerContextStorage; - private final Map commonAttributeMap = new HashMap<>(); /** * Constructor @@ -44,7 +40,6 @@ private DefaultSpanScope( this.beforeSpan = beforeSpan; this.previousSpanScope = previousSpanScope; this.tracerContextStorage = tracerContextStorage; - initializeCommonPropagationAttributes(); } /** @@ -60,13 +55,6 @@ public static SpanScope create(Span span, TracerContextStorage tra 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(); @@ -76,38 +64,9 @@ public void close() { 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) { 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 9a27eb27f4fcd..c204cb486654f 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 @@ -52,8 +52,9 @@ public Span startSpan(SpanCreationContext context) { } else { parentSpan = getCurrentSpanInternal(); } + Span span = createSpan(context, parentSpan); - addDefaultAttributes(span); + addDefaultAttributes(parentSpan, span); return span; } @@ -97,7 +98,8 @@ 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) { + addCommonParentAttributes(parentSpan, span); span.addAttribute(THREAD_NAME, Thread.currentThread().getName()); } @@ -113,4 +115,16 @@ private void addRequestAttributeToContext(SpanCreationContext spanCreationContex 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 inferredAttribute = Optional.ofNullable( + parentSpan.getAttributeBoolean(TracerContextStorage.INFERRED_SAMPLER) + ); + if (inferredAttribute.isPresent()) { + currentSpan.addAttribute(TracerContextStorage.INFERRED_SAMPLER, true); + } + } + } } 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 a0dbd1d0ec524..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 @@ -10,8 +10,6 @@ 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. @@ -96,15 +94,30 @@ public interface Span { String getSpanId(); /** - * Returns attribute. - * @param key key - * @return value + * * + * @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 */ - Object getAttribute(String key); + Long getAttributeLong(String key); /** - * Returns the attributes as map. - * @return returns the attributes map. + * * + * @param key for which we need to look for value + * @return Double attribute value */ - Map getAttributes(); + Double getAttributeDouble(String key); } diff --git a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/attributes/AttributeType.java b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/attributes/AttributeType.java deleted file mode 100644 index 134a8251e00c7..0000000000000 --- a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/attributes/AttributeType.java +++ /dev/null @@ -1,38 +0,0 @@ -/* - * 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 -} 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 30d59bb5d5d7d..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 @@ -11,8 +11,6 @@ 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} * @@ -85,24 +83,23 @@ public String getSpanId() { return "noop-span-id"; } - /** - * Returns attribute. - * - * @param key key - * @return value - */ @Override - public Object getAttribute(String key) { - return null; + public String getAttributeString(String key) { + return ""; } - /** - * Returns the attributes as map. - * - * @return returns the attributes map. - */ @Override - public Map getAttributes() { - return null; + 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/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 367713d6d5e91..d06bc702fb463 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,11 +8,10 @@ package org.opensearch.telemetry.tracing; -import java.util.HashMap; -import java.util.Map; - +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 @@ -22,8 +21,6 @@ class OTelSpan extends AbstractSpan { private final Span delegateSpan; - private final Map metadata; - /** * Constructor * @param spanName span name @@ -33,12 +30,11 @@ class OTelSpan extends AbstractSpan { 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)) { + if (getAttributeBoolean(TracerContextStorage.SAMPLED) != null && getAttributeBoolean(TracerContextStorage.SAMPLED)) { markParentForSampling(); } delegateSpan.end(); @@ -46,7 +42,7 @@ public void endSpan() { private void markParentForSampling() { org.opensearch.telemetry.tracing.Span current_parent = getParentSpan(); - while (current_parent != null && !current_parent.getAttributes().containsKey(TracerContextStorage.SAMPLED)) { + while (current_parent != null && current_parent.getAttributeBoolean(TracerContextStorage.SAMPLED) == null) { current_parent.addAttribute(TracerContextStorage.SAMPLED, true); current_parent = current_parent.getParentSpan(); } @@ -55,25 +51,21 @@ private void markParentForSampling() { @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 @@ -99,13 +91,39 @@ public String getSpanId() { } @Override - public Object getAttribute(String key) { - return metadata.get(key); + public String getAttributeString(String key) { + if (delegateSpan != null && delegateSpan instanceof ReadableSpan) return ((ReadableSpan) delegateSpan).getAttribute( + AttributeKey.stringKey(key) + ); + + return null; } @Override - public Map getAttributes() { - return metadata; + 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() { diff --git a/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelTracingTelemetry.java b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelTracingTelemetry.java index 90f471bbb8a13..01aa2c19395ab 100644 --- a/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelTracingTelemetry.java +++ b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelTracingTelemetry.java @@ -15,7 +15,6 @@ 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,23 +62,12 @@ private Span createOtelSpan(SpanCreationContext spanCreationContext, Span parent OTelAttributesConverter.convert(spanCreationContext.getAttributes()), OTelSpanKindConverter.convert(spanCreationContext.getSpanKind()) ); - Span newSpan = new OTelSpan(spanCreationContext.getSpanName(), otelSpan, parentSpan); - // This attribute is added for inferred sampling - addInferredAttribute(newSpan, spanCreationContext); + Span newSpan = new OTelSpan(spanCreationContext.getSpanName(), otelSpan, parentSpan); 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, 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 index c55594cd34a6e..3d24d5da38447 100644 --- 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 @@ -11,11 +11,11 @@ import org.opensearch.common.settings.Settings; import org.opensearch.telemetry.TelemetrySettings; import org.opensearch.telemetry.tracing.TracerContextStorage; +import org.opensearch.telemetry.tracing.samplingResult.OTelSamplingResult; import java.util.List; import java.util.Objects; -import io.opentelemetry.api.baggage.Baggage; import io.opentelemetry.api.common.Attributes; import io.opentelemetry.api.trace.SpanKind; import io.opentelemetry.context.Context; @@ -29,10 +29,8 @@ public class InferredActionSampler implements Sampler { private final Sampler fallbackSampler; - private final Sampler actionSampler; private final TelemetrySettings telemetrySettings; private final Settings settings; - private final double samplingRatio; /** * Constructor @@ -43,8 +41,6 @@ public class InferredActionSampler implements Sampler { private InferredActionSampler(TelemetrySettings telemetrySettings, Settings settings, Sampler fallbackSampler) { this.telemetrySettings = Objects.requireNonNull(telemetrySettings); this.settings = Objects.requireNonNull(settings); - this.samplingRatio = 1.0; - this.actionSampler = Sampler.traceIdRatioBased(samplingRatio); this.fallbackSampler = fallbackSampler; } @@ -71,19 +67,16 @@ public SamplingResult shouldSample( ) { boolean inferredSamplingAllowListed = telemetrySettings.getInferredSamplingAllowListed(); if (inferredSamplingAllowListed) { - // Using baggage to store the common sampler attribute for context propagation - Baggage.fromContext(parentContext) - .toBuilder() - .put(TracerContextStorage.INFERRED_SAMPLER, "true") - .build() - .storeInContext(parentContext) - .makeCurrent(); - return actionSampler.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks); - } else { - if (fallbackSampler != null) { - return fallbackSampler.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks); - } + Attributes customSampleAttributes = Attributes.builder() + .put(TracerContextStorage.INFERRED_SAMPLER, true) + .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(); } 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/server/src/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java b/server/src/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java index c20a21495b8f0..8fcd71c5da710 100644 --- a/server/src/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java +++ b/server/src/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java @@ -81,24 +81,14 @@ public Map headers(Map source) { final SpanReference current = (SpanReference) source.get(CURRENT_SPAN); if (current != null && current.getSpan() != null) { tracingTelemetry.getContextPropagator().inject(current.getSpan(), headers::put); - } - // We will be sending one more header with the response if the request is marked for sampling - if (headers.containsKey("traceparent")) { - final String currentTrace = headers.get("traceparent"); - if (currentTrace.length() > 0) { - String[] traceParent = currentTrace.split("-"); - if (traceParent.length > 2) { - String traceID = traceParent[1]; - if (current.getSpan().getTraceId().equals(traceID)) { - if (current.getSpan().getAttributes().containsKey(SAMPLED)) { - headers.put(SAMPLED, "true"); - } - if (current.getSpan().getAttributes().containsKey(INFERRED_SAMPLER)) { - headers.put(INFERRED_SAMPLER, "true"); - } - } - } + // We will be sending one more header with the response if the request is marked for sampling + if (current.getSpan().getAttributeBoolean(SAMPLED) != null && current.getSpan().getAttributeBoolean(SAMPLED)) { + headers.put(SAMPLED, "true"); + } + if (current.getSpan().getAttributeBoolean(INFERRED_SAMPLER) != null + && current.getSpan().getAttributeBoolean(INFERRED_SAMPLER)) { + headers.put(INFERRED_SAMPLER, "true"); } } } diff --git a/server/src/main/java/org/opensearch/transport/InboundHandler.java b/server/src/main/java/org/opensearch/transport/InboundHandler.java index 7f79003e5e155..af378b3b6486a 100644 --- a/server/src/main/java/org/opensearch/transport/InboundHandler.java +++ b/server/src/main/java/org/opensearch/transport/InboundHandler.java @@ -405,7 +405,7 @@ private void handleResponse( response = handler.read(stream); response.remoteAddress(new TransportAddress(remoteAddress)); if (responseSampled) { - response.markSampled(); + response.markResponseAsSampled(); } checkStreamIsFullyConsumed(requestId, handler, stream, false); } catch (Exception e) { 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 6d9721e1a3d1b..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. @@ -191,7 +211,6 @@ private static String generateTraceId() { * @param key key * @return value */ - @Override public Object getAttribute(String key) { return metadata.get(key); } @@ -200,7 +219,6 @@ public Object getAttribute(String key) { * Returns the attributes as map. * @return returns the attributes map. */ - @Override public Map getAttributes() { return metadata; } From 7504c6a0c63fb5f22e8cb9fcd3f726456709e555 Mon Sep 17 00:00:00 2001 From: Nishchay Malhotra Date: Fri, 29 Mar 2024 11:29:09 +0530 Subject: [PATCH 08/11] Fixing comments Signed-off-by: Nishchay Malhotra Adding new sampled attributes Adding new inferred sampler attributes Signed-off-by: Nishchay Malhotra --- .../core/transport/TransportMessage.java | 13 +++--- .../telemetry/tracing/DefaultTracer.java | 16 ++++---- .../tracing/TracerContextStorage.java | 10 ----- .../attributes/SamplingAttributes.java | 40 +++++++++++++++++++ .../telemetry/tracing/OTelSpan.java | 20 +++++++--- .../tracing/OTelTracingTelemetry.java | 2 - .../tracing/processor/OtelSpanProcessor.java | 11 +++-- .../sampler/InferredActionSampler.java | 4 +- ...hreadContextBasedTracerContextStorage.java | 18 ++++++--- .../TraceableTransportResponseHandler.java | 9 ++--- .../opensearch/transport/InboundHandler.java | 18 +++------ 11 files changed, 104 insertions(+), 57 deletions(-) create mode 100644 libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/attributes/SamplingAttributes.java 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 616c723a3c769..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,22 +48,22 @@ public abstract class TransportMessage implements Writeable { private TransportAddress remoteAddress; - private boolean sampled; + private Map header = Collections.emptyMap(); public void remoteAddress(TransportAddress remoteAddress) { this.remoteAddress = remoteAddress; } - public void markResponseAsSampled() { - this.sampled = true; + public void setResponseHeaders(Map header) { + this.header = header; } public TransportAddress remoteAddress() { return remoteAddress; } - public boolean sampled() { - return sampled; + public Map getResponseHeaders() { + return header; } /** 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 c204cb486654f..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; @@ -99,7 +100,7 @@ private Span createSpan(SpanCreationContext spanCreationContext, Span parentSpan * @param span the current active span */ protected void addDefaultAttributes(Span parentSpan, Span span) { - addCommonParentAttributes(parentSpan, span); + copyInheritableParentAttributes(parentSpan, span); span.addAttribute(THREAD_NAME, Thread.currentThread().getName()); } @@ -111,19 +112,18 @@ public Span startSpan(SpanCreationContext spanCreationContext, Map> headers) { - if (headers != null && headers.containsKey(TracerContextStorage.INFERRED_SAMPLER)) { - spanCreationContext.getAttributes().addAttribute(TracerContextStorage.INFERRED_SAMPLER, true); + if (headers != null && headers.containsKey(SamplingAttributes.SAMPLER.getValue())) { + spanCreationContext.getAttributes() + .addAttribute(SamplingAttributes.SAMPLER.getValue(), SamplingAttributes.INFERRED_SAMPLER.getValue()); } } - private void addCommonParentAttributes(Span parentSpan, Span currentSpan) { + 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.getAttributeBoolean(TracerContextStorage.INFERRED_SAMPLER) - ); + Optional inferredAttribute = Optional.ofNullable(parentSpan.getAttributeString(SamplingAttributes.SAMPLER.getValue())); if (inferredAttribute.isPresent()) { - currentSpan.addAttribute(TracerContextStorage.INFERRED_SAMPLER, true); + currentSpan.addAttribute(SamplingAttributes.SAMPLER.getValue(), SamplingAttributes.INFERRED_SAMPLER.getValue()); } } } diff --git a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/TracerContextStorage.java b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/TracerContextStorage.java index b08baa4be4cd1..958d054948483 100644 --- a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/TracerContextStorage.java +++ b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/TracerContextStorage.java @@ -24,16 +24,6 @@ public interface TracerContextStorage { */ 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 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..7ae1d6f945b68 --- /dev/null +++ b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/attributes/SamplingAttributes.java @@ -0,0 +1,40 @@ +/* + * 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; + + public String getValue() { + return name().toLowerCase(Locale.ROOT); + } +} 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 d06bc702fb463..cef96cebb4f8a 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,6 +8,8 @@ package org.opensearch.telemetry.tracing; +import org.opensearch.telemetry.tracing.attributes.SamplingAttributes; + import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.StatusCode; @@ -34,17 +36,25 @@ public OTelSpan(String spanName, Span span, org.opensearch.telemetry.tracing.Spa @Override public void endSpan() { - if (getAttributeBoolean(TracerContextStorage.SAMPLED) != null && getAttributeBoolean(TracerContextStorage.SAMPLED)) { + 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() { + return false; + } + 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(); + 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(); } } diff --git a/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelTracingTelemetry.java b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelTracingTelemetry.java index 01aa2c19395ab..af39617a8c744 100644 --- a/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelTracingTelemetry.java +++ b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelTracingTelemetry.java @@ -62,9 +62,7 @@ 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; } 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 index c7bdc19afb5d3..d16f41e7ddd20 100644 --- 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 @@ -8,7 +8,9 @@ package org.opensearch.telemetry.tracing.processor; -import org.opensearch.telemetry.tracing.TracerContextStorage; +import org.opensearch.telemetry.tracing.attributes.SamplingAttributes; + +import java.util.Objects; import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.trace.Span; @@ -71,8 +73,11 @@ public boolean isStartRequired() { 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) { + && 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 { 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 index 3d24d5da38447..04baef85cb14d 100644 --- 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 @@ -10,7 +10,7 @@ import org.opensearch.common.settings.Settings; import org.opensearch.telemetry.TelemetrySettings; -import org.opensearch.telemetry.tracing.TracerContextStorage; +import org.opensearch.telemetry.tracing.attributes.SamplingAttributes; import org.opensearch.telemetry.tracing.samplingResult.OTelSamplingResult; import java.util.List; @@ -68,7 +68,7 @@ public SamplingResult shouldSample( boolean inferredSamplingAllowListed = telemetrySettings.getInferredSamplingAllowListed(); if (inferredSamplingAllowListed) { Attributes customSampleAttributes = Attributes.builder() - .put(TracerContextStorage.INFERRED_SAMPLER, true) + .put(SamplingAttributes.SAMPLER.getValue(), SamplingAttributes.INFERRED_SAMPLER.getValue()) .putAll(attributes) .build(); SamplingResult result = SamplingResult.recordAndSample(); 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 8fcd71c5da710..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 @@ -83,12 +85,18 @@ public Map headers(Map source) { tracingTelemetry.getContextPropagator().inject(current.getSpan(), headers::put); // We will be sending one more header with the response if the request is marked for sampling - if (current.getSpan().getAttributeBoolean(SAMPLED) != null && current.getSpan().getAttributeBoolean(SAMPLED)) { - headers.put(SAMPLED, "true"); + Optional isSpanSampled = Optional.ofNullable( + current.getSpan().getAttributeBoolean(SamplingAttributes.SAMPLED.getValue()) + ); + if (isSpanSampled.isPresent()) { + headers.put(SamplingAttributes.SAMPLED.getValue(), "true"); } - if (current.getSpan().getAttributeBoolean(INFERRED_SAMPLER) != null - && current.getSpan().getAttributeBoolean(INFERRED_SAMPLER)) { - headers.put(INFERRED_SAMPLER, "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/handler/TraceableTransportResponseHandler.java b/server/src/main/java/org/opensearch/telemetry/tracing/handler/TraceableTransportResponseHandler.java index 3e5bdbcc4ee7b..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,7 +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.TracerContextStorage; +import org.opensearch.telemetry.tracing.attributes.SamplingAttributes; import org.opensearch.transport.TransportException; import org.opensearch.transport.TransportResponseHandler; @@ -70,8 +70,9 @@ public T read(StreamInput in) throws IOException { @Override public void handleResponse(T response) { try (SpanScope scope = tracer.withSpanInScope(span)) { - if (response.sampled()) { - span.addAttribute(TracerContextStorage.SAMPLED, true); + String sampleInformation = response.getResponseHeaders().getOrDefault(SamplingAttributes.SAMPLED.getValue(), ""); + if (sampleInformation.equals("true")) { + span.addAttribute(SamplingAttributes.SAMPLED.getValue(), true); } span.endSpan(); } finally { @@ -82,7 +83,6 @@ public void handleResponse(T response) { @Override public void handleException(TransportException exp) { try (SpanScope scope = tracer.withSpanInScope(span)) { - span.addAttribute(TracerContextStorage.SAMPLED, true); span.setError(exp); span.endSpan(); } finally { @@ -103,7 +103,6 @@ public String toString() { @Override public void handleRejection(Exception exp) { try (SpanScope scope = tracer.withSpanInScope(span)) { - span.addAttribute(TracerContextStorage.SAMPLED, true); span.setError(exp); span.endSpan(); } finally { diff --git a/server/src/main/java/org/opensearch/transport/InboundHandler.java b/server/src/main/java/org/opensearch/transport/InboundHandler.java index af378b3b6486a..2afc95f3a7198 100644 --- a/server/src/main/java/org/opensearch/transport/InboundHandler.java +++ b/server/src/main/java/org/opensearch/transport/InboundHandler.java @@ -50,7 +50,6 @@ import org.opensearch.telemetry.tracing.SpanBuilder; import org.opensearch.telemetry.tracing.SpanScope; import org.opensearch.telemetry.tracing.Tracer; -import org.opensearch.telemetry.tracing.TracerContextStorage; import org.opensearch.telemetry.tracing.channels.TraceableTcpTransportChannel; import org.opensearch.threadpool.ThreadPool; @@ -137,15 +136,12 @@ private void messageReceived(TcpChannel channel, InboundMessage message, long st final Header header = message.getHeader(); assert header.needsToReadVariableHeader() == false; ThreadContext threadContext = threadPool.getThreadContext(); - boolean responseSampled = false; + Map 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); - String sampleInformation = message.getHeader().getHeaders().v1().getOrDefault(TracerContextStorage.SAMPLED, ""); - if (sampleInformation.equals("true")) { - responseSampled = true; - } + responseHeader = message.getHeader().getHeaders().v1(); if (header.isRequest()) { handleRequest(channel, header, message); } else { @@ -175,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, responseSampled); + handleResponse(requestId, remoteAddress, streamInput, handler, responseHeader); } } else { assert header.isError() == false; - handleResponse(requestId, remoteAddress, EMPTY_STREAM_INPUT, handler, responseSampled); + handleResponse(requestId, remoteAddress, EMPTY_STREAM_INPUT, handler, responseHeader); } } @@ -398,15 +394,13 @@ private void handleResponse( InetSocketAddress remoteAddress, final StreamInput stream, final TransportResponseHandler handler, - boolean responseSampled + Map responseHeader ) { final T response; try { response = handler.read(stream); response.remoteAddress(new TransportAddress(remoteAddress)); - if (responseSampled) { - response.markResponseAsSampled(); - } + response.setResponseHeaders(responseHeader); checkStreamIsFullyConsumed(requestId, handler, stream, false); } catch (Exception e) { final Exception serializationException = new TransportSerializationException( From a1b3df136002b2cfc4d17a3dba10a07b80ee2a5d Mon Sep 17 00:00:00 2001 From: Nishchay Malhotra Date: Mon, 1 Apr 2024 10:09:30 +0530 Subject: [PATCH 09/11] Fixing spotlesschecks Signed-off-by: Nishchay Malhotra --- .../telemetry/tracing/attributes/SamplingAttributes.java | 4 ++++ 1 file changed, 4 insertions(+) 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 index 7ae1d6f945b68..08a712a408124 100644 --- 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 @@ -34,6 +34,10 @@ public enum SamplingAttributes { */ SAMPLER; + /** + * returns lower case enum value* + * @return String + */ public String getValue() { return name().toLowerCase(Locale.ROOT); } From 8bce54060321fc31204a1a251e13298143694fe5 Mon Sep 17 00:00:00 2001 From: Nishchay Malhotra Date: Mon, 1 Apr 2024 16:58:33 +0530 Subject: [PATCH 10/11] Adding unit test cases Signed-off-by: Nishchay Malhotra --- .../telemetry/tracing/DefaultTracerTests.java | 60 ++++++++++ .../tracing/OTelResourceProvider.java | 6 +- ...nProcessor.java => OTelSpanProcessor.java} | 4 +- .../telemetry/tracing/OTelSpanTests.java | 24 ++++ .../tracing/processor/MockSpanProcessor.java | 47 ++++++++ .../processor/OTelSpanProcessorTests.java | 84 ++++++++++++++ .../tracing/sampler/InferredSamplerTests.java | 109 ++++++++++++++++++ .../OTelSamplingResultTests.java | 86 ++++++++++++++ .../telemetry/TelemetrySettingsTests.java | 30 +++++ .../telemetry/tracing/SpanBuilderTests.java | 14 +++ 10 files changed, 459 insertions(+), 5 deletions(-) rename plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/processor/{OtelSpanProcessor.java => OTelSpanProcessor.java} (96%) create mode 100644 plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/processor/MockSpanProcessor.java create mode 100644 plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/processor/OTelSpanProcessorTests.java create mode 100644 plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/InferredSamplerTests.java create mode 100644 plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/samplingResult/OTelSamplingResultTests.java 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 d461a24081184..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,7 +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.processor.OTelSpanProcessor; import org.opensearch.telemetry.tracing.sampler.OTelSamplerFactory; import org.opensearch.telemetry.tracing.sampler.RequestSampler; @@ -118,8 +118,8 @@ private static SdkTracerProvider createSdkTracerProvider( .build(); } - private static OtelSpanProcessor spanProcessor(Settings settings, SpanExporter spanExporter) { - return new OtelSpanProcessor(batchSpanProcessor(settings, spanExporter)); + private static OTelSpanProcessor spanProcessor(Settings settings, SpanExporter spanExporter) { + return new OTelSpanProcessor(batchSpanProcessor(settings, spanExporter)); } private static BatchSpanProcessor batchSpanProcessor(Settings settings, SpanExporter spanExporter) { 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 similarity index 96% rename from plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/processor/OtelSpanProcessor.java rename to plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/processor/OTelSpanProcessor.java index d16f41e7ddd20..42d302458bf4b 100644 --- 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 @@ -23,7 +23,7 @@ /** * Implementation of the SpanProcessor and delegates to the configured processor. */ -public class OtelSpanProcessor implements SpanProcessor { +public class OTelSpanProcessor implements SpanProcessor { private final SpanProcessor delegateProcessor; @@ -31,7 +31,7 @@ public class OtelSpanProcessor implements SpanProcessor { * * * @param delegateProcessor the span processor to which this processor will delegate */ - public OtelSpanProcessor(SpanProcessor delegateProcessor) { + public OTelSpanProcessor(SpanProcessor delegateProcessor) { this.delegateProcessor = delegateProcessor; } 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..227c53f3ea4fb 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 @@ -85,6 +85,30 @@ 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")); + } + 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..09f1cf14fc0e2 --- /dev/null +++ b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/processor/OTelSpanProcessorTests.java @@ -0,0 +1,84 @@ +/* + * 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.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 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..6b3ca318f5572 --- /dev/null +++ b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/InferredSamplerTests.java @@ -0,0 +1,109 @@ +/* + * 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() + ); + } +} 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/test/java/org/opensearch/telemetry/TelemetrySettingsTests.java b/server/src/test/java/org/opensearch/telemetry/TelemetrySettingsTests.java index 98af6ec3ea6d9..badc8d92345ac 100644 --- a/server/src/test/java/org/opensearch/telemetry/TelemetrySettingsTests.java +++ b/server/src/test/java/org/opensearch/telemetry/TelemetrySettingsTests.java @@ -71,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 From b944c0bfc2577684207a67996b1a471916272918 Mon Sep 17 00:00:00 2001 From: Nishchay Malhotra Date: Mon, 1 Apr 2024 23:23:58 +0530 Subject: [PATCH 11/11] Addimh more unit test cases Signed-off-by: Nishchay Malhotra --- .../telemetry/tracing/OTelSpan.java | 9 +++++- .../telemetry/tracing/OTelSpanTests.java | 30 ++++++++++++++++++ .../processor/OTelSpanProcessorTests.java | 11 +++++++ .../tracing/sampler/InferredSamplerTests.java | 31 +++++++++++++++++++ 4 files changed, 80 insertions(+), 1 deletion(-) 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 cef96cebb4f8a..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 @@ -10,6 +10,8 @@ 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; @@ -47,7 +49,12 @@ public void endSpan() { * This Framework will be used to evaluate a span if that is an outlier or not. */ private boolean isSpanOutlier() { - return false; + 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() { 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 227c53f3ea4fb..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; @@ -109,6 +112,33 @@ public void testGetSpanDouble() { 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/OTelSpanProcessorTests.java b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/processor/OTelSpanProcessorTests.java index 09f1cf14fc0e2..d63f88fafd33f 100644 --- 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 @@ -15,6 +15,8 @@ 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; @@ -53,6 +55,15 @@ public void testOnEndFunction() { 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); 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 index 6b3ca318f5572..a1679ffdb0ee9 100644 --- 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 @@ -105,5 +105,36 @@ public void testGetSamplerWithAddedAttributes() { 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); } }