forked from opensearch-project/OpenSearch
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Added custom sampler support based on action in request (opensearch-p…
…roject#10136) * Added custom sampler support based on action in request Signed-off-by: Dev Agarwal <[email protected]> * UT Fix Signed-off-by: Dev Agarwal <[email protected]> * Added Transport action sampler, which will sample based on different probability for all actions Signed-off-by: Dev Agarwal <[email protected]> * Added Transport action sampler, which will sample based on different probability for all actions. Also added setting to define order of samplers Signed-off-by: Dev Agarwal <[email protected]> * Added missing java-doc Signed-off-by: Dev Agarwal <[email protected]> * Moving sampler class settings to OtelTelemetry setting Signed-off-by: Dev Agarwal <[email protected]> * Minor refactor Signed-off-by: Dev Agarwal <[email protected]> * Refactored to use chain of samplers Signed-off-by: Dev Agarwal <[email protected]> * Addressed comments Signed-off-by: Dev Agarwal <[email protected]> * Addressed comments to move action_probability to OtelTelemetrySettings Signed-off-by: Dev Agarwal <[email protected]> * Updated eror msg returned when Sampler class is not found Signed-off-by: Dev Agarwal <[email protected]> * Added UT for OTelSamplerFactory Signed-off-by: Dev Agarwal <[email protected]> * minor refactor Signed-off-by: Dev Agarwal <[email protected]> * minor refactor Signed-off-by: Dev Agarwal <[email protected]> * spotless check Signed-off-by: Dev Agarwal <[email protected]> * Updating OtelTelemetryPlugin.get() method Signed-off-by: Dev Agarwal <[email protected]> * Addressed comments Signed-off-by: Dev Agarwal <[email protected]> * minor refactor Signed-off-by: Dev Agarwal <[email protected]> * addressed comments Signed-off-by: Dev Agarwal <[email protected]> * Updated transport action sampler Signed-off-by: Dev Agarwal <[email protected]> * Empty-Commit Signed-off-by: Dev Agarwal <[email protected]> * Empty-Commit Signed-off-by: Dev Agarwal <[email protected]> --------- Signed-off-by: Dev Agarwal <[email protected]> (cherry picked from commit 445bf1f)
- Loading branch information
1 parent
ac37b38
commit 7bde40a
Showing
15 changed files
with
547 additions
and
73 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
96 changes: 96 additions & 0 deletions
96
...metry-otel/src/main/java/org/opensearch/telemetry/tracing/sampler/OTelSamplerFactory.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,96 @@ | ||
| /* | ||
| * 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.apache.logging.log4j.LogManager; | ||
| import org.apache.logging.log4j.Logger; | ||
| import org.opensearch.SpecialPermission; | ||
| import org.opensearch.common.settings.Settings; | ||
| import org.opensearch.telemetry.OTelTelemetrySettings; | ||
| import org.opensearch.telemetry.TelemetrySettings; | ||
|
|
||
| import java.lang.invoke.MethodHandles; | ||
| import java.lang.invoke.MethodType; | ||
| import java.security.AccessController; | ||
| import java.security.PrivilegedExceptionAction; | ||
| import java.util.List; | ||
| import java.util.ListIterator; | ||
|
|
||
| import io.opentelemetry.sdk.trace.samplers.Sampler; | ||
|
|
||
| /** | ||
| * Factory class to create the instance of OTelSampler | ||
| */ | ||
| public class OTelSamplerFactory { | ||
|
|
||
| /** | ||
| * Logger instance for logging messages related to the OTelSamplerFactory. | ||
| */ | ||
| private static final Logger logger = LogManager.getLogger(OTelSamplerFactory.class); | ||
|
|
||
| /** | ||
| * Base constructor. | ||
| */ | ||
| private OTelSamplerFactory() { | ||
|
|
||
| } | ||
|
|
||
| /** | ||
| * Creates the {@link Sampler} instances based on the TRACER_SPAN_SAMPLER_CLASSES value. | ||
| * | ||
| * @param telemetrySettings TelemetrySettings. | ||
| * @param settings the settings | ||
| * @return list of samplers. | ||
| */ | ||
| public static Sampler create(TelemetrySettings telemetrySettings, Settings settings) { | ||
| List<Class<Sampler>> samplersNameList = OTelTelemetrySettings.OTEL_TRACER_SPAN_SAMPLER_CLASS_SETTINGS.get(settings); | ||
| ListIterator<Class<Sampler>> li = samplersNameList.listIterator(samplersNameList.size()); | ||
|
|
||
| Sampler fallbackSampler = null; | ||
|
|
||
| // Iterating samplers list in reverse order to create chain of sampler | ||
| while (li.hasPrevious()) { | ||
| Class<Sampler> samplerName = li.previous(); | ||
| fallbackSampler = instantiateSampler(samplerName, telemetrySettings, settings, fallbackSampler); | ||
| } | ||
|
|
||
| return fallbackSampler; | ||
| } | ||
|
|
||
| private static Sampler instantiateSampler( | ||
| Class<Sampler> samplerClassName, | ||
| TelemetrySettings telemetrySettings, | ||
| Settings settings, | ||
| Sampler fallbackSampler | ||
| ) { | ||
| try { | ||
| // Check we ourselves are not being called by unprivileged code. | ||
| SpecialPermission.check(); | ||
|
|
||
| return AccessController.doPrivileged((PrivilegedExceptionAction<Sampler>) () -> { | ||
| try { | ||
| // Define the method type which receives TelemetrySettings & Sampler as arguments | ||
| MethodType methodType = MethodType.methodType(Sampler.class, TelemetrySettings.class, Settings.class, Sampler.class); | ||
|
|
||
| return (Sampler) MethodHandles.publicLookup() | ||
| .findStatic(samplerClassName, "create", methodType) | ||
| .invokeExact(telemetrySettings, settings, fallbackSampler); | ||
| } catch (Throwable e) { | ||
| if (e.getCause() instanceof NoSuchMethodException) { | ||
| throw new IllegalStateException("No create method exist in [" + samplerClassName + "]", e.getCause()); | ||
| } else { | ||
| throw new IllegalStateException("Sampler instantiation failed for class [" + samplerClassName + "]", e.getCause()); | ||
| } | ||
| } | ||
| }); | ||
| } catch (Exception e) { | ||
| throw new IllegalStateException("Sampler instantiation failed for class [" + samplerClassName + "]", e.getCause()); | ||
| } | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
99 changes: 99 additions & 0 deletions
99
...in/java/org/opensearch/telemetry/tracing/sampler/ProbabilisticTransportActionSampler.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,99 @@ | ||
| /* | ||
| * 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.OTelTelemetrySettings; | ||
| import org.opensearch.telemetry.TelemetrySettings; | ||
|
|
||
| import java.util.List; | ||
| import java.util.Objects; | ||
|
|
||
| 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.data.LinkData; | ||
| import io.opentelemetry.sdk.trace.samplers.Sampler; | ||
| import io.opentelemetry.sdk.trace.samplers.SamplingDecision; | ||
| import io.opentelemetry.sdk.trace.samplers.SamplingResult; | ||
|
|
||
| import static org.opensearch.telemetry.tracing.AttributeNames.TRANSPORT_ACTION; | ||
|
|
||
| /** | ||
| * ProbabilisticTransportActionSampler sampler samples request with action based on defined probability | ||
| */ | ||
| public class ProbabilisticTransportActionSampler implements Sampler { | ||
|
|
||
| private final Sampler fallbackSampler; | ||
| private Sampler actionSampler; | ||
| private final TelemetrySettings telemetrySettings; | ||
| private final Settings settings; | ||
| private double actionSamplingRatio; | ||
|
|
||
| /** | ||
| * Creates ProbabilisticTransportActionSampler sampler | ||
| * @param telemetrySettings TelemetrySettings | ||
| */ | ||
| private ProbabilisticTransportActionSampler(TelemetrySettings telemetrySettings, Settings settings, Sampler fallbackSampler) { | ||
| this.telemetrySettings = Objects.requireNonNull(telemetrySettings); | ||
| this.settings = Objects.requireNonNull(settings); | ||
| this.actionSamplingRatio = OTelTelemetrySettings.TRACER_SAMPLER_ACTION_PROBABILITY.get(settings); | ||
| this.actionSampler = Sampler.traceIdRatioBased(actionSamplingRatio); | ||
| this.fallbackSampler = fallbackSampler; | ||
| } | ||
|
|
||
| /** | ||
| * Create probabilistic transport action sampler. | ||
| * | ||
| * @param telemetrySettings the telemetry settings | ||
| * @param settings the settings | ||
| * @param fallbackSampler the fallback sampler | ||
| * @return the probabilistic transport action sampler | ||
| */ | ||
| public static Sampler create(TelemetrySettings telemetrySettings, Settings settings, Sampler fallbackSampler) { | ||
| return new ProbabilisticTransportActionSampler(telemetrySettings, settings, fallbackSampler); | ||
| } | ||
|
|
||
| @Override | ||
| public SamplingResult shouldSample( | ||
| Context parentContext, | ||
| String traceId, | ||
| String name, | ||
| SpanKind spanKind, | ||
| Attributes attributes, | ||
| List<LinkData> parentLinks | ||
| ) { | ||
| final String action = attributes.get(AttributeKey.stringKey(TRANSPORT_ACTION)); | ||
| if (action != null) { | ||
| final SamplingResult result = actionSampler.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks); | ||
| if (result.getDecision() != SamplingDecision.DROP && fallbackSampler != null) { | ||
| return fallbackSampler.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks); | ||
| } | ||
| return result; | ||
| } | ||
| if (fallbackSampler != null) return fallbackSampler.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks); | ||
|
|
||
| return SamplingResult.drop(); | ||
| } | ||
|
|
||
| double getSamplingRatio() { | ||
| return actionSamplingRatio; | ||
| } | ||
|
|
||
| @Override | ||
| public String getDescription() { | ||
| return "Transport Action Sampler"; | ||
| } | ||
|
|
||
| @Override | ||
| public String toString() { | ||
| return getDescription(); | ||
| } | ||
| } |
Oops, something went wrong.