Skip to content

Commit 5be018e

Browse files
devagarwal1803devagarwal1803
authored
Added custom sampler support based on action in request (opensearch-project#10136) (opensearch-project#12223)
* Added custom sampler support based on action in request Signed-off-by: Dev Agarwal <devagarwal1803@gmail.com> * UT Fix Signed-off-by: Dev Agarwal <devagarwal1803@gmail.com> * Added Transport action sampler, which will sample based on different probability for all actions Signed-off-by: Dev Agarwal <devagarwal1803@gmail.com> * 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 <devagarwal1803@gmail.com> * Added missing java-doc Signed-off-by: Dev Agarwal <devagarwal1803@gmail.com> * Moving sampler class settings to OtelTelemetry setting Signed-off-by: Dev Agarwal <devagarwal1803@gmail.com> * Minor refactor Signed-off-by: Dev Agarwal <devagarwal1803@gmail.com> * Refactored to use chain of samplers Signed-off-by: Dev Agarwal <devagarwal1803@gmail.com> * Addressed comments Signed-off-by: Dev Agarwal <devagarwal1803@gmail.com> * Addressed comments to move action_probability to OtelTelemetrySettings Signed-off-by: Dev Agarwal <devagarwal1803@gmail.com> * Updated eror msg returned when Sampler class is not found Signed-off-by: Dev Agarwal <devagarwal1803@gmail.com> * Added UT for OTelSamplerFactory Signed-off-by: Dev Agarwal <devagarwal1803@gmail.com> * minor refactor Signed-off-by: Dev Agarwal <devagarwal1803@gmail.com> * minor refactor Signed-off-by: Dev Agarwal <devagarwal1803@gmail.com> * spotless check Signed-off-by: Dev Agarwal <devagarwal1803@gmail.com> * Updating OtelTelemetryPlugin.get() method Signed-off-by: Dev Agarwal <devagarwal1803@gmail.com> * Addressed comments Signed-off-by: Dev Agarwal <devagarwal1803@gmail.com> * minor refactor Signed-off-by: Dev Agarwal <devagarwal1803@gmail.com> * addressed comments Signed-off-by: Dev Agarwal <devagarwal1803@gmail.com> * Updated transport action sampler Signed-off-by: Dev Agarwal <devagarwal1803@gmail.com> * Empty-Commit Signed-off-by: Dev Agarwal <devagarwal1803@gmail.com> * Empty-Commit Signed-off-by: Dev Agarwal <devagarwal1803@gmail.com> --------- Signed-off-by: Dev Agarwal <devagarwal1803@gmail.com> (cherry picked from commit 445bf1f)
1 parent ac37b38 commit 5be018e

15 files changed

+547
-73
lines changed

‎CHANGELOG.md

+2
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
143143
- Added Support for dynamically adding SearchRequestOperationsListeners with SearchRequestOperationsCompositeListenerFactory ([#11526](https://github.com/opensearch-project/OpenSearch/pull/11526))
144144
- [Query Insights] Query Insights Framework which currently supports retrieving the most time-consuming queries within the last configured time window ([#11903](https://github.com/opensearch-project/OpenSearch/pull/11903))
145145
- [Query Insights] Implement Top N Queries feature to collect and gather information about high latency queries in a window ([#11904](https://github.com/opensearch-project/OpenSearch/pull/11904))
146+
- Add override support for sampling based on action ([#9621](https://github.com/opensearch-project/OpenSearch/issues/9621))
147+
- Added custom sampler support based on transport action in request ([#9621](https://github.com/opensearch-project/OpenSearch/issues/9621))
146148

147149
### Deprecated
148150

‎plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/OTelTelemetryPlugin.java

+3-1
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,9 @@ public List<Setting<?>> getSettings() {
5353
OTelTelemetrySettings.TRACER_EXPORTER_DELAY_SETTING,
5454
OTelTelemetrySettings.TRACER_EXPORTER_MAX_QUEUE_SIZE_SETTING,
5555
OTelTelemetrySettings.OTEL_TRACER_SPAN_EXPORTER_CLASS_SETTING,
56-
OTelTelemetrySettings.OTEL_METRICS_EXPORTER_CLASS_SETTING
56+
OTelTelemetrySettings.OTEL_TRACER_SPAN_SAMPLER_CLASS_SETTINGS,
57+
OTelTelemetrySettings.OTEL_METRICS_EXPORTER_CLASS_SETTING,
58+
OTelTelemetrySettings.TRACER_SAMPLER_ACTION_PROBABILITY
5759
);
5860
}
5961

‎plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/OTelTelemetrySettings.java

+42
Original file line numberDiff line numberDiff line change
@@ -13,15 +13,21 @@
1313
import org.opensearch.common.unit.TimeValue;
1414
import org.opensearch.telemetry.metrics.exporter.OTelMetricsExporterFactory;
1515
import org.opensearch.telemetry.tracing.exporter.OTelSpanExporterFactory;
16+
import org.opensearch.telemetry.tracing.sampler.OTelSamplerFactory;
17+
import org.opensearch.telemetry.tracing.sampler.ProbabilisticSampler;
18+
import org.opensearch.telemetry.tracing.sampler.ProbabilisticTransportActionSampler;
1619

1720
import java.security.AccessController;
1821
import java.security.PrivilegedActionException;
1922
import java.security.PrivilegedExceptionAction;
23+
import java.util.Arrays;
24+
import java.util.List;
2025

2126
import io.opentelemetry.exporter.logging.LoggingMetricExporter;
2227
import io.opentelemetry.exporter.logging.LoggingSpanExporter;
2328
import io.opentelemetry.sdk.metrics.export.MetricExporter;
2429
import io.opentelemetry.sdk.trace.export.SpanExporter;
30+
import io.opentelemetry.sdk.trace.samplers.Sampler;
2531

2632
/**
2733
* OTel specific telemetry settings.
@@ -110,4 +116,40 @@ private OTelTelemetrySettings() {}
110116
Setting.Property.NodeScope,
111117
Setting.Property.Final
112118
);
119+
120+
/**
121+
* Samplers orders setting.
122+
*/
123+
@SuppressWarnings("unchecked")
124+
public static final Setting<List<Class<Sampler>>> OTEL_TRACER_SPAN_SAMPLER_CLASS_SETTINGS = Setting.listSetting(
125+
"telemetry.otel.tracer.span.sampler.classes",
126+
Arrays.asList(ProbabilisticTransportActionSampler.class.getName(), ProbabilisticSampler.class.getName()),
127+
sampler -> {
128+
// Check we ourselves are not being called by unprivileged code.
129+
SpecialPermission.check();
130+
try {
131+
return AccessController.doPrivileged((PrivilegedExceptionAction<Class<Sampler>>) () -> {
132+
final ClassLoader loader = OTelSamplerFactory.class.getClassLoader();
133+
return (Class<Sampler>) loader.loadClass(sampler);
134+
});
135+
} catch (PrivilegedActionException ex) {
136+
throw new IllegalStateException("Unable to load sampler class: " + sampler, ex.getCause());
137+
}
138+
},
139+
Setting.Property.NodeScope,
140+
Setting.Property.Final
141+
);
142+
143+
/**
144+
* Probability of action based sampler
145+
*/
146+
public static final Setting<Double> TRACER_SAMPLER_ACTION_PROBABILITY = Setting.doubleSetting(
147+
"telemetry.tracer.action.sampler.probability",
148+
0.001d,
149+
0.000d,
150+
1.00d,
151+
Setting.Property.NodeScope,
152+
Setting.Property.Dynamic
153+
);
154+
113155
}

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

+2-2
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
import org.opensearch.telemetry.TelemetrySettings;
1313
import org.opensearch.telemetry.metrics.exporter.OTelMetricsExporterFactory;
1414
import org.opensearch.telemetry.tracing.exporter.OTelSpanExporterFactory;
15-
import org.opensearch.telemetry.tracing.sampler.ProbabilisticSampler;
15+
import org.opensearch.telemetry.tracing.sampler.OTelSamplerFactory;
1616
import org.opensearch.telemetry.tracing.sampler.RequestSampler;
1717

1818
import java.security.AccessController;
@@ -60,7 +60,7 @@ public static OpenTelemetrySdk get(TelemetrySettings telemetrySettings, Settings
6060
settings,
6161
OTelSpanExporterFactory.create(settings),
6262
ContextPropagators.create(W3CTraceContextPropagator.getInstance()),
63-
Sampler.parentBased(new RequestSampler(new ProbabilisticSampler(telemetrySettings)))
63+
Sampler.parentBased(new RequestSampler(OTelSamplerFactory.create(telemetrySettings, settings)))
6464
)
6565
);
6666
}

‎plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/sampler/OTelSamplerFactory.java

+96
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.telemetry.tracing.sampler;
10+
11+
import org.apache.logging.log4j.LogManager;
12+
import org.apache.logging.log4j.Logger;
13+
import org.opensearch.SpecialPermission;
14+
import org.opensearch.common.settings.Settings;
15+
import org.opensearch.telemetry.OTelTelemetrySettings;
16+
import org.opensearch.telemetry.TelemetrySettings;
17+
18+
import java.lang.invoke.MethodHandles;
19+
import java.lang.invoke.MethodType;
20+
import java.security.AccessController;
21+
import java.security.PrivilegedExceptionAction;
22+
import java.util.List;
23+
import java.util.ListIterator;
24+
25+
import io.opentelemetry.sdk.trace.samplers.Sampler;
26+
27+
/**
28+
* Factory class to create the instance of OTelSampler
29+
*/
30+
public class OTelSamplerFactory {
31+
32+
/**
33+
* Logger instance for logging messages related to the OTelSamplerFactory.
34+
*/
35+
private static final Logger logger = LogManager.getLogger(OTelSamplerFactory.class);
36+
37+
/**
38+
* Base constructor.
39+
*/
40+
private OTelSamplerFactory() {
41+
42+
}
43+
44+
/**
45+
* Creates the {@link Sampler} instances based on the TRACER_SPAN_SAMPLER_CLASSES value.
46+
*
47+
* @param telemetrySettings TelemetrySettings.
48+
* @param settings the settings
49+
* @return list of samplers.
50+
*/
51+
public static Sampler create(TelemetrySettings telemetrySettings, Settings settings) {
52+
List<Class<Sampler>> samplersNameList = OTelTelemetrySettings.OTEL_TRACER_SPAN_SAMPLER_CLASS_SETTINGS.get(settings);
53+
ListIterator<Class<Sampler>> li = samplersNameList.listIterator(samplersNameList.size());
54+
55+
Sampler fallbackSampler = null;
56+
57+
// Iterating samplers list in reverse order to create chain of sampler
58+
while (li.hasPrevious()) {
59+
Class<Sampler> samplerName = li.previous();
60+
fallbackSampler = instantiateSampler(samplerName, telemetrySettings, settings, fallbackSampler);
61+
}
62+
63+
return fallbackSampler;
64+
}
65+
66+
private static Sampler instantiateSampler(
67+
Class<Sampler> samplerClassName,
68+
TelemetrySettings telemetrySettings,
69+
Settings settings,
70+
Sampler fallbackSampler
71+
) {
72+
try {
73+
// Check we ourselves are not being called by unprivileged code.
74+
SpecialPermission.check();
75+
76+
return AccessController.doPrivileged((PrivilegedExceptionAction<Sampler>) () -> {
77+
try {
78+
// Define the method type which receives TelemetrySettings & Sampler as arguments
79+
MethodType methodType = MethodType.methodType(Sampler.class, TelemetrySettings.class, Settings.class, Sampler.class);
80+
81+
return (Sampler) MethodHandles.publicLookup()
82+
.findStatic(samplerClassName, "create", methodType)
83+
.invokeExact(telemetrySettings, settings, fallbackSampler);
84+
} catch (Throwable e) {
85+
if (e.getCause() instanceof NoSuchMethodException) {
86+
throw new IllegalStateException("No create method exist in [" + samplerClassName + "]", e.getCause());
87+
} else {
88+
throw new IllegalStateException("Sampler instantiation failed for class [" + samplerClassName + "]", e.getCause());
89+
}
90+
}
91+
});
92+
} catch (Exception e) {
93+
throw new IllegalStateException("Sampler instantiation failed for class [" + samplerClassName + "]", e.getCause());
94+
}
95+
}
96+
}

‎plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/sampler/ProbabilisticSampler.java

+32-12
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88

99
package org.opensearch.telemetry.tracing.sampler;
1010

11+
import org.opensearch.common.settings.Settings;
1112
import org.opensearch.telemetry.TelemetrySettings;
1213

1314
import java.util.List;
@@ -18,36 +19,43 @@
1819
import io.opentelemetry.context.Context;
1920
import io.opentelemetry.sdk.trace.data.LinkData;
2021
import io.opentelemetry.sdk.trace.samplers.Sampler;
22+
import io.opentelemetry.sdk.trace.samplers.SamplingDecision;
2123
import io.opentelemetry.sdk.trace.samplers.SamplingResult;
2224

2325
/**
24-
* ProbabilisticSampler implements a head-based sampling strategy based on provided settings.
26+
* ProbabilisticSampler implements a probability sampling strategy based on configured sampling ratio.
2527
*/
2628
public class ProbabilisticSampler implements Sampler {
2729
private Sampler defaultSampler;
2830
private final TelemetrySettings telemetrySettings;
31+
private final Settings settings;
32+
private final Sampler fallbackSampler;
33+
2934
private double samplingRatio;
3035

3136
/**
3237
* Constructor
3338
*
3439
* @param telemetrySettings Telemetry settings.
3540
*/
36-
public ProbabilisticSampler(TelemetrySettings telemetrySettings) {
41+
private ProbabilisticSampler(TelemetrySettings telemetrySettings, Settings settings, Sampler fallbackSampler) {
3742
this.telemetrySettings = Objects.requireNonNull(telemetrySettings);
43+
this.settings = Objects.requireNonNull(settings);
3844
this.samplingRatio = telemetrySettings.getSamplingProbability();
3945
this.defaultSampler = Sampler.traceIdRatioBased(samplingRatio);
46+
this.fallbackSampler = fallbackSampler;
4047
}
4148

42-
Sampler getSampler() {
43-
double newSamplingRatio = telemetrySettings.getSamplingProbability();
44-
if (isSamplingRatioChanged(newSamplingRatio)) {
45-
synchronized (this) {
46-
this.samplingRatio = newSamplingRatio;
47-
defaultSampler = Sampler.traceIdRatioBased(samplingRatio);
48-
}
49-
}
50-
return defaultSampler;
49+
/**
50+
* Create probabilistic sampler.
51+
*
52+
* @param telemetrySettings the telemetry settings
53+
* @param settings the settings
54+
* @param fallbackSampler the fallback sampler
55+
* @return the probabilistic sampler
56+
*/
57+
public static Sampler create(TelemetrySettings telemetrySettings, Settings settings, Sampler fallbackSampler) {
58+
return new ProbabilisticSampler(telemetrySettings, settings, fallbackSampler);
5159
}
5260

5361
private boolean isSamplingRatioChanged(double newSamplingRatio) {
@@ -67,7 +75,19 @@ public SamplingResult shouldSample(
6775
Attributes attributes,
6876
List<LinkData> parentLinks
6977
) {
70-
return getSampler().shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks);
78+
double newSamplingRatio = telemetrySettings.getSamplingProbability();
79+
if (isSamplingRatioChanged(newSamplingRatio)) {
80+
synchronized (this) {
81+
this.samplingRatio = newSamplingRatio;
82+
defaultSampler = Sampler.traceIdRatioBased(samplingRatio);
83+
}
84+
}
85+
final SamplingResult result = defaultSampler.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks);
86+
if (result.getDecision() != SamplingDecision.DROP && fallbackSampler != null) {
87+
return fallbackSampler.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks);
88+
} else {
89+
return result;
90+
}
7191
}
7292

7393
@Override

‎plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/sampler/ProbabilisticTransportActionSampler.java

+99
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.telemetry.tracing.sampler;
10+
11+
import org.opensearch.common.settings.Settings;
12+
import org.opensearch.telemetry.OTelTelemetrySettings;
13+
import org.opensearch.telemetry.TelemetrySettings;
14+
15+
import java.util.List;
16+
import java.util.Objects;
17+
18+
import io.opentelemetry.api.common.AttributeKey;
19+
import io.opentelemetry.api.common.Attributes;
20+
import io.opentelemetry.api.trace.SpanKind;
21+
import io.opentelemetry.context.Context;
22+
import io.opentelemetry.sdk.trace.data.LinkData;
23+
import io.opentelemetry.sdk.trace.samplers.Sampler;
24+
import io.opentelemetry.sdk.trace.samplers.SamplingDecision;
25+
import io.opentelemetry.sdk.trace.samplers.SamplingResult;
26+
27+
import static org.opensearch.telemetry.tracing.AttributeNames.TRANSPORT_ACTION;
28+
29+
/**
30+
* ProbabilisticTransportActionSampler sampler samples request with action based on defined probability
31+
*/
32+
public class ProbabilisticTransportActionSampler implements Sampler {
33+
34+
private final Sampler fallbackSampler;
35+
private Sampler actionSampler;
36+
private final TelemetrySettings telemetrySettings;
37+
private final Settings settings;
38+
private double actionSamplingRatio;
39+
40+
/**
41+
* Creates ProbabilisticTransportActionSampler sampler
42+
* @param telemetrySettings TelemetrySettings
43+
*/
44+
private ProbabilisticTransportActionSampler(TelemetrySettings telemetrySettings, Settings settings, Sampler fallbackSampler) {
45+
this.telemetrySettings = Objects.requireNonNull(telemetrySettings);
46+
this.settings = Objects.requireNonNull(settings);
47+
this.actionSamplingRatio = OTelTelemetrySettings.TRACER_SAMPLER_ACTION_PROBABILITY.get(settings);
48+
this.actionSampler = Sampler.traceIdRatioBased(actionSamplingRatio);
49+
this.fallbackSampler = fallbackSampler;
50+
}
51+
52+
/**
53+
* Create probabilistic transport action sampler.
54+
*
55+
* @param telemetrySettings the telemetry settings
56+
* @param settings the settings
57+
* @param fallbackSampler the fallback sampler
58+
* @return the probabilistic transport action sampler
59+
*/
60+
public static Sampler create(TelemetrySettings telemetrySettings, Settings settings, Sampler fallbackSampler) {
61+
return new ProbabilisticTransportActionSampler(telemetrySettings, settings, fallbackSampler);
62+
}
63+
64+
@Override
65+
public SamplingResult shouldSample(
66+
Context parentContext,
67+
String traceId,
68+
String name,
69+
SpanKind spanKind,
70+
Attributes attributes,
71+
List<LinkData> parentLinks
72+
) {
73+
final String action = attributes.get(AttributeKey.stringKey(TRANSPORT_ACTION));
74+
if (action != null) {
75+
final SamplingResult result = actionSampler.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks);
76+
if (result.getDecision() != SamplingDecision.DROP && fallbackSampler != null) {
77+
return fallbackSampler.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks);
78+
}
79+
return result;
80+
}
81+
if (fallbackSampler != null) return fallbackSampler.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks);
82+
83+
return SamplingResult.drop();
84+
}
85+
86+
double getSamplingRatio() {
87+
return actionSamplingRatio;
88+
}
89+
90+
@Override
91+
public String getDescription() {
92+
return "Transport Action Sampler";
93+
}
94+
95+
@Override
96+
public String toString() {
97+
return getDescription();
98+
}
99+
}

0 commit comments

Comments
 (0)