Skip to content

Commit ae5df65

Browse files
kkondakakkondaka
authored
Add sanitization to metric names and labels (#6277)
* Add sanitization to metric names and labels Signed-off-by: Kondaka <krishkdk@amazon.com> * Addressed review comments Signed-off-by: Kondaka <krishkdk@amazon.com> * Fixed build issues Signed-off-by: Kondaka <krishkdk@amazon.com> * Addressed comments Signed-off-by: Kondaka <krishkdk@amazon.com> * Removed unnecessary debug statement Signed-off-by: Kondaka <krishkdk@amazon.com> --------- Signed-off-by: Kondaka <krishkdk@amazon.com>
1 parent 23eb0d9 commit ae5df65

15 files changed

Lines changed: 531 additions & 63 deletions

File tree

‎data-prepper-plugins/prometheus-sink/build.gradle‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ dependencies {
5252
implementation project(':data-prepper-plugin-framework')
5353
testImplementation 'com.fasterxml.jackson.datatype:jackson-datatype-jsr310'
5454
testImplementation project(':data-prepper-test:test-common')
55+
testImplementation project(':data-prepper-test:plugin-test-framework')
5556
testImplementation project(':data-prepper-core')
5657
testImplementation project(':data-prepper-plugin-framework')
5758
testImplementation 'com.fasterxml.jackson.datatype:jackson-datatype-jsr310'

‎data-prepper-plugins/prometheus-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/prometheus/PrometheusSinkAMPIT.java‎

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010

1111
package org.opensearch.dataprepper.plugins.sink.prometheus;
1212

13+
import org.opensearch.dataprepper.model.configuration.PipelineDescription;
1314
import org.opensearch.dataprepper.model.metric.JacksonSum;
1415
import org.opensearch.dataprepper.model.metric.JacksonSummary;
1516
import org.opensearch.dataprepper.model.metric.JacksonGauge;
@@ -94,6 +95,7 @@
9495

9596
@ExtendWith(MockitoExtension.class)
9697
public class PrometheusSinkAMPIT {
98+
private static final String TEST_PIPELINE_NAME = "testPipeline";
9799
private static final long NANO_MULTIPLIER = 1_000_000_000L;
98100
private static final int TEST_READ_BATCH_TIMEOUT = 500;
99101
private static final int TEST_PROCESSOR_THREADS = 1;
@@ -130,6 +132,8 @@ public class PrometheusSinkAMPIT {
130132
private EventHandle eventHandle;
131133
@Mock
132134
private Pipeline dlqPipeline;
135+
@Mock
136+
private PipelineDescription pipelineDescription;
133137

134138
private String awsRegion;
135139
private String awsRole;
@@ -146,6 +150,7 @@ public class PrometheusSinkAMPIT {
146150
private WebClient webClient;
147151
private PrometheusSinkThresholdConfig thresholdConfig;
148152

153+
149154
@BeforeEach
150155
void setUp() {
151156
webClient = WebClient.builder()
@@ -157,6 +162,7 @@ void setUp() {
157162
.build();
158163

159164
eventHandle = mock(EventHandle.class);
165+
pipelineDescription = mock(PipelineDescription.class);
160166
awsCredentialsProvider = DefaultCredentialsProvider.create();
161167
metricsInAMP = 0;
162168
testStartTime = Instant.now();
@@ -174,13 +180,11 @@ void setUp() {
174180
awsConfig = mock(AwsConfig.class);
175181
when(pluginSetting.getPipelineName()).thenReturn("pipeline");
176182
when(pluginSetting.getName()).thenReturn("name");
177-
178183

179184
metricsSuccessCounter = mock(Counter.class);
180185
metricsFailedCounter = mock(Counter.class);
181186
requestsSuccessCounter = mock(Counter.class);
182187
requestsFailedCounter = mock(Counter.class);
183-
184188
summary = mock(DistributionSummary.class);
185189

186190
when(pluginMetrics.counter(eq("sinkRequestsSucceeded"))).thenReturn(requestsSuccessCounter);
@@ -222,7 +226,7 @@ void setUp() {
222226
}
223227

224228
private PrometheusSink createObjectUnderTest() {
225-
return new PrometheusSink(pluginSetting, pluginMetrics, pluginFactory, prometheusSinkConfig, awsCredentialsSupplier);
229+
return new PrometheusSink(pluginSetting, pluginMetrics, pipelineDescription, prometheusSinkConfig, awsCredentialsSupplier);
226230
}
227231

228232
private void getMetricsFromAMP(final String metricName, final String qs) throws Exception {
@@ -242,7 +246,7 @@ private void getMetricsFromAMP(final String metricName, final String qs) throws
242246
queryStr = metricName+"|"+metricName+"_sum|"+metricName+"_count|attrKey1|attrKey2";
243247
query = "{__name__=~\""+queryStr+"\"}";
244248
} else if (qs.equals("sum")) {
245-
queryStr = metricName+"|"+metricName+"_sum";
249+
queryStr = metricName+"|"+metricName+"_total";
246250
query = "{__name__=~\""+queryStr+"\"}";
247251
} else {
248252
query = metricName;
@@ -263,6 +267,7 @@ private void getMetricsFromAMP(final String metricName, final String qs) throws
263267
headersBuilder.add(k, v);
264268
});
265269
});
270+
266271
HttpRequest request = HttpRequest.of(headersBuilder.build(), HttpData.ofAscii(getUrlQuery));
267272
webClient.execute(request).aggregate()
268273
.thenAccept(response -> {
@@ -328,6 +333,7 @@ void TestSumMetrics() throws Exception {
328333

329334
@Test
330335
void TestSumMetricsFailuresWithDLQ() throws Exception {
336+
when(pipelineDescription.getPipelineName()).thenReturn(TEST_PIPELINE_NAME);
331337
dlqPipeline = mock(Pipeline.class);
332338
doAnswer(a -> {
333339
Collection<Record<Event>> records = (Collection<Record<Event>>)a.getArgument(0);
@@ -443,7 +449,7 @@ private Collection<Record<Event>> getSumRecordList(int numberOfRecords, final St
443449
.withStartTime(convertUnixNanosToISO8601(startTimeNanos))
444450
.withIsMonotonic(true)
445451
.withUnit("1")
446-
.withAggregationTemporality("delta")
452+
.withAggregationTemporality("AGGREGATION_TEMPORALITY_CUMULATIVE")
447453
.withResource(Map.of("attributes", Map.of("attrKey1", 1, "attrKey2", Map.of("attrKey3", "attrValue3"))))
448454
.withValue((double)metricValue+i)
449455
.withEventHandle(eventHandle)

‎data-prepper-plugins/prometheus-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/prometheus/PrometheusHttpSender.java‎

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
import org.opensearch.dataprepper.model.codec.CompressionEngine;
3434
import com.linecorp.armeria.client.ClientFactory;
3535
import com.linecorp.armeria.client.ClientOptions;
36+
import com.google.common.annotations.VisibleForTesting;
3637

3738
import javax.annotation.Nonnull;
3839
import java.nio.charset.StandardCharsets;
@@ -46,6 +47,9 @@ public class PrometheusHttpSender {
4647
private static final Logger LOG = LoggerFactory.getLogger(PrometheusHttpSender.class);
4748
private static final int DEFAULT_MAX_REQUEST_SIZE = 1024*1024; // 1MB
4849
private static final Set<Integer> RETRYABLE_STATUS_CODES = Set.of(429, 502, 503, 504);
50+
private static final int BACKOFF_INITIAL_DELAY_MS = 100;
51+
private static final int BACKOFF_MAX_DELAY_MS = 10_000;
52+
private static final double BACKOFF_DEFAULT_JITTER = 0.2;
4953

5054
private final PrometheusSigV4Signer signer;
5155
private final WebClient webClient;
@@ -61,14 +65,20 @@ public class PrometheusHttpSender {
6165
* @param awsCredentialsSupplier the AWS credentials supplier
6266
* @param config The configuration for the Prometheus sink plugin.
6367
*/
64-
public PrometheusHttpSender(@Nonnull final AwsCredentialsSupplier awsCredentialsSupplier, @Nonnull final PrometheusSinkConfiguration config, @Nonnull final SinkMetrics sinkMetrics, final long connectionTimeoutMillis, final long idleTimeoutMillis) {
68+
public PrometheusHttpSender(@Nonnull final AwsCredentialsSupplier awsCredentialsSupplier, @Nonnull final PrometheusSinkConfiguration config, @Nonnull final SinkMetrics sinkMetrics) {
69+
this(awsCredentialsSupplier, buildWebClient(config), config, sinkMetrics);
70+
}
71+
72+
@VisibleForTesting
73+
public PrometheusHttpSender(@Nonnull final AwsCredentialsSupplier awsCredentialsSupplier, @Nonnull WebClient webClient, @Nonnull final PrometheusSinkConfiguration config, @Nonnull final SinkMetrics sinkMetrics) {
6574
this.signer = config.getAwsConfig() != null ? new PrometheusSigV4Signer(awsCredentialsSupplier, config) : null;
66-
this.webClient = buildWebClient(config);
75+
this.webClient = webClient;
6776
this.compressionEngine = config.getEncoding().getCompressionEngine();
6877
this.sinkMetrics = sinkMetrics;
6978
this.config = config;
70-
this.connectionTimeoutMillis = connectionTimeoutMillis;
71-
this.idleTimeoutMillis = idleTimeoutMillis;
79+
this.connectionTimeoutMillis = config.getConnectionTimeout().toMillis();
80+
this.idleTimeoutMillis = config.getIdleTimeout().toMillis();
81+
7282
}
7383

7484
/**
@@ -86,7 +96,7 @@ public PrometheusHttpSender(@Nonnull final AwsCredentialsSupplier awsCredentials
8696
private static WebClient buildWebClient(final PrometheusSinkConfiguration config) {
8797
final RetryRuleWithContent<HttpResponse> retryRule = RetryRuleWithContent.<HttpResponse>builder()
8898
.onStatus((ctx, status) -> RETRYABLE_STATUS_CODES.contains(status.code()))
89-
.thenBackoff(Backoff.exponential(100, 10_000).withJitter(0.2));
99+
.thenBackoff(Backoff.exponential(BACKOFF_INITIAL_DELAY_MS, BACKOFF_MAX_DELAY_MS).withJitter(BACKOFF_DEFAULT_JITTER));
90100

91101
final long estimatedContentLimit = Math.max(1, DEFAULT_MAX_REQUEST_SIZE) * (config.getMaxRetries() + 1);
92102
final int safeContentLimit = (int) Math.min(estimatedContentLimit, Integer.MAX_VALUE);
@@ -125,7 +135,7 @@ public PrometheusPushResult pushToEndpoint(final byte[] payload) {
125135
.thenApply(response -> {
126136
final long latency = System.currentTimeMillis() - startTime;
127137
sinkMetrics.recordRequestSize(compressedBufferData.length);
128-
LOG.error("Response received in {}ms. Status: {}", latency, response.status());
138+
LOG.debug("Response received in {}ms. Status: {}", latency, response.status());
129139

130140
int statusCode = response.status().code();
131141
final byte[] responseBytes = response.content().array();
@@ -168,6 +178,7 @@ private HttpRequest buildHttpRequest(final byte[] payload) {
168178
.path(sdkHttpRequest.getUri().getRawPath())
169179
.authority(sdkHttpRequest.getUri().getAuthority());
170180

181+
171182
// Preserve all original headers from the signed request without modification
172183
sdkHttpRequest.headers().forEach((k, vList) -> {
173184
// Add each header value individually to preserve exact format

‎data-prepper-plugins/prometheus-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/prometheus/PrometheusSigV4Signer.java‎

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
*/
1010
package org.opensearch.dataprepper.plugins.sink.prometheus;
1111

12+
import org.opensearch.dataprepper.aws.api.AwsConfig;
1213
import org.opensearch.dataprepper.plugins.sink.prometheus.configuration.PrometheusSinkConfiguration;
1314
import org.opensearch.dataprepper.aws.api.AwsCredentialsOptions;
1415
import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier;
@@ -41,15 +42,24 @@ class PrometheusSigV4Signer {
4142
this.region = config.getAwsConfig().getAwsRegion();
4243

4344
this.config = config;
44-
this.credentialsProvider = awsCredentialsSupplier.getProvider(AwsCredentialsOptions.builder()
45-
.withRegion(region)
46-
.withStsRoleArn(config.getAwsConfig().getAwsStsRoleArn())
47-
.withStsExternalId(config.getAwsConfig().getAwsStsExternalId())
48-
.build());
45+
this.credentialsProvider = awsCredentialsSupplier.getProvider(convertToCredentialOptions(config.getAwsConfig()));
4946

5047
this.endpointUri = URI.create(url);
5148
}
5249

50+
private static AwsCredentialsOptions convertToCredentialOptions(final AwsConfig awsConfig) {
51+
if (awsConfig == null) {
52+
return AwsCredentialsOptions.builder().build();
53+
}
54+
return AwsCredentialsOptions.builder()
55+
.withRegion(awsConfig.getAwsRegion())
56+
.withStsRoleArn(awsConfig.getAwsStsRoleArn())
57+
.withStsExternalId(awsConfig.getAwsStsExternalId())
58+
.withStsHeaderOverrides(awsConfig.getAwsStsHeaderOverrides())
59+
.build();
60+
}
61+
62+
5363
/**
5464
* Constructs a SigV4 signer helper.
5565
*

‎data-prepper-plugins/prometheus-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/prometheus/PrometheusSink.java‎

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,9 @@
2121
import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin;
2222
import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor;
2323
import org.opensearch.dataprepper.model.configuration.PluginSetting;
24+
import org.opensearch.dataprepper.model.configuration.PipelineDescription;
2425
import org.opensearch.dataprepper.model.event.Event;
2526
import org.opensearch.dataprepper.model.pipeline.HeadlessPipeline;
26-
import org.opensearch.dataprepper.model.plugin.PluginFactory;
2727
import org.opensearch.dataprepper.metrics.PluginMetrics;
2828
import org.opensearch.dataprepper.model.record.Record;
2929
import org.opensearch.dataprepper.model.sink.AbstractSink;
@@ -48,7 +48,7 @@ public class PrometheusSink extends AbstractSink<Record<Event>> {
4848
@DataPrepperPluginConstructor
4949
public PrometheusSink(final PluginSetting pluginSetting,
5050
final PluginMetrics pluginMetrics,
51-
final PluginFactory pluginFactory,
51+
final PipelineDescription pipelineDescription,
5252
final PrometheusSinkConfiguration prometheusSinkConfiguration,
5353
final AwsCredentialsSupplier awsCredentialsSupplier) {
5454
super(pluginSetting);
@@ -58,9 +58,7 @@ public PrometheusSink(final PluginSetting pluginSetting,
5858
Region region = (awsConfig != null) ? awsConfig.getAwsRegion() : awsCredentialsSupplier.getDefaultRegion().get();
5959

6060
sinkMetrics = new DefaultSinkMetrics(pluginMetrics, "Metric");
61-
httpSender = new PrometheusHttpSender(awsCredentialsSupplier, prometheusSinkConfiguration, sinkMetrics,
62-
prometheusSinkConfiguration.getConnectionTimeout().toMillis(),
63-
prometheusSinkConfiguration.getIdleTimeout().toMillis());
61+
httpSender = new PrometheusHttpSender(awsCredentialsSupplier, prometheusSinkConfiguration, sinkMetrics);
6462

6563
PrometheusSinkThresholdConfig thresholdConfig = prometheusSinkConfiguration.getThresholdConfig();
6664

@@ -69,8 +67,7 @@ public PrometheusSink(final PluginSetting pluginSetting,
6967
sinkMetrics,
7068
httpSender,
7169
getFailurePipeline(),
72-
pluginMetrics,
73-
pluginSetting);
70+
pipelineDescription);
7471
}
7572

7673
private static AwsCredentialsOptions convertToCredentialOptions(final AwsConfig awsConfig) {

‎data-prepper-plugins/prometheus-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/prometheus/configuration/PrometheusSinkThresholdConfig.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ public class PrometheusSinkThresholdConfig {
3232
private int maxEvents = DEFAULT_MAX_EVENTS;
3333

3434
@JsonProperty("max_request_size")
35-
private String maxRequestSize = DEFAULT_MAX_REQUEST_SIZE;
35+
private ByteCount maxRequestSize = ByteCount.parse(DEFAULT_MAX_REQUEST_SIZE);
3636

3737
@JsonProperty("flush_interval")
3838
@DurationMin(seconds = 1)
@@ -44,7 +44,7 @@ public int getMaxEvents() {
4444
}
4545

4646
public long getMaxRequestSizeBytes() {
47-
return ByteCount.parse(maxRequestSize).getBytes();
47+
return maxRequestSize.getBytes();
4848
}
4949

5050
public long getFlushInterval() {

‎data-prepper-plugins/prometheus-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/prometheus/service/PrometheusSinkService.java‎

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,7 @@
1717
import org.opensearch.dataprepper.common.sink.SinkMetrics;
1818
import org.opensearch.dataprepper.common.sink.SinkBufferEntry;
1919
import org.opensearch.dataprepper.common.sink.ReentrantLockStrategy;
20-
import org.opensearch.dataprepper.metrics.PluginMetrics;
21-
import org.opensearch.dataprepper.model.configuration.PluginSetting;
20+
import org.opensearch.dataprepper.model.configuration.PipelineDescription;
2221
import org.opensearch.dataprepper.model.event.Event;
2322
import org.opensearch.dataprepper.model.record.Record;
2423
import org.opensearch.dataprepper.model.pipeline.HeadlessPipeline;
@@ -33,24 +32,24 @@
3332

3433

3534
public class PrometheusSinkService extends DefaultSinkOutputStrategy {
36-
35+
static final String PLUGIN_NAME = "prometheus";
3736
private static final Logger LOG = LoggerFactory.getLogger(PrometheusSinkService.class);
3837
public static final String PROMETHEUS_SINK_RECORDS_SUCCESS_COUNTER = "prometheusSinkRecordsNumberOfSuccessful";
3938

4039
public static final String PROMETHEUS_SINK_RECORDS_FAILED_COUNTER = "prometheusSinkRecordsNumberOfFailed";
4140
private final PrometheusHttpSender httpSender;
42-
private final PluginSetting pluginSetting;
41+
private final PipelineDescription pipelineDescription;
4342
private final List<Record<Event>> dlqRecords;
4443
private final boolean sanitizeNames;
4544
private HeadlessPipeline dlqPipeline;
4645
private boolean dropIfNoDLQConfigured;
46+
private String pluginName;
4747

4848
public PrometheusSinkService(final PrometheusSinkConfiguration prometheusSinkConfiguration,
4949
final SinkMetrics sinkMetrics,
5050
final PrometheusHttpSender httpSender,
5151
final HeadlessPipeline dlqPipeline,
52-
final PluginMetrics pluginMetrics,
53-
final PluginSetting pluginSetting) {
52+
final PipelineDescription pipelineDescription) {
5453
super(new ReentrantLockStrategy(),
5554
new DefaultSinkBuffer(prometheusSinkConfiguration.getThresholdConfig().getMaxEvents(),
5655
prometheusSinkConfiguration.getThresholdConfig().getMaxRequestSizeBytes(),
@@ -63,7 +62,7 @@ public PrometheusSinkService(final PrometheusSinkConfiguration prometheusSinkCon
6362
this.dlqPipeline = dlqPipeline;
6463
this.dlqRecords = new ArrayList<>();
6564
this.httpSender = httpSender;
66-
this.pluginSetting = pluginSetting;
65+
this.pipelineDescription = pipelineDescription;
6766
}
6867

6968
public void addFailedEventsToDLQ(final List<Event> events, final Throwable ex) {
@@ -104,8 +103,8 @@ public void addFailedEventsToDlq(final List<Event> failedEvents, final Throwable
104103
}
105104
event.updateFailureMetadata()
106105
.with("statusCode", statusCode)
107-
.with("pluginName", pluginSetting.getName())
108-
.with("pipelineName", pluginSetting.getPipelineName());
106+
.with("pluginName", PLUGIN_NAME)
107+
.with("pipelineName", pipelineDescription.getPipelineName());
109108
if (ex != null) {
110109
event.updateFailureMetadata()
111110
.with("message", ex.getMessage());

0 commit comments

Comments
 (0)