From 1fc9a608cc72e25323dc6405110e25f50abc79f7 Mon Sep 17 00:00:00 2001 From: huyPham Date: Thu, 3 Apr 2025 11:42:38 -0700 Subject: [PATCH 01/23] Feature/xray init (#1) * feat(xray-otlp-sink): Add X-Ray OTLP Sink plugin skeleton - Added test resources and support for Span records - Added sample pipeline config and OTLP test span JSON under `src/test/resources` - Verified local pipeline ingest and logging using `grpcurl` - Added README with developer instructions for running and testing locally These changes establish the foundation for local testing and future X-Ray e2e testing. Signed-off-by: huy pham --- data-prepper-plugins/xray-otlp-sink/README.md | 54 +++++++++++ .../xray-otlp-sink/build.gradle | 20 ++++ .../plugins/sink/xrayotlp/XRayOTLPSink.java | 93 +++++++++++++++++++ .../sink/xrayotlp/XRayOTLPSinkTest.java | 48 ++++++++++ .../test/resources/data-prepper-config.yaml | 1 + .../src/test/resources/pipelines.yaml | 13 +++ .../src/test/resources/sample-trace.json | 28 ++++++ settings.gradle | 1 + 8 files changed, 258 insertions(+) create mode 100644 data-prepper-plugins/xray-otlp-sink/README.md create mode 100644 data-prepper-plugins/xray-otlp-sink/build.gradle create mode 100644 data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSink.java create mode 100644 data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkTest.java create mode 100644 data-prepper-plugins/xray-otlp-sink/src/test/resources/data-prepper-config.yaml create mode 100644 data-prepper-plugins/xray-otlp-sink/src/test/resources/pipelines.yaml create mode 100644 data-prepper-plugins/xray-otlp-sink/src/test/resources/sample-trace.json diff --git a/data-prepper-plugins/xray-otlp-sink/README.md b/data-prepper-plugins/xray-otlp-sink/README.md new file mode 100644 index 0000000000..9bd72e77e9 --- /dev/null +++ b/data-prepper-plugins/xray-otlp-sink/README.md @@ -0,0 +1,54 @@ +# X-Ray OTLP Sink + +The `xray_otlp_sink` plugin sends span data to [AWS X-Ray](https://docs.aws.amazon.com/xray/) using the OTLP (OpenTelemetry Protocol) format. + +## Usage + +For information on usage, see the forthcoming documentation in the [Data Prepper Sink Plugins section](https://opensearch.org/docs/latest/data-prepper/pipelines/configuration/sinks/). + +A sample pipeline configuration will be added once the plugin is ready for testing. + +## Developer Guide + +See the [CONTRIBUTING](https://github.com/opensearch-project/data-prepper/blob/main/CONTRIBUTING.md) guide for general information on contributions. + +The integration tests for this plugin do not run as part of the main Data Prepper build. + +#### Run unit tests locally + +```bash +./gradlew :data-prepper-plugins:xray-otlp-sink:test +``` + +#### Run a local pipeline that uses this sink + +1. Install `grpcurl` – Used to send OTLP span data to the running pipeline. +2. Build the plugin and Data Prepper: +``` +./gradlew build` +``` +3. Start the pipeline: +``` +cd release/archives/linux/build/install/opensearch-data-prepper-2.11.0-SNAPSHOT-linux-x64 + +bin/data-prepper \ + /path/to/data-prepper-plugins/xray-otlp-sink/src/test/resources/pipelines.yaml \ + /path/to/data-prepper-plugins/xray-otlp-sink/src/test/resources/data-prepper-config.yaml +``` +4. Send test spans to the local pipeline: +``` +cd /path/to/opentelemetry-proto + +grpcurl -plaintext \ + -import-path . \ + -proto opentelemetry/proto/collector/trace/v1/trace_service.proto \ + -proto opentelemetry/proto/common/v1/common.proto \ + -proto opentelemetry/proto/resource/v1/resource.proto \ + -proto opentelemetry/proto/trace/v1/trace.proto \ + -d @ \ + localhost:21890 \ + opentelemetry.proto.collector.trace.v1.TraceService/Export \ + < /path/to/data-prepper-plugins/xray-otlp-sink/src/test/resources/sample-trace.json +``` + +You should see log output from XRayOTLPSink that confirms the span data was received and parsed correctly. diff --git a/data-prepper-plugins/xray-otlp-sink/build.gradle b/data-prepper-plugins/xray-otlp-sink/build.gradle new file mode 100644 index 0000000000..b3611457ee --- /dev/null +++ b/data-prepper-plugins/xray-otlp-sink/build.gradle @@ -0,0 +1,20 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +plugins { + id 'java-library' +} + +dependencies { + implementation project(':data-prepper-api') + implementation 'com.fasterxml.jackson.core:jackson-databind' + + testImplementation 'org.junit.jupiter:junit-jupiter-api:5.9.1' + testRuntimeOnly 'org.junit.jupiter:junit-jupiter-engine:5.9.1' +} + +test { + useJUnitPlatform() +} \ No newline at end of file diff --git a/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSink.java b/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSink.java new file mode 100644 index 0000000000..c6a78e27d4 --- /dev/null +++ b/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSink.java @@ -0,0 +1,93 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.xrayotlp; + +import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; +import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.model.sink.Sink; +import org.opensearch.dataprepper.model.trace.Span; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Collection; + +/** + * A Data Prepper Sink plugin that forwards traces to AWS X-Ray's OTLP endpoint. + */ +@DataPrepperPlugin( + name = "xray_otlp_sink", + pluginType = Sink.class +) +public class XRayOTLPSink implements Sink> { + + private static final Logger LOG = LoggerFactory.getLogger(XRayOTLPSink.class); + + /** + * Constructs the OTLP X-Ray Sink. + * Configuration loading will be added in a later iteration. + */ + public XRayOTLPSink() { + // TODO: Inject config or plugin setting. + } + + /** + * Lifecycle hook invoked during pipeline startup. + * Initialize AWS clients or other resources here. + */ + @Override + public void initialize() { + // TODO: Initialize AWS X-Ray client + } + + /** + * Called each time a batch of records is emitted to the sink. + * This method is responsible for handling delivery to AWS X-Ray. + * + * @param records Collection of OTLP log records to process. + */ + @Override + public void output(final Collection> records) { + for (Record record : records) { + final Span span = record.getData(); + + LOG.info("===> Span name: {}", span.getName()); + LOG.info("===> Trace ID: {}", span.getTraceId()); + LOG.info("===> Span ID: {}", span.getSpanId()); + LOG.info("===> Parent ID: {}", span.getParentSpanId()); + LOG.info("===> Start time (epoch nanos): {}", span.getStartTime()); + LOG.info("===> End time (epoch nanos): {}", span.getEndTime()); + LOG.info("===> Attributes: {}", span.getAttributes()); + } + } + + /** + * Indicates whether this sink is ready to receive data. + * + * @return true if the sink is ready + */ + @Override + public boolean isReady() { + // TODO: Implement readiness logic + return true; + } + + /** + * Hook called during pipeline shutdown. + */ + @Override + public void shutdown() { + // TODO: Clean up resources + } + + /** + * Updates internal latency metrics using the received records. + * + * @param events Collection of records used for latency tracking. + */ + @Override + public void updateLatencyMetrics(final Collection> events) { + // TODO: Implement latency tracking with PluginMetrics + } +} diff --git a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkTest.java b/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkTest.java new file mode 100644 index 0000000000..293a07e6e8 --- /dev/null +++ b/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkTest.java @@ -0,0 +1,48 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.sink.xrayotlp; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.model.trace.Span; + +import java.util.Collections; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; + +class XRayOTLPSinkTest { + private XRayOTLPSink sink; + + @BeforeEach + void setUp() { + sink = new XRayOTLPSink(); + } + + @Test + void testInitialize_doesNotThrow() { + assertDoesNotThrow(() -> sink.initialize()); + } + + @Test + void testOutput_printsRecordData() { + final Span mockSpan = mock(Span.class); + final Record record = new Record<>(mockSpan); + assertDoesNotThrow(() -> sink.output(Collections.singletonList(record))); + } + + @Test + void testIsReady_returnsTrue() { + assertTrue(sink.isReady()); + } + + @Test + void testShutdown_doesNotThrow() { + assertDoesNotThrow(() -> sink.shutdown()); + } +} diff --git a/data-prepper-plugins/xray-otlp-sink/src/test/resources/data-prepper-config.yaml b/data-prepper-plugins/xray-otlp-sink/src/test/resources/data-prepper-config.yaml new file mode 100644 index 0000000000..7462c0a70e --- /dev/null +++ b/data-prepper-plugins/xray-otlp-sink/src/test/resources/data-prepper-config.yaml @@ -0,0 +1 @@ +ssl: false diff --git a/data-prepper-plugins/xray-otlp-sink/src/test/resources/pipelines.yaml b/data-prepper-plugins/xray-otlp-sink/src/test/resources/pipelines.yaml new file mode 100644 index 0000000000..dc0015e3de --- /dev/null +++ b/data-prepper-plugins/xray-otlp-sink/src/test/resources/pipelines.yaml @@ -0,0 +1,13 @@ +otel_trace_source_to_xray: + source: + otel_trace_source: + ssl: false + + sink: + - xray_otlp_sink: + region: us-west-2 + + buffer: + bounded_blocking: + buffer_size: 10 + batch_size: 5 diff --git a/data-prepper-plugins/xray-otlp-sink/src/test/resources/sample-trace.json b/data-prepper-plugins/xray-otlp-sink/src/test/resources/sample-trace.json new file mode 100644 index 0000000000..e03778ed6b --- /dev/null +++ b/data-prepper-plugins/xray-otlp-sink/src/test/resources/sample-trace.json @@ -0,0 +1,28 @@ +{ + "resourceSpans": [ + { + "resource": { + "attributes": [ + { + "key": "service.name", + "value": { "stringValue": "test-service" } + } + ] + }, + "scopeSpans": [ + { + "spans": [ + { + "traceId": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "spanId": "bbbbbbbbbbbbbbbb", + "name": "test-span", + "kind": "SPAN_KIND_INTERNAL", + "startTimeUnixNano": "1697051838000000000", + "endTimeUnixNano": "1697051839000000000" + } + ] + } + ] + } + ] +} \ No newline at end of file diff --git a/settings.gradle b/settings.gradle index 744852c78e..5b36fa86d8 100644 --- a/settings.gradle +++ b/settings.gradle @@ -196,3 +196,4 @@ include 'data-prepper-plugins:saas-source-plugins:jira-source' include 'data-prepper-plugins:saas-source-plugins:confluence-source' include 'data-prepper-plugins:saas-source-plugins:atlassian-commons' include 'data-prepper-plugins:saas-source-plugins:crowdstrike-source' +include 'data-prepper-plugins:xray-otlp-sink' From 5978b82a2c127c1a59ea48707342356bd7b454e0 Mon Sep 17 00:00:00 2001 From: huyPham Date: Thu, 3 Apr 2025 13:24:03 -0700 Subject: [PATCH 02/23] Add a simple integration-test (#2) Signed-off-by: huy pham --- data-prepper-plugins/xray-otlp-sink/README.md | 6 +++ .../xray-otlp-sink/build.gradle | 33 +++++++++++++++- .../plugins/sink/xrayotlp/XRayOTLPSinkIT.java | 38 +++++++++++++++++++ 3 files changed, 76 insertions(+), 1 deletion(-) create mode 100644 data-prepper-plugins/xray-otlp-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkIT.java diff --git a/data-prepper-plugins/xray-otlp-sink/README.md b/data-prepper-plugins/xray-otlp-sink/README.md index 9bd72e77e9..ad0bb60c9f 100644 --- a/data-prepper-plugins/xray-otlp-sink/README.md +++ b/data-prepper-plugins/xray-otlp-sink/README.md @@ -20,6 +20,12 @@ The integration tests for this plugin do not run as part of the main Data Preppe ./gradlew :data-prepper-plugins:xray-otlp-sink:test ``` +#### Run integration tests locally + +``` +./gradlew :data-prepper-plugins:xray-otlp-sink:integrationTest +``` + #### Run a local pipeline that uses this sink 1. Install `grpcurl` – Used to send OTLP span data to the running pipeline. diff --git a/data-prepper-plugins/xray-otlp-sink/build.gradle b/data-prepper-plugins/xray-otlp-sink/build.gradle index b3611457ee..c76e8dd290 100644 --- a/data-prepper-plugins/xray-otlp-sink/build.gradle +++ b/data-prepper-plugins/xray-otlp-sink/build.gradle @@ -17,4 +17,35 @@ dependencies { test { useJUnitPlatform() -} \ No newline at end of file +} + +sourceSets { + integrationTest { + java.srcDir file('src/integrationTest/java') + resources.srcDir file('src/integrationTest/resources') + compileClasspath += sourceSets.main.output + runtimeClasspath += sourceSets.main.output + } +} + +configurations { + integrationTestImplementation.extendsFrom testImplementation + integrationTestRuntimeOnly.extendsFrom testRuntimeOnly +} + +dependencies { + integrationTestImplementation project(':data-prepper-test-common') + integrationTestImplementation 'org.junit.jupiter:junit-jupiter-api:5.9.2' + integrationTestRuntimeOnly 'org.junit.jupiter:junit-jupiter-engine:5.9.2' +} + +tasks.register('integrationTest', Test) { + description = 'Runs integration tests.' + group = 'verification' + testClassesDirs = sourceSets.integrationTest.output.classesDirs + classpath = sourceSets.integrationTest.runtimeClasspath + shouldRunAfter test + useJUnitPlatform() +} + +check.dependsOn integrationTest \ No newline at end of file diff --git a/data-prepper-plugins/xray-otlp-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkIT.java b/data-prepper-plugins/xray-otlp-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkIT.java new file mode 100644 index 0000000000..421635c4d9 --- /dev/null +++ b/data-prepper-plugins/xray-otlp-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkIT.java @@ -0,0 +1,38 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.sink.xrayotlp; + +import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.model.trace.JacksonStandardSpan; +import org.opensearch.dataprepper.model.trace.Span; + +import java.time.Instant; +import java.util.Collections; +import java.util.List; + +class XRayOTLPSinkIT { + @Test + void testSinkProcessesHardcodedSpan() { + final Span testSpan = JacksonStandardSpan.builder() + .withTraceId("abc123") + .withSpanId("def456") + .withParentSpanId("parent-testSpan-id") + .withName("my-test-testSpan") + .withStartTime(String.valueOf(Instant.now())) + .withEndTime(String.valueOf(Instant.now().plusMillis(10))) + .withAttributes(Collections.emptyMap()) + .withKind("test") + .build(); + + final Record record = new Record<>(testSpan); + final XRayOTLPSink sink = new XRayOTLPSink(); + + sink.initialize(); + sink.output(List.of(record)); + sink.shutdown(); + } +} From e57e57410d621ea0e70d854e3c4f895109b8859f Mon Sep 17 00:00:00 2001 From: Heli Date: Thu, 3 Apr 2025 15:34:31 -0700 Subject: [PATCH 03/23] Add X-Ray OTLP sink configuration and tests (#5) Signed-off-by: Heli --- data-prepper-plugins/xray-otlp-sink/README.md | 36 ++++++++++ .../xray-otlp-sink/build.gradle | 3 + .../sink/xrayotlp/XRayOTLPSinkConfig.java | 33 +++++++++ .../AwsAuthenticationConfiguration.java | 57 +++++++++++++++ .../AwsAuthenticationConfigurationTest.java | 43 ++++++++++++ .../sink/xrayotlp/XRayOTLPSinkConfigTest.java | 70 +++++++++++++++++++ 6 files changed, 242 insertions(+) create mode 100644 data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkConfig.java create mode 100644 data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/AwsAuthenticationConfiguration.java create mode 100644 data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/AwsAuthenticationConfigurationTest.java create mode 100644 data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkConfigTest.java diff --git a/data-prepper-plugins/xray-otlp-sink/README.md b/data-prepper-plugins/xray-otlp-sink/README.md index ad0bb60c9f..0fffacda9a 100644 --- a/data-prepper-plugins/xray-otlp-sink/README.md +++ b/data-prepper-plugins/xray-otlp-sink/README.md @@ -8,6 +8,42 @@ For information on usage, see the forthcoming documentation in the [Data Prepper A sample pipeline configuration will be added once the plugin is ready for testing. +### Configuration Options + +#### aws (Required) +Configuration options for AWS authentication and region settings. + +* `region` (Required): The AWS region where X-Ray service is located + * Must be a valid AWS region identifier (e.g., us-east-1, us-west-2) + * Cannot be empty + +* `sts_role_arn` (Required): AWS STS Role ARN for assuming role-based access + * Format: arn:aws:iam::{account}:role/{role-name} + * Length must be between 20 and 2048 characters + +* `sts_external_id` (Optional): External ID for additional security when assuming an IAM role + * Required only if the trust policy requires an external ID + * Length must be between 2 and 1224 characters + +### Sample Pipeline Configuration + +```yaml +pipeline: + source: + otel_trace_source: + ssl: true + + buffer: + bounded_blocking: + buffer_size: 10 + batch_size: 5 + + sink: + - xray-otlp-sink: + aws: + region: us-west-2 + sts_role_arn: arn:aws:iam::123456789012:role/XrayRole + ## Developer Guide See the [CONTRIBUTING](https://github.com/opensearch-project/data-prepper/blob/main/CONTRIBUTING.md) guide for general information on contributions. diff --git a/data-prepper-plugins/xray-otlp-sink/build.gradle b/data-prepper-plugins/xray-otlp-sink/build.gradle index c76e8dd290..2c6a38d4c3 100644 --- a/data-prepper-plugins/xray-otlp-sink/build.gradle +++ b/data-prepper-plugins/xray-otlp-sink/build.gradle @@ -8,8 +8,11 @@ plugins { } dependencies { + compileOnly 'org.projectlombok:lombok:1.18.30' + annotationProcessor 'org.projectlombok:lombok:1.18.30' implementation project(':data-prepper-api') implementation 'com.fasterxml.jackson.core:jackson-databind' + implementation 'software.amazon.awssdk:regions' testImplementation 'org.junit.jupiter:junit-jupiter-api:5.9.1' testRuntimeOnly 'org.junit.jupiter:junit-jupiter-engine:5.9.1' diff --git a/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkConfig.java b/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkConfig.java new file mode 100644 index 0000000000..3abcbb6427 --- /dev/null +++ b/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkConfig.java @@ -0,0 +1,33 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.xrayotlp; + +import com.fasterxml.jackson.annotation.JsonProperty; +import jakarta.validation.Valid; +import jakarta.validation.constraints.NotNull; +import lombok.Getter; +import org.opensearch.dataprepper.plugins.sink.xrayotlp.configuration.AwsAuthenticationConfiguration; + +/** + * Configuration class for the X-Ray OTLP sink plugin. + * This class defines the configuration options available when setting up + * the X-Ray OTLP sink in Data Prepper pipelines. + * + * @since 2.6 + */ +public class XRayOTLPSinkConfig { + public static final String DEFAULT_AWS_REGION = "us-east-1"; + + /** + * AWS configuration for X-Ray access. + * Contains authentication and region settings required for AWS X-Ray service. + * This is a required configuration and must be valid. + */ + @Getter + @JsonProperty("aws") + @NotNull + @Valid + AwsAuthenticationConfiguration awsAuthenticationConfiguration; +} diff --git a/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/AwsAuthenticationConfiguration.java b/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/AwsAuthenticationConfiguration.java new file mode 100644 index 0000000000..75ef5e92aa --- /dev/null +++ b/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/AwsAuthenticationConfiguration.java @@ -0,0 +1,57 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.xrayotlp.configuration; + +import com.fasterxml.jackson.annotation.JsonProperty; +import jakarta.validation.constraints.Size; +import lombok.Getter; +import software.amazon.awssdk.regions.Region; + +import static org.opensearch.dataprepper.plugins.sink.xrayotlp.XRayOTLPSinkConfig.DEFAULT_AWS_REGION; + +/** + * Configuration class for AWS authentication settings. + * Handles region, STS role ARN, and external ID configurations required for AWS service access. + * + * @since 2.6 + */ +public class AwsAuthenticationConfiguration { + /** + * AWS region for X-Ray service. + * Must be a valid AWS region identifier (e.g., us-east-1, us-west-2). + */ + @JsonProperty("region") + @Size(min = 1, message = "Region cannot be empty string") + public String awsRegion = DEFAULT_AWS_REGION; + + /** + * AWS STS Role ARN for assuming role-based access. + * Format: arn:aws:iam::{account}:role/{role-name} + * Length must be between 20 and 2048 characters. + */ + @Getter + @JsonProperty("sts_role_arn") + @Size(min = 20, max = 2048, message = "awsStsRoleArn length should be between 1 and 2048 characters") + public String awsStsRoleArn; + + /** + * External ID for additional security when assuming an IAM role. + * Required only if the trust policy requires an external ID. + * Length must be between 2 and 1224 characters. + */ + @Getter + @JsonProperty("sts_external_id") + @Size(min = 2, max = 1224, message = "awsStsExternalId length should be between 2 and 1224 characters") + public String awsStsExternalId; + + /** + * Gets the AWS Region object corresponding to the configured region string. + * + * @return Region object if awsRegion is set, null otherwise + */ + public Region getAwsRegion() { + return awsRegion != null ? Region.of(awsRegion) : null; + } +} diff --git a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/AwsAuthenticationConfigurationTest.java b/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/AwsAuthenticationConfigurationTest.java new file mode 100644 index 0000000000..4b99b71b15 --- /dev/null +++ b/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/AwsAuthenticationConfigurationTest.java @@ -0,0 +1,43 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.xrayotlp; + +import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.plugins.sink.xrayotlp.configuration.AwsAuthenticationConfiguration; +import software.amazon.awssdk.regions.Region; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.CoreMatchers.notNullValue; +import static org.hamcrest.MatcherAssert.assertThat; + +class AwsAuthenticationConfigurationTest { + + @Test + void testGetRegion() { + final AwsAuthenticationConfiguration config = new AwsAuthenticationConfiguration(); + config.awsRegion = "us-west-2"; + + assertThat(config.getAwsRegion(), notNullValue()); + assertThat(config.getAwsRegion(), equalTo(Region.US_WEST_2)); + } + + @Test + void testGetStsRoleArn() { + final AwsAuthenticationConfiguration config = new AwsAuthenticationConfiguration(); + final String roleArn = "arn:aws:iam::123456789012:role/MyRole"; + config.awsStsRoleArn = roleArn; + + assertThat(config.awsStsRoleArn, equalTo(roleArn)); + } + + @Test + void testGetStsExternalId() { + final AwsAuthenticationConfiguration config = new AwsAuthenticationConfiguration(); + final String externalId = "myExternalId"; + config.awsStsExternalId = externalId; + + assertThat(config.awsStsExternalId, equalTo(externalId)); + } +} diff --git a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkConfigTest.java b/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkConfigTest.java new file mode 100644 index 0000000000..cda6ffaa66 --- /dev/null +++ b/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkConfigTest.java @@ -0,0 +1,70 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.xrayotlp; + +import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.plugins.sink.xrayotlp.configuration.AwsAuthenticationConfiguration; +import software.amazon.awssdk.regions.Region; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.CoreMatchers.notNullValue; +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.CoreMatchers.nullValue; +import static org.hamcrest.MatcherAssert.assertThat; + +class XRayOTLPSinkConfigTest { + + @Test + void testGetAwsConfiguration() { + final XRayOTLPSinkConfig config = new XRayOTLPSinkConfig(); + final AwsAuthenticationConfiguration awsConfig = new AwsAuthenticationConfiguration(); + config.awsAuthenticationConfiguration = awsConfig; + + assertThat(config.awsAuthenticationConfiguration, notNullValue()); + assertThat(config.awsAuthenticationConfiguration, equalTo(awsConfig)); + } + + @Test + void testGetAwsConfiguration_whenNull() { + final XRayOTLPSinkConfig config = new XRayOTLPSinkConfig(); + assertThat(config.awsAuthenticationConfiguration, is(nullValue())); + } + + @Test + void testAwsConfiguration_withRegion() { + final XRayOTLPSinkConfig config = new XRayOTLPSinkConfig(); + final AwsAuthenticationConfiguration awsConfig = new AwsAuthenticationConfiguration(); + awsConfig.awsRegion = "us-west-2"; + config.awsAuthenticationConfiguration = awsConfig; + + assertThat(config.awsAuthenticationConfiguration.getAwsRegion(), equalTo(Region.US_WEST_2)); + } + + @Test + void testAwsConfiguration_withStsRole() { + final XRayOTLPSinkConfig config = new XRayOTLPSinkConfig(); + final AwsAuthenticationConfiguration awsConfig = new AwsAuthenticationConfiguration(); + awsConfig.awsStsRoleArn = "arn:aws:iam::123456789012:role/MyRole"; + config.awsAuthenticationConfiguration = awsConfig; + + assertThat(config.awsAuthenticationConfiguration.awsStsRoleArn, + equalTo("arn:aws:iam::123456789012:role/MyRole")); + } + + @Test + void testAwsConfiguration_withCompleteConfig() { + final XRayOTLPSinkConfig config = new XRayOTLPSinkConfig(); + final AwsAuthenticationConfiguration awsConfig = new AwsAuthenticationConfiguration(); + awsConfig.awsRegion = "us-west-2"; + awsConfig.awsStsRoleArn = "arn:aws:iam::123456789012:role/MyRole"; + awsConfig.awsStsExternalId = "MyExternalId"; + config.awsAuthenticationConfiguration = awsConfig; + + assertThat(config.awsAuthenticationConfiguration.getAwsRegion(), equalTo(Region.US_WEST_2)); + assertThat(config.awsAuthenticationConfiguration.awsStsRoleArn, + equalTo("arn:aws:iam::123456789012:role/MyRole")); + assertThat(config.awsAuthenticationConfiguration.awsStsExternalId, equalTo("MyExternalId")); + } +} From e992cf61b2839b9886b4038b0623e388cb3c2490 Mon Sep 17 00:00:00 2001 From: huyPham Date: Tue, 8 Apr 2025 21:45:40 -0700 Subject: [PATCH 04/23] feature/sigv4: Add SigV4 signer and encapsulate AWS config (#6) - Introduced SigV4Signer helper to sign requests for AWS OTLP endpoint. - Encapsulated AwsAuthenticationConfiguration and exposed only via XRayOTLPSinkConfig. - Added unit tests with 90%+ coverage. Signed-off-by: huy pham --- .../xray-otlp-sink/build.gradle | 66 +++++-- .../sink/xrayotlp/XRayOTLPSinkConfig.java | 33 ---- .../AwsAuthenticationConfiguration.java | 19 +- .../configuration/XRayOTLPSinkConfig.java | 44 +++++ .../sink/xrayotlp/http/SigV4Signer.java | 101 +++++++++++ .../AwsAuthenticationConfigurationTest.java | 43 ----- .../sink/xrayotlp/XRayOTLPSinkConfigTest.java | 70 -------- .../AwsAuthenticationConfigurationTest.java | 41 +++++ .../configuration/XRayOTLPSinkConfigTest.java | 34 ++++ .../sink/xrayotlp/http/SigV4SignerTest.java | 166 ++++++++++++++++++ 10 files changed, 450 insertions(+), 167 deletions(-) delete mode 100644 data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkConfig.java create mode 100644 data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/XRayOTLPSinkConfig.java create mode 100644 data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/SigV4Signer.java delete mode 100644 data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/AwsAuthenticationConfigurationTest.java delete mode 100644 data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkConfigTest.java create mode 100644 data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/AwsAuthenticationConfigurationTest.java create mode 100644 data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/XRayOTLPSinkConfigTest.java create mode 100644 data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/SigV4SignerTest.java diff --git a/data-prepper-plugins/xray-otlp-sink/build.gradle b/data-prepper-plugins/xray-otlp-sink/build.gradle index 2c6a38d4c3..e61e5d6eea 100644 --- a/data-prepper-plugins/xray-otlp-sink/build.gradle +++ b/data-prepper-plugins/xray-otlp-sink/build.gradle @@ -5,21 +5,74 @@ plugins { id 'java-library' + id 'jacoco' +} + +configurations { + integrationTestImplementation.extendsFrom testImplementation + integrationTestRuntimeOnly.extendsFrom testRuntimeOnly } dependencies { + implementation 'software.amazon.awssdk:core:2.28.23' + implementation 'software.amazon.awssdk:auth:2.28.23' + implementation 'software.amazon.awssdk:sts:2.28.23' + implementation 'software.amazon.awssdk:regions:2.28.23' + compileOnly 'org.projectlombok:lombok:1.18.30' annotationProcessor 'org.projectlombok:lombok:1.18.30' + implementation project(':data-prepper-api') implementation 'com.fasterxml.jackson.core:jackson-databind' - implementation 'software.amazon.awssdk:regions' testImplementation 'org.junit.jupiter:junit-jupiter-api:5.9.1' testRuntimeOnly 'org.junit.jupiter:junit-jupiter-engine:5.9.1' + testImplementation 'org.slf4j:slf4j-simple:2.0.7' + testImplementation 'org.mockito:mockito-core:5.10.0' + + integrationTestImplementation project(':data-prepper-test-common') + integrationTestImplementation 'org.junit.jupiter:junit-jupiter-api:5.9.2' + integrationTestRuntimeOnly 'org.junit.jupiter:junit-jupiter-engine:5.9.2' +} + +jacoco { + toolVersion = "0.8.11" +} + +jacocoTestReport { + dependsOn test + reports { + xml.required = true + html.required = true + } +} + +jacocoTestCoverageVerification { + violationRules { + rule { + enabled = true + element = 'CLASS' + limit { + counter = 'LINE' + value = 'COVEREDRATIO' + minimum = 0.9 + } + + includes = ['org.opensearch.dataprepper.plugins.sink.xrayotlp.*'] + } + } +} + +tasks.named('jacocoTestReport') { + doLast { + def reportPath = layout.buildDirectory.file("reports/jacoco/test/html/index.html").get().asFile.toURI() + println "\nView test coverage report here:\n ${reportPath}\n" + } } test { useJUnitPlatform() + finalizedBy jacocoTestReport, jacocoTestCoverageVerification } sourceSets { @@ -31,17 +84,6 @@ sourceSets { } } -configurations { - integrationTestImplementation.extendsFrom testImplementation - integrationTestRuntimeOnly.extendsFrom testRuntimeOnly -} - -dependencies { - integrationTestImplementation project(':data-prepper-test-common') - integrationTestImplementation 'org.junit.jupiter:junit-jupiter-api:5.9.2' - integrationTestRuntimeOnly 'org.junit.jupiter:junit-jupiter-engine:5.9.2' -} - tasks.register('integrationTest', Test) { description = 'Runs integration tests.' group = 'verification' diff --git a/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkConfig.java b/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkConfig.java deleted file mode 100644 index 3abcbb6427..0000000000 --- a/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkConfig.java +++ /dev/null @@ -1,33 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ -package org.opensearch.dataprepper.plugins.sink.xrayotlp; - -import com.fasterxml.jackson.annotation.JsonProperty; -import jakarta.validation.Valid; -import jakarta.validation.constraints.NotNull; -import lombok.Getter; -import org.opensearch.dataprepper.plugins.sink.xrayotlp.configuration.AwsAuthenticationConfiguration; - -/** - * Configuration class for the X-Ray OTLP sink plugin. - * This class defines the configuration options available when setting up - * the X-Ray OTLP sink in Data Prepper pipelines. - * - * @since 2.6 - */ -public class XRayOTLPSinkConfig { - public static final String DEFAULT_AWS_REGION = "us-east-1"; - - /** - * AWS configuration for X-Ray access. - * Contains authentication and region settings required for AWS X-Ray service. - * This is a required configuration and must be valid. - */ - @Getter - @JsonProperty("aws") - @NotNull - @Valid - AwsAuthenticationConfiguration awsAuthenticationConfiguration; -} diff --git a/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/AwsAuthenticationConfiguration.java b/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/AwsAuthenticationConfiguration.java index 75ef5e92aa..347136b81c 100644 --- a/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/AwsAuthenticationConfiguration.java +++ b/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/AwsAuthenticationConfiguration.java @@ -6,50 +6,51 @@ import com.fasterxml.jackson.annotation.JsonProperty; import jakarta.validation.constraints.Size; +import lombok.Builder; import lombok.Getter; import software.amazon.awssdk.regions.Region; -import static org.opensearch.dataprepper.plugins.sink.xrayotlp.XRayOTLPSinkConfig.DEFAULT_AWS_REGION; - /** * Configuration class for AWS authentication settings. * Handles region, STS role ARN, and external ID configurations required for AWS service access. + * This class will be automatically wired by Data-Prepper; the Builder is only for testing. * * @since 2.6 */ -public class AwsAuthenticationConfiguration { +@Getter +@Builder +class AwsAuthenticationConfiguration { /** * AWS region for X-Ray service. * Must be a valid AWS region identifier (e.g., us-east-1, us-west-2). */ @JsonProperty("region") @Size(min = 1, message = "Region cannot be empty string") - public String awsRegion = DEFAULT_AWS_REGION; + private String awsRegion; /** * AWS STS Role ARN for assuming role-based access. * Format: arn:aws:iam::{account}:role/{role-name} * Length must be between 20 and 2048 characters. */ - @Getter @JsonProperty("sts_role_arn") @Size(min = 20, max = 2048, message = "awsStsRoleArn length should be between 1 and 2048 characters") - public String awsStsRoleArn; + private String awsStsRoleArn; /** * External ID for additional security when assuming an IAM role. * Required only if the trust policy requires an external ID. * Length must be between 2 and 1224 characters. */ - @Getter @JsonProperty("sts_external_id") @Size(min = 2, max = 1224, message = "awsStsExternalId length should be between 2 and 1224 characters") - public String awsStsExternalId; + private String awsStsExternalId; /** * Gets the AWS Region object corresponding to the configured region string. * - * @return Region object if awsRegion is set, null otherwise + * @return Region object if awsRegion is set, otherwise returns null. + * Note: Default region fallback is handled externally by the caller. */ public Region getAwsRegion() { return awsRegion != null ? Region.of(awsRegion) : null; diff --git a/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/XRayOTLPSinkConfig.java b/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/XRayOTLPSinkConfig.java new file mode 100644 index 0000000000..d5e1376b62 --- /dev/null +++ b/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/XRayOTLPSinkConfig.java @@ -0,0 +1,44 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.xrayotlp.configuration; + +import com.fasterxml.jackson.annotation.JsonProperty; +import jakarta.validation.Valid; +import jakarta.validation.constraints.NotNull; +import lombok.Builder; +import software.amazon.awssdk.regions.Region; + +/** + * Configuration class for the X-Ray OTLP sink plugin. + * This class defines the configuration options available when setting up + * the X-Ray OTLP sink in Data Prepper pipelines. + * This class will be automatically wired by Data-Prepper; the Builder is only for testing. + * + * @since 2.6 + */ +@Builder +public class XRayOTLPSinkConfig { + /** + * AWS configuration for X-Ray access. + * Contains authentication and region settings required for AWS X-Ray service. + * This is a required configuration and must be valid. + */ + @JsonProperty("aws") + @NotNull + @Valid + private AwsAuthenticationConfiguration awsAuthenticationConfiguration; + + public Region getAwsRegion() { + return awsAuthenticationConfiguration.getAwsRegion(); + } + + public String getStsRoleArn() { + return awsAuthenticationConfiguration.getAwsStsRoleArn(); + } + + public String getStsExternalId() { + return awsAuthenticationConfiguration.getAwsStsExternalId(); + } +} diff --git a/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/SigV4Signer.java b/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/SigV4Signer.java new file mode 100644 index 0000000000..460148558e --- /dev/null +++ b/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/SigV4Signer.java @@ -0,0 +1,101 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.sink.xrayotlp.http; + +import com.google.common.annotations.VisibleForTesting; +import org.opensearch.dataprepper.plugins.sink.xrayotlp.configuration.XRayOTLPSinkConfig; +import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider; +import software.amazon.awssdk.auth.signer.Aws4Signer; +import software.amazon.awssdk.auth.signer.params.Aws4SignerParams; +import software.amazon.awssdk.core.SdkBytes; +import software.amazon.awssdk.http.SdkHttpFullRequest; +import software.amazon.awssdk.http.SdkHttpMethod; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.regions.providers.DefaultAwsRegionProviderChain; +import software.amazon.awssdk.services.sts.StsClient; +import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider; +import java.net.URI; + +/** + * Helper class to apply AWS SigV4 signing to outgoing HTTP requests + * before sending them to the AWS X-Ray OTLP endpoint. + */ +public class SigV4Signer { + private static final String SERVICE_NAME = "xray"; + private static final String OTLP_PATH = "/v1/traces"; + private final Aws4Signer signer = Aws4Signer.create(); + + private final AwsCredentialsProvider credentialsProvider; + private final Region region; + private final URI endpointUri; + + /** + * Constructs a SigV4 signer helper based on the AWS authentication configuration. + * Supports optional STS role assumption. + * + * @param config Configuration for region and optional STS role + */ + public SigV4Signer(final XRayOTLPSinkConfig config) { + this(config, null); + } + + /** + * Package-private constructor for unit testing with mocked STS client. + */ + @VisibleForTesting + SigV4Signer(final XRayOTLPSinkConfig config, final StsClient stsClient) { + this.region = resolveRegion(config.getAwsRegion()); + this.credentialsProvider = initCredentialsProvider(config.getAwsRegion(), config.getStsRoleArn(), config.getStsExternalId(), stsClient); + this.endpointUri = URI.create(String.format("https://xray.%s.amazonaws.com%s", region.id(), OTLP_PATH)); + } + + private Region resolveRegion(final Region region) { + return region != null ? region : DefaultAwsRegionProviderChain.builder().build().getRegion(); + } + + private AwsCredentialsProvider initCredentialsProvider( + final Region region, + final String stsRoleArn, + final String stsExternalId, + final StsClient stsClient + ) { + if (region != null && stsRoleArn != null) { + return StsAssumeRoleCredentialsProvider.builder() + .refreshRequest(r -> { + r.roleArn(stsRoleArn); + if (stsExternalId != null) { + r.externalId(stsExternalId); + } + }) + .stsClient(stsClient != null ? stsClient : StsClient.builder().region(region).build()) + .build(); + } + + return DefaultCredentialsProvider.create(); + } + + /** + * Signs a request payload using AWS SigV4 and returns a fully signed request. + * + * @param payload The OTLP Protobuf-encoded request body to be sent + * @return A signed {@link SdkHttpFullRequest} ready for transmission to the AWS OTLP endpoint + */ + public SdkHttpFullRequest signRequest(final byte[] payload) { + final SdkHttpFullRequest unsignedRequest = SdkHttpFullRequest.builder() + .method(SdkHttpMethod.POST) + .uri(endpointUri) + .putHeader("Content-Type", "application/x-protobuf") + .contentStreamProvider(() -> SdkBytes.fromByteArray(payload).asInputStream()) + .build(); + + return signer.sign(unsignedRequest, Aws4SignerParams.builder() + .signingRegion(region) + .signingName(SERVICE_NAME) + .awsCredentials(credentialsProvider.resolveCredentials()) + .build()); + } +} diff --git a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/AwsAuthenticationConfigurationTest.java b/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/AwsAuthenticationConfigurationTest.java deleted file mode 100644 index 4b99b71b15..0000000000 --- a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/AwsAuthenticationConfigurationTest.java +++ /dev/null @@ -1,43 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ -package org.opensearch.dataprepper.plugins.sink.xrayotlp; - -import org.junit.jupiter.api.Test; -import org.opensearch.dataprepper.plugins.sink.xrayotlp.configuration.AwsAuthenticationConfiguration; -import software.amazon.awssdk.regions.Region; - -import static org.hamcrest.CoreMatchers.equalTo; -import static org.hamcrest.CoreMatchers.notNullValue; -import static org.hamcrest.MatcherAssert.assertThat; - -class AwsAuthenticationConfigurationTest { - - @Test - void testGetRegion() { - final AwsAuthenticationConfiguration config = new AwsAuthenticationConfiguration(); - config.awsRegion = "us-west-2"; - - assertThat(config.getAwsRegion(), notNullValue()); - assertThat(config.getAwsRegion(), equalTo(Region.US_WEST_2)); - } - - @Test - void testGetStsRoleArn() { - final AwsAuthenticationConfiguration config = new AwsAuthenticationConfiguration(); - final String roleArn = "arn:aws:iam::123456789012:role/MyRole"; - config.awsStsRoleArn = roleArn; - - assertThat(config.awsStsRoleArn, equalTo(roleArn)); - } - - @Test - void testGetStsExternalId() { - final AwsAuthenticationConfiguration config = new AwsAuthenticationConfiguration(); - final String externalId = "myExternalId"; - config.awsStsExternalId = externalId; - - assertThat(config.awsStsExternalId, equalTo(externalId)); - } -} diff --git a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkConfigTest.java b/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkConfigTest.java deleted file mode 100644 index cda6ffaa66..0000000000 --- a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkConfigTest.java +++ /dev/null @@ -1,70 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ -package org.opensearch.dataprepper.plugins.sink.xrayotlp; - -import org.junit.jupiter.api.Test; -import org.opensearch.dataprepper.plugins.sink.xrayotlp.configuration.AwsAuthenticationConfiguration; -import software.amazon.awssdk.regions.Region; - -import static org.hamcrest.CoreMatchers.equalTo; -import static org.hamcrest.CoreMatchers.notNullValue; -import static org.hamcrest.CoreMatchers.is; -import static org.hamcrest.CoreMatchers.nullValue; -import static org.hamcrest.MatcherAssert.assertThat; - -class XRayOTLPSinkConfigTest { - - @Test - void testGetAwsConfiguration() { - final XRayOTLPSinkConfig config = new XRayOTLPSinkConfig(); - final AwsAuthenticationConfiguration awsConfig = new AwsAuthenticationConfiguration(); - config.awsAuthenticationConfiguration = awsConfig; - - assertThat(config.awsAuthenticationConfiguration, notNullValue()); - assertThat(config.awsAuthenticationConfiguration, equalTo(awsConfig)); - } - - @Test - void testGetAwsConfiguration_whenNull() { - final XRayOTLPSinkConfig config = new XRayOTLPSinkConfig(); - assertThat(config.awsAuthenticationConfiguration, is(nullValue())); - } - - @Test - void testAwsConfiguration_withRegion() { - final XRayOTLPSinkConfig config = new XRayOTLPSinkConfig(); - final AwsAuthenticationConfiguration awsConfig = new AwsAuthenticationConfiguration(); - awsConfig.awsRegion = "us-west-2"; - config.awsAuthenticationConfiguration = awsConfig; - - assertThat(config.awsAuthenticationConfiguration.getAwsRegion(), equalTo(Region.US_WEST_2)); - } - - @Test - void testAwsConfiguration_withStsRole() { - final XRayOTLPSinkConfig config = new XRayOTLPSinkConfig(); - final AwsAuthenticationConfiguration awsConfig = new AwsAuthenticationConfiguration(); - awsConfig.awsStsRoleArn = "arn:aws:iam::123456789012:role/MyRole"; - config.awsAuthenticationConfiguration = awsConfig; - - assertThat(config.awsAuthenticationConfiguration.awsStsRoleArn, - equalTo("arn:aws:iam::123456789012:role/MyRole")); - } - - @Test - void testAwsConfiguration_withCompleteConfig() { - final XRayOTLPSinkConfig config = new XRayOTLPSinkConfig(); - final AwsAuthenticationConfiguration awsConfig = new AwsAuthenticationConfiguration(); - awsConfig.awsRegion = "us-west-2"; - awsConfig.awsStsRoleArn = "arn:aws:iam::123456789012:role/MyRole"; - awsConfig.awsStsExternalId = "MyExternalId"; - config.awsAuthenticationConfiguration = awsConfig; - - assertThat(config.awsAuthenticationConfiguration.getAwsRegion(), equalTo(Region.US_WEST_2)); - assertThat(config.awsAuthenticationConfiguration.awsStsRoleArn, - equalTo("arn:aws:iam::123456789012:role/MyRole")); - assertThat(config.awsAuthenticationConfiguration.awsStsExternalId, equalTo("MyExternalId")); - } -} diff --git a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/AwsAuthenticationConfigurationTest.java b/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/AwsAuthenticationConfigurationTest.java new file mode 100644 index 0000000000..2f46d4de03 --- /dev/null +++ b/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/AwsAuthenticationConfigurationTest.java @@ -0,0 +1,41 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.xrayotlp.configuration; + +import org.junit.jupiter.api.Test; +import software.amazon.awssdk.regions.Region; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.CoreMatchers.notNullValue; +import static org.hamcrest.CoreMatchers.nullValue; +import static org.hamcrest.MatcherAssert.assertThat; + +class AwsAuthenticationConfigurationTest { + + @Test + void testGetRegion_whenRegionIsNull_returnNull() { + final AwsAuthenticationConfiguration config = AwsAuthenticationConfiguration.builder().build(); + + assertThat(config.getAwsRegion(), nullValue()); + } + + @Test + void testAwsAuthenticationConfigurationFields() { + final String expectedRegion = "us-west-2"; + final String expectedRoleArn = "arn:aws:iam::123456789012:role/MyRole"; + final String expectedExternalId = "myExternalId"; + + final AwsAuthenticationConfiguration config = AwsAuthenticationConfiguration.builder() + .awsRegion(expectedRegion) + .awsStsRoleArn(expectedRoleArn) + .awsStsExternalId(expectedExternalId) + .build(); + + assertThat(config.getAwsRegion(), notNullValue()); + assertThat(config.getAwsRegion(), equalTo(Region.US_WEST_2)); + assertThat(config.getAwsStsRoleArn(), equalTo(expectedRoleArn)); + assertThat(config.getAwsStsExternalId(), equalTo(expectedExternalId)); + } +} diff --git a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/XRayOTLPSinkConfigTest.java b/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/XRayOTLPSinkConfigTest.java new file mode 100644 index 0000000000..d9d0fe0cd3 --- /dev/null +++ b/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/XRayOTLPSinkConfigTest.java @@ -0,0 +1,34 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.xrayotlp.configuration; + +import org.junit.jupiter.api.Test; +import software.amazon.awssdk.regions.Region; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; + +class XRayOTLPSinkConfigTest { + @Test + void testAwsAuthenticationConfiguration_withAllFields() { + final String expectedRegion = "us-west-2"; + final String expectedRoleArn = "arn:aws:iam::123456789012:role/MyRole"; + final String expectedExternalId = "external-id-123"; + + AwsAuthenticationConfiguration awsConfig = AwsAuthenticationConfiguration.builder() + .awsRegion(expectedRegion) + .awsStsRoleArn(expectedRoleArn) + .awsStsExternalId(expectedExternalId) + .build(); + + XRayOTLPSinkConfig config = XRayOTLPSinkConfig.builder() + .awsAuthenticationConfiguration(awsConfig) + .build(); + + assertThat(config.getAwsRegion(), equalTo(Region.of(expectedRegion))); + assertThat(config.getStsRoleArn(), equalTo(expectedRoleArn)); + assertThat(config.getStsExternalId(), equalTo(expectedExternalId)); + } +} diff --git a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/SigV4SignerTest.java b/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/SigV4SignerTest.java new file mode 100644 index 0000000000..62285ed6f0 --- /dev/null +++ b/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/SigV4SignerTest.java @@ -0,0 +1,166 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.xrayotlp.http; + +import java.nio.charset.StandardCharsets; +import java.time.Instant; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentMatcher; +import org.mockito.MockedStatic; +import org.opensearch.dataprepper.plugins.sink.xrayotlp.configuration.XRayOTLPSinkConfig; +import software.amazon.awssdk.http.SdkHttpFullRequest; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.regions.providers.DefaultAwsRegionProviderChain; +import software.amazon.awssdk.services.sts.StsClient; +import software.amazon.awssdk.services.sts.StsClientBuilder; +import software.amazon.awssdk.services.sts.model.AssumeRoleRequest; +import software.amazon.awssdk.services.sts.model.AssumeRoleResponse; +import software.amazon.awssdk.services.sts.model.Credentials; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.argThat; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +class SigV4SignerTest { + + private static final byte[] PAYLOAD = "test-payload".getBytes(StandardCharsets.UTF_8); + private XRayOTLPSinkConfig mockXrayConfig; + private SigV4Signer target; + + @BeforeEach + void setUp(){ + mockXrayConfig = mock(XRayOTLPSinkConfig.class); + } + + @AfterEach + void cleanUp() { + System.clearProperty("aws.region"); + } + + @Test + void testSignRequest_withFallbackRegion_whenRegionNotSet() { + // setup + System.setProperty("aws.region", Region.US_WEST_2.toString()); + when(mockXrayConfig.getAwsRegion()).thenReturn(null); + target = new SigV4Signer(mockXrayConfig); + + // run + final SdkHttpFullRequest signedRequest = target.signRequest(PAYLOAD); + + // assert + assertNotNull(signedRequest); + assertEquals("POST", signedRequest.method().name()); + assertTrue(signedRequest.headers().containsKey("Authorization")); + assertEquals("application/x-protobuf", signedRequest.firstMatchingHeader("Content-Type").orElse(null)); + + final String expectedRegion = DefaultAwsRegionProviderChain.builder().build().getRegion().id(); + assertTrue(signedRequest.getUri().toString().contains(String.format("https://xray.%s.amazonaws.com/v1/traces", expectedRegion))); + } + + @Test + void testSignRequest_withFallbackStsRole_whenStsRoleNotSet() { + // setup + when(mockXrayConfig.getAwsRegion()).thenReturn(Region.US_WEST_2); + target = new SigV4Signer(mockXrayConfig); + + // run + final SdkHttpFullRequest signedRequest = target.signRequest(PAYLOAD); + + // assert + assertNotNull(signedRequest); + assertEquals("POST", signedRequest.method().name()); + assertTrue(signedRequest.headers().containsKey("Authorization")); + assertEquals("application/x-protobuf", signedRequest.firstMatchingHeader("Content-Type").orElse(null)); + assertTrue(signedRequest.getUri().toString().contains(String.format("https://xray.%s.amazonaws.com/v1/traces", Region.US_WEST_2.toString()))); + } + + @Test + void testSignRequest_withCustomCredentials_usingDefaultStsClientFallback() { + when(mockXrayConfig.getAwsRegion()).thenReturn(Region.US_WEST_2); + when(mockXrayConfig.getStsRoleArn()).thenReturn("arn:aws:iam::123456789012:role/test-role"); + + // Use mocked static builder to simulate StsClient.builder() + try (MockedStatic mockedStsClientStatic = mockStatic(StsClient.class)) { + final StsClient mockStsClient = mock(StsClient.class); + + // Setup fake STS response + when(mockStsClient.assumeRole(any(AssumeRoleRequest.class))) + .thenReturn(AssumeRoleResponse.builder() + .credentials(Credentials.builder() + .accessKeyId("fake-access") + .secretAccessKey("fake-secret") + .sessionToken("fake-token") + .expiration(Instant.now().plusSeconds(3600)) + .build()) + .build()); + + // Mock StsClient builder + final StsClientBuilder mockBuilder = mock(StsClientBuilder.class); + when(mockBuilder.region(Region.US_WEST_2)).thenReturn(mockBuilder); + when(mockBuilder.build()).thenReturn(mockStsClient); + + // Mock the static method StsClient.builder() + mockedStsClientStatic.when(StsClient::builder).thenReturn(mockBuilder); + + // run + target = new SigV4Signer(mockXrayConfig, null); + SdkHttpFullRequest signedRequest = target.signRequest(PAYLOAD); + + // assert + assertNotNull(signedRequest); + assertTrue(signedRequest.headers().containsKey("Authorization")); + assertTrue(signedRequest.getUri().toString().contains(String.format("https://xray.%s.amazonaws.com/v1/traces", Region.US_WEST_2.toString()))); + } + } + + @Test + void testSignRequest_withCustomCredentials_usingMockedSts() { + // setup + final String expectedRoleArn = "arn:aws:iam::123456789012:role/test-role"; + final String expectedExternalId = "external-id-123"; + when(mockXrayConfig.getAwsRegion()).thenReturn(Region.US_WEST_2); + when(mockXrayConfig.getStsRoleArn()).thenReturn(expectedRoleArn); + when(mockXrayConfig.getStsExternalId()).thenReturn(expectedExternalId); + + final StsClient mockStsClient = mock(StsClient.class); + when(mockStsClient.assumeRole(any(AssumeRoleRequest.class))).thenReturn( + AssumeRoleResponse.builder() + .credentials(Credentials.builder() + .accessKeyId("fake-access-key") + .secretAccessKey("fake-secret-key") + .sessionToken("fake-session-token") + .expiration(Instant.now().plusSeconds(3600)) + .build()) + .build() + ); + + target = new SigV4Signer(mockXrayConfig, mockStsClient); + + // run + final SdkHttpFullRequest signedRequest = target.signRequest(PAYLOAD); + + // assert + assertNotNull(signedRequest); + assertEquals("POST", signedRequest.method().name()); + assertTrue(signedRequest.headers().containsKey("Authorization")); + assertEquals("application/x-protobuf", signedRequest.firstMatchingHeader("Content-Type").orElse(null)); + assertTrue(signedRequest.getUri().toString().contains("https://xray.us-west-2.amazonaws.com/v1/traces")); + + ArgumentMatcher matcher = request -> + expectedRoleArn.equals(request.roleArn()) && + expectedExternalId.equals(request.externalId()); + + verify(mockStsClient).assumeRole(argThat(matcher)); + } +} From 64df871cc5f31310910c317fbc80b2f37d6fb0b6 Mon Sep 17 00:00:00 2001 From: Heli Date: Fri, 11 Apr 2025 11:53:52 -0700 Subject: [PATCH 05/23] Add otlp proto trace output codec (#8) Signed-off-by: Heli --- .../otel-proto-common/build.gradle | 2 + .../otel/codec/OtlpTraceOutputCodec.java | 96 +++++++++++++ .../otel/codec/OtlpTraceOutputCodecTest.java | 129 ++++++++++++++++++ 3 files changed, 227 insertions(+) create mode 100644 data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OtlpTraceOutputCodec.java create mode 100644 data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OtlpTraceOutputCodecTest.java diff --git a/data-prepper-plugins/otel-proto-common/build.gradle b/data-prepper-plugins/otel-proto-common/build.gradle index 40636b049c..ab5ea220e1 100644 --- a/data-prepper-plugins/otel-proto-common/build.gradle +++ b/data-prepper-plugins/otel-proto-common/build.gradle @@ -12,6 +12,8 @@ test { } dependencies { + compileOnly 'org.projectlombok:lombok:1.18.30' + annotationProcessor 'org.projectlombok:lombok:1.18.30' implementation project(':data-prepper-api') implementation libs.opentelemetry.proto implementation libs.protobuf.util diff --git a/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OtlpTraceOutputCodec.java b/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OtlpTraceOutputCodec.java new file mode 100644 index 0000000000..901b6cc7c4 --- /dev/null +++ b/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OtlpTraceOutputCodec.java @@ -0,0 +1,96 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.otel.codec; + +import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceRequest; +import io.opentelemetry.proto.trace.v1.ResourceSpans; +import lombok.NonNull; +import org.apache.commons.codec.DecoderException; +import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; +import org.opensearch.dataprepper.model.codec.OutputCodec; +import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.sink.OutputCodecContext; +import org.opensearch.dataprepper.model.trace.Span; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.OutputStream; + +/** + * An implementation of {@link OutputCodec} that converts {@link Span} events + * into OpenTelemetry Protocol (OTLP) binary format using protobuf. + *

+ * This codec is primarily intended for use with trace data that will be + * forwarded to systems such as AWS X-Ray via OTLP. + */ +@DataPrepperPlugin(name = "otlp_trace", pluginType = OutputCodec.class) +public class OtlpTraceOutputCodec implements OutputCodec { + + private static final Logger LOG = LoggerFactory.getLogger(OtlpTraceOutputCodec.class); + private static final String OTLP_EXTENSION = "pb"; + + private final OTelProtoStandardCodec.OTelProtoEncoder encoder = new OTelProtoStandardCodec.OTelProtoEncoder(); + + /** + * Initializes the output stream. No-op for OTLP format. + * + * @param outputStream The output stream to write to. + * @param event The event (not used). + * @param context The codec context (not used). + */ + @Override + public void start(final OutputStream outputStream, final Event event, final OutputCodecContext context) { + // No-op for OTLP format + } + + /** + * Writes a single {@link Span} event to the output stream in OTLP binary format. + * + * @param event The event to encode. Must be of type {@link Span}. + * @param outputStream The stream to which the encoded bytes will be written. + */ + @Override + public void writeEvent(@NonNull final Event event, @NonNull final OutputStream outputStream) { + if (!(event instanceof Span)) { + throw new IllegalArgumentException("OtlpTraceOutputCodec only supports Span events"); + } + + final Span span = (Span) event; + try { + final ResourceSpans resourceSpans = encoder.convertToResourceSpans(span); + final ExportTraceServiceRequest request = ExportTraceServiceRequest.newBuilder() + .addResourceSpans(resourceSpans) + .build(); + + outputStream.write(request.toByteArray()); + } catch (final DecoderException e) { + LOG.warn("Skipping invalid span with ID [{}] due to decoding error.", span.getSpanId(), e); + } catch (final Exception e) { + LOG.error("Unexpected error while writing span with ID [{}] to OTLP output.", span.getSpanId(), e); + } + } + + /** + * Finalizes the output stream. No-op for OTLP format. + * + * @param outputStream The output stream to finalize. + */ + @Override + public void complete(final OutputStream outputStream) { + // No-op for OTLP format + } + + /** + * Returns the file extension used by this codec. + * In this case, "pb" for protobuf binary format. + * + * @return The string "pb". + */ + @Override + public String getExtension() { + return OTLP_EXTENSION; + } +} diff --git a/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OtlpTraceOutputCodecTest.java b/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OtlpTraceOutputCodecTest.java new file mode 100644 index 0000000000..9b881ebca3 --- /dev/null +++ b/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OtlpTraceOutputCodecTest.java @@ -0,0 +1,129 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.otel.codec; + +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.ObjectMapper; +import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceRequest; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.model.event.JacksonEvent; +import org.opensearch.dataprepper.model.sink.OutputCodecContext; +import org.opensearch.dataprepper.model.trace.DefaultTraceGroupFields; +import org.opensearch.dataprepper.model.trace.JacksonSpan; +import org.opensearch.dataprepper.model.trace.Span; + +import java.io.ByteArrayOutputStream; +import java.io.InputStream; +import java.util.Map; +import java.util.Objects; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +class OtlpTraceOutputCodecTest { + + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + private static final String TEST_SPAN_EVENT_JSON_FILE = "test-span-event.json"; + + private OtlpTraceOutputCodec codec; + + @BeforeEach + void setup() { + codec = new OtlpTraceOutputCodec(); + } + + private Span buildSpanFromTestFile(String fileName, String traceIdOverride) { + try (InputStream inputStream = Objects.requireNonNull( + getClass().getClassLoader().getResourceAsStream(fileName))) { + + final Map spanMap = OBJECT_MAPPER.readValue(inputStream, new TypeReference<>() {}); + final JacksonSpan.Builder builder = JacksonSpan.builder() + .withTraceId(traceIdOverride != null ? traceIdOverride : (String) spanMap.get("traceId")) + .withSpanId((String) spanMap.get("spanId")) + .withParentSpanId((String) spanMap.get("parentSpanId")) + .withTraceState((String) spanMap.get("traceState")) + .withName((String) spanMap.get("name")) + .withKind((String) spanMap.get("kind")) + .withDurationInNanos(((Number) spanMap.get("durationInNanos")).longValue()) + .withStartTime((String) spanMap.get("startTime")) + .withEndTime((String) spanMap.get("endTime")) + .withTraceGroup((String) spanMap.get("traceGroup")); + + final Map traceGroupFieldsMap = (Map) spanMap.get("traceGroupFields"); + if (traceGroupFieldsMap != null) { + builder.withTraceGroupFields(DefaultTraceGroupFields.builder() + .withStatusCode((Integer) traceGroupFieldsMap.getOrDefault("statusCode", 0)) + .withEndTime((String) spanMap.get("endTime")) + .withDurationInNanos(((Number) spanMap.get("durationInNanos")).longValue()) + .build()); + } + + return builder.build(); + } catch (Exception e) { + throw new RuntimeException("Failed to load span from file", e); + } + } + + @Test + void testWriteEvent_withValidSpanFromTestFile_writesSuccessfully() throws Exception { + final Span span = buildSpanFromTestFile(TEST_SPAN_EVENT_JSON_FILE, null); + final ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); + + codec.start(outputStream, span, new OutputCodecContext()); + codec.writeEvent(span, outputStream); + codec.complete(outputStream); + + final byte[] bytes = outputStream.toByteArray(); + assertThat(bytes).isNotEmpty(); + + final ExportTraceServiceRequest request = ExportTraceServiceRequest.parseFrom(bytes); + assertThat(request.getResourceSpansCount()).isGreaterThan(0); + } + + @Test + void testWriteEvent_withBadTraceId_logsAndSkips() throws Exception { + final Span span = buildSpanFromTestFile(TEST_SPAN_EVENT_JSON_FILE,"bad-trace-id" ); + + final ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); + + codec.writeEvent(span, outputStream); + + // Since traceId is invalid hex, it will trigger DecoderException + assertThat(outputStream.toByteArray()).isEmpty(); // nothing written + } + + @Test + void testWriteEvent_withNonSpanEvent_throwsException() { + final JacksonEvent nonSpanEvent = JacksonEvent.builder() + .withEventType("fake") + .withData(Map.of("key", "value")) + .build(); + final ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); + + assertThatThrownBy(() -> codec.writeEvent(nonSpanEvent, outputStream)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("OtlpTraceOutputCodec only supports Span events"); + } + + @Test + void testWriteEvent_withNullEvent_throwsNullPointerException() { + final ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); + + assertThatThrownBy(() -> codec.writeEvent(null, outputStream)) + .isInstanceOf(NullPointerException.class) + .hasMessage("event is marked non-null but is null"); + } + + @Test + void testWriteEvent_withNullOutputStream_throwsNullPointerException() { + final Span span = buildSpanFromTestFile(TEST_SPAN_EVENT_JSON_FILE, null); + + assertThatThrownBy(() -> codec.writeEvent(span, null)) + .isInstanceOf(NullPointerException.class) + .hasMessage("outputStream is marked non-null but is null"); + } +} From 4cf85ee01d1c2c435c64a0de62af11c9eee27e3c Mon Sep 17 00:00:00 2001 From: Heli Date: Fri, 11 Apr 2025 19:31:56 -0700 Subject: [PATCH 06/23] Add XRayOTLPSink and XRayOtlpHttpSender for sending OTLP trace data to AWS X-Ray (#9) Signed-off-by: Heli --- .../xray-otlp-sink/build.gradle | 3 + .../plugins/sink/xrayotlp/XRayOTLPSink.java | 63 ++++++--- .../configuration/XRayOTLPSinkConfig.java | 6 + .../xrayotlp/http/XRayOtlpHttpSender.java | 73 +++++++++++ .../sink/xrayotlp/XRayOTLPSinkTest.java | 124 ++++++++++++++++-- .../configuration/XRayOTLPSinkConfigTest.java | 16 +++ .../xrayotlp/http/XRayOtlpHttpSenderTest.java | 91 +++++++++++++ .../src/test/resources/test-span-event.json | 23 ++++ 8 files changed, 372 insertions(+), 27 deletions(-) create mode 100644 data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/XRayOtlpHttpSender.java create mode 100644 data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/XRayOtlpHttpSenderTest.java create mode 100644 data-prepper-plugins/xray-otlp-sink/src/test/resources/test-span-event.json diff --git a/data-prepper-plugins/xray-otlp-sink/build.gradle b/data-prepper-plugins/xray-otlp-sink/build.gradle index e61e5d6eea..11d6872dad 100644 --- a/data-prepper-plugins/xray-otlp-sink/build.gradle +++ b/data-prepper-plugins/xray-otlp-sink/build.gradle @@ -18,11 +18,14 @@ dependencies { implementation 'software.amazon.awssdk:auth:2.28.23' implementation 'software.amazon.awssdk:sts:2.28.23' implementation 'software.amazon.awssdk:regions:2.28.23' + implementation("software.amazon.awssdk:http-client-spi") + implementation("software.amazon.awssdk:apache-client") compileOnly 'org.projectlombok:lombok:1.18.30' annotationProcessor 'org.projectlombok:lombok:1.18.30' implementation project(':data-prepper-api') + implementation project(':data-prepper-plugins:otel-proto-common') implementation 'com.fasterxml.jackson.core:jackson-databind' testImplementation 'org.junit.jupiter:junit-jupiter-api:5.9.1' diff --git a/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSink.java b/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSink.java index c6a78e27d4..a5e86f1197 100644 --- a/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSink.java +++ b/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSink.java @@ -2,19 +2,28 @@ * Copyright OpenSearch Contributors * SPDX-License-Identifier: Apache-2.0 */ + package org.opensearch.dataprepper.plugins.sink.xrayotlp; import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; +import org.opensearch.dataprepper.model.codec.OutputCodec; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.model.sink.Sink; import org.opensearch.dataprepper.model.trace.Span; +import org.opensearch.dataprepper.plugins.otel.codec.OtlpTraceOutputCodec; +import org.opensearch.dataprepper.plugins.sink.xrayotlp.configuration.XRayOTLPSinkConfig; +import org.opensearch.dataprepper.plugins.sink.xrayotlp.http.SigV4Signer; +import org.opensearch.dataprepper.plugins.sink.xrayotlp.http.XRayOtlpHttpSender; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import software.amazon.awssdk.http.apache.ApacheHttpClient; +import java.io.ByteArrayOutputStream; import java.util.Collection; /** - * A Data Prepper Sink plugin that forwards traces to AWS X-Ray's OTLP endpoint. + * A Data Prepper Sink plugin that forwards spans to AWS X-Ray OTLP endpoint using + * OTLP Protobuf encoding and AWS SigV4 authentication. */ @DataPrepperPlugin( name = "xray_otlp_sink", @@ -24,41 +33,58 @@ public class XRayOTLPSink implements Sink> { private static final Logger LOG = LoggerFactory.getLogger(XRayOTLPSink.class); + private final OutputCodec codec; + private final XRayOtlpHttpSender httpSender; + /** - * Constructs the OTLP X-Ray Sink. - * Configuration loading will be added in a later iteration. + * Default constructor used by Data Prepper. Initializes codec and HTTP sender + * using default ApacheHttpClient and basic configuration. */ public XRayOTLPSink() { - // TODO: Inject config or plugin setting. + this.codec = new OtlpTraceOutputCodec(); + final XRayOTLPSinkConfig config = new XRayOTLPSinkConfig(); // TODO: Load real config + final SigV4Signer signer = new SigV4Signer(config); + this.httpSender = new XRayOtlpHttpSender(signer, ApacheHttpClient.builder().build()); + } + + /** + * Constructor for unit testing with injected dependencies. + * + * @param codec the OutputCodec to encode spans + * @param httpSender the HTTP sender to transmit OTLP data + */ + public XRayOTLPSink(final OutputCodec codec, final XRayOtlpHttpSender httpSender) { + this.codec = codec; + this.httpSender = httpSender; } /** - * Lifecycle hook invoked during pipeline startup. - * Initialize AWS clients or other resources here. + * Initializes the sink. Called once during pipeline startup. */ @Override public void initialize() { // TODO: Initialize AWS X-Ray client + LOG.info("Initialized XRay OTLP Sink"); } /** - * Called each time a batch of records is emitted to the sink. - * This method is responsible for handling delivery to AWS X-Ray. + * Processes a batch of spans and sends them to the AWS X-Ray OTLP endpoint. * - * @param records Collection of OTLP log records to process. + * @param records a collection of span records */ @Override public void output(final Collection> records) { - for (Record record : records) { - final Span span = record.getData(); + if (records == null || records.isEmpty()) { + return; + } - LOG.info("===> Span name: {}", span.getName()); - LOG.info("===> Trace ID: {}", span.getTraceId()); - LOG.info("===> Span ID: {}", span.getSpanId()); - LOG.info("===> Parent ID: {}", span.getParentSpanId()); - LOG.info("===> Start time (epoch nanos): {}", span.getStartTime()); - LOG.info("===> End time (epoch nanos): {}", span.getEndTime()); - LOG.info("===> Attributes: {}", span.getAttributes()); + try (ByteArrayOutputStream out = new ByteArrayOutputStream()) { + for (final Record record : records) { + codec.writeEvent(record.getData(), out); + } + httpSender.send(out.toByteArray()); + } catch (final Exception e) { + LOG.error("Failed to process span records", e); } } @@ -79,6 +105,7 @@ public boolean isReady() { @Override public void shutdown() { // TODO: Clean up resources + httpSender.close(); } /** diff --git a/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/XRayOTLPSinkConfig.java b/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/XRayOTLPSinkConfig.java index d5e1376b62..541e630699 100644 --- a/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/XRayOTLPSinkConfig.java +++ b/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/XRayOTLPSinkConfig.java @@ -7,7 +7,10 @@ import com.fasterxml.jackson.annotation.JsonProperty; import jakarta.validation.Valid; import jakarta.validation.constraints.NotNull; +import lombok.AllArgsConstructor; import lombok.Builder; +import lombok.Getter; +import lombok.NoArgsConstructor; import software.amazon.awssdk.regions.Region; /** @@ -18,7 +21,10 @@ * * @since 2.6 */ +@Getter @Builder +@NoArgsConstructor +@AllArgsConstructor public class XRayOTLPSinkConfig { /** * AWS configuration for X-Ray access. diff --git a/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/XRayOtlpHttpSender.java b/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/XRayOtlpHttpSender.java new file mode 100644 index 0000000000..8388616abd --- /dev/null +++ b/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/XRayOtlpHttpSender.java @@ -0,0 +1,73 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.sink.xrayotlp.http; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import software.amazon.awssdk.http.HttpExecuteRequest; +import software.amazon.awssdk.http.HttpExecuteResponse; +import software.amazon.awssdk.http.SdkHttpClient; +import software.amazon.awssdk.http.SdkHttpFullRequest; + +import java.io.ByteArrayInputStream; + +/** + * Responsible for sending signed OTLP Protobuf trace data to AWS X-Ray OTLP endpoint. + * This class uses AWS SDK's HTTP client to send the request signed using AWS SigV4. + */ +public class XRayOtlpHttpSender implements AutoCloseable { + + private static final Logger LOG = LoggerFactory.getLogger(XRayOtlpHttpSender.class); + + private final SigV4Signer signer; + private final SdkHttpClient httpClient; + + /** + * Constructs a new OtlpHttpSender. + * + * @param signer The SigV4Signer used to sign HTTP requests. + * @param httpClient The AWS SDK HTTP client used to send signed requests. + */ + public XRayOtlpHttpSender(final SigV4Signer signer, final SdkHttpClient httpClient) { + this.signer = signer; + this.httpClient = httpClient; + } + + /** + * Signs and sends the given OTLP-encoded span payload to the X-Ray OTLP endpoint. + * + * @param payload The OTLP Protobuf payload to send. + */ + public void send(final byte[] payload) { + try { + final SdkHttpFullRequest signedRequest = signer.signRequest(payload); + + final HttpExecuteRequest httpRequest = HttpExecuteRequest.builder() + .request(signedRequest) + .contentStreamProvider(() -> new ByteArrayInputStream(payload)) + .build(); + + final HttpExecuteResponse response = httpClient.prepareRequest(httpRequest).call(); + final int status = response.httpResponse().statusCode(); + + if (status >= 200 && status < 300) { + LOG.info("Successfully sent OTLP data to AWS X-Ray. Status: {}", status); + } else { + LOG.warn("Failed to send OTLP data to AWS X-Ray. Status: {}", status); + } + } catch (final Exception e) { + LOG.error("Error sending OTLP data to AWS X-Ray", e); + } + } + + /** + * Closes the underlying HTTP client to release resources. + */ + @Override + public void close() { + httpClient.close(); + } +} diff --git a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkTest.java b/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkTest.java index 293a07e6e8..5f3ec01ea8 100644 --- a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkTest.java +++ b/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkTest.java @@ -5,35 +5,141 @@ package org.opensearch.dataprepper.plugins.sink.xrayotlp; +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.ObjectMapper; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.model.codec.OutputCodec; import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.model.trace.DefaultTraceGroupFields; +import org.opensearch.dataprepper.model.trace.JacksonSpan; import org.opensearch.dataprepper.model.trace.Span; +import org.opensearch.dataprepper.plugins.sink.xrayotlp.http.XRayOtlpHttpSender; +import java.io.InputStream; +import java.io.OutputStream; import java.util.Collections; +import java.util.Map; +import java.util.Objects; -import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; -import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.mockito.Mockito.mock; +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.*; class XRayOTLPSinkTest { + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + private static final String TEST_SPAN_EVENT_JSON_FILE = "test-span-event.json"; + + private OutputCodec mockCodec; + private XRayOtlpHttpSender mockSender; private XRayOTLPSink sink; @BeforeEach void setUp() { - sink = new XRayOTLPSink(); + mockCodec = mock(OutputCodec.class); + mockSender = mock(XRayOtlpHttpSender.class); + sink = new XRayOTLPSink(mockCodec, mockSender); + } + + private Span buildSpanFromTestFile(String fileName, String traceIdOverride) { + try (InputStream inputStream = Objects.requireNonNull( + getClass().getClassLoader().getResourceAsStream(fileName))) { + + final Map spanMap = OBJECT_MAPPER.readValue(inputStream, new TypeReference<>() {}); + final JacksonSpan.Builder builder = JacksonSpan.builder() + .withTraceId(traceIdOverride != null ? traceIdOverride : (String) spanMap.get("traceId")) + .withSpanId((String) spanMap.get("spanId")) + .withParentSpanId((String) spanMap.get("parentSpanId")) + .withTraceState((String) spanMap.get("traceState")) + .withName((String) spanMap.get("name")) + .withKind((String) spanMap.get("kind")) + .withDurationInNanos(((Number) spanMap.get("durationInNanos")).longValue()) + .withStartTime((String) spanMap.get("startTime")) + .withEndTime((String) spanMap.get("endTime")) + .withTraceGroup((String) spanMap.get("traceGroup")); + + final Map traceGroupFieldsMap = (Map) spanMap.get("traceGroupFields"); + if (traceGroupFieldsMap != null) { + builder.withTraceGroupFields(DefaultTraceGroupFields.builder() + .withStatusCode((Integer) traceGroupFieldsMap.getOrDefault("statusCode", 0)) + .withEndTime((String) spanMap.get("endTime")) + .withDurationInNanos(((Number) spanMap.get("durationInNanos")).longValue()) + .build()); + } + + return builder.build(); + } catch (Exception e) { + throw new RuntimeException("Failed to load span from file", e); + } } @Test - void testInitialize_doesNotThrow() { - assertDoesNotThrow(() -> sink.initialize()); + void testOutput_sendsDataToXRay_onSuccessResponse() throws Exception { + final Span span = buildSpanFromTestFile(TEST_SPAN_EVENT_JSON_FILE, "bad-trace-id"); + Record record = new Record<>(span); + + doAnswer(invocation -> { + OutputStream out = invocation.getArgument(1); + out.write("dummy-otlp-payload".getBytes()); + return null; + }).when(mockCodec).writeEvent(eq(span), any()); + + doNothing().when(mockSender).send(eq("dummy-otlp-payload".getBytes())); + + sink.output(Collections.singletonList(record)); + + verify(mockCodec).writeEvent(eq(span), any()); + verify(mockSender).send(eq("dummy-otlp-payload".getBytes())); } @Test - void testOutput_printsRecordData() { - final Span mockSpan = mock(Span.class); - final Record record = new Record<>(mockSpan); + void testOutput_handlesException_gracefully() throws Exception { + Span span = mock(Span.class); + Record record = new Record<>(span); + + doThrow(new RuntimeException("codec error")).when(mockCodec).writeEvent(eq(span), any()); + + sink.output(Collections.singletonList(record)); + + verify(mockCodec).writeEvent(eq(span), any()); + verifyNoInteractions(mockSender); + } + + @Test + void testOutput_senderThrowsException_isHandledGracefully() throws Exception { + final Span span = buildSpanFromTestFile(TEST_SPAN_EVENT_JSON_FILE, "bad-trace-id"); + Record record = new Record<>(span); + + doAnswer(invocation -> { + OutputStream out = invocation.getArgument(1); + out.write("dummy-otlp-payload".getBytes()); + return null; + }).when(mockCodec).writeEvent(eq(span), any()); + + doThrow(new RuntimeException("Send failed")).when(mockSender).send(any()); + assertDoesNotThrow(() -> sink.output(Collections.singletonList(record))); + + verify(mockCodec).writeEvent(eq(span), any()); + verify(mockSender).send(any()); + } + + @Test + void testOutput_withNullRecordList_doesNothing() { + assertDoesNotThrow(() -> sink.output(null)); + verifyNoInteractions(mockCodec, mockSender); + } + + @Test + void testOutput_withEmptyRecordList_doesNothing() { + assertDoesNotThrow(() -> sink.output(Collections.emptyList())); + verifyNoInteractions(mockCodec, mockSender); + } + + @Test + void testInitialize_doesNotThrow() { + assertDoesNotThrow(() -> sink.initialize()); } @Test diff --git a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/XRayOTLPSinkConfigTest.java b/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/XRayOTLPSinkConfigTest.java index d9d0fe0cd3..1993a88cd4 100644 --- a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/XRayOTLPSinkConfigTest.java +++ b/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/XRayOTLPSinkConfigTest.java @@ -9,6 +9,8 @@ import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.notNullValue; +import static org.hamcrest.Matchers.nullValue; class XRayOTLPSinkConfigTest { @Test @@ -31,4 +33,18 @@ void testAwsAuthenticationConfiguration_withAllFields() { assertThat(config.getStsRoleArn(), equalTo(expectedRoleArn)); assertThat(config.getStsExternalId(), equalTo(expectedExternalId)); } + + @Test + void testDefaultConstructorAndSetters() { + final XRayOTLPSinkConfig config = new XRayOTLPSinkConfig(); + assertThat(config, notNullValue()); + } + + @Test + void testBuilder_withNullAwsAuthConfig() { + XRayOTLPSinkConfig config = XRayOTLPSinkConfig.builder() + .awsAuthenticationConfiguration(null) + .build(); + assertThat(config.getAwsAuthenticationConfiguration(), nullValue()); + } } diff --git a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/XRayOtlpHttpSenderTest.java b/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/XRayOtlpHttpSenderTest.java new file mode 100644 index 0000000000..c7bc03e9ed --- /dev/null +++ b/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/XRayOtlpHttpSenderTest.java @@ -0,0 +1,91 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.sink.xrayotlp.http; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import software.amazon.awssdk.http.ExecutableHttpRequest; +import software.amazon.awssdk.http.HttpExecuteRequest; +import software.amazon.awssdk.http.HttpExecuteResponse; +import software.amazon.awssdk.http.SdkHttpClient; +import software.amazon.awssdk.http.SdkHttpFullRequest; +import software.amazon.awssdk.http.SdkHttpResponse; + +import java.io.IOException; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.*; + +class XRayOtlpHttpSenderTest { + + private SigV4Signer mockSigner; + private SdkHttpClient mockHttpClient; + private XRayOtlpHttpSender sender; + + @BeforeEach + void setUp() { + mockSigner = mock(SigV4Signer.class); + mockHttpClient = mock(SdkHttpClient.class); + sender = new XRayOtlpHttpSender(mockSigner, mockHttpClient); + } + + @Test + void testSend_successfulRequest_logsInfo() throws Exception { + byte[] payload = "test-payload".getBytes(); + + SdkHttpFullRequest signedRequest = mock(SdkHttpFullRequest.class); + when(mockSigner.signRequest(payload)).thenReturn(signedRequest); + + HttpExecuteResponse mockResponse = mock(HttpExecuteResponse.class); + when(mockResponse.httpResponse()).thenReturn(SdkHttpResponse.builder().statusCode(200).build()); + + ExecutableHttpRequest executableRequest = mock(ExecutableHttpRequest.class); + when(executableRequest.call()).thenReturn(mockResponse); + when(mockHttpClient.prepareRequest(any(HttpExecuteRequest.class))).thenReturn(executableRequest); + + assertDoesNotThrow(() -> sender.send(payload)); + + verify(mockSigner).signRequest(payload); + verify(mockHttpClient).prepareRequest(any(HttpExecuteRequest.class)); + } + + @Test + void testSend_httpError_logsWarning() throws Exception { + byte[] payload = "test-payload".getBytes(); + + SdkHttpFullRequest signedRequest = mock(SdkHttpFullRequest.class); + when(mockSigner.signRequest(payload)).thenReturn(signedRequest); + + HttpExecuteResponse mockResponse = mock(HttpExecuteResponse.class); + when(mockResponse.httpResponse()).thenReturn(SdkHttpResponse.builder().statusCode(500).build()); + + ExecutableHttpRequest executableRequest = mock(ExecutableHttpRequest.class); + when(executableRequest.call()).thenReturn(mockResponse); + when(mockHttpClient.prepareRequest(any(HttpExecuteRequest.class))).thenReturn(executableRequest); + + sender.send(payload); + + verify(mockSigner).signRequest(payload); + verify(mockHttpClient).prepareRequest(any(HttpExecuteRequest.class)); + } + + @Test + void testSend_exceptionDuringSend_logsError() throws Exception { + byte[] payload = "test-payload".getBytes(); + + when(mockSigner.signRequest(payload)).thenThrow(new RuntimeException("signing failed")); + + assertDoesNotThrow(() -> sender.send(payload)); + verify(mockSigner).signRequest(payload); + } + + @Test + void testClose_closesHttpClient() throws IOException { + sender.close(); + verify(mockHttpClient).close(); + } +} diff --git a/data-prepper-plugins/xray-otlp-sink/src/test/resources/test-span-event.json b/data-prepper-plugins/xray-otlp-sink/src/test/resources/test-span-event.json new file mode 100644 index 0000000000..eae0ae4460 --- /dev/null +++ b/data-prepper-plugins/xray-otlp-sink/src/test/resources/test-span-event.json @@ -0,0 +1,23 @@ +{ + "traceId": "ffa576d321173ac6cef3601c8f4bde75", + "spanId": "085ac082ffcfbf8d", + "traceState": "TRACE_STATE_1", + "kind": "SPAN_KIND_CLIENT", + "traceGroupFields": { + "endTime": "2020-08-20T05:40:43.217170200Z", + "durationInNanos": 49160000, + "statusCode": 1 + }, + "name": "TRACE_1_ROOT_SPAN", + "traceGroup": "TRACE_1_ROOT_SPAN", + "startTime": "2020-08-20T05:40:43.168010200Z", + "durationInNanos": 49160000, + "endTime": "2020-08-20T05:40:43.217170200Z", + "parentSpanId": "", + "attributes": {}, + "droppedAttributesCount": 0, + "links": [], + "droppedLinksCount": 0, + "events": [], + "droppedEventsCount": 0 +} \ No newline at end of file From fb080cca8e43516c3132e7ee322d1e174a234338 Mon Sep 17 00:00:00 2001 From: huyPham Date: Mon, 14 Apr 2025 12:30:12 -0700 Subject: [PATCH 07/23] feat(otlp-sink): Add retry, batching, OkHttp client, and encoder refactor (#10) This update introduces several enhancements to the OTLP sink: Retry Mechanism: Implements exponential backoff with jitter to gracefully handle transient failures (e.g. 429, 5xx). Batching Support: Adds configurable batching for efficient trace delivery. OkHttp Integration: Replaces basic HTTP logic with OkHttp client, enabling future async and connection reuse support. Encoder Refactor: Replaces OutputCodec usage with OTelProtoStandardCodec.OTelProtoEncoder for direct protobuf encoding of ExportTraceServiceRequest. Test Coverage: Improves unit test coverage to ~95% with additional edge case validation. E2E Testing: Performed end-to-end testing with real OTLP trace ingestion into AWS X-Ray. Sink Renaming: Renamed plugin to otlp sink for clarity and consistency. Signed-off-by: huy pham --- .../otel/codec/OtlpTraceOutputCodec.java | 13 +- .../otel/codec/OtlpTraceOutputCodecTest.java | 8 +- .../{xray-otlp-sink => otlp-sink}/README.md | 15 +- .../build.gradle | 85 ++-- .../plugins/sink/otlp/OtlpSinkIT.java | 84 ++++ .../plugins/sink/otlp/OtlpSink.java | 164 +++++++ .../AwsAuthenticationConfiguration.java | 12 +- .../otlp/configuration/OtlpSinkConfig.java | 80 ++++ .../sink/otlp/http/OtlpHttpSender.java | 196 +++++++++ .../plugins/sink/otlp}/http/SigV4Signer.java | 29 +- .../plugins/sink/otlp/http/Sleeper.java | 20 + .../plugins/sink/otlp/http/ThreadSleeper.java | 16 + .../plugins/sink/otlp/OtlpSinkTest.java | 161 +++++++ .../AwsAuthenticationConfigurationTest.java | 45 ++ .../configuration/OtlpSinkConfigTest.java | 57 +++ .../sink/otlp/http/OtlpHttpSenderTest.java | 412 ++++++++++++++++++ .../sink/otlp}/http/SigV4SignerTest.java | 39 +- .../sink/otlp/http/ThreadSleeperTest.java | 50 +++ .../test/resources/data-prepper-config.yaml | 0 .../src/test/resources/pipelines.yaml | 12 +- .../src/test/resources/sample-trace.json | 0 .../src/test/resources/test-span-event.json | 0 .../plugins/sink/xrayotlp/XRayOTLPSinkIT.java | 38 -- .../plugins/sink/xrayotlp/XRayOTLPSink.java | 120 ----- .../configuration/XRayOTLPSinkConfig.java | 50 --- .../xrayotlp/http/XRayOtlpHttpSender.java | 73 ---- .../sink/xrayotlp/XRayOTLPSinkTest.java | 154 ------- .../AwsAuthenticationConfigurationTest.java | 41 -- .../configuration/XRayOTLPSinkConfigTest.java | 50 --- .../xrayotlp/http/XRayOtlpHttpSenderTest.java | 91 ---- settings.gradle | 2 +- 31 files changed, 1422 insertions(+), 695 deletions(-) rename data-prepper-plugins/{xray-otlp-sink => otlp-sink}/README.md (82%) rename data-prepper-plugins/{xray-otlp-sink => otlp-sink}/build.gradle (77%) create mode 100644 data-prepper-plugins/otlp-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkIT.java create mode 100644 data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java rename data-prepper-plugins/{xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp => otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp}/configuration/AwsAuthenticationConfiguration.java (86%) create mode 100644 data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfig.java create mode 100644 data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java rename data-prepper-plugins/{xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp => otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp}/http/SigV4Signer.java (77%) create mode 100644 data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/Sleeper.java create mode 100644 data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeper.java create mode 100644 data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java create mode 100644 data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfigurationTest.java create mode 100644 data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfigTest.java create mode 100644 data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java rename data-prepper-plugins/{xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp => otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp}/http/SigV4SignerTest.java (83%) create mode 100644 data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeperTest.java rename data-prepper-plugins/{xray-otlp-sink => otlp-sink}/src/test/resources/data-prepper-config.yaml (100%) rename data-prepper-plugins/{xray-otlp-sink => otlp-sink}/src/test/resources/pipelines.yaml (56%) rename data-prepper-plugins/{xray-otlp-sink => otlp-sink}/src/test/resources/sample-trace.json (100%) rename data-prepper-plugins/{xray-otlp-sink => otlp-sink}/src/test/resources/test-span-event.json (100%) delete mode 100644 data-prepper-plugins/xray-otlp-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkIT.java delete mode 100644 data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSink.java delete mode 100644 data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/XRayOTLPSinkConfig.java delete mode 100644 data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/XRayOtlpHttpSender.java delete mode 100644 data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkTest.java delete mode 100644 data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/AwsAuthenticationConfigurationTest.java delete mode 100644 data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/XRayOTLPSinkConfigTest.java delete mode 100644 data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/XRayOtlpHttpSenderTest.java diff --git a/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OtlpTraceOutputCodec.java b/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OtlpTraceOutputCodec.java index 901b6cc7c4..b6d26fd373 100644 --- a/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OtlpTraceOutputCodec.java +++ b/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OtlpTraceOutputCodec.java @@ -24,7 +24,7 @@ * into OpenTelemetry Protocol (OTLP) binary format using protobuf. *

* This codec is primarily intended for use with trace data that will be - * forwarded to systems such as AWS X-Ray via OTLP. + * forwarded to OTLP endpoint. */ @DataPrepperPlugin(name = "otlp_trace", pluginType = OutputCodec.class) public class OtlpTraceOutputCodec implements OutputCodec { @@ -49,8 +49,17 @@ public void start(final OutputStream outputStream, final Event event, final Outp /** * Writes a single {@link Span} event to the output stream in OTLP binary format. * + *

This method throws a {@link RuntimeException} (e.g., wrapping {@link DecoderException}) + * if the span is malformed or cannot be encoded. This allows upstream components (such as Sink plugins) + * to track and report failed span encoding attempts via plugin metrics or perform custom error handling. + * + *

Failing fast here instead of silently logging ensures invalid spans are not silently dropped and + * gives pipeline developers better visibility into pipeline health and data loss. + * * @param event The event to encode. Must be of type {@link Span}. * @param outputStream The stream to which the encoded bytes will be written. + * @throws IllegalArgumentException If the event is not a {@link Span}. + * @throws RuntimeException If encoding the span fails. */ @Override public void writeEvent(@NonNull final Event event, @NonNull final OutputStream outputStream) { @@ -68,8 +77,10 @@ public void writeEvent(@NonNull final Event event, @NonNull final OutputStream o outputStream.write(request.toByteArray()); } catch (final DecoderException e) { LOG.warn("Skipping invalid span with ID [{}] due to decoding error.", span.getSpanId(), e); + throw new RuntimeException(e); } catch (final Exception e) { LOG.error("Unexpected error while writing span with ID [{}] to OTLP output.", span.getSpanId(), e); + throw new RuntimeException(e); } } diff --git a/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OtlpTraceOutputCodecTest.java b/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OtlpTraceOutputCodecTest.java index 9b881ebca3..d6ec2e94d5 100644 --- a/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OtlpTraceOutputCodecTest.java +++ b/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OtlpTraceOutputCodecTest.java @@ -85,15 +85,13 @@ void testWriteEvent_withValidSpanFromTestFile_writesSuccessfully() throws Except } @Test - void testWriteEvent_withBadTraceId_logsAndSkips() throws Exception { + void testWriteEvent_withBadTraceId_throwsException() throws Exception { final Span span = buildSpanFromTestFile(TEST_SPAN_EVENT_JSON_FILE,"bad-trace-id" ); final ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); - codec.writeEvent(span, outputStream); - - // Since traceId is invalid hex, it will trigger DecoderException - assertThat(outputStream.toByteArray()).isEmpty(); // nothing written + assertThatThrownBy(() -> codec.writeEvent(span, outputStream)) + .isInstanceOf(RuntimeException.class); } @Test diff --git a/data-prepper-plugins/xray-otlp-sink/README.md b/data-prepper-plugins/otlp-sink/README.md similarity index 82% rename from data-prepper-plugins/xray-otlp-sink/README.md rename to data-prepper-plugins/otlp-sink/README.md index 0fffacda9a..4cc7a4fde8 100644 --- a/data-prepper-plugins/xray-otlp-sink/README.md +++ b/data-prepper-plugins/otlp-sink/README.md @@ -1,6 +1,6 @@ # X-Ray OTLP Sink -The `xray_otlp_sink` plugin sends span data to [AWS X-Ray](https://docs.aws.amazon.com/xray/) using the OTLP (OpenTelemetry Protocol) format. +The `otlp_sink` plugin sends span data to [AWS X-Ray](https://docs.aws.amazon.com/xray/) using the OTLP (OpenTelemetry Protocol) format. ## Usage @@ -39,10 +39,11 @@ pipeline: batch_size: 5 sink: - - xray-otlp-sink: + - otlp: aws: region: us-west-2 sts_role_arn: arn:aws:iam::123456789012:role/XrayRole +``` ## Developer Guide @@ -53,13 +54,13 @@ The integration tests for this plugin do not run as part of the main Data Preppe #### Run unit tests locally ```bash -./gradlew :data-prepper-plugins:xray-otlp-sink:test +./gradlew :data-prepper-plugins:otlp-sink:test ``` #### Run integration tests locally ``` -./gradlew :data-prepper-plugins:xray-otlp-sink:integrationTest +./gradlew :data-prepper-plugins:otlp-sink:integrationTest ``` #### Run a local pipeline that uses this sink @@ -74,8 +75,8 @@ The integration tests for this plugin do not run as part of the main Data Preppe cd release/archives/linux/build/install/opensearch-data-prepper-2.11.0-SNAPSHOT-linux-x64 bin/data-prepper \ - /path/to/data-prepper-plugins/xray-otlp-sink/src/test/resources/pipelines.yaml \ - /path/to/data-prepper-plugins/xray-otlp-sink/src/test/resources/data-prepper-config.yaml + /path/to/data-prepper-plugins/otlp-sink/src/test/resources/pipelines.yaml \ + /path/to/data-prepper-plugins/otlp-sink/src/test/resources/data-prepper-config.yaml ``` 4. Send test spans to the local pipeline: ``` @@ -90,7 +91,7 @@ grpcurl -plaintext \ -d @ \ localhost:21890 \ opentelemetry.proto.collector.trace.v1.TraceService/Export \ - < /path/to/data-prepper-plugins/xray-otlp-sink/src/test/resources/sample-trace.json + < /path/to/data-prepper-plugins/otlp-sink/src/test/resources/sample-trace.json ``` You should see log output from XRayOTLPSink that confirms the span data was received and parsed correctly. diff --git a/data-prepper-plugins/xray-otlp-sink/build.gradle b/data-prepper-plugins/otlp-sink/build.gradle similarity index 77% rename from data-prepper-plugins/xray-otlp-sink/build.gradle rename to data-prepper-plugins/otlp-sink/build.gradle index 11d6872dad..49a94a05d4 100644 --- a/data-prepper-plugins/xray-otlp-sink/build.gradle +++ b/data-prepper-plugins/otlp-sink/build.gradle @@ -8,40 +8,83 @@ plugins { id 'jacoco' } +jacoco { + toolVersion = "0.8.11" +} + +sourceSets { + integrationTest { + java.srcDir file('src/integrationTest/java') + resources.srcDir file('src/integrationTest/resources') + compileClasspath += sourceSets.main.output + runtimeClasspath += sourceSets.main.output + } +} + configurations { integrationTestImplementation.extendsFrom testImplementation integrationTestRuntimeOnly.extendsFrom testRuntimeOnly } dependencies { + // AWS SDK implementation 'software.amazon.awssdk:core:2.28.23' implementation 'software.amazon.awssdk:auth:2.28.23' implementation 'software.amazon.awssdk:sts:2.28.23' implementation 'software.amazon.awssdk:regions:2.28.23' - implementation("software.amazon.awssdk:http-client-spi") - implementation("software.amazon.awssdk:apache-client") + implementation 'software.amazon.awssdk:http-client-spi' + implementation 'software.amazon.awssdk:apache-client' + + // OkHttp + implementation 'com.squareup.okhttp3:okhttp:4.12.0' + + // OpenTelemetry Protobuf + implementation libs.opentelemetry.proto + implementation libs.protobuf.util + testImplementation libs.opentelemetry.proto + testImplementation libs.protobuf.util + // Jackson + implementation 'com.fasterxml.jackson.core:jackson-databind' + implementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml:2.15.3' + testImplementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml:2.15.3' + + // Lombok compileOnly 'org.projectlombok:lombok:1.18.30' annotationProcessor 'org.projectlombok:lombok:1.18.30' + // Data Prepper Projects implementation project(':data-prepper-api') implementation project(':data-prepper-plugins:otel-proto-common') - implementation 'com.fasterxml.jackson.core:jackson-databind' + // Unit Testing testImplementation 'org.junit.jupiter:junit-jupiter-api:5.9.1' testRuntimeOnly 'org.junit.jupiter:junit-jupiter-engine:5.9.1' testImplementation 'org.slf4j:slf4j-simple:2.0.7' testImplementation 'org.mockito:mockito-core:5.10.0' + // Integration Testing integrationTestImplementation project(':data-prepper-test-common') integrationTestImplementation 'org.junit.jupiter:junit-jupiter-api:5.9.2' integrationTestRuntimeOnly 'org.junit.jupiter:junit-jupiter-engine:5.9.2' } -jacoco { - toolVersion = "0.8.11" +test { + useJUnitPlatform() + finalizedBy jacocoTestReport, jacocoTestCoverageVerification } +tasks.register('integrationTest', Test) { + description = 'Runs integration tests.' + group = 'verification' + testClassesDirs = sourceSets.integrationTest.output.classesDirs + classpath = sourceSets.integrationTest.runtimeClasspath + shouldRunAfter test + useJUnitPlatform() +} + +check.dependsOn integrationTest + jacocoTestReport { dependsOn test reports { @@ -55,13 +98,12 @@ jacocoTestCoverageVerification { rule { enabled = true element = 'CLASS' + includes = ['org.opensearch.dataprepper.plugins.sink.otlp.*'] limit { counter = 'LINE' value = 'COVEREDRATIO' - minimum = 0.9 + minimum = 0.95 } - - includes = ['org.opensearch.dataprepper.plugins.sink.xrayotlp.*'] } } } @@ -71,29 +113,4 @@ tasks.named('jacocoTestReport') { def reportPath = layout.buildDirectory.file("reports/jacoco/test/html/index.html").get().asFile.toURI() println "\nView test coverage report here:\n ${reportPath}\n" } -} - -test { - useJUnitPlatform() - finalizedBy jacocoTestReport, jacocoTestCoverageVerification -} - -sourceSets { - integrationTest { - java.srcDir file('src/integrationTest/java') - resources.srcDir file('src/integrationTest/resources') - compileClasspath += sourceSets.main.output - runtimeClasspath += sourceSets.main.output - } -} - -tasks.register('integrationTest', Test) { - description = 'Runs integration tests.' - group = 'verification' - testClassesDirs = sourceSets.integrationTest.output.classesDirs - classpath = sourceSets.integrationTest.runtimeClasspath - shouldRunAfter test - useJUnitPlatform() -} - -check.dependsOn integrationTest \ No newline at end of file +} \ No newline at end of file diff --git a/data-prepper-plugins/otlp-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkIT.java b/data-prepper-plugins/otlp-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkIT.java new file mode 100644 index 0000000000..b3c3e66155 --- /dev/null +++ b/data-prepper-plugins/otlp-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkIT.java @@ -0,0 +1,84 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.sink.otlp; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.model.trace.JacksonStandardSpan; +import org.opensearch.dataprepper.model.trace.Span; +import org.opensearch.dataprepper.plugins.otel.codec.OTelProtoStandardCodec; +import org.opensearch.dataprepper.plugins.sink.otlp.configuration.OtlpSinkConfig; +import org.opensearch.dataprepper.plugins.sink.otlp.http.OtlpHttpSender; +import software.amazon.awssdk.regions.Region; + +import java.time.Instant; +import java.util.Collections; +import java.util.List; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** + * Integration test for OtlpSink. Requires AWS credentials to be set up in the environment. + * This will not run as part of the Data Prepper build. + *

+ * ./gradlew :data-prepper-plugins:otlp-sink:integrationTest \ + * -Dtests.xray.region=us-west-2 \ + * -Dtests.xray.profile=dev + */ +class OtlpSinkIT { + + private OtlpSinkConfig mockConfig; + private OtlpSink target; + + @BeforeEach + void setUp() { + System.setProperty("aws.accessKeyId", "dummy"); + System.setProperty("aws.secretAccessKey", "dummy"); + System.setProperty("aws.region", Region.US_WEST_2.toString()); + + mockConfig = mock(OtlpSinkConfig.class); + when(mockConfig.getMaxRetries()).thenReturn(3); + when(mockConfig.getBatchSize()).thenReturn(100); + + target = new OtlpSink(mockConfig); + } + + @AfterEach + void cleanUp() { + System.clearProperty("aws.accessKeyId"); + System.clearProperty("aws.secretAccessKey"); + System.clearProperty("aws.region"); + } + + /** + * This test is not part of the Data Prepper build. It requires AWS credentials to be set up in the environment. + * + * @throws InterruptedException + */ + @Test + void testSinkProcessesHardcodedSpan() { + final Span testSpan = JacksonStandardSpan.builder() + .withTraceId("0123456789abcdef0123456789abcdef") + .withSpanId("0123456789abcdef") + .withParentSpanId("1111111111111111") + .withName("my-test-span") + .withStartTime(Instant.now().toString()) + .withEndTime(Instant.now().plusMillis(10).toString()) + .withAttributes(Collections.emptyMap()) + .withKind("INTERNAL") + .build(); + + final Record record = new Record<>(testSpan); + final OtlpSink sink = new OtlpSink(mockConfig, mock(OTelProtoStandardCodec.OTelProtoEncoder.class), mock(OtlpHttpSender.class)); + + sink.initialize(); + sink.output(List.of(record)); + sink.shutdown(); + } +} diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java new file mode 100644 index 0000000000..007c861282 --- /dev/null +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java @@ -0,0 +1,164 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.sink.otlp; + +import com.google.common.annotations.VisibleForTesting; +import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceRequest; +import io.opentelemetry.proto.trace.v1.ResourceSpans; +import org.jetbrains.annotations.Nullable; +import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; +import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor; +import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.model.sink.Sink; +import org.opensearch.dataprepper.model.trace.Span; +import org.opensearch.dataprepper.plugins.otel.codec.OTelProtoStandardCodec; +import org.opensearch.dataprepper.plugins.sink.otlp.configuration.OtlpSinkConfig; +import org.opensearch.dataprepper.plugins.sink.otlp.http.OtlpHttpSender; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nonnull; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.Objects; +import java.util.stream.Collectors; + +/** + * A Data Prepper Sink plugin that forwards spans to OTLP endpoint using + */ +@DataPrepperPlugin( + name = "otlp", + pluginType = Sink.class, + pluginConfigurationType = OtlpSinkConfig.class +) +public class OtlpSink implements Sink> { + + private static final Logger LOG = LoggerFactory.getLogger(OtlpSink.class); + + private final int batchSize; + private final OtlpHttpSender httpSender; + private final OTelProtoStandardCodec.OTelProtoEncoder encoder; + + /** + * Constructor for the OTLPSink plugin. + * + * @param config the configuration for the OTLP sink + */ + @DataPrepperPluginConstructor + public OtlpSink(@Nonnull final OtlpSinkConfig config) { + this(config, null, null); + } + + /** + * Constructor for testing purposes. + * + * @param config the configuration for the OTLP sink + * @param encoder OTEL Protobuf encoder to use + * @param httpSender the HTTP sender to use + */ + @VisibleForTesting + OtlpSink(@Nonnull final OtlpSinkConfig config, final OTelProtoStandardCodec.OTelProtoEncoder encoder, final OtlpHttpSender httpSender) { + this.batchSize = config.getBatchSize(); + + if (encoder == null && httpSender == null) { + this.encoder = new OTelProtoStandardCodec.OTelProtoEncoder(); + this.httpSender = new OtlpHttpSender(config); + } else { + this.encoder = encoder; + this.httpSender = httpSender; + } + + LOG.info("Config setting: endpoint = {}", config.getEndpoint()); + LOG.info("Config setting: batch_size = {}", config.getBatchSize()); + LOG.info("Config setting: max_retries = {}", config.getMaxRetries()); + LOG.info("Config setting: aws_region = {}", config.getAwsRegion()); + LOG.info("Config setting: aws_sts_role_arn = {}", config.getStsRoleArn()); + LOG.info("Config setting: aws_sts_external_id = {}", config.getStsExternalId()); + } + + /** + * Initializes the sink. Called once during pipeline startup. + */ + @Override + public void initialize() { + LOG.info("Initialized OTLP Sink"); + } + + /** + * Processes a batch of spans and sends them to the OTLP endpoint. + * + * @param records a collection of span records + */ + @Override + public void output(@Nonnull final Collection> records) { + final List> recordList = new ArrayList<>(records); + for (int i = 0; i < recordList.size(); i += this.batchSize) { + final int end = Math.min(i + this.batchSize, recordList.size()); + final List> batch = recordList.subList(i, end); + + try { + final List resourceSpans = batch.stream() + .map(Record::getData) + .map(this::getResourceSpans) + .filter(Objects::nonNull) + .collect(Collectors.toList()); + + if (resourceSpans.isEmpty()) { + LOG.debug("Skipping empty span batch, nothing to send."); + continue; + } + + final ExportTraceServiceRequest request = ExportTraceServiceRequest.newBuilder() + .addAllResourceSpans(resourceSpans) + .build(); + httpSender.send(request.toByteArray()); + LOG.info("Finished processing {} spans.", resourceSpans.size()); + } catch (final RuntimeException e) { + LOG.error("Unexpected error processing OTLP span batch: {}", e.getMessage(), e); + } + } + } + + @Nullable + private ResourceSpans getResourceSpans(final Span span) { + try { + return encoder.convertToResourceSpans(span); + } catch (final Exception e) { + LOG.warn("Failed to encode span with ID [{}], skipping.", span.getSpanId(), e); + return null; + } + } + + /** + * Indicates whether this sink is ready to receive data. + * + * @return true if the sink is ready + */ + @Override + public boolean isReady() { + return true; + } + + /** + * Hook called during pipeline shutdown. + */ + @Override + public void shutdown() { + // OkHttpClient is shared and self-managed; no need to explicitly close + LOG.info("OTLP Sink shutdown complete"); + } + + /** + * Updates internal latency metrics using the received records. + * + * @param events Collection of records used for latency tracking. + */ + @Override + public void updateLatencyMetrics(@Nonnull final Collection> events) { + // TODO: Implement latency tracking with PluginMetrics + } +} diff --git a/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/AwsAuthenticationConfiguration.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfiguration.java similarity index 86% rename from data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/AwsAuthenticationConfiguration.java rename to data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfiguration.java index 347136b81c..8ee088d2a4 100644 --- a/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/AwsAuthenticationConfiguration.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfiguration.java @@ -2,26 +2,26 @@ * Copyright OpenSearch Contributors * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.dataprepper.plugins.sink.xrayotlp.configuration; +package org.opensearch.dataprepper.plugins.sink.otlp.configuration; import com.fasterxml.jackson.annotation.JsonProperty; import jakarta.validation.constraints.Size; -import lombok.Builder; import lombok.Getter; +import lombok.NoArgsConstructor; import software.amazon.awssdk.regions.Region; /** * Configuration class for AWS authentication settings. * Handles region, STS role ARN, and external ID configurations required for AWS service access. - * This class will be automatically wired by Data-Prepper; the Builder is only for testing. + * This class will be automatically wired by Data-Prepper. * * @since 2.6 */ @Getter -@Builder +@NoArgsConstructor class AwsAuthenticationConfiguration { /** - * AWS region for X-Ray service. + * AWS region. * Must be a valid AWS region identifier (e.g., us-east-1, us-west-2). */ @JsonProperty("region") @@ -52,7 +52,7 @@ class AwsAuthenticationConfiguration { * @return Region object if awsRegion is set, otherwise returns null. * Note: Default region fallback is handled externally by the caller. */ - public Region getAwsRegion() { + Region getAwsRegion() { return awsRegion != null ? Region.of(awsRegion) : null; } } diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfig.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfig.java new file mode 100644 index 0000000000..35da8a01e4 --- /dev/null +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfig.java @@ -0,0 +1,80 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.otlp.configuration; + +import com.fasterxml.jackson.annotation.JsonProperty; +import jakarta.validation.Valid; +import jakarta.validation.constraints.Max; +import jakarta.validation.constraints.Min; +import jakarta.validation.constraints.Size; +import lombok.Getter; +import lombok.NoArgsConstructor; +import software.amazon.awssdk.regions.Region; + +/** + * Configuration class for the OTLP sink plugin. + * This class defines the configuration options available when setting up + * the OTLP sink in Data Prepper pipelines. + *

+ * Note that {@code @Getter} is applied at the field level (not the class level) + * to preserve encapsulation and maintain control over exposed configuration data. + *

+ * This class is automatically wired by the Data Prepper framework during pipeline initialization. + * + * @since 2.6 + */ +@NoArgsConstructor +public class OtlpSinkConfig { + + @Getter + @JsonProperty("endpoint") + @Size(min = 1, message = "endpoint cannot be empty string") + private String endpoint; + + @Getter + @JsonProperty("batch_size") + @Min(value = 10, message = "batch_size must be at least 10") + @Max(value = 512, message = "batch_size must be at most 512") + private int batchSize = 100; + + @Getter + @JsonProperty("max_retries") + @Min(value = 1, message = "max_retries must be at least 1") + @Max(value = 5, message = "max_retries must be at most 5") + private int maxRetries = 3; + + /** + * AWS authentication configuration. + * This object contains the AWS region and STS role ARN (if applicable). + * This field is kept private and its contents should be accessed via the generated getter methods. + */ + @JsonProperty("aws") + @Valid + private AwsAuthenticationConfiguration awsAuthenticationConfiguration; + + public Region getAwsRegion() { + if (awsAuthenticationConfiguration == null) { + return null; + } + + return awsAuthenticationConfiguration.getAwsRegion(); + } + + public String getStsRoleArn() { + if (awsAuthenticationConfiguration == null) { + return null; + } + + return awsAuthenticationConfiguration.getAwsStsRoleArn(); + } + + public String getStsExternalId() { + if (awsAuthenticationConfiguration == null) { + return null; + } + + return awsAuthenticationConfiguration.getAwsStsExternalId(); + } +} diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java new file mode 100644 index 0000000000..32f3f501ca --- /dev/null +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java @@ -0,0 +1,196 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.otlp.http; + +import com.google.common.annotations.VisibleForTesting; +import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceResponse; +import okhttp3.MediaType; +import okhttp3.OkHttpClient; +import okhttp3.Request; +import okhttp3.RequestBody; +import okhttp3.Response; +import org.opensearch.dataprepper.plugins.sink.otlp.configuration.OtlpSinkConfig; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import software.amazon.awssdk.http.SdkHttpFullRequest; + +import javax.annotation.Nonnull; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.security.SecureRandom; +import java.util.ArrayList; +import java.util.List; +import java.util.Set; + +/** + * Responsible for sending signed OTLP Protobuf trace data to AWS OTLP endpoint using OkHttp. + */ +public class OtlpHttpSender implements AutoCloseable { + @VisibleForTesting + static final Set NON_RETRYABLE_STATUS_CODES = Set.of(400, 401, 403, 422); + + private static final int BASE_RETRY_DELAY_MS = 100; + private static final Logger LOG = LoggerFactory.getLogger(OtlpHttpSender.class); + private static final MediaType PROTOBUF = MediaType.get("application/x-protobuf"); + private final SecureRandom random = new SecureRandom(); + + private final SigV4Signer signer; + private final OkHttpClient httpClient; + private final Sleeper sleeper; + private final int maxRetries; + private final List retryDelaysMs; + + /** + * Constructor for the OtlpHttpSender. + * Initializes the signer and HTTP client. + * @param config The configuration for the OTLP sink plugin. + */ + public OtlpHttpSender(@Nonnull final OtlpSinkConfig config) { + this(config, null, null, null); + } + + /** + * Constructor for unit testing with injected dependencies. + * + * @param config The configuration for the OTLP sink plugin. + * @param signer The SigV4Signer instance for signing requests. + * @param httpClient The OkHttpClient instance for making HTTP requests. + */ + @VisibleForTesting + OtlpHttpSender(@Nonnull final OtlpSinkConfig config, final SigV4Signer signer, final OkHttpClient httpClient, final Sleeper sleeper) { + this.signer = signer != null ? signer : new SigV4Signer(config); + this.httpClient = httpClient != null ? httpClient : new OkHttpClient(); + this.sleeper = sleeper != null ? sleeper : new ThreadSleeper(); + + this.retryDelaysMs = generateExponentialBackoffDelays(config.getMaxRetries()); + this.maxRetries = config.getMaxRetries(); + } + + /** + * Generates exponential backoff delays with jitter. + * + * @param retries Number of retries. + * @return List of delay durations in milliseconds. + */ + private List generateExponentialBackoffDelays(final int retries) { + List delays = new ArrayList<>(); + + for (int i = 0; i < retries; i++) { + // Exponential backoff: 100ms, 200ms, 400ms, ... + delays.add(BASE_RETRY_DELAY_MS * (1 << i)); + } + + return delays; + } + + /** + * Sends the provided OTLP Protobuf trace data to the OTLP endpoint. + * Retries with exponential backoff and jitter on failure. + * + * @param payload The OTLP Protobuf-encoded trace data to be sent. + */ + public void send(@Nonnull final byte[] payload) { + for (int attempt = 0; attempt <= maxRetries; attempt++) { + try { + final SdkHttpFullRequest signedRequest = signer.signRequest(payload); + + final Request.Builder requestBuilder = new Request.Builder() + .url(signedRequest.getUri().toString()) + .post(RequestBody.create(payload, PROTOBUF)); + + signedRequest.headers().forEach((key, values) -> { + for (final String value : values) { + requestBuilder.addHeader(key, value); + } + }); + + final Request request = requestBuilder.build(); + + try (final Response response = httpClient.newCall(request).execute()) { + handleResponse(response); + return; + } + } catch (final Exception e) { + if (attempt < maxRetries) { + final int jitter = random.nextInt(100); + final int retryIndex = Math.min(attempt, retryDelaysMs.size() - 1); + final int delay = retryDelaysMs.get(retryIndex) + jitter; + LOG.warn("Retrying after failure in attempt {}. Sleeping {}ms.", attempt + 1, delay, e); + try { + sleeper.sleep(delay); + } catch (final InterruptedException ie) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Retry interrupted while sending OTLP data", ie); + } + } else { + LOG.error("All retry attempts failed while signing or sending OTLP data", e); + throw new RuntimeException("Failed to sign/send data after retries", e); + } + } + } + } + + /** + * Handles the response from the OTLP endpoint. + * Logs the response status and body, and throws an exception for retryable errors. + * + * @param response The HTTP response from the OTLP endpoint. + * @throws IOException If the response status is not successful and retryable. + */ + private void handleResponse(@Nonnull final Response response) throws IOException { + final int status = response.code(); + final byte[] responseBytes = response.body() != null + ? response.body().bytes() + : null; + + if (status >= 200 && status < 300) { + handleSuccessfulResponse(responseBytes); + return; + } + + final String responseBody = responseBytes != null ? new String(responseBytes, StandardCharsets.UTF_8) : ""; + if (NON_RETRYABLE_STATUS_CODES.contains(status)) { + LOG.error("Non-retryable client error. Status: {}, Response: {}", status, responseBody); + return; + } + + final String errorMsg = String.format("Failed to send OTLP data. Status: %d, Response: %s", status, responseBody); + LOG.error(errorMsg); + throw new IOException(errorMsg); + } + + private void handleSuccessfulResponse(final byte[] responseBytes) { + if (responseBytes == null || responseBytes.length == 0) { + LOG.info("OTLP export successful. No response body."); + return; + } + + try { + final ExportTraceServiceResponse otlpResponse = ExportTraceServiceResponse.parseFrom(responseBytes); + + if (otlpResponse.hasPartialSuccess()) { + final var partial = otlpResponse.getPartialSuccess(); + final long rejectedSpans = partial.getRejectedSpans(); + final String errorMessage = partial.getErrorMessage(); + + if (rejectedSpans > 0 || !errorMessage.isEmpty()) { + LOG.warn("OTLP Partial Success: rejectedSpans={}, message={}", rejectedSpans, errorMessage); + } else { + LOG.info("OTLP export successful with no rejections."); + } + } else { + LOG.info("OTLP export successful with no partial success field."); + } + } catch (final Exception e) { + LOG.error("Could not parse OTLP response as ExportTraceServiceResponse: {}", e.getMessage()); + } + } + + @Override + public void close() { + // No explicit shutdown required for OkHttpClient unless using dispatcher or connection pool tuning. + // https://square.github.io/okhttp/4.x/okhttp/okhttp3/-ok-http-client/dispatchers/ + } +} \ No newline at end of file diff --git a/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/SigV4Signer.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4Signer.java similarity index 77% rename from data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/SigV4Signer.java rename to data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4Signer.java index 460148558e..92f3500294 100644 --- a/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/SigV4Signer.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4Signer.java @@ -3,10 +3,10 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.dataprepper.plugins.sink.xrayotlp.http; +package org.opensearch.dataprepper.plugins.sink.otlp.http; import com.google.common.annotations.VisibleForTesting; -import org.opensearch.dataprepper.plugins.sink.xrayotlp.configuration.XRayOTLPSinkConfig; +import org.opensearch.dataprepper.plugins.sink.otlp.configuration.OtlpSinkConfig; import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider; import software.amazon.awssdk.auth.signer.Aws4Signer; @@ -18,15 +18,18 @@ import software.amazon.awssdk.regions.providers.DefaultAwsRegionProviderChain; import software.amazon.awssdk.services.sts.StsClient; import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider; + +import javax.annotation.Nonnull; import java.net.URI; /** * Helper class to apply AWS SigV4 signing to outgoing HTTP requests - * before sending them to the AWS X-Ray OTLP endpoint. + * before sending them to the AWS OTLP endpoint. */ -public class SigV4Signer { +class SigV4Signer { private static final String SERVICE_NAME = "xray"; private static final String OTLP_PATH = "/v1/traces"; + private static final String OTLP_SINK_SESSION = "otlp-sink-session"; private final Aws4Signer signer = Aws4Signer.create(); private final AwsCredentialsProvider credentialsProvider; @@ -39,7 +42,7 @@ public class SigV4Signer { * * @param config Configuration for region and optional STS role */ - public SigV4Signer(final XRayOTLPSinkConfig config) { + SigV4Signer(@Nonnull final OtlpSinkConfig config) { this(config, null); } @@ -47,10 +50,13 @@ public SigV4Signer(final XRayOTLPSinkConfig config) { * Package-private constructor for unit testing with mocked STS client. */ @VisibleForTesting - SigV4Signer(final XRayOTLPSinkConfig config, final StsClient stsClient) { + SigV4Signer(@Nonnull final OtlpSinkConfig config, final StsClient stsClient) { this.region = resolveRegion(config.getAwsRegion()); - this.credentialsProvider = initCredentialsProvider(config.getAwsRegion(), config.getStsRoleArn(), config.getStsExternalId(), stsClient); - this.endpointUri = URI.create(String.format("https://xray.%s.amazonaws.com%s", region.id(), OTLP_PATH)); + this.credentialsProvider = initCredentialsProvider(region, config.getStsRoleArn(), config.getStsExternalId(), stsClient); + this.endpointUri = config.getEndpoint() != null + ? URI.create(config.getEndpoint()) + : URI.create(String.format("https://xray.%s.amazonaws.com%s", region.id(), OTLP_PATH)); + } private Region resolveRegion(final Region region) { @@ -58,15 +64,16 @@ private Region resolveRegion(final Region region) { } private AwsCredentialsProvider initCredentialsProvider( - final Region region, + @Nonnull final Region region, final String stsRoleArn, final String stsExternalId, final StsClient stsClient ) { - if (region != null && stsRoleArn != null) { + if (stsRoleArn != null) { return StsAssumeRoleCredentialsProvider.builder() .refreshRequest(r -> { r.roleArn(stsRoleArn); + r.roleSessionName(OTLP_SINK_SESSION); if (stsExternalId != null) { r.externalId(stsExternalId); } @@ -84,7 +91,7 @@ private AwsCredentialsProvider initCredentialsProvider( * @param payload The OTLP Protobuf-encoded request body to be sent * @return A signed {@link SdkHttpFullRequest} ready for transmission to the AWS OTLP endpoint */ - public SdkHttpFullRequest signRequest(final byte[] payload) { + SdkHttpFullRequest signRequest(@Nonnull final byte[] payload) { final SdkHttpFullRequest unsignedRequest = SdkHttpFullRequest.builder() .method(SdkHttpMethod.POST) .uri(endpointUri) diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/Sleeper.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/Sleeper.java new file mode 100644 index 0000000000..1ba9f34021 --- /dev/null +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/Sleeper.java @@ -0,0 +1,20 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.sink.otlp.http; + +/** + * Interface for sleeping in tests. + */ +interface Sleeper { + + /** + * Sleeps for the specified number of milliseconds. + * + * @param millis the number of milliseconds to sleep + * @throws InterruptedException if the sleep is interrupted + */ + void sleep(int millis) throws InterruptedException; +} \ No newline at end of file diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeper.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeper.java new file mode 100644 index 0000000000..61f8d23d36 --- /dev/null +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeper.java @@ -0,0 +1,16 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.sink.otlp.http; + +/** + * Implementation of {@link Sleeper} that sleeps using {@link Thread#sleep(long)}. + */ +class ThreadSleeper implements Sleeper { + @Override + public void sleep(int millis) throws InterruptedException { + Thread.sleep(millis); + } +} diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java new file mode 100644 index 0000000000..8207a8e72f --- /dev/null +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java @@ -0,0 +1,161 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.sink.otlp; + +import io.opentelemetry.proto.trace.v1.ResourceSpans; +import org.apache.commons.codec.DecoderException; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.model.trace.Span; +import org.opensearch.dataprepper.plugins.otel.codec.OTelProtoStandardCodec; +import org.opensearch.dataprepper.plugins.sink.otlp.configuration.OtlpSinkConfig; +import org.opensearch.dataprepper.plugins.sink.otlp.http.OtlpHttpSender; +import software.amazon.awssdk.regions.Region; + +import java.io.UnsupportedEncodingException; +import java.util.ArrayList; +import java.util.List; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +class OtlpSinkTest { + + private OtlpSinkConfig mockConfig; + private OTelProtoStandardCodec.OTelProtoEncoder mockEncoder; + private OtlpHttpSender mockSender; + private OtlpSink target; + + @BeforeEach + void setUp() { + System.setProperty("aws.accessKeyId", "dummy"); + System.setProperty("aws.secretAccessKey", "dummy"); + + mockConfig = mock(OtlpSinkConfig.class); + when(mockConfig.getAwsRegion()).thenReturn(Region.US_WEST_2); + when(mockConfig.getBatchSize()).thenReturn(100); + when(mockConfig.getMaxRetries()).thenReturn(3); + + mockEncoder = mock(OTelProtoStandardCodec.OTelProtoEncoder.class); + mockSender = mock(OtlpHttpSender.class); + + target = new OtlpSink(mockConfig, mockEncoder, mockSender); + } + + @AfterEach + void cleanUp() { + System.clearProperty("aws.accessKeyId"); + System.clearProperty("aws.secretAccessKey"); + System.clearProperty("aws.region"); + } + + @Test + void testOutput_shouldSendAllBatches() throws Exception { + final int recordCount = 250; + final List> records = new ArrayList<>(); + for (int i = 0; i < recordCount; i++) { + Record mockRecord = mock(Record.class); + Span span = mock(Span.class); + ResourceSpans resourceSpans = ResourceSpans.getDefaultInstance(); + when(mockEncoder.convertToResourceSpans(span)).thenReturn(resourceSpans); + when(mockRecord.getData()).thenReturn(span); + records.add(mockRecord); + } + + target.output(records); + + // 250 total / 100 batch size = 3 calls to httpSender + verify(mockSender, times(3)).send(any(byte[].class)); + verify(mockEncoder, times(recordCount)).convertToResourceSpans(any()); + } + + @Test + void testOutput_shouldHandleRuntimeException() throws Exception { + Span span = mock(Span.class); + when(span.getSpanId()).thenReturn("span123"); + Record record = mock(Record.class); + when(record.getData()).thenReturn(span); + + when(mockEncoder.convertToResourceSpans(span)).thenReturn(ResourceSpans.getDefaultInstance()); + doThrow(new RuntimeException("sender failure")).when(mockSender).send(any()); + + target.output(List.of(record)); + + verify(mockEncoder).convertToResourceSpans(eq(span)); + verify(mockSender).send(any()); + } + + @Test + void testOutput_shouldSendPartialBatchWhenSomeSpansSucceed() throws Exception { + // Good span + Span goodSpan = mock(Span.class); + when(goodSpan.getSpanId()).thenReturn("good-span"); + Record goodRecord = mock(Record.class); + when(goodRecord.getData()).thenReturn(goodSpan); + + // Bad span (encoder throws) + Span badSpan = mock(Span.class); + when(badSpan.getSpanId()).thenReturn("bad-span"); + Record badRecord = mock(Record.class); + when(badRecord.getData()).thenReturn(badSpan); + + // Good span gets encoded properly + ResourceSpans goodResourceSpans = ResourceSpans.getDefaultInstance(); + when(mockEncoder.convertToResourceSpans(goodSpan)).thenReturn(goodResourceSpans); + + // Bad span causes exception during encode + when(mockEncoder.convertToResourceSpans(badSpan)).thenThrow(new RuntimeException("bad span")); + + target.output(List.of(badRecord, goodRecord)); + + // Encoder is called on both + verify(mockEncoder).convertToResourceSpans(badSpan); + verify(mockEncoder).convertToResourceSpans(goodSpan); + + // Sender should still be called with only the good span + verify(mockSender, times(1)).send(any(byte[].class)); + } + + @Test + void testOutput_shouldNotSendEmptyBatch() throws DecoderException, UnsupportedEncodingException { + Span span = mock(Span.class); + when(span.getSpanId()).thenReturn("bad-span"); + Record record = mock(Record.class); + when(record.getData()).thenReturn(span); + + when(mockEncoder.convertToResourceSpans(span)).thenReturn(null); // simulate invalid span + + target.output(List.of(record)); + + verify(mockSender, never()).send(any()); + } + + @Test + void testConstructor_withOnlyConfig_shouldInitializeWithoutException() { + OtlpSink sink = new OtlpSink(mockConfig); + + sink.initialize(); + sink.shutdown(); + } + + @Test + void testIsReady_returnsTrue() { + assert target.isReady(); + } + + @Test + void testShutdown_doesNotThrow() { + target.shutdown(); + } +} diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfigurationTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfigurationTest.java new file mode 100644 index 0000000000..94fe1ae98e --- /dev/null +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfigurationTest.java @@ -0,0 +1,45 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.otlp.configuration; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.dataformat.yaml.YAMLFactory; +import org.junit.jupiter.api.Test; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.nullValue; +import static org.junit.jupiter.api.Assertions.assertEquals; + +class AwsAuthenticationConfigurationTest { + + private final String expectedRoleArn = "arn:aws:iam::123456789012:role/MyRole"; + private final String expectedExternalId = "external-id-123"; + private final String expectedRegion = "us-west-2"; + private final ObjectMapper mapper = new ObjectMapper(new YAMLFactory()); + + @Test + void testDeserializationFromYaml() throws Exception { + final String yaml = String.join("\n", + "region: " + expectedRegion, + "sts_role_arn: " + expectedRoleArn, + "sts_external_id: " + expectedExternalId + ); + + AwsAuthenticationConfiguration config = mapper.readValue(yaml, AwsAuthenticationConfiguration.class); + + assertEquals(expectedRegion, config.getAwsRegion().toString()); + assertEquals(expectedRoleArn, config.getAwsStsRoleArn()); + assertEquals(expectedExternalId, config.getAwsStsExternalId()); + } + + @Test + void testGetRegion_whenAllIsNull_returnNull() throws JsonProcessingException { + final String yaml = "{}"; + AwsAuthenticationConfiguration config = mapper.readValue(yaml, AwsAuthenticationConfiguration.class); + + assertThat(config.getAwsRegion(), nullValue()); + } +} diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfigTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfigTest.java new file mode 100644 index 0000000000..bacf23487f --- /dev/null +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfigTest.java @@ -0,0 +1,57 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.otlp.configuration; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.dataformat.yaml.YAMLFactory; +import org.junit.jupiter.api.Test; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.nullValue; + +class OtlpSinkConfigTest { + + private static final String EXPECTED_ENDPOINT = "https://example.com/otlp"; + private static final int EXPECTED_BATCH_SIZE = 128; + private static final int EXPECTED_MAX_RETRIES = 4; + private static final String EXPECTED_REGION = "us-west-2"; + private static final String EXPECTED_ROLE_ARN = "arn:aws:iam::123456789012:role/OtlpRole"; + private static final String EXPECTED_EXTERNAL_ID = "my-ext-id"; + + @Test + void testDeserializationFromYaml() throws Exception { + final String yaml = + "endpoint: \"" + EXPECTED_ENDPOINT + "\"\n" + + "batch_size: " + EXPECTED_BATCH_SIZE + "\n" + + "max_retries: " + EXPECTED_MAX_RETRIES + "\n" + + "aws:\n" + + " region: \"" + EXPECTED_REGION + "\"\n" + + " sts_role_arn: \"" + EXPECTED_ROLE_ARN + "\"\n" + + " sts_external_id: \"" + EXPECTED_EXTERNAL_ID + "\"\n"; + + final ObjectMapper objectMapper = new ObjectMapper(new YAMLFactory()); + final OtlpSinkConfig config = objectMapper.readValue(yaml, OtlpSinkConfig.class); + + assertThat(config.getEndpoint(), equalTo(EXPECTED_ENDPOINT)); + assertThat(config.getBatchSize(), equalTo(EXPECTED_BATCH_SIZE)); + assertThat(config.getMaxRetries(), equalTo(EXPECTED_MAX_RETRIES)); + + assertThat(config.getAwsRegion().toString(), equalTo(EXPECTED_REGION)); + assertThat(config.getStsRoleArn(), equalTo(EXPECTED_ROLE_ARN)); + assertThat(config.getStsExternalId(), equalTo(EXPECTED_EXTERNAL_ID)); + } + + @Test + void testDeserializationFromYaml_withNullAwsValues() throws Exception { + final String yaml = "{}"; + final ObjectMapper objectMapper = new ObjectMapper(new YAMLFactory()); + final OtlpSinkConfig config = objectMapper.readValue(yaml, OtlpSinkConfig.class); + + assertThat(config.getAwsRegion(), nullValue()); + assertThat(config.getStsRoleArn(), nullValue()); + assertThat(config.getStsExternalId(), nullValue()); + } +} \ No newline at end of file diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java new file mode 100644 index 0000000000..1d63eb1d58 --- /dev/null +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java @@ -0,0 +1,412 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.sink.otlp.http; + +import com.google.protobuf.UnknownFieldSet; +import com.google.protobuf.UnknownFieldSet.Field; +import io.opentelemetry.proto.collector.trace.v1.ExportTracePartialSuccess; +import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceResponse; +import okhttp3.Call; +import okhttp3.HttpUrl; +import okhttp3.MediaType; +import okhttp3.OkHttpClient; +import okhttp3.Protocol; +import okhttp3.Request; +import okhttp3.Response; +import okhttp3.ResponseBody; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.plugins.sink.otlp.configuration.OtlpSinkConfig; +import software.amazon.awssdk.http.SdkHttpFullRequest; +import software.amazon.awssdk.regions.Region; + +import java.io.IOException; +import java.net.URI; +import java.nio.charset.StandardCharsets; +import java.util.Collections; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.fail; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.reset; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.opensearch.dataprepper.plugins.sink.otlp.http.OtlpHttpSender.NON_RETRYABLE_STATUS_CODES; + +class OtlpHttpSenderTest { + + private static final byte[] PAYLOAD = "test-otlp-payload".getBytes(StandardCharsets.UTF_8); + private static final String ERROR_BODY = "{\"error\": \"Something went wrong\"}"; + + private OtlpSinkConfig mockConfig; + private SigV4Signer mockSigner; + private OkHttpClient mockHttpClient; + private Sleeper mockSleeper; + private OtlpHttpSender target; + + @BeforeEach + void setUp() { + System.setProperty("aws.accessKeyId", "dummy"); + System.setProperty("aws.secretAccessKey", "dummy"); + + mockConfig = mock(OtlpSinkConfig.class); + when(mockConfig.getAwsRegion()).thenReturn(Region.US_WEST_2); + when(mockConfig.getMaxRetries()).thenReturn(3); + + mockSigner = mock(SigV4Signer.class); + mockHttpClient = mock(OkHttpClient.class); + mockSleeper = mock(Sleeper.class); + + target = new OtlpHttpSender(mockConfig, mockSigner, mockHttpClient, mockSleeper); + } + + @AfterEach + void cleanUp() { + System.clearProperty("aws.accessKeyId"); + System.clearProperty("aws.secretAccessKey"); + System.clearProperty("aws.region"); + } + + @Test + void testSend_successfulResponse() throws IOException { + SdkHttpFullRequest mockSignedRequest = mock(SdkHttpFullRequest.class); + when(mockSignedRequest.getUri()).thenReturn(HttpUrl.get("https://xray.us-west-2.amazonaws.com/v1/traces").uri()); + when(mockSignedRequest.headers()).thenReturn(Map.of("Authorization", Collections.singletonList("signed-header"))); + + when(mockSigner.signRequest(PAYLOAD)).thenReturn(mockSignedRequest); + + Call mockCall = mock(Call.class); + Response mockResponse = new Response.Builder() + .request(new Request.Builder().url("https://xray.us-west-2.amazonaws.com/v1/traces").build()) + .protocol(Protocol.HTTP_1_1) + .code(200) + .message("OK") + .body(ResponseBody.create("", MediaType.get("application/x-protobuf"))) + .build(); + + when(mockHttpClient.newCall(any())).thenReturn(mockCall); + when(mockCall.execute()).thenReturn(mockResponse); + + assertDoesNotThrow(() -> target.send(PAYLOAD)); + } + + @Test + void testSend_doesNotRetryOnNonRetryable4xxResponses() throws IOException { + final SdkHttpFullRequest mockSignedRequest = mock(SdkHttpFullRequest.class); + when(mockSignedRequest.getUri()).thenReturn(HttpUrl.get("https://xray.us-west-2.amazonaws.com/v1/traces").uri()); + when(mockSignedRequest.headers()).thenReturn(Map.of()); + when(mockSigner.signRequest(PAYLOAD)).thenReturn(mockSignedRequest); + + final okhttp3.Request okHttpRequest = new Request.Builder() + .url("https://xray.us-west-2.amazonaws.com/v1/traces") + .build(); + + for (int statusCode : NON_RETRYABLE_STATUS_CODES) { + final String responseBodyText = "Non-retryable error from server"; + + final Response mockResponse = new Response.Builder() + .request(okHttpRequest) + .protocol(Protocol.HTTP_1_1) + .code(statusCode) + .message("Client Error") + .body(ResponseBody.create(responseBodyText, MediaType.get("application/json"))) + .build(); + + final Call mockCall = mock(Call.class); + when(mockCall.execute()).thenReturn(mockResponse); + when(mockHttpClient.newCall(any())).thenReturn(mockCall); + + assertDoesNotThrow(() -> target.send(PAYLOAD), "Should not throw on non-retryable status " + statusCode); + verify(mockHttpClient, times(1)).newCall(any()); + + reset(mockHttpClient); // reset between iterations + } + } + + @Test + void testSend_retryOnFailure_thenSuccess() throws IOException { + SdkHttpFullRequest mockSignedRequest = mock(SdkHttpFullRequest.class); + when(mockSignedRequest.getUri()).thenReturn(HttpUrl.get("https://xray.us-west-2.amazonaws.com/v1/traces").uri()); + when(mockSignedRequest.headers()).thenReturn(Map.of()); + + when(mockSigner.signRequest(PAYLOAD)).thenReturn(mockSignedRequest); + + Call mockCall1 = mock(Call.class); + Call mockCall2 = mock(Call.class); + + when(mockHttpClient.newCall(any())) + .thenReturn(mockCall1) + .thenReturn(mockCall2); + + when(mockCall1.execute()).thenThrow(new IOException("first attempt failed")); + Response successResponse = new Response.Builder() + .request(new Request.Builder().url("https://xray.us-west-2.amazonaws.com/v1/traces").build()) + .protocol(Protocol.HTTP_1_1) + .code(200) + .message("OK") + .body(ResponseBody.create("", MediaType.get("application/x-protobuf"))) + .build(); + when(mockCall2.execute()).thenReturn(successResponse); + + assertDoesNotThrow(() -> target.send(PAYLOAD)); + } + + @Test + void testSend_failsAfterAllRetries() throws IOException { + SdkHttpFullRequest mockSignedRequest = mock(SdkHttpFullRequest.class); + when(mockSignedRequest.getUri()).thenReturn(HttpUrl.get("https://xray.us-west-2.amazonaws.com/v1/traces").uri()); + when(mockSignedRequest.headers()).thenReturn(Map.of()); + + when(mockSigner.signRequest(PAYLOAD)).thenReturn(mockSignedRequest); + + Call mockCall = mock(Call.class); + when(mockHttpClient.newCall(any())).thenReturn(mockCall); + when(mockCall.execute()).thenThrow(new IOException("always fail")); + + assertThrows(RuntimeException.class, () -> target.send(PAYLOAD)); + } + + @Test + void testSend_throwsIOException_on500ResponseWithBody() throws IOException { + // Mock signed request + SdkHttpFullRequest sdkRequest = SdkHttpFullRequest.builder() + .method(software.amazon.awssdk.http.SdkHttpMethod.POST) + .uri(URI.create("https://xray.us-west-2.amazonaws.com/v1/traces")) + .putHeader("Content-Type", "application/x-protobuf") + .build(); + + when(mockSigner.signRequest(PAYLOAD)).thenReturn(sdkRequest); + + // Build actual OkHttp request (we need this to inject into the mocked response) + okhttp3.Request okHttpRequest = new Request.Builder() + .url(sdkRequest.getUri().toString()) + .build(); + + // Mock 500 response with body + Response mockResponse = new Response.Builder() + .code(500) + .message("Internal Server Error") + .request(okHttpRequest) + .protocol(Protocol.HTTP_1_1) + .body(ResponseBody.create(ERROR_BODY, MediaType.get("application/json"))) + .build(); + + var call = mock(okhttp3.Call.class); + when(mockHttpClient.newCall(any(Request.class))).thenReturn(call); + when(call.execute()).thenReturn(mockResponse); + + // Run test + assertThrows(RuntimeException.class, () -> target.send(PAYLOAD)); + } + + @Test + void testInterruptedExceptionDuringRetryThrowsRuntimeException() throws IOException, InterruptedException { + final SdkHttpFullRequest signedRequest = mock(SdkHttpFullRequest.class); + + when(mockSigner.signRequest(any())).thenReturn(signedRequest); + when(signedRequest.getUri()).thenReturn(URI.create("https://example.com")); + when(signedRequest.headers()).thenReturn(Map.of()); + + final Call mockCall = mock(Call.class); + when(mockCall.execute()).thenThrow(new IOException("boom")); + when(mockHttpClient.newCall(any())).thenReturn(mockCall); + + doThrow(new InterruptedException("interrupted")).when(mockSleeper).sleep(anyInt()); + + target = new OtlpHttpSender(mockConfig, mockSigner, mockHttpClient, mockSleeper); + + final RuntimeException thrown = assertThrows(RuntimeException.class, () -> + target.send(PAYLOAD) + ); + + assertTrue(thrown.getMessage().contains("Retry interrupted")); + assertTrue(Thread.currentThread().isInterrupted()); + } + + @Test + void testSend_partialSuccessResponse_logsWarning() throws IOException { + final ExportTraceServiceResponse responseProto = ExportTraceServiceResponse.newBuilder() + .setPartialSuccess(ExportTracePartialSuccess.newBuilder() + .setRejectedSpans(5) + .setErrorMessage("Some spans were rejected due to invalid format") + .build()) + .build(); + + final byte[] responseBytes = responseProto.toByteArray(); + + final SdkHttpFullRequest sdkRequest = SdkHttpFullRequest.builder() + .method(software.amazon.awssdk.http.SdkHttpMethod.POST) + .uri(URI.create("https://xray.us-west-2.amazonaws.com/v1/traces")) + .putHeader("Content-Type", "application/x-protobuf") + .build(); + + when(mockSigner.signRequest(PAYLOAD)).thenReturn(sdkRequest); + + final okhttp3.Request okHttpRequest = new Request.Builder() + .url(sdkRequest.getUri().toString()) + .build(); + + final Response mockResponse = new Response.Builder() + .request(okHttpRequest) + .protocol(Protocol.HTTP_1_1) + .code(200) + .message("OK") + .body(ResponseBody.create(responseBytes, MediaType.get("application/x-protobuf"))) + .build(); + + final Call mockCall = mock(Call.class); + when(mockHttpClient.newCall(any())).thenReturn(mockCall); + when(mockCall.execute()).thenReturn(mockResponse); + + assertDoesNotThrow(() -> target.send(PAYLOAD)); + } + + @Test + void testSend_partialSuccessEmpty_logsInfo() throws IOException { + final ExportTraceServiceResponse responseProto = ExportTraceServiceResponse.newBuilder() + .setPartialSuccess(ExportTracePartialSuccess.newBuilder() + .setRejectedSpans(0) + .setErrorMessage("") // Empty = no warning + .build()) + .build(); + + final byte[] responseBytes = responseProto.toByteArray(); + + final SdkHttpFullRequest sdkRequest = SdkHttpFullRequest.builder() + .method(software.amazon.awssdk.http.SdkHttpMethod.POST) + .uri(URI.create("https://xray.us-west-2.amazonaws.com/v1/traces")) + .putHeader("Content-Type", "application/x-protobuf") + .build(); + + when(mockSigner.signRequest(PAYLOAD)).thenReturn(sdkRequest); + + final okhttp3.Request okHttpRequest = new Request.Builder() + .url(sdkRequest.getUri().toString()) + .build(); + + final Response mockResponse = new Response.Builder() + .request(okHttpRequest) + .protocol(Protocol.HTTP_1_1) + .code(200) + .message("OK") + .body(ResponseBody.create(responseBytes, MediaType.get("application/x-protobuf"))) + .build(); + + final Call mockCall = mock(Call.class); + when(mockHttpClient.newCall(any())).thenReturn(mockCall); + when(mockCall.execute()).thenReturn(mockResponse); + + assertDoesNotThrow(() -> target.send(PAYLOAD)); + } + + @Test + void testSend_successResponseWithNonEmptyBodyNoPartialSuccess_logsInfo() throws IOException { + // Manually add an unknown field to make the serialized body non-empty + final UnknownFieldSet unknownFields = UnknownFieldSet.newBuilder() + .addField(123, Field.newBuilder().addVarint(42).build()) // dummy varint + .build(); + + final ExportTraceServiceResponse responseProto = ExportTraceServiceResponse.newBuilder() + .mergeUnknownFields(unknownFields) + .build(); + + final byte[] responseBytes = responseProto.toByteArray(); + assertTrue(responseBytes.length > 0); // sanity check + + final SdkHttpFullRequest sdkRequest = SdkHttpFullRequest.builder() + .method(software.amazon.awssdk.http.SdkHttpMethod.POST) + .uri(URI.create("https://xray.us-west-2.amazonaws.com/v1/traces")) + .putHeader("Content-Type", "application/x-protobuf") + .build(); + + when(mockSigner.signRequest(PAYLOAD)).thenReturn(sdkRequest); + + final okhttp3.Request okHttpRequest = new Request.Builder() + .url(sdkRequest.getUri().toString()) + .build(); + + final Response mockResponse = new Response.Builder() + .request(okHttpRequest) + .protocol(Protocol.HTTP_1_1) + .code(200) + .message("OK") + .body(ResponseBody.create(responseBytes, MediaType.get("application/x-protobuf"))) + .build(); + + final Call mockCall = mock(Call.class); + when(mockHttpClient.newCall(any())).thenReturn(mockCall); + when(mockCall.execute()).thenReturn(mockResponse); + + assertDoesNotThrow(() -> target.send(PAYLOAD)); + } + + @Test + void testSend_invalidProtoResponse_logsError() throws IOException { + // Build a response with invalid OTLP proto data (random bytes) + byte[] invalidProtoBytes = "this-is-not-valid-proto".getBytes(StandardCharsets.UTF_8); + + // Mock the signed request + final SdkHttpFullRequest sdkRequest = SdkHttpFullRequest.builder() + .method(software.amazon.awssdk.http.SdkHttpMethod.POST) + .uri(URI.create("https://xray.us-west-2.amazonaws.com/v1/traces")) + .putHeader("Content-Type", "application/x-protobuf") + .build(); + + when(mockSigner.signRequest(PAYLOAD)).thenReturn(sdkRequest); + + // Build OkHttp request to satisfy response builder + final okhttp3.Request okHttpRequest = new Request.Builder() + .url(sdkRequest.getUri().toString()) + .build(); + + // Create mock response with invalid proto payload + final Response mockResponse = new Response.Builder() + .request(okHttpRequest) + .protocol(Protocol.HTTP_1_1) + .code(200) + .message("OK") + .body(ResponseBody.create(invalidProtoBytes, MediaType.get("application/x-protobuf"))) + .build(); + + final Call mockCall = mock(Call.class); + when(mockHttpClient.newCall(any())).thenReturn(mockCall); + when(mockCall.execute()).thenReturn(mockResponse); + + // Should not throw, just logs error + assertDoesNotThrow(() -> target.send(PAYLOAD)); + } + + @Test + void testDefaultConstructorInitializesDefaults() { + target = new OtlpHttpSender(mockConfig); + + // Reflection to assert internal fields (not great, but useful for unit validation) + assertNotNull(getPrivateField(target, "signer")); + assertNotNull(getPrivateField(target, "httpClient")); + assertNotNull(getPrivateField(target, "sleeper")); + } + + private Object getPrivateField(Object instance, String fieldName) { + try { + var field = instance.getClass().getDeclaredField(fieldName); + field.setAccessible(true); + return field.get(instance); + } catch (Exception e) { + fail("Could not access field: " + fieldName); + return null; + } + } + +} diff --git a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/SigV4SignerTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4SignerTest.java similarity index 83% rename from data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/SigV4SignerTest.java rename to data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4SignerTest.java index 62285ed6f0..8fb69e754e 100644 --- a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/SigV4SignerTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4SignerTest.java @@ -2,17 +2,14 @@ * Copyright OpenSearch Contributors * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.dataprepper.plugins.sink.xrayotlp.http; - -import java.nio.charset.StandardCharsets; -import java.time.Instant; +package org.opensearch.dataprepper.plugins.sink.otlp.http; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.ArgumentMatcher; import org.mockito.MockedStatic; -import org.opensearch.dataprepper.plugins.sink.xrayotlp.configuration.XRayOTLPSinkConfig; +import org.opensearch.dataprepper.plugins.sink.otlp.configuration.OtlpSinkConfig; import software.amazon.awssdk.http.SdkHttpFullRequest; import software.amazon.awssdk.regions.Region; import software.amazon.awssdk.regions.providers.DefaultAwsRegionProviderChain; @@ -22,6 +19,9 @@ import software.amazon.awssdk.services.sts.model.AssumeRoleResponse; import software.amazon.awssdk.services.sts.model.Credentials; +import java.nio.charset.StandardCharsets; +import java.time.Instant; + import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -35,19 +35,44 @@ class SigV4SignerTest { private static final byte[] PAYLOAD = "test-payload".getBytes(StandardCharsets.UTF_8); - private XRayOTLPSinkConfig mockXrayConfig; + private OtlpSinkConfig mockXrayConfig; private SigV4Signer target; @BeforeEach void setUp(){ - mockXrayConfig = mock(XRayOTLPSinkConfig.class); + System.setProperty("aws.accessKeyId", "dummy"); + System.setProperty("aws.secretAccessKey", "dummy"); + + mockXrayConfig = mock(OtlpSinkConfig.class); } @AfterEach void cleanUp() { + System.clearProperty("aws.accessKeyId"); + System.clearProperty("aws.secretAccessKey"); System.clearProperty("aws.region"); } + @Test + void testSignRequest_withInputEndpoint_whenEndpointIsSet() { + // setup + final String endpoint = "https://performance.us-west-2.xray.cloudwatch.aws.dev/v1/traces"; + System.setProperty("aws.region", Region.US_WEST_2.toString()); + when(mockXrayConfig.getAwsRegion()).thenReturn(null); + when(mockXrayConfig.getEndpoint()).thenReturn(endpoint); + target = new SigV4Signer(mockXrayConfig); + + // run + final SdkHttpFullRequest signedRequest = target.signRequest(PAYLOAD); + + // assert + assertNotNull(signedRequest); + assertEquals("POST", signedRequest.method().name()); + assertTrue(signedRequest.headers().containsKey("Authorization")); + assertEquals("application/x-protobuf", signedRequest.firstMatchingHeader("Content-Type").orElse(null)); + assertTrue(signedRequest.getUri().toString().contains(endpoint)); + } + @Test void testSignRequest_withFallbackRegion_whenRegionNotSet() { // setup diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeperTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeperTest.java new file mode 100644 index 0000000000..9da7148267 --- /dev/null +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeperTest.java @@ -0,0 +1,50 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.sink.otlp.http; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.fail; + +class ThreadSleeperTest { + private ThreadSleeper target; + + @BeforeEach + void setUp() { + target = new ThreadSleeper(); + } + + @Test + void testSleepDoesNotThrowWhenNotInterrupted() { + try { + target.sleep(1); + } catch (InterruptedException e) { + fail("Sleep was interrupted unexpectedly"); + } + } + + @Test + void testSleepThrowsInterruptedExceptionIfThreadInterrupted() { + Thread thread = new Thread(() -> { + try { + Thread.currentThread().interrupt(); + target.sleep(10); + fail("Expected InterruptedException"); + } catch (InterruptedException e) { + assertTrue(Thread.currentThread().isInterrupted(), "Thread should remain interrupted"); + } + }); + + thread.start(); + try { + thread.join(); + } catch (InterruptedException e) { + fail("Test thread was interrupted"); + } + } +} diff --git a/data-prepper-plugins/xray-otlp-sink/src/test/resources/data-prepper-config.yaml b/data-prepper-plugins/otlp-sink/src/test/resources/data-prepper-config.yaml similarity index 100% rename from data-prepper-plugins/xray-otlp-sink/src/test/resources/data-prepper-config.yaml rename to data-prepper-plugins/otlp-sink/src/test/resources/data-prepper-config.yaml diff --git a/data-prepper-plugins/xray-otlp-sink/src/test/resources/pipelines.yaml b/data-prepper-plugins/otlp-sink/src/test/resources/pipelines.yaml similarity index 56% rename from data-prepper-plugins/xray-otlp-sink/src/test/resources/pipelines.yaml rename to data-prepper-plugins/otlp-sink/src/test/resources/pipelines.yaml index dc0015e3de..51c1135323 100644 --- a/data-prepper-plugins/xray-otlp-sink/src/test/resources/pipelines.yaml +++ b/data-prepper-plugins/otlp-sink/src/test/resources/pipelines.yaml @@ -2,12 +2,12 @@ otel_trace_source_to_xray: source: otel_trace_source: ssl: false - - sink: - - xray_otlp_sink: - region: us-west-2 + port: 21890 buffer: bounded_blocking: - buffer_size: 10 - batch_size: 5 + buffer_size: 1000000 + batch_size: 125000 + + sink: + - otlp: \ No newline at end of file diff --git a/data-prepper-plugins/xray-otlp-sink/src/test/resources/sample-trace.json b/data-prepper-plugins/otlp-sink/src/test/resources/sample-trace.json similarity index 100% rename from data-prepper-plugins/xray-otlp-sink/src/test/resources/sample-trace.json rename to data-prepper-plugins/otlp-sink/src/test/resources/sample-trace.json diff --git a/data-prepper-plugins/xray-otlp-sink/src/test/resources/test-span-event.json b/data-prepper-plugins/otlp-sink/src/test/resources/test-span-event.json similarity index 100% rename from data-prepper-plugins/xray-otlp-sink/src/test/resources/test-span-event.json rename to data-prepper-plugins/otlp-sink/src/test/resources/test-span-event.json diff --git a/data-prepper-plugins/xray-otlp-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkIT.java b/data-prepper-plugins/xray-otlp-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkIT.java deleted file mode 100644 index 421635c4d9..0000000000 --- a/data-prepper-plugins/xray-otlp-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkIT.java +++ /dev/null @@ -1,38 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.sink.xrayotlp; - -import org.junit.jupiter.api.Test; -import org.opensearch.dataprepper.model.record.Record; -import org.opensearch.dataprepper.model.trace.JacksonStandardSpan; -import org.opensearch.dataprepper.model.trace.Span; - -import java.time.Instant; -import java.util.Collections; -import java.util.List; - -class XRayOTLPSinkIT { - @Test - void testSinkProcessesHardcodedSpan() { - final Span testSpan = JacksonStandardSpan.builder() - .withTraceId("abc123") - .withSpanId("def456") - .withParentSpanId("parent-testSpan-id") - .withName("my-test-testSpan") - .withStartTime(String.valueOf(Instant.now())) - .withEndTime(String.valueOf(Instant.now().plusMillis(10))) - .withAttributes(Collections.emptyMap()) - .withKind("test") - .build(); - - final Record record = new Record<>(testSpan); - final XRayOTLPSink sink = new XRayOTLPSink(); - - sink.initialize(); - sink.output(List.of(record)); - sink.shutdown(); - } -} diff --git a/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSink.java b/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSink.java deleted file mode 100644 index a5e86f1197..0000000000 --- a/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSink.java +++ /dev/null @@ -1,120 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.sink.xrayotlp; - -import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; -import org.opensearch.dataprepper.model.codec.OutputCodec; -import org.opensearch.dataprepper.model.record.Record; -import org.opensearch.dataprepper.model.sink.Sink; -import org.opensearch.dataprepper.model.trace.Span; -import org.opensearch.dataprepper.plugins.otel.codec.OtlpTraceOutputCodec; -import org.opensearch.dataprepper.plugins.sink.xrayotlp.configuration.XRayOTLPSinkConfig; -import org.opensearch.dataprepper.plugins.sink.xrayotlp.http.SigV4Signer; -import org.opensearch.dataprepper.plugins.sink.xrayotlp.http.XRayOtlpHttpSender; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; -import software.amazon.awssdk.http.apache.ApacheHttpClient; - -import java.io.ByteArrayOutputStream; -import java.util.Collection; - -/** - * A Data Prepper Sink plugin that forwards spans to AWS X-Ray OTLP endpoint using - * OTLP Protobuf encoding and AWS SigV4 authentication. - */ -@DataPrepperPlugin( - name = "xray_otlp_sink", - pluginType = Sink.class -) -public class XRayOTLPSink implements Sink> { - - private static final Logger LOG = LoggerFactory.getLogger(XRayOTLPSink.class); - - private final OutputCodec codec; - private final XRayOtlpHttpSender httpSender; - - /** - * Default constructor used by Data Prepper. Initializes codec and HTTP sender - * using default ApacheHttpClient and basic configuration. - */ - public XRayOTLPSink() { - this.codec = new OtlpTraceOutputCodec(); - final XRayOTLPSinkConfig config = new XRayOTLPSinkConfig(); // TODO: Load real config - final SigV4Signer signer = new SigV4Signer(config); - this.httpSender = new XRayOtlpHttpSender(signer, ApacheHttpClient.builder().build()); - } - - /** - * Constructor for unit testing with injected dependencies. - * - * @param codec the OutputCodec to encode spans - * @param httpSender the HTTP sender to transmit OTLP data - */ - public XRayOTLPSink(final OutputCodec codec, final XRayOtlpHttpSender httpSender) { - this.codec = codec; - this.httpSender = httpSender; - } - - /** - * Initializes the sink. Called once during pipeline startup. - */ - @Override - public void initialize() { - // TODO: Initialize AWS X-Ray client - LOG.info("Initialized XRay OTLP Sink"); - } - - /** - * Processes a batch of spans and sends them to the AWS X-Ray OTLP endpoint. - * - * @param records a collection of span records - */ - @Override - public void output(final Collection> records) { - if (records == null || records.isEmpty()) { - return; - } - - try (ByteArrayOutputStream out = new ByteArrayOutputStream()) { - for (final Record record : records) { - codec.writeEvent(record.getData(), out); - } - httpSender.send(out.toByteArray()); - } catch (final Exception e) { - LOG.error("Failed to process span records", e); - } - } - - /** - * Indicates whether this sink is ready to receive data. - * - * @return true if the sink is ready - */ - @Override - public boolean isReady() { - // TODO: Implement readiness logic - return true; - } - - /** - * Hook called during pipeline shutdown. - */ - @Override - public void shutdown() { - // TODO: Clean up resources - httpSender.close(); - } - - /** - * Updates internal latency metrics using the received records. - * - * @param events Collection of records used for latency tracking. - */ - @Override - public void updateLatencyMetrics(final Collection> events) { - // TODO: Implement latency tracking with PluginMetrics - } -} diff --git a/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/XRayOTLPSinkConfig.java b/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/XRayOTLPSinkConfig.java deleted file mode 100644 index 541e630699..0000000000 --- a/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/XRayOTLPSinkConfig.java +++ /dev/null @@ -1,50 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ -package org.opensearch.dataprepper.plugins.sink.xrayotlp.configuration; - -import com.fasterxml.jackson.annotation.JsonProperty; -import jakarta.validation.Valid; -import jakarta.validation.constraints.NotNull; -import lombok.AllArgsConstructor; -import lombok.Builder; -import lombok.Getter; -import lombok.NoArgsConstructor; -import software.amazon.awssdk.regions.Region; - -/** - * Configuration class for the X-Ray OTLP sink plugin. - * This class defines the configuration options available when setting up - * the X-Ray OTLP sink in Data Prepper pipelines. - * This class will be automatically wired by Data-Prepper; the Builder is only for testing. - * - * @since 2.6 - */ -@Getter -@Builder -@NoArgsConstructor -@AllArgsConstructor -public class XRayOTLPSinkConfig { - /** - * AWS configuration for X-Ray access. - * Contains authentication and region settings required for AWS X-Ray service. - * This is a required configuration and must be valid. - */ - @JsonProperty("aws") - @NotNull - @Valid - private AwsAuthenticationConfiguration awsAuthenticationConfiguration; - - public Region getAwsRegion() { - return awsAuthenticationConfiguration.getAwsRegion(); - } - - public String getStsRoleArn() { - return awsAuthenticationConfiguration.getAwsStsRoleArn(); - } - - public String getStsExternalId() { - return awsAuthenticationConfiguration.getAwsStsExternalId(); - } -} diff --git a/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/XRayOtlpHttpSender.java b/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/XRayOtlpHttpSender.java deleted file mode 100644 index 8388616abd..0000000000 --- a/data-prepper-plugins/xray-otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/XRayOtlpHttpSender.java +++ /dev/null @@ -1,73 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.sink.xrayotlp.http; - -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; -import software.amazon.awssdk.http.HttpExecuteRequest; -import software.amazon.awssdk.http.HttpExecuteResponse; -import software.amazon.awssdk.http.SdkHttpClient; -import software.amazon.awssdk.http.SdkHttpFullRequest; - -import java.io.ByteArrayInputStream; - -/** - * Responsible for sending signed OTLP Protobuf trace data to AWS X-Ray OTLP endpoint. - * This class uses AWS SDK's HTTP client to send the request signed using AWS SigV4. - */ -public class XRayOtlpHttpSender implements AutoCloseable { - - private static final Logger LOG = LoggerFactory.getLogger(XRayOtlpHttpSender.class); - - private final SigV4Signer signer; - private final SdkHttpClient httpClient; - - /** - * Constructs a new OtlpHttpSender. - * - * @param signer The SigV4Signer used to sign HTTP requests. - * @param httpClient The AWS SDK HTTP client used to send signed requests. - */ - public XRayOtlpHttpSender(final SigV4Signer signer, final SdkHttpClient httpClient) { - this.signer = signer; - this.httpClient = httpClient; - } - - /** - * Signs and sends the given OTLP-encoded span payload to the X-Ray OTLP endpoint. - * - * @param payload The OTLP Protobuf payload to send. - */ - public void send(final byte[] payload) { - try { - final SdkHttpFullRequest signedRequest = signer.signRequest(payload); - - final HttpExecuteRequest httpRequest = HttpExecuteRequest.builder() - .request(signedRequest) - .contentStreamProvider(() -> new ByteArrayInputStream(payload)) - .build(); - - final HttpExecuteResponse response = httpClient.prepareRequest(httpRequest).call(); - final int status = response.httpResponse().statusCode(); - - if (status >= 200 && status < 300) { - LOG.info("Successfully sent OTLP data to AWS X-Ray. Status: {}", status); - } else { - LOG.warn("Failed to send OTLP data to AWS X-Ray. Status: {}", status); - } - } catch (final Exception e) { - LOG.error("Error sending OTLP data to AWS X-Ray", e); - } - } - - /** - * Closes the underlying HTTP client to release resources. - */ - @Override - public void close() { - httpClient.close(); - } -} diff --git a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkTest.java b/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkTest.java deleted file mode 100644 index 5f3ec01ea8..0000000000 --- a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/XRayOTLPSinkTest.java +++ /dev/null @@ -1,154 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.sink.xrayotlp; - -import com.fasterxml.jackson.core.type.TypeReference; -import com.fasterxml.jackson.databind.ObjectMapper; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.opensearch.dataprepper.model.codec.OutputCodec; -import org.opensearch.dataprepper.model.record.Record; -import org.opensearch.dataprepper.model.trace.DefaultTraceGroupFields; -import org.opensearch.dataprepper.model.trace.JacksonSpan; -import org.opensearch.dataprepper.model.trace.Span; -import org.opensearch.dataprepper.plugins.sink.xrayotlp.http.XRayOtlpHttpSender; - -import java.io.InputStream; -import java.io.OutputStream; -import java.util.Collections; -import java.util.Map; -import java.util.Objects; - -import static org.junit.jupiter.api.Assertions.*; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.*; - -class XRayOTLPSinkTest { - private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); - private static final String TEST_SPAN_EVENT_JSON_FILE = "test-span-event.json"; - - private OutputCodec mockCodec; - private XRayOtlpHttpSender mockSender; - private XRayOTLPSink sink; - - @BeforeEach - void setUp() { - mockCodec = mock(OutputCodec.class); - mockSender = mock(XRayOtlpHttpSender.class); - sink = new XRayOTLPSink(mockCodec, mockSender); - } - - private Span buildSpanFromTestFile(String fileName, String traceIdOverride) { - try (InputStream inputStream = Objects.requireNonNull( - getClass().getClassLoader().getResourceAsStream(fileName))) { - - final Map spanMap = OBJECT_MAPPER.readValue(inputStream, new TypeReference<>() {}); - final JacksonSpan.Builder builder = JacksonSpan.builder() - .withTraceId(traceIdOverride != null ? traceIdOverride : (String) spanMap.get("traceId")) - .withSpanId((String) spanMap.get("spanId")) - .withParentSpanId((String) spanMap.get("parentSpanId")) - .withTraceState((String) spanMap.get("traceState")) - .withName((String) spanMap.get("name")) - .withKind((String) spanMap.get("kind")) - .withDurationInNanos(((Number) spanMap.get("durationInNanos")).longValue()) - .withStartTime((String) spanMap.get("startTime")) - .withEndTime((String) spanMap.get("endTime")) - .withTraceGroup((String) spanMap.get("traceGroup")); - - final Map traceGroupFieldsMap = (Map) spanMap.get("traceGroupFields"); - if (traceGroupFieldsMap != null) { - builder.withTraceGroupFields(DefaultTraceGroupFields.builder() - .withStatusCode((Integer) traceGroupFieldsMap.getOrDefault("statusCode", 0)) - .withEndTime((String) spanMap.get("endTime")) - .withDurationInNanos(((Number) spanMap.get("durationInNanos")).longValue()) - .build()); - } - - return builder.build(); - } catch (Exception e) { - throw new RuntimeException("Failed to load span from file", e); - } - } - - @Test - void testOutput_sendsDataToXRay_onSuccessResponse() throws Exception { - final Span span = buildSpanFromTestFile(TEST_SPAN_EVENT_JSON_FILE, "bad-trace-id"); - Record record = new Record<>(span); - - doAnswer(invocation -> { - OutputStream out = invocation.getArgument(1); - out.write("dummy-otlp-payload".getBytes()); - return null; - }).when(mockCodec).writeEvent(eq(span), any()); - - doNothing().when(mockSender).send(eq("dummy-otlp-payload".getBytes())); - - sink.output(Collections.singletonList(record)); - - verify(mockCodec).writeEvent(eq(span), any()); - verify(mockSender).send(eq("dummy-otlp-payload".getBytes())); - } - - @Test - void testOutput_handlesException_gracefully() throws Exception { - Span span = mock(Span.class); - Record record = new Record<>(span); - - doThrow(new RuntimeException("codec error")).when(mockCodec).writeEvent(eq(span), any()); - - sink.output(Collections.singletonList(record)); - - verify(mockCodec).writeEvent(eq(span), any()); - verifyNoInteractions(mockSender); - } - - @Test - void testOutput_senderThrowsException_isHandledGracefully() throws Exception { - final Span span = buildSpanFromTestFile(TEST_SPAN_EVENT_JSON_FILE, "bad-trace-id"); - Record record = new Record<>(span); - - doAnswer(invocation -> { - OutputStream out = invocation.getArgument(1); - out.write("dummy-otlp-payload".getBytes()); - return null; - }).when(mockCodec).writeEvent(eq(span), any()); - - doThrow(new RuntimeException("Send failed")).when(mockSender).send(any()); - - assertDoesNotThrow(() -> sink.output(Collections.singletonList(record))); - - verify(mockCodec).writeEvent(eq(span), any()); - verify(mockSender).send(any()); - } - - @Test - void testOutput_withNullRecordList_doesNothing() { - assertDoesNotThrow(() -> sink.output(null)); - verifyNoInteractions(mockCodec, mockSender); - } - - @Test - void testOutput_withEmptyRecordList_doesNothing() { - assertDoesNotThrow(() -> sink.output(Collections.emptyList())); - verifyNoInteractions(mockCodec, mockSender); - } - - @Test - void testInitialize_doesNotThrow() { - assertDoesNotThrow(() -> sink.initialize()); - } - - @Test - void testIsReady_returnsTrue() { - assertTrue(sink.isReady()); - } - - @Test - void testShutdown_doesNotThrow() { - assertDoesNotThrow(() -> sink.shutdown()); - } -} diff --git a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/AwsAuthenticationConfigurationTest.java b/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/AwsAuthenticationConfigurationTest.java deleted file mode 100644 index 2f46d4de03..0000000000 --- a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/AwsAuthenticationConfigurationTest.java +++ /dev/null @@ -1,41 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ -package org.opensearch.dataprepper.plugins.sink.xrayotlp.configuration; - -import org.junit.jupiter.api.Test; -import software.amazon.awssdk.regions.Region; - -import static org.hamcrest.CoreMatchers.equalTo; -import static org.hamcrest.CoreMatchers.notNullValue; -import static org.hamcrest.CoreMatchers.nullValue; -import static org.hamcrest.MatcherAssert.assertThat; - -class AwsAuthenticationConfigurationTest { - - @Test - void testGetRegion_whenRegionIsNull_returnNull() { - final AwsAuthenticationConfiguration config = AwsAuthenticationConfiguration.builder().build(); - - assertThat(config.getAwsRegion(), nullValue()); - } - - @Test - void testAwsAuthenticationConfigurationFields() { - final String expectedRegion = "us-west-2"; - final String expectedRoleArn = "arn:aws:iam::123456789012:role/MyRole"; - final String expectedExternalId = "myExternalId"; - - final AwsAuthenticationConfiguration config = AwsAuthenticationConfiguration.builder() - .awsRegion(expectedRegion) - .awsStsRoleArn(expectedRoleArn) - .awsStsExternalId(expectedExternalId) - .build(); - - assertThat(config.getAwsRegion(), notNullValue()); - assertThat(config.getAwsRegion(), equalTo(Region.US_WEST_2)); - assertThat(config.getAwsStsRoleArn(), equalTo(expectedRoleArn)); - assertThat(config.getAwsStsExternalId(), equalTo(expectedExternalId)); - } -} diff --git a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/XRayOTLPSinkConfigTest.java b/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/XRayOTLPSinkConfigTest.java deleted file mode 100644 index 1993a88cd4..0000000000 --- a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/configuration/XRayOTLPSinkConfigTest.java +++ /dev/null @@ -1,50 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ -package org.opensearch.dataprepper.plugins.sink.xrayotlp.configuration; - -import org.junit.jupiter.api.Test; -import software.amazon.awssdk.regions.Region; - -import static org.hamcrest.CoreMatchers.equalTo; -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.notNullValue; -import static org.hamcrest.Matchers.nullValue; - -class XRayOTLPSinkConfigTest { - @Test - void testAwsAuthenticationConfiguration_withAllFields() { - final String expectedRegion = "us-west-2"; - final String expectedRoleArn = "arn:aws:iam::123456789012:role/MyRole"; - final String expectedExternalId = "external-id-123"; - - AwsAuthenticationConfiguration awsConfig = AwsAuthenticationConfiguration.builder() - .awsRegion(expectedRegion) - .awsStsRoleArn(expectedRoleArn) - .awsStsExternalId(expectedExternalId) - .build(); - - XRayOTLPSinkConfig config = XRayOTLPSinkConfig.builder() - .awsAuthenticationConfiguration(awsConfig) - .build(); - - assertThat(config.getAwsRegion(), equalTo(Region.of(expectedRegion))); - assertThat(config.getStsRoleArn(), equalTo(expectedRoleArn)); - assertThat(config.getStsExternalId(), equalTo(expectedExternalId)); - } - - @Test - void testDefaultConstructorAndSetters() { - final XRayOTLPSinkConfig config = new XRayOTLPSinkConfig(); - assertThat(config, notNullValue()); - } - - @Test - void testBuilder_withNullAwsAuthConfig() { - XRayOTLPSinkConfig config = XRayOTLPSinkConfig.builder() - .awsAuthenticationConfiguration(null) - .build(); - assertThat(config.getAwsAuthenticationConfiguration(), nullValue()); - } -} diff --git a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/XRayOtlpHttpSenderTest.java b/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/XRayOtlpHttpSenderTest.java deleted file mode 100644 index c7bc03e9ed..0000000000 --- a/data-prepper-plugins/xray-otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/xrayotlp/http/XRayOtlpHttpSenderTest.java +++ /dev/null @@ -1,91 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.sink.xrayotlp.http; - -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import software.amazon.awssdk.http.ExecutableHttpRequest; -import software.amazon.awssdk.http.HttpExecuteRequest; -import software.amazon.awssdk.http.HttpExecuteResponse; -import software.amazon.awssdk.http.SdkHttpClient; -import software.amazon.awssdk.http.SdkHttpFullRequest; -import software.amazon.awssdk.http.SdkHttpResponse; - -import java.io.IOException; - -import static org.junit.jupiter.api.Assertions.*; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.*; - -class XRayOtlpHttpSenderTest { - - private SigV4Signer mockSigner; - private SdkHttpClient mockHttpClient; - private XRayOtlpHttpSender sender; - - @BeforeEach - void setUp() { - mockSigner = mock(SigV4Signer.class); - mockHttpClient = mock(SdkHttpClient.class); - sender = new XRayOtlpHttpSender(mockSigner, mockHttpClient); - } - - @Test - void testSend_successfulRequest_logsInfo() throws Exception { - byte[] payload = "test-payload".getBytes(); - - SdkHttpFullRequest signedRequest = mock(SdkHttpFullRequest.class); - when(mockSigner.signRequest(payload)).thenReturn(signedRequest); - - HttpExecuteResponse mockResponse = mock(HttpExecuteResponse.class); - when(mockResponse.httpResponse()).thenReturn(SdkHttpResponse.builder().statusCode(200).build()); - - ExecutableHttpRequest executableRequest = mock(ExecutableHttpRequest.class); - when(executableRequest.call()).thenReturn(mockResponse); - when(mockHttpClient.prepareRequest(any(HttpExecuteRequest.class))).thenReturn(executableRequest); - - assertDoesNotThrow(() -> sender.send(payload)); - - verify(mockSigner).signRequest(payload); - verify(mockHttpClient).prepareRequest(any(HttpExecuteRequest.class)); - } - - @Test - void testSend_httpError_logsWarning() throws Exception { - byte[] payload = "test-payload".getBytes(); - - SdkHttpFullRequest signedRequest = mock(SdkHttpFullRequest.class); - when(mockSigner.signRequest(payload)).thenReturn(signedRequest); - - HttpExecuteResponse mockResponse = mock(HttpExecuteResponse.class); - when(mockResponse.httpResponse()).thenReturn(SdkHttpResponse.builder().statusCode(500).build()); - - ExecutableHttpRequest executableRequest = mock(ExecutableHttpRequest.class); - when(executableRequest.call()).thenReturn(mockResponse); - when(mockHttpClient.prepareRequest(any(HttpExecuteRequest.class))).thenReturn(executableRequest); - - sender.send(payload); - - verify(mockSigner).signRequest(payload); - verify(mockHttpClient).prepareRequest(any(HttpExecuteRequest.class)); - } - - @Test - void testSend_exceptionDuringSend_logsError() throws Exception { - byte[] payload = "test-payload".getBytes(); - - when(mockSigner.signRequest(payload)).thenThrow(new RuntimeException("signing failed")); - - assertDoesNotThrow(() -> sender.send(payload)); - verify(mockSigner).signRequest(payload); - } - - @Test - void testClose_closesHttpClient() throws IOException { - sender.close(); - verify(mockHttpClient).close(); - } -} diff --git a/settings.gradle b/settings.gradle index 5b36fa86d8..a9b70cfbec 100644 --- a/settings.gradle +++ b/settings.gradle @@ -196,4 +196,4 @@ include 'data-prepper-plugins:saas-source-plugins:jira-source' include 'data-prepper-plugins:saas-source-plugins:confluence-source' include 'data-prepper-plugins:saas-source-plugins:atlassian-commons' include 'data-prepper-plugins:saas-source-plugins:crowdstrike-source' -include 'data-prepper-plugins:xray-otlp-sink' +include 'data-prepper-plugins:otlp-sink' From 6e732ba2759f973eeb8b3e072d837baa5cdea883 Mon Sep 17 00:00:00 2001 From: huyPham Date: Wed, 16 Apr 2025 08:16:56 -0700 Subject: [PATCH 08/23] feature: Add PluginMetrics instrumentation for OTLP delivery and throughput (#11) - Introduced OtlpSinkMetrics class to track key metrics: - recordsIn, recordsOut, droppedRecords - payloadSize, deliveryLatency - retriesCount, rejectedSpansCount, errorsCount - response code categorization (e.g., http_2xx_responses) - Added instrumentation in OtlpSink and OtlpHttpSender - Metrics follow the required format: pipeline-name_sink-name_METRIC_NAME - Added unit tests Signed-off-by: huy pham --- data-prepper-plugins/otlp-sink/build.gradle | 1 + .../plugins/sink/otlp/OtlpSinkIT.java | 14 ++- .../plugins/sink/otlp/OtlpSink.java | 48 +++++++-- .../sink/otlp/http/OtlpHttpSender.java | 19 +++- .../sink/otlp/metrics/OtlpSinkMetrics.java | 74 ++++++++++++++ .../plugins/sink/otlp/OtlpSinkTest.java | 51 +++++++++- .../sink/otlp/http/OtlpHttpSenderTest.java | 9 +- .../otlp/metrics/OtlpSinkMetricsTest.java | 97 +++++++++++++++++++ .../test/resources/data-prepper-config.yaml | 2 + .../src/test/resources/pipelines.yaml | 2 +- 10 files changed, 298 insertions(+), 19 deletions(-) create mode 100644 data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetrics.java create mode 100644 data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java diff --git a/data-prepper-plugins/otlp-sink/build.gradle b/data-prepper-plugins/otlp-sink/build.gradle index 49a94a05d4..6255275fa6 100644 --- a/data-prepper-plugins/otlp-sink/build.gradle +++ b/data-prepper-plugins/otlp-sink/build.gradle @@ -56,6 +56,7 @@ dependencies { // Data Prepper Projects implementation project(':data-prepper-api') implementation project(':data-prepper-plugins:otel-proto-common') + implementation 'io.micrometer:micrometer-core' // Unit Testing testImplementation 'org.junit.jupiter:junit-jupiter-api:5.9.1' diff --git a/data-prepper-plugins/otlp-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkIT.java b/data-prepper-plugins/otlp-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkIT.java index b3c3e66155..dbd6a6e2ad 100644 --- a/data-prepper-plugins/otlp-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkIT.java +++ b/data-prepper-plugins/otlp-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkIT.java @@ -5,9 +5,12 @@ package org.opensearch.dataprepper.plugins.sink.otlp; +import io.micrometer.core.instrument.Counter; +import io.micrometer.core.instrument.DistributionSummary; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.model.trace.JacksonStandardSpan; import org.opensearch.dataprepper.model.trace.Span; @@ -20,6 +23,7 @@ import java.util.Collections; import java.util.List; +import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -34,6 +38,7 @@ class OtlpSinkIT { private OtlpSinkConfig mockConfig; + private PluginMetrics mockPluginMetrics; private OtlpSink target; @BeforeEach @@ -46,7 +51,12 @@ void setUp() { when(mockConfig.getMaxRetries()).thenReturn(3); when(mockConfig.getBatchSize()).thenReturn(100); - target = new OtlpSink(mockConfig); + mockPluginMetrics = mock(PluginMetrics.class); + mockPluginMetrics = mock(PluginMetrics.class); + when(mockPluginMetrics.counter(anyString())).thenReturn(mock(Counter.class)); + when(mockPluginMetrics.summary(anyString())).thenReturn(mock(DistributionSummary.class)); + + target = new OtlpSink(mockConfig, mockPluginMetrics); } @AfterEach @@ -75,7 +85,7 @@ void testSinkProcessesHardcodedSpan() { .build(); final Record record = new Record<>(testSpan); - final OtlpSink sink = new OtlpSink(mockConfig, mock(OTelProtoStandardCodec.OTelProtoEncoder.class), mock(OtlpHttpSender.class)); + final OtlpSink sink = new OtlpSink(mockConfig, mockPluginMetrics, mock(OTelProtoStandardCodec.OTelProtoEncoder.class), mock(OtlpHttpSender.class)); sink.initialize(); sink.output(List.of(record)); diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java index 007c861282..bb840a859a 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java @@ -9,6 +9,7 @@ import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceRequest; import io.opentelemetry.proto.trace.v1.ResourceSpans; import org.jetbrains.annotations.Nullable; +import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor; import org.opensearch.dataprepper.model.record.Record; @@ -17,10 +18,13 @@ import org.opensearch.dataprepper.plugins.otel.codec.OTelProtoStandardCodec; import org.opensearch.dataprepper.plugins.sink.otlp.configuration.OtlpSinkConfig; import org.opensearch.dataprepper.plugins.sink.otlp.http.OtlpHttpSender; +import org.opensearch.dataprepper.plugins.sink.otlp.metrics.OtlpSinkMetrics; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import javax.annotation.Nonnull; +import java.time.Duration; +import java.time.Instant; import java.util.ArrayList; import java.util.Collection; import java.util.List; @@ -42,6 +46,7 @@ public class OtlpSink implements Sink> { private final int batchSize; private final OtlpHttpSender httpSender; private final OTelProtoStandardCodec.OTelProtoEncoder encoder; + private final OtlpSinkMetrics sinkMetrics; /** * Constructor for the OTLPSink plugin. @@ -49,8 +54,8 @@ public class OtlpSink implements Sink> { * @param config the configuration for the OTLP sink */ @DataPrepperPluginConstructor - public OtlpSink(@Nonnull final OtlpSinkConfig config) { - this(config, null, null); + public OtlpSink(@Nonnull final OtlpSinkConfig config, @Nonnull final PluginMetrics pluginMetrics) { + this(config, pluginMetrics, null, null); } /** @@ -61,12 +66,13 @@ public OtlpSink(@Nonnull final OtlpSinkConfig config) { * @param httpSender the HTTP sender to use */ @VisibleForTesting - OtlpSink(@Nonnull final OtlpSinkConfig config, final OTelProtoStandardCodec.OTelProtoEncoder encoder, final OtlpHttpSender httpSender) { + OtlpSink(@Nonnull final OtlpSinkConfig config, @Nonnull final PluginMetrics pluginMetrics, final OTelProtoStandardCodec.OTelProtoEncoder encoder, final OtlpHttpSender httpSender) { this.batchSize = config.getBatchSize(); + this.sinkMetrics = new OtlpSinkMetrics(pluginMetrics); if (encoder == null && httpSender == null) { this.encoder = new OTelProtoStandardCodec.OTelProtoEncoder(); - this.httpSender = new OtlpHttpSender(config); + this.httpSender = new OtlpHttpSender(config, sinkMetrics); } else { this.encoder = encoder; this.httpSender = httpSender; @@ -95,6 +101,8 @@ public void initialize() { */ @Override public void output(@Nonnull final Collection> records) { + sinkMetrics.incrementRecordsIn(records.size()); + final List> recordList = new ArrayList<>(records); for (int i = 0; i < recordList.size(); i += this.batchSize) { final int end = Math.min(i + this.batchSize, recordList.size()); @@ -115,10 +123,18 @@ public void output(@Nonnull final Collection> records) { final ExportTraceServiceRequest request = ExportTraceServiceRequest.newBuilder() .addAllResourceSpans(resourceSpans) .build(); - httpSender.send(request.toByteArray()); + final byte[] payload = request.toByteArray(); + sinkMetrics.incrementPayloadSize(payload.length); + + httpSender.send(payload); LOG.info("Finished processing {} spans.", resourceSpans.size()); + + sinkMetrics.incrementRecordsOut(batch.size()); } catch (final RuntimeException e) { LOG.error("Unexpected error processing OTLP span batch: {}", e.getMessage(), e); + + sinkMetrics.incrementDroppedRecords(batch.size()); + sinkMetrics.incrementErrorsCount(); } } } @@ -129,6 +145,7 @@ private ResourceSpans getResourceSpans(final Span span) { return encoder.convertToResourceSpans(span); } catch (final Exception e) { LOG.warn("Failed to encode span with ID [{}], skipping.", span.getSpanId(), e); + sinkMetrics.incrementErrorsCount(); return null; } } @@ -153,12 +170,27 @@ public void shutdown() { } /** - * Updates internal latency metrics using the received records. + * Records the latency between when each span originally started (as specified in the span's start time) + * and when it was received by the sink (i.e., when this method is called). + *

+ * This measures end-to-end ingestion latency from the span's source to the sink. + * It does not include any processing or export latency within the sink itself. * - * @param events Collection of records used for latency tracking. + * @param events A collection of spans received by the sink. */ @Override public void updateLatencyMetrics(@Nonnull final Collection> events) { - // TODO: Implement latency tracking with PluginMetrics + final Instant now = Instant.now(); + + for (final Record record : events) { + try { + final Instant startTime = Instant.parse(record.getData().getStartTime()); + final long durationMillis = Duration.between(startTime, now).toMillis(); + sinkMetrics.recordDeliveryLatency(durationMillis); + } catch (final Exception e) { + LOG.warn("Failed to parse startTime: {}", record.getData().getStartTime(), e); + sinkMetrics.incrementErrorsCount(); + } + } } } diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java index 32f3f501ca..2e0bcf3248 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java @@ -12,6 +12,7 @@ import okhttp3.RequestBody; import okhttp3.Response; import org.opensearch.dataprepper.plugins.sink.otlp.configuration.OtlpSinkConfig; +import org.opensearch.dataprepper.plugins.sink.otlp.metrics.OtlpSinkMetrics; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import software.amazon.awssdk.http.SdkHttpFullRequest; @@ -41,14 +42,15 @@ public class OtlpHttpSender implements AutoCloseable { private final Sleeper sleeper; private final int maxRetries; private final List retryDelaysMs; + private final OtlpSinkMetrics sinkMetrics; /** * Constructor for the OtlpHttpSender. * Initializes the signer and HTTP client. * @param config The configuration for the OTLP sink plugin. */ - public OtlpHttpSender(@Nonnull final OtlpSinkConfig config) { - this(config, null, null, null); + public OtlpHttpSender(@Nonnull final OtlpSinkConfig config, @Nonnull final OtlpSinkMetrics sinkMetrics) { + this(config, sinkMetrics, null, null, null); } /** @@ -59,7 +61,8 @@ public OtlpHttpSender(@Nonnull final OtlpSinkConfig config) { * @param httpClient The OkHttpClient instance for making HTTP requests. */ @VisibleForTesting - OtlpHttpSender(@Nonnull final OtlpSinkConfig config, final SigV4Signer signer, final OkHttpClient httpClient, final Sleeper sleeper) { + OtlpHttpSender(@Nonnull final OtlpSinkConfig config, @Nonnull final OtlpSinkMetrics sinkMetrics, final SigV4Signer signer, final OkHttpClient httpClient, final Sleeper sleeper) { + this.sinkMetrics = sinkMetrics; this.signer = signer != null ? signer : new SigV4Signer(config); this.httpClient = httpClient != null ? httpClient : new OkHttpClient(); this.sleeper = sleeper != null ? sleeper : new ThreadSleeper(); @@ -108,7 +111,11 @@ public void send(@Nonnull final byte[] payload) { final Request request = requestBuilder.build(); + final long startTime = System.currentTimeMillis(); try (final Response response = httpClient.newCall(request).execute()) { + final long duration = System.currentTimeMillis() - startTime; + sinkMetrics.recordHttpLatency(duration); + handleResponse(response); return; } @@ -118,6 +125,7 @@ public void send(@Nonnull final byte[] payload) { final int retryIndex = Math.min(attempt, retryDelaysMs.size() - 1); final int delay = retryDelaysMs.get(retryIndex) + jitter; LOG.warn("Retrying after failure in attempt {}. Sleeping {}ms.", attempt + 1, delay, e); + sinkMetrics.incrementRetriesCount(); try { sleeper.sleep(delay); } catch (final InterruptedException ie) { @@ -141,6 +149,8 @@ public void send(@Nonnull final byte[] payload) { */ private void handleResponse(@Nonnull final Response response) throws IOException { final int status = response.code(); + sinkMetrics.recordResponseCode(status); + final byte[] responseBytes = response.body() != null ? response.body().bytes() : null; @@ -173,6 +183,8 @@ private void handleSuccessfulResponse(final byte[] responseBytes) { if (otlpResponse.hasPartialSuccess()) { final var partial = otlpResponse.getPartialSuccess(); final long rejectedSpans = partial.getRejectedSpans(); + sinkMetrics.incrementRejectedSpansCount(rejectedSpans); + final String errorMessage = partial.getErrorMessage(); if (rejectedSpans > 0 || !errorMessage.isEmpty()) { @@ -185,6 +197,7 @@ private void handleSuccessfulResponse(final byte[] responseBytes) { } } catch (final Exception e) { LOG.error("Could not parse OTLP response as ExportTraceServiceResponse: {}", e.getMessage()); + sinkMetrics.incrementErrorsCount(); } } diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetrics.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetrics.java new file mode 100644 index 0000000000..fb68066a80 --- /dev/null +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetrics.java @@ -0,0 +1,74 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.sink.otlp.metrics; + +import org.opensearch.dataprepper.metrics.PluginMetrics; + +import javax.annotation.Nonnull; + +/** + * Metrics class for the otlp-sink + */ +public class OtlpSinkMetrics { + + private final PluginMetrics pluginMetrics; + + /** + * Constructor for OtlpSinkMetrics + * + * @param pluginMetrics The plugin metrics instance + */ + public OtlpSinkMetrics(@Nonnull final PluginMetrics pluginMetrics) { + this.pluginMetrics = pluginMetrics; + } + + public void incrementRecordsIn(long count) { + pluginMetrics.counter("recordsIn").increment(count); + } + + public void incrementRecordsOut(long count) { + pluginMetrics.counter("recordsOut").increment(count); + } + + public void incrementDroppedRecords(long count) { + pluginMetrics.counter("droppedRecords").increment(count); + } + + public void incrementErrorsCount() { + pluginMetrics.counter("errorsCount").increment(1); + } + + public void incrementPayloadSize(long bytes) { + pluginMetrics.summary("payloadSize").record(bytes); + } + + public void recordDeliveryLatency(long durationMillis) { + pluginMetrics.summary("deliveryLatency").record(durationMillis); + } + + public void recordHttpLatency(long durationMillis) { + pluginMetrics.summary("httpLatency").record(durationMillis); + } + + public void incrementRetriesCount() { + pluginMetrics.counter("retriesCount").increment(1); + } + + public void incrementRejectedSpansCount(long count) { + pluginMetrics.counter("rejectedSpansCount").increment(count); + } + + /** + * Records the response code in the metrics. + * Group response codes by category: 2xx, 4xx, 5xx, etc. + * + * @param statusCode The HTTP response code. + */ + public void recordResponseCode(final int statusCode) { + String codeCategory = (statusCode / 100) + "xx"; + pluginMetrics.counter("http_" + codeCategory + "_responses").increment(); + } +} diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java index 8207a8e72f..753c91c6a2 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java @@ -5,11 +5,14 @@ package org.opensearch.dataprepper.plugins.sink.otlp; +import io.micrometer.core.instrument.Counter; +import io.micrometer.core.instrument.DistributionSummary; import io.opentelemetry.proto.trace.v1.ResourceSpans; import org.apache.commons.codec.DecoderException; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.model.trace.Span; import org.opensearch.dataprepper.plugins.otel.codec.OTelProtoStandardCodec; @@ -18,10 +21,14 @@ import software.amazon.awssdk.regions.Region; import java.io.UnsupportedEncodingException; +import java.time.Instant; import java.util.ArrayList; +import java.util.Collection; import java.util.List; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyDouble; +import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; @@ -35,6 +42,7 @@ class OtlpSinkTest { private OtlpSinkConfig mockConfig; private OTelProtoStandardCodec.OTelProtoEncoder mockEncoder; private OtlpHttpSender mockSender; + private PluginMetrics mockPluginMetrics; private OtlpSink target; @BeforeEach @@ -50,7 +58,12 @@ void setUp() { mockEncoder = mock(OTelProtoStandardCodec.OTelProtoEncoder.class); mockSender = mock(OtlpHttpSender.class); - target = new OtlpSink(mockConfig, mockEncoder, mockSender); + mockPluginMetrics = mock(PluginMetrics.class); + mockPluginMetrics = mock(PluginMetrics.class); + when(mockPluginMetrics.counter(anyString())).thenReturn(mock(Counter.class)); + when(mockPluginMetrics.summary(anyString())).thenReturn(mock(DistributionSummary.class)); + + target = new OtlpSink(mockConfig, mockPluginMetrics, mockEncoder, mockSender); } @AfterEach @@ -141,9 +154,43 @@ void testOutput_shouldNotSendEmptyBatch() throws DecoderException, UnsupportedEn verify(mockSender, never()).send(any()); } + @Test + void testUpdateLatencyMetrics_shouldRecordLatency() { + Span mockSpan = mock(Span.class); + String startTime = Instant.now().minusSeconds(5).toString(); + when(mockSpan.getStartTime()).thenReturn(startTime); + + Record mockRecord = mock(Record.class); + when(mockRecord.getData()).thenReturn(mockSpan); + + Collection> events = List.of(mockRecord); + + target.updateLatencyMetrics(events); + + verify(mockPluginMetrics.summary("deliveryLatency"), times(1)).record(any(Double.class)); + } + + @Test + void testUpdateLatencyMetrics_shouldHandleInvalidStartTime() { + // Mock a Span with an invalid start time string + final Span badSpan = mock(Span.class); + when(badSpan.getStartTime()).thenReturn("invalid-timestamp"); + + final Record badRecord = mock(Record.class); + when(badRecord.getData()).thenReturn(badSpan); + + final List> records = List.of(badRecord); + + target.updateLatencyMetrics(records); + + // Ensure it attempted to parse and failed gracefully + verify(mockPluginMetrics.summary("deliveryLatency"), never()).record(anyDouble()); + verify(mockPluginMetrics.counter("errorsCount")).increment(1); + } + @Test void testConstructor_withOnlyConfig_shouldInitializeWithoutException() { - OtlpSink sink = new OtlpSink(mockConfig); + OtlpSink sink = new OtlpSink(mockConfig, mockPluginMetrics); sink.initialize(); sink.shutdown(); diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java index 1d63eb1d58..29047bd4f1 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java @@ -21,6 +21,7 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.opensearch.dataprepper.plugins.sink.otlp.configuration.OtlpSinkConfig; +import org.opensearch.dataprepper.plugins.sink.otlp.metrics.OtlpSinkMetrics; import software.amazon.awssdk.http.SdkHttpFullRequest; import software.amazon.awssdk.regions.Region; @@ -54,6 +55,7 @@ class OtlpHttpSenderTest { private SigV4Signer mockSigner; private OkHttpClient mockHttpClient; private Sleeper mockSleeper; + private OtlpSinkMetrics mockSinkMetrics; private OtlpHttpSender target; @BeforeEach @@ -68,8 +70,9 @@ void setUp() { mockSigner = mock(SigV4Signer.class); mockHttpClient = mock(OkHttpClient.class); mockSleeper = mock(Sleeper.class); + mockSinkMetrics = mock(OtlpSinkMetrics.class); - target = new OtlpHttpSender(mockConfig, mockSigner, mockHttpClient, mockSleeper); + target = new OtlpHttpSender(mockConfig, mockSinkMetrics, mockSigner, mockHttpClient, mockSleeper); } @AfterEach @@ -225,7 +228,7 @@ void testInterruptedExceptionDuringRetryThrowsRuntimeException() throws IOExcept doThrow(new InterruptedException("interrupted")).when(mockSleeper).sleep(anyInt()); - target = new OtlpHttpSender(mockConfig, mockSigner, mockHttpClient, mockSleeper); + target = new OtlpHttpSender(mockConfig, mockSinkMetrics, mockSigner, mockHttpClient, mockSleeper); final RuntimeException thrown = assertThrows(RuntimeException.class, () -> target.send(PAYLOAD) @@ -390,7 +393,7 @@ void testSend_invalidProtoResponse_logsError() throws IOException { @Test void testDefaultConstructorInitializesDefaults() { - target = new OtlpHttpSender(mockConfig); + target = new OtlpHttpSender(mockConfig, mockSinkMetrics); // Reflection to assert internal fields (not great, but useful for unit validation) assertNotNull(getPrivateField(target, "signer")); diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java new file mode 100644 index 0000000000..8fbda2c909 --- /dev/null +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java @@ -0,0 +1,97 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.sink.otlp.metrics; + +import io.micrometer.core.instrument.Counter; +import io.micrometer.core.instrument.DistributionSummary; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.metrics.PluginMetrics; + +import static org.mockito.Mockito.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +class OtlpSinkMetricsTest { + + private PluginMetrics pluginMetrics; + private Counter counterMock; + private DistributionSummary summaryMock; + private OtlpSinkMetrics sinkMetrics; + + @BeforeEach + void setUp() { + pluginMetrics = mock(PluginMetrics.class); + counterMock = mock(Counter.class); + summaryMock = mock(DistributionSummary.class); + + when(pluginMetrics.counter(anyString())).thenReturn(counterMock); + when(pluginMetrics.summary(anyString())).thenReturn(summaryMock); + + sinkMetrics = new OtlpSinkMetrics(pluginMetrics); + } + + @Test + void testIncrementRecordsIn() { + sinkMetrics.incrementRecordsIn(3); + verify(counterMock).increment(3); + } + + @Test + void testIncrementRecordsOut() { + sinkMetrics.incrementRecordsOut(2); + verify(counterMock).increment(2); + } + + @Test + void testIncrementDroppedRecords() { + sinkMetrics.incrementDroppedRecords(1); + verify(counterMock).increment(1); + } + + @Test + void testIncrementErrorsCount() { + sinkMetrics.incrementErrorsCount(); + verify(counterMock).increment(1); + } + + @Test + void testIncrementPayloadSize() { + sinkMetrics.incrementPayloadSize(1024); + verify(summaryMock).record(1024); + } + + @Test + void testRecordDeliveryLatency() { + sinkMetrics.recordDeliveryLatency(150); + verify(summaryMock).record(150); + } + + @Test + void testRecordHttpLatency() { + sinkMetrics.recordHttpLatency(150); + verify(summaryMock).record(150); + } + + @Test + void testIncrementRetriesCount() { + sinkMetrics.incrementRetriesCount(); + verify(counterMock).increment(1); + } + + @Test + void testIncrementRejectedSpansCount() { + sinkMetrics.incrementRejectedSpansCount(5); + verify(counterMock).increment(5); + } + + @Test + void testRecordResponseCode() { + sinkMetrics.recordResponseCode(200); + verify(pluginMetrics).counter("http_2xx_responses"); + } +} \ No newline at end of file diff --git a/data-prepper-plugins/otlp-sink/src/test/resources/data-prepper-config.yaml b/data-prepper-plugins/otlp-sink/src/test/resources/data-prepper-config.yaml index 7462c0a70e..544c112df1 100644 --- a/data-prepper-plugins/otlp-sink/src/test/resources/data-prepper-config.yaml +++ b/data-prepper-plugins/otlp-sink/src/test/resources/data-prepper-config.yaml @@ -1 +1,3 @@ ssl: false +metricRegistries: + - CloudWatch diff --git a/data-prepper-plugins/otlp-sink/src/test/resources/pipelines.yaml b/data-prepper-plugins/otlp-sink/src/test/resources/pipelines.yaml index 51c1135323..449c21032e 100644 --- a/data-prepper-plugins/otlp-sink/src/test/resources/pipelines.yaml +++ b/data-prepper-plugins/otlp-sink/src/test/resources/pipelines.yaml @@ -1,4 +1,4 @@ -otel_trace_source_to_xray: +otlp_pipeline: source: otel_trace_source: ssl: false From 7dadd574a8a706ef96a2f6f0dd9a4cd932405190 Mon Sep 17 00:00:00 2001 From: huyPham Date: Thu, 17 Apr 2025 15:56:18 -0700 Subject: [PATCH 09/23] fix: use Timer to track latency metrics and correcting latency metric semantics (#12) - Updated unit tests Signed-off-by: huy pham --- .../plugins/sink/otlp/metrics/OtlpSinkMetrics.java | 5 +++-- .../dataprepper/plugins/sink/otlp/OtlpSinkTest.java | 5 ++++- .../plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java | 10 ++++++++-- 3 files changed, 15 insertions(+), 5 deletions(-) diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetrics.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetrics.java index fb68066a80..f4b5926f78 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetrics.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetrics.java @@ -8,6 +8,7 @@ import org.opensearch.dataprepper.metrics.PluginMetrics; import javax.annotation.Nonnull; +import java.time.Duration; /** * Metrics class for the otlp-sink @@ -46,11 +47,11 @@ public void incrementPayloadSize(long bytes) { } public void recordDeliveryLatency(long durationMillis) { - pluginMetrics.summary("deliveryLatency").record(durationMillis); + pluginMetrics.timer("deliveryLatency").record(Duration.ofMillis(durationMillis)); } public void recordHttpLatency(long durationMillis) { - pluginMetrics.summary("httpLatency").record(durationMillis); + pluginMetrics.timer("httpLatency").record(Duration.ofMillis(durationMillis)); } public void incrementRetriesCount() { diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java index 753c91c6a2..66715f4cc9 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java @@ -7,6 +7,7 @@ import io.micrometer.core.instrument.Counter; import io.micrometer.core.instrument.DistributionSummary; +import io.micrometer.core.instrument.Timer; import io.opentelemetry.proto.trace.v1.ResourceSpans; import org.apache.commons.codec.DecoderException; import org.junit.jupiter.api.AfterEach; @@ -21,6 +22,7 @@ import software.amazon.awssdk.regions.Region; import java.io.UnsupportedEncodingException; +import java.time.Duration; import java.time.Instant; import java.util.ArrayList; import java.util.Collection; @@ -62,6 +64,7 @@ void setUp() { mockPluginMetrics = mock(PluginMetrics.class); when(mockPluginMetrics.counter(anyString())).thenReturn(mock(Counter.class)); when(mockPluginMetrics.summary(anyString())).thenReturn(mock(DistributionSummary.class)); + when(mockPluginMetrics.timer(anyString())).thenReturn(mock(Timer.class)); target = new OtlpSink(mockConfig, mockPluginMetrics, mockEncoder, mockSender); } @@ -167,7 +170,7 @@ void testUpdateLatencyMetrics_shouldRecordLatency() { target.updateLatencyMetrics(events); - verify(mockPluginMetrics.summary("deliveryLatency"), times(1)).record(any(Double.class)); + verify(mockPluginMetrics.timer("deliveryLatency"), times(1)).record(any(Duration.class)); } @Test diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java index 8fbda2c909..7568d703e7 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java @@ -7,10 +7,13 @@ import io.micrometer.core.instrument.Counter; import io.micrometer.core.instrument.DistributionSummary; +import io.micrometer.core.instrument.Timer; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.opensearch.dataprepper.metrics.PluginMetrics; +import java.time.Duration; + import static org.mockito.Mockito.anyString; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; @@ -20,6 +23,7 @@ class OtlpSinkMetricsTest { private PluginMetrics pluginMetrics; private Counter counterMock; + private Timer timerMock; private DistributionSummary summaryMock; private OtlpSinkMetrics sinkMetrics; @@ -28,9 +32,11 @@ void setUp() { pluginMetrics = mock(PluginMetrics.class); counterMock = mock(Counter.class); summaryMock = mock(DistributionSummary.class); + timerMock = mock(Timer.class); when(pluginMetrics.counter(anyString())).thenReturn(counterMock); when(pluginMetrics.summary(anyString())).thenReturn(summaryMock); + when(pluginMetrics.timer(anyString())).thenReturn(timerMock); sinkMetrics = new OtlpSinkMetrics(pluginMetrics); } @@ -68,13 +74,13 @@ void testIncrementPayloadSize() { @Test void testRecordDeliveryLatency() { sinkMetrics.recordDeliveryLatency(150); - verify(summaryMock).record(150); + verify(timerMock).record(Duration.ofMillis(150)); } @Test void testRecordHttpLatency() { sinkMetrics.recordHttpLatency(150); - verify(summaryMock).record(150); + verify(timerMock).record(Duration.ofMillis(150)); } @Test From 4608a77c1db3caea637f913185376bb06c16b278 Mon Sep 17 00:00:00 2001 From: huyPham Date: Wed, 23 Apr 2025 07:23:50 -0700 Subject: [PATCH 10/23] feature: add gzip and percentile metrics (#13) Signed-off-by: huy pham --- .../plugins/sink/otlp/OtlpSinkIT.java | 12 ++- .../plugins/sink/otlp/OtlpSink.java | 77 +++++++------------ .../sink/otlp/http/GzipCompressor.java | 63 +++++++++++++++ .../sink/otlp/http/OtlpHttpSender.java | 64 ++++++++------- .../plugins/sink/otlp/http/Sleeper.java | 20 ----- .../plugins/sink/otlp/http/ThreadSleeper.java | 24 ++++-- .../sink/otlp/metrics/OtlpSinkMetrics.java | 71 ++++++++++++++--- .../plugins/sink/otlp/OtlpSinkTest.java | 50 +++++------- .../sink/otlp/http/GzipCompressorTest.java | 64 +++++++++++++++ .../sink/otlp/http/OtlpHttpSenderTest.java | 47 +++++++---- .../sink/otlp/http/ThreadSleeperTest.java | 12 +-- .../otlp/metrics/OtlpSinkMetricsTest.java | 47 +++++++---- 12 files changed, 368 insertions(+), 183 deletions(-) create mode 100644 data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressor.java delete mode 100644 data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/Sleeper.java create mode 100644 data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressorTest.java diff --git a/data-prepper-plugins/otlp-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkIT.java b/data-prepper-plugins/otlp-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkIT.java index dbd6a6e2ad..5630f4675f 100644 --- a/data-prepper-plugins/otlp-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkIT.java +++ b/data-prepper-plugins/otlp-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkIT.java @@ -11,6 +11,7 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.model.trace.JacksonStandardSpan; import org.opensearch.dataprepper.model.trace.Span; @@ -39,6 +40,7 @@ class OtlpSinkIT { private OtlpSinkConfig mockConfig; private PluginMetrics mockPluginMetrics; + private PluginSetting mockPluginSetting; private OtlpSink target; @BeforeEach @@ -56,7 +58,11 @@ void setUp() { when(mockPluginMetrics.counter(anyString())).thenReturn(mock(Counter.class)); when(mockPluginMetrics.summary(anyString())).thenReturn(mock(DistributionSummary.class)); - target = new OtlpSink(mockConfig, mockPluginMetrics); + mockPluginSetting = mock(PluginSetting.class); + when(mockPluginSetting.getPipelineName()).thenReturn("otlp_pipeline"); + when(mockPluginSetting.getName()).thenReturn("otlp"); + + target = new OtlpSink(mockConfig, mockPluginMetrics, mockPluginSetting); } @AfterEach @@ -68,8 +74,6 @@ void cleanUp() { /** * This test is not part of the Data Prepper build. It requires AWS credentials to be set up in the environment. - * - * @throws InterruptedException */ @Test void testSinkProcessesHardcodedSpan() { @@ -85,7 +89,7 @@ void testSinkProcessesHardcodedSpan() { .build(); final Record record = new Record<>(testSpan); - final OtlpSink sink = new OtlpSink(mockConfig, mockPluginMetrics, mock(OTelProtoStandardCodec.OTelProtoEncoder.class), mock(OtlpHttpSender.class)); + final OtlpSink sink = new OtlpSink(mockConfig, mockPluginMetrics, mockPluginSetting, mock(OTelProtoStandardCodec.OTelProtoEncoder.class), mock(OtlpHttpSender.class)); sink.initialize(); sink.output(List.of(record)); diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java index bb840a859a..494fa1bc8e 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java @@ -12,6 +12,7 @@ import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor; +import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.model.sink.Sink; import org.opensearch.dataprepper.model.trace.Span; @@ -49,26 +50,24 @@ public class OtlpSink implements Sink> { private final OtlpSinkMetrics sinkMetrics; /** - * Constructor for the OTLPSink plugin. + * Constructor for the OTLP sink plugin. * - * @param config the configuration for the OTLP sink + * @param config the configuration for the sink + * @param pluginMetrics the plugin metrics to use + * @param pluginSetting the plugin setting to use */ @DataPrepperPluginConstructor - public OtlpSink(@Nonnull final OtlpSinkConfig config, @Nonnull final PluginMetrics pluginMetrics) { - this(config, pluginMetrics, null, null); + public OtlpSink(@Nonnull final OtlpSinkConfig config, @Nonnull final PluginMetrics pluginMetrics, @Nonnull final PluginSetting pluginSetting) { + this(config, pluginMetrics, pluginSetting, null, null); } /** - * Constructor for testing purposes. - * - * @param config the configuration for the OTLP sink - * @param encoder OTEL Protobuf encoder to use - * @param httpSender the HTTP sender to use + * Constructor for the OTLP sink plugin. Used for testing ONLY. */ @VisibleForTesting - OtlpSink(@Nonnull final OtlpSinkConfig config, @Nonnull final PluginMetrics pluginMetrics, final OTelProtoStandardCodec.OTelProtoEncoder encoder, final OtlpHttpSender httpSender) { + OtlpSink(@Nonnull final OtlpSinkConfig config, @Nonnull final PluginMetrics pluginMetrics, @Nonnull final PluginSetting pluginSetting, final OTelProtoStandardCodec.OTelProtoEncoder encoder, final OtlpHttpSender httpSender) { this.batchSize = config.getBatchSize(); - this.sinkMetrics = new OtlpSinkMetrics(pluginMetrics); + this.sinkMetrics = new OtlpSinkMetrics(pluginMetrics, pluginSetting); if (encoder == null && httpSender == null) { this.encoder = new OTelProtoStandardCodec.OTelProtoEncoder(); @@ -77,13 +76,6 @@ public OtlpSink(@Nonnull final OtlpSinkConfig config, @Nonnull final PluginMetri this.encoder = encoder; this.httpSender = httpSender; } - - LOG.info("Config setting: endpoint = {}", config.getEndpoint()); - LOG.info("Config setting: batch_size = {}", config.getBatchSize()); - LOG.info("Config setting: max_retries = {}", config.getMaxRetries()); - LOG.info("Config setting: aws_region = {}", config.getAwsRegion()); - LOG.info("Config setting: aws_sts_role_arn = {}", config.getStsRoleArn()); - LOG.info("Config setting: aws_sts_external_id = {}", config.getStsExternalId()); } /** @@ -91,7 +83,7 @@ public OtlpSink(@Nonnull final OtlpSinkConfig config, @Nonnull final PluginMetri */ @Override public void initialize() { - LOG.info("Initialized OTLP Sink"); + LOG.debug("Initialized OTLP Sink"); } /** @@ -107,35 +99,23 @@ public void output(@Nonnull final Collection> records) { for (int i = 0; i < recordList.size(); i += this.batchSize) { final int end = Math.min(i + this.batchSize, recordList.size()); final List> batch = recordList.subList(i, end); - - try { - final List resourceSpans = batch.stream() - .map(Record::getData) - .map(this::getResourceSpans) - .filter(Objects::nonNull) - .collect(Collectors.toList()); - - if (resourceSpans.isEmpty()) { - LOG.debug("Skipping empty span batch, nothing to send."); - continue; - } - - final ExportTraceServiceRequest request = ExportTraceServiceRequest.newBuilder() - .addAllResourceSpans(resourceSpans) - .build(); - final byte[] payload = request.toByteArray(); - sinkMetrics.incrementPayloadSize(payload.length); - - httpSender.send(payload); - LOG.info("Finished processing {} spans.", resourceSpans.size()); - - sinkMetrics.incrementRecordsOut(batch.size()); - } catch (final RuntimeException e) { - LOG.error("Unexpected error processing OTLP span batch: {}", e.getMessage(), e); - - sinkMetrics.incrementDroppedRecords(batch.size()); - sinkMetrics.incrementErrorsCount(); + final List resourceSpans = batch.stream() + .map(Record::getData) + .map(this::getResourceSpans) + .filter(Objects::nonNull) + .collect(Collectors.toList()); + + if (resourceSpans.isEmpty()) { + LOG.debug("Skipping empty span batch, nothing to send."); + continue; } + + final ExportTraceServiceRequest request = ExportTraceServiceRequest.newBuilder() + .addAllResourceSpans(resourceSpans) + .build(); + final byte[] payload = request.toByteArray(); + httpSender.send(payload); + sinkMetrics.incrementRecordsOut(resourceSpans.size()); } } @@ -165,8 +145,7 @@ public boolean isReady() { */ @Override public void shutdown() { - // OkHttpClient is shared and self-managed; no need to explicitly close - LOG.info("OTLP Sink shutdown complete"); + httpSender.close(); } /** diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressor.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressor.java new file mode 100644 index 0000000000..9b0b9a3f70 --- /dev/null +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressor.java @@ -0,0 +1,63 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.sink.otlp.http; + +import com.google.common.annotations.VisibleForTesting; +import org.opensearch.dataprepper.plugins.sink.otlp.metrics.OtlpSinkMetrics; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.util.Optional; +import java.util.function.Function; +import java.util.zip.GZIPOutputStream; + +class GzipCompressor implements Function> { + private static final Logger LOG = LoggerFactory.getLogger(GzipCompressor.class); + private final OtlpSinkMetrics sinkMetrics; + + /** + * Constructor for the GzipCompressor. + * + * @param sinkMetrics The metrics for the OTLP sink plugin. + */ + GzipCompressor(final OtlpSinkMetrics sinkMetrics) { + this.sinkMetrics = sinkMetrics; + } + + /** + * Compresses the provided payload using GZIP compression. + * Logs an error if compression fails. + * + * @param payload The payload to be compressed. + * @return Optional containing the compressed payload, or empty if compression failed. + */ + @Override + public Optional apply(final byte[] payload) { + try { + return Optional.of(compressInternal(payload)); + } catch (final IOException e) { + LOG.error("Failed to compress payload", e); + sinkMetrics.incrementErrorsCount(); + return Optional.empty(); + } + } + + /** + * Internal method to enable mocked-testing. + */ + @VisibleForTesting + byte[] compressInternal(final byte[] payload) throws IOException { + try (final ByteArrayOutputStream out = new ByteArrayOutputStream(); + final GZIPOutputStream gzip = new GZIPOutputStream(out)) { + + gzip.write(payload); + gzip.finish(); + return out.toByteArray(); + } + } +} diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java index 2e0bcf3248..c7cda9de62 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java @@ -23,7 +23,10 @@ import java.security.SecureRandom; import java.util.ArrayList; import java.util.List; +import java.util.Optional; import java.util.Set; +import java.util.function.Consumer; +import java.util.function.Function; /** * Responsible for sending signed OTLP Protobuf trace data to AWS OTLP endpoint using OkHttp. @@ -37,32 +40,33 @@ public class OtlpHttpSender implements AutoCloseable { private static final MediaType PROTOBUF = MediaType.get("application/x-protobuf"); private final SecureRandom random = new SecureRandom(); + private final int maxRetries; private final SigV4Signer signer; private final OkHttpClient httpClient; - private final Sleeper sleeper; - private final int maxRetries; + private final Consumer sleeper; private final List retryDelaysMs; private final OtlpSinkMetrics sinkMetrics; + private final Function> gzipCompressor; /** * Constructor for the OtlpHttpSender. * Initializes the signer and HTTP client. * @param config The configuration for the OTLP sink plugin. + * @param sinkMetrics The metrics for the OTLP sink plugin. */ public OtlpHttpSender(@Nonnull final OtlpSinkConfig config, @Nonnull final OtlpSinkMetrics sinkMetrics) { - this(config, sinkMetrics, null, null, null); + this(config, sinkMetrics, new GzipCompressor(sinkMetrics), null, null, null); } /** * Constructor for unit testing with injected dependencies. - * - * @param config The configuration for the OTLP sink plugin. - * @param signer The SigV4Signer instance for signing requests. - * @param httpClient The OkHttpClient instance for making HTTP requests. */ @VisibleForTesting - OtlpHttpSender(@Nonnull final OtlpSinkConfig config, @Nonnull final OtlpSinkMetrics sinkMetrics, final SigV4Signer signer, final OkHttpClient httpClient, final Sleeper sleeper) { + OtlpHttpSender(@Nonnull final OtlpSinkConfig config, @Nonnull final OtlpSinkMetrics sinkMetrics, @Nonnull final Function> gzipCompressor, + final SigV4Signer signer, final OkHttpClient httpClient, final Consumer sleeper) { + this.sinkMetrics = sinkMetrics; + this.gzipCompressor = gzipCompressor; this.signer = signer != null ? signer : new SigV4Signer(config); this.httpClient = httpClient != null ? httpClient : new OkHttpClient(); this.sleeper = sleeper != null ? sleeper : new ThreadSleeper(); @@ -79,7 +83,6 @@ public OtlpHttpSender(@Nonnull final OtlpSinkConfig config, @Nonnull final OtlpS */ private List generateExponentialBackoffDelays(final int retries) { List delays = new ArrayList<>(); - for (int i = 0; i < retries; i++) { // Exponential backoff: 100ms, 200ms, 400ms, ... delays.add(BASE_RETRY_DELAY_MS * (1 << i)); @@ -97,11 +100,16 @@ private List generateExponentialBackoffDelays(final int retries) { public void send(@Nonnull final byte[] payload) { for (int attempt = 0; attempt <= maxRetries; attempt++) { try { - final SdkHttpFullRequest signedRequest = signer.signRequest(payload); + final Optional compressedPayload = gzipCompressor.apply(payload); + if (compressedPayload.isEmpty()) { + return; + } + final SdkHttpFullRequest signedRequest = signer.signRequest(compressedPayload.get()); final Request.Builder requestBuilder = new Request.Builder() .url(signedRequest.getUri().toString()) - .post(RequestBody.create(payload, PROTOBUF)); + .post(RequestBody.create(compressedPayload.get(), PROTOBUF)) + .addHeader("Content-Encoding", "gzip"); signedRequest.headers().forEach((key, values) -> { for (final String value : values) { @@ -117,24 +125,30 @@ public void send(@Nonnull final byte[] payload) { sinkMetrics.recordHttpLatency(duration); handleResponse(response); + + sinkMetrics.incrementPayloadSize(payload.length); + sinkMetrics.incrementPayloadGzipSize(compressedPayload.get().length); return; } - } catch (final Exception e) { + } catch (final IOException ioException) { if (attempt < maxRetries) { final int jitter = random.nextInt(100); final int retryIndex = Math.min(attempt, retryDelaysMs.size() - 1); final int delay = retryDelaysMs.get(retryIndex) + jitter; - LOG.warn("Retrying after failure in attempt {}. Sleeping {}ms.", attempt + 1, delay, e); + LOG.debug("Retrying after failure in attempt {}. Sleeping {}ms.", attempt + 1, delay, ioException); sinkMetrics.incrementRetriesCount(); try { - sleeper.sleep(delay); - } catch (final InterruptedException ie) { + sleeper.accept(delay); + } catch (final RuntimeException runtimeException) { Thread.currentThread().interrupt(); - throw new RuntimeException("Retry interrupted while sending OTLP data", ie); + LOG.error("Interrupted while sleeping between retries", runtimeException); + sinkMetrics.incrementErrorsCount(); + return; } } else { - LOG.error("All retry attempts failed while signing or sending OTLP data", e); - throw new RuntimeException("Failed to sign/send data after retries", e); + LOG.error("Failed to sign/send data after all retries", ioException); + sinkMetrics.incrementErrorsCount(); + return; } } } @@ -173,7 +187,7 @@ private void handleResponse(@Nonnull final Response response) throws IOException private void handleSuccessfulResponse(final byte[] responseBytes) { if (responseBytes == null || responseBytes.length == 0) { - LOG.info("OTLP export successful. No response body."); + LOG.debug("OTLP export successful. No response body."); return; } @@ -183,17 +197,11 @@ private void handleSuccessfulResponse(final byte[] responseBytes) { if (otlpResponse.hasPartialSuccess()) { final var partial = otlpResponse.getPartialSuccess(); final long rejectedSpans = partial.getRejectedSpans(); - sinkMetrics.incrementRejectedSpansCount(rejectedSpans); - final String errorMessage = partial.getErrorMessage(); - if (rejectedSpans > 0 || !errorMessage.isEmpty()) { LOG.warn("OTLP Partial Success: rejectedSpans={}, message={}", rejectedSpans, errorMessage); - } else { - LOG.info("OTLP export successful with no rejections."); + sinkMetrics.incrementRejectedSpansCount(rejectedSpans); } - } else { - LOG.info("OTLP export successful with no partial success field."); } } catch (final Exception e) { LOG.error("Could not parse OTLP response as ExportTraceServiceResponse: {}", e.getMessage()); @@ -203,7 +211,7 @@ private void handleSuccessfulResponse(final byte[] responseBytes) { @Override public void close() { - // No explicit shutdown required for OkHttpClient unless using dispatcher or connection pool tuning. - // https://square.github.io/okhttp/4.x/okhttp/okhttp3/-ok-http-client/dispatchers/ + httpClient.connectionPool().evictAll(); + httpClient.dispatcher().executorService().shutdown(); } } \ No newline at end of file diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/Sleeper.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/Sleeper.java deleted file mode 100644 index 1ba9f34021..0000000000 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/Sleeper.java +++ /dev/null @@ -1,20 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.sink.otlp.http; - -/** - * Interface for sleeping in tests. - */ -interface Sleeper { - - /** - * Sleeps for the specified number of milliseconds. - * - * @param millis the number of milliseconds to sleep - * @throws InterruptedException if the sleep is interrupted - */ - void sleep(int millis) throws InterruptedException; -} \ No newline at end of file diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeper.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeper.java index 61f8d23d36..8355d40290 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeper.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeper.java @@ -5,12 +5,24 @@ package org.opensearch.dataprepper.plugins.sink.otlp.http; -/** - * Implementation of {@link Sleeper} that sleeps using {@link Thread#sleep(long)}. - */ -class ThreadSleeper implements Sleeper { +import javax.annotation.Nonnull; +import java.util.function.Consumer; + +class ThreadSleeper implements Consumer { + + /** + * Sleeps for the specified duration in milliseconds. + * If the thread is interrupted while sleeping, the interrupted status is cleared. + * + * @param millis the input argument + */ @Override - public void sleep(int millis) throws InterruptedException { - Thread.sleep(millis); + public void accept(final @Nonnull Integer millis) { + try { + Thread.sleep(millis); + } catch (final InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } } } diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetrics.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetrics.java index f4b5926f78..b492b79e8e 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetrics.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetrics.java @@ -5,7 +5,11 @@ package org.opensearch.dataprepper.plugins.sink.otlp.metrics; +import io.micrometer.core.instrument.DistributionSummary; +import io.micrometer.core.instrument.Metrics; +import io.micrometer.core.instrument.Timer; import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.configuration.PluginSetting; import javax.annotation.Nonnull; import java.time.Duration; @@ -16,14 +20,63 @@ public class OtlpSinkMetrics { private final PluginMetrics pluginMetrics; + private final Timer httpLatency; + private final Timer deliveryLatency; + private final DistributionSummary payloadSize; + private final DistributionSummary payloadGzipSize; /** * Constructor for OtlpSinkMetrics * - * @param pluginMetrics The plugin metrics instance + * @param pluginMetrics The plugin metrics + * @param pluginSetting The plugin setting */ - public OtlpSinkMetrics(@Nonnull final PluginMetrics pluginMetrics) { + public OtlpSinkMetrics(@Nonnull final PluginMetrics pluginMetrics, @Nonnull final PluginSetting pluginSetting) { this.pluginMetrics = pluginMetrics; + + final String pipelineName = pluginSetting.getPipelineName(); + final String pluginName = pluginSetting.getName(); + + httpLatency = buildLatencyTimer(pipelineName, pluginName, "httpLatency"); + deliveryLatency = buildLatencyTimer(pipelineName, pluginName, "deliveryLatency"); + + payloadSize = buildDistributionSummary(pipelineName, pluginName, "payloadSize"); + payloadGzipSize = buildDistributionSummary(pipelineName, pluginName, "payloadGzipSize"); + } + + /** + * Builds a timer for latency metrics with percentiles + * + * @param pipelineName The pipeline name + * @param pluginName The plugin name + * @param metricName The metric name + * @return The timer + */ + private static Timer buildLatencyTimer(@Nonnull final String pipelineName, @Nonnull final String pluginName, @Nonnull final String metricName) { + return Timer.builder(String.format("%s_%s_%s", pipelineName, pluginName, metricName)) + .publishPercentiles(0.5, 0.9, 0.95, 1.0) + .publishPercentileHistogram(true) + .distributionStatisticBufferLength(1024) + .distributionStatisticExpiry(Duration.ofMinutes(10)) + .register(Metrics.globalRegistry); + } + + /** + * Builds a distribution summary for payload size metrics with percentiles + * + * @param pipelineName The pipeline name + * @param pluginName The plugin name + * @param metricName The metric name + * @return The distribution summary + */ + private static DistributionSummary buildDistributionSummary(@Nonnull final String pipelineName, @Nonnull final String pluginName, @Nonnull final String metricName) { + return DistributionSummary.builder(String.format("%s_%s_%s", pipelineName, pluginName, metricName)) + .baseUnit("bytes") + .publishPercentiles(0.5, 0.9, 0.95, 1.0) + .publishPercentileHistogram(true) + .distributionStatisticBufferLength(1024) + .distributionStatisticExpiry(Duration.ofMinutes(10)) + .register(Metrics.globalRegistry); } public void incrementRecordsIn(long count) { @@ -34,24 +87,24 @@ public void incrementRecordsOut(long count) { pluginMetrics.counter("recordsOut").increment(count); } - public void incrementDroppedRecords(long count) { - pluginMetrics.counter("droppedRecords").increment(count); - } - public void incrementErrorsCount() { pluginMetrics.counter("errorsCount").increment(1); } public void incrementPayloadSize(long bytes) { - pluginMetrics.summary("payloadSize").record(bytes); + payloadSize.record(bytes); + } + + public void incrementPayloadGzipSize(long bytes) { + payloadGzipSize.record(bytes); } public void recordDeliveryLatency(long durationMillis) { - pluginMetrics.timer("deliveryLatency").record(Duration.ofMillis(durationMillis)); + deliveryLatency.record(Duration.ofMillis(durationMillis)); } public void recordHttpLatency(long durationMillis) { - pluginMetrics.timer("httpLatency").record(Duration.ofMillis(durationMillis)); + httpLatency.record(Duration.ofMillis(durationMillis)); } public void incrementRetriesCount() { diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java index 66715f4cc9..68483d9da7 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java @@ -14,6 +14,7 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.opensearch.dataprepper.metrics.PluginMetrics; +import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.model.trace.Span; import org.opensearch.dataprepper.plugins.otel.codec.OTelProtoStandardCodec; @@ -22,7 +23,6 @@ import software.amazon.awssdk.regions.Region; import java.io.UnsupportedEncodingException; -import java.time.Duration; import java.time.Instant; import java.util.ArrayList; import java.util.Collection; @@ -31,8 +31,6 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyDouble; import static org.mockito.ArgumentMatchers.anyString; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; @@ -45,6 +43,7 @@ class OtlpSinkTest { private OTelProtoStandardCodec.OTelProtoEncoder mockEncoder; private OtlpHttpSender mockSender; private PluginMetrics mockPluginMetrics; + private PluginSetting mockPluginSetting; private OtlpSink target; @BeforeEach @@ -66,7 +65,11 @@ void setUp() { when(mockPluginMetrics.summary(anyString())).thenReturn(mock(DistributionSummary.class)); when(mockPluginMetrics.timer(anyString())).thenReturn(mock(Timer.class)); - target = new OtlpSink(mockConfig, mockPluginMetrics, mockEncoder, mockSender); + mockPluginSetting = mock(PluginSetting.class); + when(mockPluginSetting.getPipelineName()).thenReturn("otlp_pipeline"); + when(mockPluginSetting.getName()).thenReturn("otlp"); + + target = new OtlpSink(mockConfig, mockPluginMetrics, mockPluginSetting, mockEncoder, mockSender); } @AfterEach @@ -96,22 +99,6 @@ void testOutput_shouldSendAllBatches() throws Exception { verify(mockEncoder, times(recordCount)).convertToResourceSpans(any()); } - @Test - void testOutput_shouldHandleRuntimeException() throws Exception { - Span span = mock(Span.class); - when(span.getSpanId()).thenReturn("span123"); - Record record = mock(Record.class); - when(record.getData()).thenReturn(span); - - when(mockEncoder.convertToResourceSpans(span)).thenReturn(ResourceSpans.getDefaultInstance()); - doThrow(new RuntimeException("sender failure")).when(mockSender).send(any()); - - target.output(List.of(record)); - - verify(mockEncoder).convertToResourceSpans(eq(span)); - verify(mockSender).send(any()); - } - @Test void testOutput_shouldSendPartialBatchWhenSomeSpansSucceed() throws Exception { // Good span @@ -145,20 +132,25 @@ void testOutput_shouldSendPartialBatchWhenSomeSpansSucceed() throws Exception { @Test void testOutput_shouldNotSendEmptyBatch() throws DecoderException, UnsupportedEncodingException { - Span span = mock(Span.class); + // Given: a span that fails encoding + final Span span = mock(Span.class); when(span.getSpanId()).thenReturn("bad-span"); - Record record = mock(Record.class); + final Record record = mock(Record.class); when(record.getData()).thenReturn(span); - when(mockEncoder.convertToResourceSpans(span)).thenReturn(null); // simulate invalid span + // Simulate encoding failure (e.g., due to missing fields) + when(mockEncoder.convertToResourceSpans(span)).thenReturn(null); + // When: the output method is called with that bad span target.output(List.of(record)); + // Then: it should not send anything to the OTLP endpoint verify(mockSender, never()).send(any()); } + @Test - void testUpdateLatencyMetrics_shouldRecordLatency() { + void testUpdateLatencyMetrics_runsWithoutException_givenValidSpanStartTime() { Span mockSpan = mock(Span.class); String startTime = Instant.now().minusSeconds(5).toString(); when(mockSpan.getStartTime()).thenReturn(startTime); @@ -166,16 +158,16 @@ void testUpdateLatencyMetrics_shouldRecordLatency() { Record mockRecord = mock(Record.class); when(mockRecord.getData()).thenReturn(mockSpan); - Collection> events = List.of(mockRecord); + Collection> records = List.of(mockRecord); - target.updateLatencyMetrics(events); + target.updateLatencyMetrics(records); - verify(mockPluginMetrics.timer("deliveryLatency"), times(1)).record(any(Duration.class)); + // verify that no error counter incremented + verify(mockPluginMetrics.counter("errorsCount"), never()).increment(); } @Test void testUpdateLatencyMetrics_shouldHandleInvalidStartTime() { - // Mock a Span with an invalid start time string final Span badSpan = mock(Span.class); when(badSpan.getStartTime()).thenReturn("invalid-timestamp"); @@ -193,7 +185,7 @@ void testUpdateLatencyMetrics_shouldHandleInvalidStartTime() { @Test void testConstructor_withOnlyConfig_shouldInitializeWithoutException() { - OtlpSink sink = new OtlpSink(mockConfig, mockPluginMetrics); + OtlpSink sink = new OtlpSink(mockConfig, mockPluginMetrics, mockPluginSetting); sink.initialize(); sink.shutdown(); diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressorTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressorTest.java new file mode 100644 index 0000000000..d05b47c217 --- /dev/null +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressorTest.java @@ -0,0 +1,64 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.sink.otlp.http; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.plugins.sink.otlp.metrics.OtlpSinkMetrics; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.Optional; +import java.util.zip.GZIPInputStream; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; + +class GzipCompressorTest { + + private OtlpSinkMetrics sinkMetrics; + + @BeforeEach + void setUp() { + sinkMetrics = mock(OtlpSinkMetrics.class); + } + + @Test + void apply_returnsCompressedPayload() throws IOException { + byte[] input = "test-payload".getBytes(); + GzipCompressor gzipCompressor = new GzipCompressor(sinkMetrics); + Optional compressed = gzipCompressor.apply(input); + + assertTrue(compressed.isPresent(), "Expected compressed payload to be present"); + + // Validate decompression gives original input + byte[] decompressed = decompress(compressed.get()); + assertArrayEquals(input, decompressed); + } + + @Test + void apply_handlesIOException_andIncrementsErrorMetric() throws IOException { + GzipCompressor gzipCompressor = spy(new GzipCompressor(sinkMetrics)); + doThrow(new IOException("boom")).when(gzipCompressor).compressInternal(any()); + + Optional result = gzipCompressor.apply("payload".getBytes(StandardCharsets.UTF_8)); + + assertTrue(result.isEmpty()); + verify(sinkMetrics).incrementErrorsCount(); + } + + private byte[] decompress(byte[] compressed) throws IOException { + try (GZIPInputStream gzipStream = new GZIPInputStream(new ByteArrayInputStream(compressed))) { + return gzipStream.readAllBytes(); + } + } +} \ No newline at end of file diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java index 29047bd4f1..74c1a378d0 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java @@ -30,10 +30,12 @@ import java.nio.charset.StandardCharsets; import java.util.Collections; import java.util.Map; +import java.util.Optional; +import java.util.function.Consumer; +import java.util.function.Function; import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; import static org.mockito.ArgumentMatchers.any; @@ -54,7 +56,8 @@ class OtlpHttpSenderTest { private OtlpSinkConfig mockConfig; private SigV4Signer mockSigner; private OkHttpClient mockHttpClient; - private Sleeper mockSleeper; + private Consumer mockSleeper; + private Function> mockGzipCompressor; private OtlpSinkMetrics mockSinkMetrics; private OtlpHttpSender target; @@ -69,10 +72,16 @@ void setUp() { mockSigner = mock(SigV4Signer.class); mockHttpClient = mock(OkHttpClient.class); - mockSleeper = mock(Sleeper.class); + mockSleeper = mock(ThreadSleeper.class); mockSinkMetrics = mock(OtlpSinkMetrics.class); - target = new OtlpHttpSender(mockConfig, mockSinkMetrics, mockSigner, mockHttpClient, mockSleeper); + mockGzipCompressor = mock(GzipCompressor.class); + when(mockGzipCompressor.apply(any())).thenAnswer(invocation -> { + byte[] input = invocation.getArgument(0); + return input == null ? Optional.empty() : Optional.of(input); + }); + + target = new OtlpHttpSender(mockConfig, mockSinkMetrics, mockGzipCompressor, mockSigner, mockHttpClient, mockSleeper); } @AfterEach @@ -167,7 +176,7 @@ void testSend_retryOnFailure_thenSuccess() throws IOException { } @Test - void testSend_failsAfterAllRetries() throws IOException { + void testSend_doesNotThrowException_failsAfterAllRetries() throws IOException { SdkHttpFullRequest mockSignedRequest = mock(SdkHttpFullRequest.class); when(mockSignedRequest.getUri()).thenReturn(HttpUrl.get("https://xray.us-west-2.amazonaws.com/v1/traces").uri()); when(mockSignedRequest.headers()).thenReturn(Map.of()); @@ -178,11 +187,11 @@ void testSend_failsAfterAllRetries() throws IOException { when(mockHttpClient.newCall(any())).thenReturn(mockCall); when(mockCall.execute()).thenThrow(new IOException("always fail")); - assertThrows(RuntimeException.class, () -> target.send(PAYLOAD)); + assertDoesNotThrow(() -> target.send(PAYLOAD)); } @Test - void testSend_throwsIOException_on500ResponseWithBody() throws IOException { + void testSend_doesNotThrowsException_on500ResponseWithBody() throws IOException { // Mock signed request SdkHttpFullRequest sdkRequest = SdkHttpFullRequest.builder() .method(software.amazon.awssdk.http.SdkHttpMethod.POST) @@ -211,11 +220,11 @@ void testSend_throwsIOException_on500ResponseWithBody() throws IOException { when(call.execute()).thenReturn(mockResponse); // Run test - assertThrows(RuntimeException.class, () -> target.send(PAYLOAD)); + assertDoesNotThrow(() -> target.send(PAYLOAD)); } @Test - void testInterruptedExceptionDuringRetryThrowsRuntimeException() throws IOException, InterruptedException { + void testSend_doesNotThrowsException_onInterruptedExceptionDuringRetry() throws IOException { final SdkHttpFullRequest signedRequest = mock(SdkHttpFullRequest.class); when(mockSigner.signRequest(any())).thenReturn(signedRequest); @@ -226,15 +235,11 @@ void testInterruptedExceptionDuringRetryThrowsRuntimeException() throws IOExcept when(mockCall.execute()).thenThrow(new IOException("boom")); when(mockHttpClient.newCall(any())).thenReturn(mockCall); - doThrow(new InterruptedException("interrupted")).when(mockSleeper).sleep(anyInt()); + doThrow(new RuntimeException("")).when(mockSleeper).accept(anyInt()); - target = new OtlpHttpSender(mockConfig, mockSinkMetrics, mockSigner, mockHttpClient, mockSleeper); + target = new OtlpHttpSender(mockConfig, mockSinkMetrics, mockGzipCompressor, mockSigner, mockHttpClient, mockSleeper); - final RuntimeException thrown = assertThrows(RuntimeException.class, () -> - target.send(PAYLOAD) - ); - - assertTrue(thrown.getMessage().contains("Retry interrupted")); + assertDoesNotThrow(() -> target.send(PAYLOAD)); assertTrue(Thread.currentThread().isInterrupted()); } @@ -391,6 +396,16 @@ void testSend_invalidProtoResponse_logsError() throws IOException { assertDoesNotThrow(() -> target.send(PAYLOAD)); } + @Test + void testSend_skipsSend_whenGzipCompressionFails() { + when(mockGzipCompressor.apply(any())).thenReturn(Optional.empty()); + + assertDoesNotThrow(() -> target.send(PAYLOAD)); + + verify(mockSigner, times(0)).signRequest(any()); + verify(mockHttpClient, times(0)).newCall(any()); + } + @Test void testDefaultConstructorInitializesDefaults() { target = new OtlpHttpSender(mockConfig, mockSinkMetrics); diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeperTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeperTest.java index 9da7148267..df368566e4 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeperTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeperTest.java @@ -8,11 +8,13 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import java.util.function.Consumer; + import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; class ThreadSleeperTest { - private ThreadSleeper target; + private Consumer target; @BeforeEach void setUp() { @@ -22,8 +24,8 @@ void setUp() { @Test void testSleepDoesNotThrowWhenNotInterrupted() { try { - target.sleep(1); - } catch (InterruptedException e) { + target.accept(1); + } catch (RuntimeException e) { fail("Sleep was interrupted unexpectedly"); } } @@ -33,9 +35,9 @@ void testSleepThrowsInterruptedExceptionIfThreadInterrupted() { Thread thread = new Thread(() -> { try { Thread.currentThread().interrupt(); - target.sleep(10); + target.accept(10); fail("Expected InterruptedException"); - } catch (InterruptedException e) { + } catch (RuntimeException e) { assertTrue(Thread.currentThread().isInterrupted(), "Thread should remain interrupted"); } }); diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java index 7568d703e7..6a53fb48e1 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java @@ -11,8 +11,7 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.opensearch.dataprepper.metrics.PluginMetrics; - -import java.time.Duration; +import org.opensearch.dataprepper.model.configuration.PluginSetting; import static org.mockito.Mockito.anyString; import static org.mockito.Mockito.mock; @@ -22,23 +21,25 @@ class OtlpSinkMetricsTest { private PluginMetrics pluginMetrics; + private PluginSetting pluginSetting; private Counter counterMock; - private Timer timerMock; private DistributionSummary summaryMock; + private Timer timerMock; private OtlpSinkMetrics sinkMetrics; @BeforeEach void setUp() { pluginMetrics = mock(PluginMetrics.class); + pluginSetting = mock(PluginSetting.class); counterMock = mock(Counter.class); summaryMock = mock(DistributionSummary.class); timerMock = mock(Timer.class); when(pluginMetrics.counter(anyString())).thenReturn(counterMock); - when(pluginMetrics.summary(anyString())).thenReturn(summaryMock); - when(pluginMetrics.timer(anyString())).thenReturn(timerMock); + when(pluginSetting.getPipelineName()).thenReturn("otlp_pipeline"); + when(pluginSetting.getName()).thenReturn("otlp"); - sinkMetrics = new OtlpSinkMetrics(pluginMetrics); + sinkMetrics = new OtlpSinkMetrics(pluginMetrics, pluginSetting); } @Test @@ -53,12 +54,6 @@ void testIncrementRecordsOut() { verify(counterMock).increment(2); } - @Test - void testIncrementDroppedRecords() { - sinkMetrics.incrementDroppedRecords(1); - verify(counterMock).increment(1); - } - @Test void testIncrementErrorsCount() { sinkMetrics.incrementErrorsCount(); @@ -68,19 +63,25 @@ void testIncrementErrorsCount() { @Test void testIncrementPayloadSize() { sinkMetrics.incrementPayloadSize(1024); - verify(summaryMock).record(1024); + // Cannot verify summaryMock as DistributionSummary is built statically inside constructor + } + + @Test + void testIncrementPayloadGzipSize() { + sinkMetrics.incrementPayloadGzipSize(2048); + // Cannot verify summaryMock without injecting mock } @Test void testRecordDeliveryLatency() { sinkMetrics.recordDeliveryLatency(150); - verify(timerMock).record(Duration.ofMillis(150)); + // Would require refactoring to inject mock Timer for verification } @Test void testRecordHttpLatency() { - sinkMetrics.recordHttpLatency(150); - verify(timerMock).record(Duration.ofMillis(150)); + sinkMetrics.recordHttpLatency(100); + // Would require refactoring to inject mock Timer for verification } @Test @@ -96,7 +97,19 @@ void testIncrementRejectedSpansCount() { } @Test - void testRecordResponseCode() { + void testRecordResponseCode_5xx() { + sinkMetrics.recordResponseCode(503); + verify(pluginMetrics).counter("http_5xx_responses"); + } + + @Test + void testRecordResponseCode_4xx() { + sinkMetrics.recordResponseCode(404); + verify(pluginMetrics).counter("http_4xx_responses"); + } + + @Test + void testRecordResponseCode_2xx() { sinkMetrics.recordResponseCode(200); verify(pluginMetrics).counter("http_2xx_responses"); } From e876183413fcf09cbcbb9fc4f8a2401fa4edaa28 Mon Sep 17 00:00:00 2001 From: huyPham Date: Mon, 28 Apr 2025 21:41:43 -0700 Subject: [PATCH 11/23] feature(threshold-config): add threshold configuration support (#14) - Introduce `ThresholdConfig` with `max_events`, `max_batch_size`, and `flush_timeout` - Implement `OtlpSinkBuffer` using a bounded `LinkedBlockingQueue` driven by threshold settings - Enhance `OtlpSinkMetrics` to register queue gauges and cover new payload/latency metrics - Update README with configuration options and note on AWS region derivation - Add comprehensive unit tests for config, buffer, metrics, and HTTP sender - Removed unused integ test package and resources Signed-off-by: huy pham --- data-prepper-plugins/otlp-sink/README.md | 117 ++--- data-prepper-plugins/otlp-sink/build.gradle | 3 + .../plugins/sink/otlp/OtlpSinkIT.java | 98 ---- .../plugins/sink/otlp/OtlpSink.java | 122 +---- .../sink/otlp/buffer/OtlpSinkBuffer.java | 206 ++++++++ ...tion.java => AwsAuthenticationConfig.java} | 5 +- .../otlp/configuration/OtlpSinkConfig.java | 53 +- .../otlp/configuration/ThresholdConfig.java | 35 ++ .../sink/otlp/http/GzipCompressor.java | 4 + .../sink/otlp/http/OtlpHttpSender.java | 53 +- .../plugins/sink/otlp/http/SigV4Signer.java | 4 +- .../plugins/sink/otlp/http/ThreadSleeper.java | 11 +- .../sink/otlp/metrics/OtlpSinkMetrics.java | 34 +- .../plugins/sink/otlp/OtlpSinkTest.java | 212 +++----- .../sink/otlp/buffer/OtlpSinkBufferTest.java | 276 +++++++++++ ....java => AwsAuthenticationConfigTest.java} | 6 +- .../configuration/OtlpSinkConfigTest.java | 147 ++++-- .../configuration/ThresholdConfigTest.java | 73 +++ .../sink/otlp/http/GzipCompressorTest.java | 1 + .../sink/otlp/http/OtlpHttpSenderTest.java | 451 ++++++++---------- .../otlp/metrics/OtlpSinkMetricsTest.java | 146 ++++-- .../test/resources/data-prepper-config.yaml | 3 - .../src/test/resources/pipelines.yaml | 13 - .../src/test/resources/sample-trace.json | 28 -- .../src/test/resources/test-span-event.json | 23 - 25 files changed, 1249 insertions(+), 875 deletions(-) delete mode 100644 data-prepper-plugins/otlp-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkIT.java create mode 100644 data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBuffer.java rename data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/{AwsAuthenticationConfiguration.java => AwsAuthenticationConfig.java} (92%) create mode 100644 data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/ThresholdConfig.java create mode 100644 data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBufferTest.java rename data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/{AwsAuthenticationConfigurationTest.java => AwsAuthenticationConfigTest.java} (85%) create mode 100644 data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/ThresholdConfigTest.java delete mode 100644 data-prepper-plugins/otlp-sink/src/test/resources/data-prepper-config.yaml delete mode 100644 data-prepper-plugins/otlp-sink/src/test/resources/pipelines.yaml delete mode 100644 data-prepper-plugins/otlp-sink/src/test/resources/sample-trace.json delete mode 100644 data-prepper-plugins/otlp-sink/src/test/resources/test-span-event.json diff --git a/data-prepper-plugins/otlp-sink/README.md b/data-prepper-plugins/otlp-sink/README.md index 4cc7a4fde8..63e26e35df 100644 --- a/data-prepper-plugins/otlp-sink/README.md +++ b/data-prepper-plugins/otlp-sink/README.md @@ -1,97 +1,76 @@ -# X-Ray OTLP Sink +# OTLP Sink Plugin -The `otlp_sink` plugin sends span data to [AWS X-Ray](https://docs.aws.amazon.com/xray/) using the OTLP (OpenTelemetry Protocol) format. +In the first release, the otlp sink plugin sends span data to AWS X-Ray using the OTLP (OpenTelemetry Protocol) format. +Future releases will enhance the sink to send spans, metrics, and traces to any OTLP Protobuf endpoints. -## Usage - -For information on usage, see the forthcoming documentation in the [Data Prepper Sink Plugins section](https://opensearch.org/docs/latest/data-prepper/pipelines/configuration/sinks/). - -A sample pipeline configuration will be added once the plugin is ready for testing. - -### Configuration Options +--- -#### aws (Required) -Configuration options for AWS authentication and region settings. - -* `region` (Required): The AWS region where X-Ray service is located - * Must be a valid AWS region identifier (e.g., us-east-1, us-west-2) - * Cannot be empty - -* `sts_role_arn` (Required): AWS STS Role ARN for assuming role-based access - * Format: arn:aws:iam::{account}:role/{role-name} - * Length must be between 20 and 2048 characters - -* `sts_external_id` (Optional): External ID for additional security when assuming an IAM role - * Required only if the trust policy requires an external ID - * Length must be between 2 and 1224 characters +## Usage ### Sample Pipeline Configuration ```yaml -pipeline: +otlp_pipeline: + workers: 2 + source: otel_trace_source: - ssl: true - + ssl: false + port: 21890 + buffer: bounded_blocking: - buffer_size: 10 - batch_size: 5 - + buffer_size: 1000000 + batch_size: 125000 + sink: - otlp: + endpoint: "https://performance.us-west-2.xray.cloudwatch.aws.dev/v1/traces" + max_retries: 5 # Optional, default: 5 + threshold: + max_events: 512 # Optional, default: 512 + max_batch_size: 1mb # Optional, default: 1mb + flush_timeout: 200ms # Optional, default: 200ms aws: - region: us-west-2 - sts_role_arn: arn:aws:iam::123456789012:role/XrayRole + sts_role_arn: arn:aws:iam::123456789012:role/MyRole # Optional STS Role ARN + sts_external_id: external-id-value # Optional external ID for STS ``` - + +--- + +## Configuration Options + +| Property | Type | Required | Default | Description | +|----------------------------|----------|----------|---------|---------------------------------------------------------------------------------------| +| `endpoint` | `String` | Yes | — | OTLP gRPC or HTTP endpoint where spans will be sent. | +| `max_retries` | `int` | No | `5` | Maximum number of retry attempts on HTTP send failures. | +| **threshold** | `Object` | No | — | Controls batching behavior. See below for sub-properties. | +| `threshold.max_events` | `int` | No | `512` | Maximum number of spans per batch. | +| `threshold.max_batch_size` | `String` | No | `1mb` | Maximum total payload bytes per batch. Supports human-readable suffixes (`kb`, `mb`). | +| `threshold.flush_timeout` | `String` | No | `200ms` | Time to wait (in milliseconds) before flushing a non-empty batch. | +| **aws** | `Object` | No | — | AWS authentication settings. See below. | +| `aws.sts_role_arn` | `String` | No | — | Amazon Resource Name of the IAM role to assume. | +| `aws.sts_external_id` | `String` | No | — | Optional external ID for assuming IAM roles with STS. | + +**Note:** `aws.region` will be derived from the provided AWS endpoint. + +--- + ## Developer Guide See the [CONTRIBUTING](https://github.com/opensearch-project/data-prepper/blob/main/CONTRIBUTING.md) guide for general information on contributions. -The integration tests for this plugin do not run as part of the main Data Prepper build. +The integration tests for this plugin do **not** run as part of the main Data Prepper build and will be added in future +releases. -#### Run unit tests locally +### Run unit tests locally ```bash ./gradlew :data-prepper-plugins:otlp-sink:test ``` -#### Run integration tests locally +### Run integration tests locally -``` +```bash ./gradlew :data-prepper-plugins:otlp-sink:integrationTest ``` - -#### Run a local pipeline that uses this sink - -1. Install `grpcurl` – Used to send OTLP span data to the running pipeline. -2. Build the plugin and Data Prepper: -``` -./gradlew build` -``` -3. Start the pipeline: -``` -cd release/archives/linux/build/install/opensearch-data-prepper-2.11.0-SNAPSHOT-linux-x64 - -bin/data-prepper \ - /path/to/data-prepper-plugins/otlp-sink/src/test/resources/pipelines.yaml \ - /path/to/data-prepper-plugins/otlp-sink/src/test/resources/data-prepper-config.yaml -``` -4. Send test spans to the local pipeline: -``` -cd /path/to/opentelemetry-proto - -grpcurl -plaintext \ - -import-path . \ - -proto opentelemetry/proto/collector/trace/v1/trace_service.proto \ - -proto opentelemetry/proto/common/v1/common.proto \ - -proto opentelemetry/proto/resource/v1/resource.proto \ - -proto opentelemetry/proto/trace/v1/trace.proto \ - -d @ \ - localhost:21890 \ - opentelemetry.proto.collector.trace.v1.TraceService/Export \ - < /path/to/data-prepper-plugins/otlp-sink/src/test/resources/sample-trace.json -``` - -You should see log output from XRayOTLPSink that confirms the span data was received and parsed correctly. diff --git a/data-prepper-plugins/otlp-sink/build.gradle b/data-prepper-plugins/otlp-sink/build.gradle index 6255275fa6..2e41ea1961 100644 --- a/data-prepper-plugins/otlp-sink/build.gradle +++ b/data-prepper-plugins/otlp-sink/build.gradle @@ -35,6 +35,9 @@ dependencies { implementation 'software.amazon.awssdk:http-client-spi' implementation 'software.amazon.awssdk:apache-client' + // Hibernate + implementation 'org.hibernate.validator:hibernate-validator:8.0.1.Final' + // OkHttp implementation 'com.squareup.okhttp3:okhttp:4.12.0' diff --git a/data-prepper-plugins/otlp-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkIT.java b/data-prepper-plugins/otlp-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkIT.java deleted file mode 100644 index 5630f4675f..0000000000 --- a/data-prepper-plugins/otlp-sink/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkIT.java +++ /dev/null @@ -1,98 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.sink.otlp; - -import io.micrometer.core.instrument.Counter; -import io.micrometer.core.instrument.DistributionSummary; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.opensearch.dataprepper.metrics.PluginMetrics; -import org.opensearch.dataprepper.model.configuration.PluginSetting; -import org.opensearch.dataprepper.model.record.Record; -import org.opensearch.dataprepper.model.trace.JacksonStandardSpan; -import org.opensearch.dataprepper.model.trace.Span; -import org.opensearch.dataprepper.plugins.otel.codec.OTelProtoStandardCodec; -import org.opensearch.dataprepper.plugins.sink.otlp.configuration.OtlpSinkConfig; -import org.opensearch.dataprepper.plugins.sink.otlp.http.OtlpHttpSender; -import software.amazon.awssdk.regions.Region; - -import java.time.Instant; -import java.util.Collections; -import java.util.List; - -import static org.mockito.ArgumentMatchers.anyString; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; - -/** - * Integration test for OtlpSink. Requires AWS credentials to be set up in the environment. - * This will not run as part of the Data Prepper build. - *

- * ./gradlew :data-prepper-plugins:otlp-sink:integrationTest \ - * -Dtests.xray.region=us-west-2 \ - * -Dtests.xray.profile=dev - */ -class OtlpSinkIT { - - private OtlpSinkConfig mockConfig; - private PluginMetrics mockPluginMetrics; - private PluginSetting mockPluginSetting; - private OtlpSink target; - - @BeforeEach - void setUp() { - System.setProperty("aws.accessKeyId", "dummy"); - System.setProperty("aws.secretAccessKey", "dummy"); - System.setProperty("aws.region", Region.US_WEST_2.toString()); - - mockConfig = mock(OtlpSinkConfig.class); - when(mockConfig.getMaxRetries()).thenReturn(3); - when(mockConfig.getBatchSize()).thenReturn(100); - - mockPluginMetrics = mock(PluginMetrics.class); - mockPluginMetrics = mock(PluginMetrics.class); - when(mockPluginMetrics.counter(anyString())).thenReturn(mock(Counter.class)); - when(mockPluginMetrics.summary(anyString())).thenReturn(mock(DistributionSummary.class)); - - mockPluginSetting = mock(PluginSetting.class); - when(mockPluginSetting.getPipelineName()).thenReturn("otlp_pipeline"); - when(mockPluginSetting.getName()).thenReturn("otlp"); - - target = new OtlpSink(mockConfig, mockPluginMetrics, mockPluginSetting); - } - - @AfterEach - void cleanUp() { - System.clearProperty("aws.accessKeyId"); - System.clearProperty("aws.secretAccessKey"); - System.clearProperty("aws.region"); - } - - /** - * This test is not part of the Data Prepper build. It requires AWS credentials to be set up in the environment. - */ - @Test - void testSinkProcessesHardcodedSpan() { - final Span testSpan = JacksonStandardSpan.builder() - .withTraceId("0123456789abcdef0123456789abcdef") - .withSpanId("0123456789abcdef") - .withParentSpanId("1111111111111111") - .withName("my-test-span") - .withStartTime(Instant.now().toString()) - .withEndTime(Instant.now().plusMillis(10).toString()) - .withAttributes(Collections.emptyMap()) - .withKind("INTERNAL") - .build(); - - final Record record = new Record<>(testSpan); - final OtlpSink sink = new OtlpSink(mockConfig, mockPluginMetrics, mockPluginSetting, mock(OTelProtoStandardCodec.OTelProtoEncoder.class), mock(OtlpHttpSender.class)); - - sink.initialize(); - sink.output(List.of(record)); - sink.shutdown(); - } -} diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java index 494fa1bc8e..f3c756d387 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java @@ -5,48 +5,32 @@ package org.opensearch.dataprepper.plugins.sink.otlp; -import com.google.common.annotations.VisibleForTesting; -import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceRequest; -import io.opentelemetry.proto.trace.v1.ResourceSpans; -import org.jetbrains.annotations.Nullable; import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor; import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.model.sink.AbstractSink; import org.opensearch.dataprepper.model.sink.Sink; import org.opensearch.dataprepper.model.trace.Span; -import org.opensearch.dataprepper.plugins.otel.codec.OTelProtoStandardCodec; +import org.opensearch.dataprepper.plugins.sink.otlp.buffer.OtlpSinkBuffer; import org.opensearch.dataprepper.plugins.sink.otlp.configuration.OtlpSinkConfig; -import org.opensearch.dataprepper.plugins.sink.otlp.http.OtlpHttpSender; import org.opensearch.dataprepper.plugins.sink.otlp.metrics.OtlpSinkMetrics; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; import javax.annotation.Nonnull; -import java.time.Duration; -import java.time.Instant; -import java.util.ArrayList; import java.util.Collection; -import java.util.List; -import java.util.Objects; -import java.util.stream.Collectors; /** - * A Data Prepper Sink plugin that forwards spans to OTLP endpoint using + * OTLP Sink Plugin for Data Prepper. */ @DataPrepperPlugin( name = "otlp", pluginType = Sink.class, pluginConfigurationType = OtlpSinkConfig.class ) -public class OtlpSink implements Sink> { +public class OtlpSink extends AbstractSink> { - private static final Logger LOG = LoggerFactory.getLogger(OtlpSink.class); - - private final int batchSize; - private final OtlpHttpSender httpSender; - private final OTelProtoStandardCodec.OTelProtoEncoder encoder; + private final OtlpSinkBuffer buffer; private final OtlpSinkMetrics sinkMetrics; /** @@ -58,75 +42,29 @@ public class OtlpSink implements Sink> { */ @DataPrepperPluginConstructor public OtlpSink(@Nonnull final OtlpSinkConfig config, @Nonnull final PluginMetrics pluginMetrics, @Nonnull final PluginSetting pluginSetting) { - this(config, pluginMetrics, pluginSetting, null, null); - } + super(pluginSetting); - /** - * Constructor for the OTLP sink plugin. Used for testing ONLY. - */ - @VisibleForTesting - OtlpSink(@Nonnull final OtlpSinkConfig config, @Nonnull final PluginMetrics pluginMetrics, @Nonnull final PluginSetting pluginSetting, final OTelProtoStandardCodec.OTelProtoEncoder encoder, final OtlpHttpSender httpSender) { - this.batchSize = config.getBatchSize(); this.sinkMetrics = new OtlpSinkMetrics(pluginMetrics, pluginSetting); - - if (encoder == null && httpSender == null) { - this.encoder = new OTelProtoStandardCodec.OTelProtoEncoder(); - this.httpSender = new OtlpHttpSender(config, sinkMetrics); - } else { - this.encoder = encoder; - this.httpSender = httpSender; - } + this.buffer = new OtlpSinkBuffer(config, sinkMetrics); } /** - * Initializes the sink. Called once during pipeline startup. + * Initialize the buffer */ @Override - public void initialize() { - LOG.debug("Initialized OTLP Sink"); + public void doInitialize() { + buffer.start(); } /** - * Processes a batch of spans and sends them to the OTLP endpoint. + * Implement the sink's output logic * - * @param records a collection of span records + * @param records Records to be output */ @Override - public void output(@Nonnull final Collection> records) { - sinkMetrics.incrementRecordsIn(records.size()); - - final List> recordList = new ArrayList<>(records); - for (int i = 0; i < recordList.size(); i += this.batchSize) { - final int end = Math.min(i + this.batchSize, recordList.size()); - final List> batch = recordList.subList(i, end); - final List resourceSpans = batch.stream() - .map(Record::getData) - .map(this::getResourceSpans) - .filter(Objects::nonNull) - .collect(Collectors.toList()); - - if (resourceSpans.isEmpty()) { - LOG.debug("Skipping empty span batch, nothing to send."); - continue; - } - - final ExportTraceServiceRequest request = ExportTraceServiceRequest.newBuilder() - .addAllResourceSpans(resourceSpans) - .build(); - final byte[] payload = request.toByteArray(); - httpSender.send(payload); - sinkMetrics.incrementRecordsOut(resourceSpans.size()); - } - } - - @Nullable - private ResourceSpans getResourceSpans(final Span span) { - try { - return encoder.convertToResourceSpans(span); - } catch (final Exception e) { - LOG.warn("Failed to encode span with ID [{}], skipping.", span.getSpanId(), e); - sinkMetrics.incrementErrorsCount(); - return null; + public void doOutput(@Nonnull final Collection> records) { + for (final Record record : records) { + buffer.add(record); } } @@ -137,7 +75,7 @@ private ResourceSpans getResourceSpans(final Span span) { */ @Override public boolean isReady() { - return true; + return buffer.isRunning(); } /** @@ -145,31 +83,7 @@ public boolean isReady() { */ @Override public void shutdown() { - httpSender.close(); - } - - /** - * Records the latency between when each span originally started (as specified in the span's start time) - * and when it was received by the sink (i.e., when this method is called). - *

- * This measures end-to-end ingestion latency from the span's source to the sink. - * It does not include any processing or export latency within the sink itself. - * - * @param events A collection of spans received by the sink. - */ - @Override - public void updateLatencyMetrics(@Nonnull final Collection> events) { - final Instant now = Instant.now(); - - for (final Record record : events) { - try { - final Instant startTime = Instant.parse(record.getData().getStartTime()); - final long durationMillis = Duration.between(startTime, now).toMillis(); - sinkMetrics.recordDeliveryLatency(durationMillis); - } catch (final Exception e) { - LOG.warn("Failed to parse startTime: {}", record.getData().getStartTime(), e); - sinkMetrics.incrementErrorsCount(); - } - } + super.shutdown(); + buffer.stop(); } } diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBuffer.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBuffer.java new file mode 100644 index 0000000000..2a2a435a31 --- /dev/null +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBuffer.java @@ -0,0 +1,206 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.sink.otlp.buffer; + +import com.google.common.annotations.VisibleForTesting; +import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceRequest; +import io.opentelemetry.proto.trace.v1.ResourceSpans; +import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.model.trace.Span; +import org.opensearch.dataprepper.plugins.otel.codec.OTelProtoStandardCodec; +import org.opensearch.dataprepper.plugins.sink.otlp.configuration.OtlpSinkConfig; +import org.opensearch.dataprepper.plugins.sink.otlp.http.OtlpHttpSender; +import org.opensearch.dataprepper.plugins.sink.otlp.metrics.OtlpSinkMetrics; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nonnull; +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.TimeUnit; + +/** + * A lossless, back-pressure aware buffer for OTLP sink. + *

+ * Spans submitted via {@link #add(Record)} are enqueued, batched by count, size or time, + * encoded to ResourceSpans, and flushed asynchronously over HTTP. + */ +public class OtlpSinkBuffer { + private static final Logger LOG = LoggerFactory.getLogger(OtlpSinkBuffer.class); + private static final int SAFETY_FACTOR = 10; + private static final int MIN_QUEUE_CAPACITY = 2000; + + private final BlockingQueue> queue; + private final OTelProtoStandardCodec.OTelProtoEncoder encoder; + private final OtlpHttpSender sender; + private final OtlpSinkMetrics sinkMetrics; + + private final int maxEvents; + private final long maxBatchBytes; + private final long flushTimeoutMillis; + + private final Thread workerThread; + private volatile boolean running = true; + + /** + * Creates a new OTLP sink buffer using default encoder and HTTP sender. + * + * @param config the OTLP sink configuration + * @param sinkMetrics the metrics collector to use + */ + public OtlpSinkBuffer(@Nonnull final OtlpSinkConfig config, @Nonnull final OtlpSinkMetrics sinkMetrics) { + this(config, sinkMetrics, null, null); + } + + /** + * Visible for testing only: constructs an OTLP sink buffer with injected encoder and sender. + * + * @param config the OTLP sink configuration + * @param sinkMetrics the metrics collector + * @param encoder custom OTLP encoder (or null to use default) + * @param sender custom HTTP sender (or null to use default) + */ + @VisibleForTesting + OtlpSinkBuffer(@Nonnull final OtlpSinkConfig config, + @Nonnull final OtlpSinkMetrics sinkMetrics, + final OTelProtoStandardCodec.OTelProtoEncoder encoder, + final OtlpHttpSender sender) { + + this.sinkMetrics = sinkMetrics; + this.encoder = encoder != null ? encoder : new OTelProtoStandardCodec.OTelProtoEncoder(); + this.sender = sender != null ? sender : new OtlpHttpSender(config, sinkMetrics); + + this.maxEvents = config.getMaxEvents(); + this.maxBatchBytes = config.getMaxBatchSize(); + this.flushTimeoutMillis = config.getFlushTimeoutMillis(); + + this.queue = new LinkedBlockingQueue<>(getQueueCapacity()); + sinkMetrics.registerQueueGauges(queue); + + this.workerThread = new Thread(this::run, "otlp-sink-buffer-thread"); + } + + private int getQueueCapacity() { + return Math.max(maxEvents * SAFETY_FACTOR, MIN_QUEUE_CAPACITY); + } + + public boolean isRunning() { + return running && workerThread.isAlive(); + } + + public void start() { + running = true; + workerThread.start(); + } + + public void stop() { + running = false; + workerThread.interrupt(); + } + + /** + * Enqueues a span record for later batching and sending. + *

+ * This will block if the internal queue is full, guaranteeing + * lossless delivery. On interruption, the span is rejected and + * error metrics are incremented. + * + * @param record the span record to enqueue + */ + public void add(final Record record) { + try { + queue.put(record); // block until space available; guaranteeing lossless delivery + } catch (final InterruptedException e) { + Thread.currentThread().interrupt(); + LOG.error("Interrupted while enqueuing span", e); + sinkMetrics.incrementRejectedSpansCount(1); + sinkMetrics.incrementErrorsCount(); + } + } + + /** + * Worker loop that batches spans by count, size, or time and then flushes them. + *

+ * Continues running as long as {@link #running} is true or the queue is not empty. + * Handles encoding failures, timeout-based flush, and final flush on shutdown. + */ + private void run() { + final List batch = new ArrayList<>(); + long batchSize = 0; + long lastFlush = System.currentTimeMillis(); + + while (running || !queue.isEmpty()) { + try { + final long now = System.currentTimeMillis(); + final Record record = queue.poll(100, TimeUnit.MILLISECONDS); + + if (record != null) { + final ResourceSpans resourceSpans; + try { + resourceSpans = encoder.convertToResourceSpans(record.getData()); + } catch (final Exception e) { + LOG.error("Failed to encode span, skipping", e); + sinkMetrics.incrementRejectedSpansCount(1); + sinkMetrics.incrementErrorsCount(); + continue; + } + + batch.add(resourceSpans); + batchSize += resourceSpans.getSerializedSize(); + } + + final boolean flushBySize = batch.size() >= maxEvents || batchSize >= maxBatchBytes; + final boolean flushByTime = !batch.isEmpty() && (now - lastFlush >= flushTimeoutMillis); + + if (flushBySize || flushByTime) { + send(batch); + + batchSize = 0; + lastFlush = now; + } + + } catch (final InterruptedException e) { + Thread.currentThread().interrupt(); + LOG.error("Interrupted while polling span", e); + sinkMetrics.incrementErrorsCount(); + } + } + + // Final flush on shutdown + if (!batch.isEmpty()) { + send(batch); + } + } + + /** + * Builds an ExportTraceServiceRequest from the given batch, sends it over HTTP, + * and updates metrics on success or failure. + *

+ * The batch is cleared in all cases to prepare for the next batch. + * + * @param batch the list of ResourceSpans to send + */ + private void send(final List batch) { + final ExportTraceServiceRequest request = ExportTraceServiceRequest.newBuilder() + .addAllResourceSpans(batch) + .build(); + final byte[] payload = request.toByteArray(); + + try { + sender.send(payload); + sinkMetrics.incrementRecordsOut(batch.size()); + } catch (final IOException e) { + LOG.error("Failed to send payload.", e); + sinkMetrics.incrementRejectedSpansCount(batch.size()); + sinkMetrics.incrementErrorsCount(); + } finally { + batch.clear(); + } + } +} diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfiguration.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfig.java similarity index 92% rename from data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfiguration.java rename to data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfig.java index 8ee088d2a4..fc94f5aca6 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfiguration.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfig.java @@ -12,14 +12,11 @@ /** * Configuration class for AWS authentication settings. - * Handles region, STS role ARN, and external ID configurations required for AWS service access. * This class will be automatically wired by Data-Prepper. - * - * @since 2.6 */ @Getter @NoArgsConstructor -class AwsAuthenticationConfiguration { +class AwsAuthenticationConfig { /** * AWS region. * Must be a valid AWS region identifier (e.g., us-east-1, us-west-2). diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfig.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfig.java index 35da8a01e4..f0b6e68cfd 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfig.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfig.java @@ -6,9 +6,8 @@ import com.fasterxml.jackson.annotation.JsonProperty; import jakarta.validation.Valid; -import jakarta.validation.constraints.Max; import jakarta.validation.constraints.Min; -import jakarta.validation.constraints.Size; +import jakarta.validation.constraints.NotBlank; import lombok.Getter; import lombok.NoArgsConstructor; import software.amazon.awssdk.regions.Region; @@ -22,28 +21,40 @@ * to preserve encapsulation and maintain control over exposed configuration data. *

* This class is automatically wired by the Data Prepper framework during pipeline initialization. - * - * @since 2.6 */ @NoArgsConstructor public class OtlpSinkConfig { @Getter @JsonProperty("endpoint") - @Size(min = 1, message = "endpoint cannot be empty string") + @NotBlank(message = "endpoint is required") private String endpoint; - @Getter - @JsonProperty("batch_size") - @Min(value = 10, message = "batch_size must be at least 10") - @Max(value = 512, message = "batch_size must be at most 512") - private int batchSize = 100; - @Getter @JsonProperty("max_retries") - @Min(value = 1, message = "max_retries must be at least 1") - @Max(value = 5, message = "max_retries must be at most 5") - private int maxRetries = 3; + @Min(value = 0) + private int maxRetries = 5; + + /** + * The threshold configuration for sending spans to the OTLP endpoint. + * This field is kept private and its contents should be accessed via the generated getter methods. + * Using eager-default values and allows the configuration to be optional in the pipeline configuration. + */ + @JsonProperty("threshold") + @Valid + private ThresholdConfig thresholdConfig = new ThresholdConfig(); + + public int getMaxEvents() { + return thresholdConfig.getMaxEvents(); + } + + public long getMaxBatchSize() { + return thresholdConfig.getMaxBatchSize().getBytes(); + } + + public long getFlushTimeoutMillis() { + return thresholdConfig.getFlushTimeout().toMillis(); + } /** * AWS authentication configuration. @@ -52,29 +63,29 @@ public class OtlpSinkConfig { */ @JsonProperty("aws") @Valid - private AwsAuthenticationConfiguration awsAuthenticationConfiguration; + private AwsAuthenticationConfig awsAuthenticationConfig; public Region getAwsRegion() { - if (awsAuthenticationConfiguration == null) { + if (awsAuthenticationConfig == null) { return null; } - return awsAuthenticationConfiguration.getAwsRegion(); + return awsAuthenticationConfig.getAwsRegion(); } public String getStsRoleArn() { - if (awsAuthenticationConfiguration == null) { + if (awsAuthenticationConfig == null) { return null; } - return awsAuthenticationConfiguration.getAwsStsRoleArn(); + return awsAuthenticationConfig.getAwsStsRoleArn(); } public String getStsExternalId() { - if (awsAuthenticationConfiguration == null) { + if (awsAuthenticationConfig == null) { return null; } - return awsAuthenticationConfiguration.getAwsStsExternalId(); + return awsAuthenticationConfig.getAwsStsExternalId(); } } diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/ThresholdConfig.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/ThresholdConfig.java new file mode 100644 index 0000000000..0375566a9d --- /dev/null +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/ThresholdConfig.java @@ -0,0 +1,35 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.sink.otlp.configuration; + +import com.fasterxml.jackson.annotation.JsonProperty; +import jakarta.validation.constraints.Min; +import lombok.Getter; +import lombok.NoArgsConstructor; +import org.hibernate.validator.constraints.time.DurationMin; +import org.opensearch.dataprepper.model.types.ByteCount; + +import java.time.Duration; + +/** + * Configuration class for threshold settings. + * This class will be automatically wired by Data-Prepper. + */ +@NoArgsConstructor +@Getter +class ThresholdConfig { + + @JsonProperty("max_events") + @Min(value = 1, message = "max_events must be at least 1") + private int maxEvents = 512; + + @JsonProperty("max_batch_size") + private ByteCount maxBatchSize = ByteCount.parse("1mb"); + + @JsonProperty("flush_timeout") + @DurationMin(millis = 1, message = "flush_timeout must be at least 1ms") + private Duration flushTimeout = Duration.ofMillis(200); +} diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressor.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressor.java index 9b0b9a3f70..d548e2d7db 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressor.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressor.java @@ -16,6 +16,9 @@ import java.util.function.Function; import java.util.zip.GZIPOutputStream; +/** + * Perform GZIP-compression on OTLP byte payloads. + */ class GzipCompressor implements Function> { private static final Logger LOG = LoggerFactory.getLogger(GzipCompressor.class); private final OtlpSinkMetrics sinkMetrics; @@ -42,6 +45,7 @@ public Optional apply(final byte[] payload) { return Optional.of(compressInternal(payload)); } catch (final IOException e) { LOG.error("Failed to compress payload", e); + sinkMetrics.incrementRejectedSpansCount(1); sinkMetrics.incrementErrorsCount(); return Optional.empty(); } diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java index c7cda9de62..67d1d47f03 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java @@ -25,11 +25,12 @@ import java.util.List; import java.util.Optional; import java.util.Set; +import java.util.concurrent.TimeUnit; import java.util.function.Consumer; import java.util.function.Function; /** - * Responsible for sending signed OTLP Protobuf trace data to AWS OTLP endpoint using OkHttp. + * Responsible for sending signed OTLP Protobuf requests to OTLP endpoint using OkHttp. */ public class OtlpHttpSender implements AutoCloseable { @VisibleForTesting @@ -51,6 +52,7 @@ public class OtlpHttpSender implements AutoCloseable { /** * Constructor for the OtlpHttpSender. * Initializes the signer and HTTP client. + * * @param config The configuration for the OTLP sink plugin. * @param sinkMetrics The metrics for the OTLP sink plugin. */ @@ -68,21 +70,32 @@ public OtlpHttpSender(@Nonnull final OtlpSinkConfig config, @Nonnull final OtlpS this.sinkMetrics = sinkMetrics; this.gzipCompressor = gzipCompressor; this.signer = signer != null ? signer : new SigV4Signer(config); - this.httpClient = httpClient != null ? httpClient : new OkHttpClient(); this.sleeper = sleeper != null ? sleeper : new ThreadSleeper(); + this.httpClient = httpClient != null ? httpClient : buildOkHttpClient(config.getFlushTimeoutMillis()); this.retryDelaysMs = generateExponentialBackoffDelays(config.getMaxRetries()); this.maxRetries = config.getMaxRetries(); } + private static OkHttpClient buildOkHttpClient(final long flushTimeoutMs) { + final long httpTimeoutMs = Math.min( + Math.max(flushTimeoutMs * 2, 3_000), + 10_000 + ); + + return new OkHttpClient.Builder() + .callTimeout(httpTimeoutMs, TimeUnit.MILLISECONDS) + .build(); + } + /** * Generates exponential backoff delays with jitter. * * @param retries Number of retries. * @return List of delay durations in milliseconds. */ - private List generateExponentialBackoffDelays(final int retries) { - List delays = new ArrayList<>(); + private static List generateExponentialBackoffDelays(final int retries) { + final List delays = new ArrayList<>(); for (int i = 0; i < retries; i++) { // Exponential backoff: 100ms, 200ms, 400ms, ... delays.add(BASE_RETRY_DELAY_MS * (1 << i)); @@ -92,12 +105,13 @@ private List generateExponentialBackoffDelays(final int retries) { } /** - * Sends the provided OTLP Protobuf trace data to the OTLP endpoint. + * Sends the provided OTLP Protobuf payload to the OTLP endpoint. * Retries with exponential backoff and jitter on failure. * - * @param payload The OTLP Protobuf-encoded trace data to be sent. + * @param payload The OTLP Protobuf-encoded data to be sent. + * @throws IOException when failed to send the payload. */ - public void send(@Nonnull final byte[] payload) { + public void send(@Nonnull final byte[] payload) throws IOException { for (int attempt = 0; attempt <= maxRetries; attempt++) { try { final Optional compressedPayload = gzipCompressor.apply(payload); @@ -135,20 +149,16 @@ public void send(@Nonnull final byte[] payload) { final int jitter = random.nextInt(100); final int retryIndex = Math.min(attempt, retryDelaysMs.size() - 1); final int delay = retryDelaysMs.get(retryIndex) + jitter; - LOG.debug("Retrying after failure in attempt {}. Sleeping {}ms.", attempt + 1, delay, ioException); - sinkMetrics.incrementRetriesCount(); try { sleeper.accept(delay); + + LOG.info("Retrying after failure in attempt {}. Sleeping {}ms.", attempt + 1, delay, ioException); + sinkMetrics.incrementRetriesCount(); } catch (final RuntimeException runtimeException) { - Thread.currentThread().interrupt(); - LOG.error("Interrupted while sleeping between retries", runtimeException); - sinkMetrics.incrementErrorsCount(); - return; + throw new IOException("Sender failed to sleep before retrying.", runtimeException); } } else { - LOG.error("Failed to sign/send data after all retries", ioException); - sinkMetrics.incrementErrorsCount(); - return; + throw ioException; } } } @@ -176,18 +186,16 @@ private void handleResponse(@Nonnull final Response response) throws IOException final String responseBody = responseBytes != null ? new String(responseBytes, StandardCharsets.UTF_8) : ""; if (NON_RETRYABLE_STATUS_CODES.contains(status)) { - LOG.error("Non-retryable client error. Status: {}, Response: {}", status, responseBody); + LOG.error("Non-retryable error. Status: {}, Response: {}", status, responseBody); return; } - final String errorMsg = String.format("Failed to send OTLP data. Status: %d, Response: %s", status, responseBody); - LOG.error(errorMsg); - throw new IOException(errorMsg); + throw new IOException(String.format("Failed to send OTLP data. Status: %d, Response: %s", status, responseBody)); } private void handleSuccessfulResponse(final byte[] responseBytes) { if (responseBytes == null || responseBytes.length == 0) { - LOG.debug("OTLP export successful. No response body."); + LOG.info("OTLP export successful. No response body."); return; } @@ -199,8 +207,9 @@ private void handleSuccessfulResponse(final byte[] responseBytes) { final long rejectedSpans = partial.getRejectedSpans(); final String errorMessage = partial.getErrorMessage(); if (rejectedSpans > 0 || !errorMessage.isEmpty()) { - LOG.warn("OTLP Partial Success: rejectedSpans={}, message={}", rejectedSpans, errorMessage); + LOG.error("OTLP Partial Success: rejectedSpans={}, message={}", rejectedSpans, errorMessage); sinkMetrics.incrementRejectedSpansCount(rejectedSpans); + sinkMetrics.incrementErrorsCount(); } } } catch (final Exception e) { diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4Signer.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4Signer.java index 92f3500294..264aa1c62e 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4Signer.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4Signer.java @@ -59,11 +59,11 @@ class SigV4Signer { } - private Region resolveRegion(final Region region) { + private static Region resolveRegion(final Region region) { return region != null ? region : DefaultAwsRegionProviderChain.builder().build().getRegion(); } - private AwsCredentialsProvider initCredentialsProvider( + private static AwsCredentialsProvider initCredentialsProvider( @Nonnull final Region region, final String stsRoleArn, final String stsExternalId, diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeper.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeper.java index 8355d40290..bd9727919f 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeper.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeper.java @@ -8,16 +8,21 @@ import javax.annotation.Nonnull; import java.util.function.Consumer; +/** + * A simple {@link Consumer} that pauses the current thread for a given number + * of milliseconds. + */ class ThreadSleeper implements Consumer { /** * Sleeps for the specified duration in milliseconds. - * If the thread is interrupted while sleeping, the interrupted status is cleared. + * If the thread is interrupted while sleeping, the interrupted status is cleared + * and a {@link RuntimeException} is thrown to signal the failure. * - * @param millis the input argument + * @param millis the number of milliseconds to sleep */ @Override - public void accept(final @Nonnull Integer millis) { + public void accept(@Nonnull final Integer millis) { try { Thread.sleep(millis); } catch (final InterruptedException e) { diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetrics.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetrics.java index b492b79e8e..81d2b38c70 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetrics.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetrics.java @@ -13,9 +13,14 @@ import javax.annotation.Nonnull; import java.time.Duration; +import java.util.concurrent.BlockingQueue; /** - * Metrics class for the otlp-sink + * A central metrics facade for the OTLP sink plugin. + *

+ * All OTLP sink components should use this class to record and expose common metrics + * such as record counts, latencies, payload sizes, and queue statistics. + *

*/ public class OtlpSinkMetrics { @@ -54,7 +59,7 @@ public OtlpSinkMetrics(@Nonnull final PluginMetrics pluginMetrics, @Nonnull fina */ private static Timer buildLatencyTimer(@Nonnull final String pipelineName, @Nonnull final String pluginName, @Nonnull final String metricName) { return Timer.builder(String.format("%s_%s_%s", pipelineName, pluginName, metricName)) - .publishPercentiles(0.5, 0.9, 0.95, 1.0) + .publishPercentiles(0.5, 0.9, 0.95, 0.99, 1.0) .publishPercentileHistogram(true) .distributionStatisticBufferLength(1024) .distributionStatisticExpiry(Duration.ofMinutes(10)) @@ -72,18 +77,14 @@ private static Timer buildLatencyTimer(@Nonnull final String pipelineName, @Nonn private static DistributionSummary buildDistributionSummary(@Nonnull final String pipelineName, @Nonnull final String pluginName, @Nonnull final String metricName) { return DistributionSummary.builder(String.format("%s_%s_%s", pipelineName, pluginName, metricName)) .baseUnit("bytes") - .publishPercentiles(0.5, 0.9, 0.95, 1.0) + .publishPercentiles(0.5, 0.9, 0.95, 0.99, 1.0) .publishPercentileHistogram(true) .distributionStatisticBufferLength(1024) .distributionStatisticExpiry(Duration.ofMinutes(10)) .register(Metrics.globalRegistry); } - public void incrementRecordsIn(long count) { - pluginMetrics.counter("recordsIn").increment(count); - } - - public void incrementRecordsOut(long count) { + public void incrementRecordsOut(final long count) { pluginMetrics.counter("recordsOut").increment(count); } @@ -91,19 +92,19 @@ public void incrementErrorsCount() { pluginMetrics.counter("errorsCount").increment(1); } - public void incrementPayloadSize(long bytes) { + public void incrementPayloadSize(final long bytes) { payloadSize.record(bytes); } - public void incrementPayloadGzipSize(long bytes) { + public void incrementPayloadGzipSize(final long bytes) { payloadGzipSize.record(bytes); } - public void recordDeliveryLatency(long durationMillis) { + public void recordDeliveryLatency(final long durationMillis) { deliveryLatency.record(Duration.ofMillis(durationMillis)); } - public void recordHttpLatency(long durationMillis) { + public void recordHttpLatency(final long durationMillis) { httpLatency.record(Duration.ofMillis(durationMillis)); } @@ -111,10 +112,15 @@ public void incrementRetriesCount() { pluginMetrics.counter("retriesCount").increment(1); } - public void incrementRejectedSpansCount(long count) { + public void incrementRejectedSpansCount(final long count) { pluginMetrics.counter("rejectedSpansCount").increment(count); } + public void registerQueueGauges(final BlockingQueue queue) { + pluginMetrics.gauge("queueSize", queue, BlockingQueue::size); + pluginMetrics.gauge("queueCapacity", queue, q -> q.remainingCapacity() + q.size()); + } + /** * Records the response code in the metrics. * Group response codes by category: 2xx, 4xx, 5xx, etc. @@ -122,7 +128,7 @@ public void incrementRejectedSpansCount(long count) { * @param statusCode The HTTP response code. */ public void recordResponseCode(final int statusCode) { - String codeCategory = (statusCode / 100) + "xx"; + final String codeCategory = (statusCode / 100) + "xx"; pluginMetrics.counter("http_" + codeCategory + "_responses").increment(); } } diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java index 68483d9da7..5c758650bb 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java @@ -2,14 +2,8 @@ * Copyright OpenSearch Contributors * SPDX-License-Identifier: Apache-2.0 */ - package org.opensearch.dataprepper.plugins.sink.otlp; -import io.micrometer.core.instrument.Counter; -import io.micrometer.core.instrument.DistributionSummary; -import io.micrometer.core.instrument.Timer; -import io.opentelemetry.proto.trace.v1.ResourceSpans; -import org.apache.commons.codec.DecoderException; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -17,187 +11,101 @@ import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.model.trace.Span; -import org.opensearch.dataprepper.plugins.otel.codec.OTelProtoStandardCodec; +import org.opensearch.dataprepper.plugins.sink.otlp.buffer.OtlpSinkBuffer; import org.opensearch.dataprepper.plugins.sink.otlp.configuration.OtlpSinkConfig; -import org.opensearch.dataprepper.plugins.sink.otlp.http.OtlpHttpSender; import software.amazon.awssdk.regions.Region; -import java.io.UnsupportedEncodingException; -import java.time.Instant; -import java.util.ArrayList; -import java.util.Collection; +import java.lang.reflect.Field; import java.util.List; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyDouble; -import static org.mockito.ArgumentMatchers.anyString; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; class OtlpSinkTest { - - private OtlpSinkConfig mockConfig; - private OTelProtoStandardCodec.OTelProtoEncoder mockEncoder; - private OtlpHttpSender mockSender; - private PluginMetrics mockPluginMetrics; - private PluginSetting mockPluginSetting; private OtlpSink target; + private OtlpSinkBuffer mockBuffer; + private OtlpSinkConfig mockConfig; + private PluginMetrics mockMetrics; + private PluginSetting mockSetting; @BeforeEach - void setUp() { - System.setProperty("aws.accessKeyId", "dummy"); - System.setProperty("aws.secretAccessKey", "dummy"); + void setUp() throws Exception { + System.setProperty("aws.region", Region.US_WEST_2.id()); + // Arrange: stub out config, metrics, setting mockConfig = mock(OtlpSinkConfig.class); - when(mockConfig.getAwsRegion()).thenReturn(Region.US_WEST_2); - when(mockConfig.getBatchSize()).thenReturn(100); - when(mockConfig.getMaxRetries()).thenReturn(3); - - mockEncoder = mock(OTelProtoStandardCodec.OTelProtoEncoder.class); - mockSender = mock(OtlpHttpSender.class); - - mockPluginMetrics = mock(PluginMetrics.class); - mockPluginMetrics = mock(PluginMetrics.class); - when(mockPluginMetrics.counter(anyString())).thenReturn(mock(Counter.class)); - when(mockPluginMetrics.summary(anyString())).thenReturn(mock(DistributionSummary.class)); - when(mockPluginMetrics.timer(anyString())).thenReturn(mock(Timer.class)); - - mockPluginSetting = mock(PluginSetting.class); - when(mockPluginSetting.getPipelineName()).thenReturn("otlp_pipeline"); - when(mockPluginSetting.getName()).thenReturn("otlp"); - - target = new OtlpSink(mockConfig, mockPluginMetrics, mockPluginSetting, mockEncoder, mockSender); + mockMetrics = mock(PluginMetrics.class); + mockSetting = mock(PluginSetting.class); + when(mockSetting.getPipelineName()).thenReturn("pipeline"); + when(mockSetting.getName()).thenReturn("otlp"); + + // Create the real sink + target = new OtlpSink(mockConfig, mockMetrics, mockSetting); + + // Replace its private buffer with a mock + mockBuffer = mock(OtlpSinkBuffer.class); + final Field bufferField = OtlpSink.class.getDeclaredField("buffer"); + bufferField.setAccessible(true); + bufferField.set(target, mockBuffer); } @AfterEach - void cleanUp() { - System.clearProperty("aws.accessKeyId"); - System.clearProperty("aws.secretAccessKey"); + void tearDown() { System.clearProperty("aws.region"); } @Test - void testOutput_shouldSendAllBatches() throws Exception { - final int recordCount = 250; - final List> records = new ArrayList<>(); - for (int i = 0; i < recordCount; i++) { - Record mockRecord = mock(Record.class); - Span span = mock(Span.class); - ResourceSpans resourceSpans = ResourceSpans.getDefaultInstance(); - when(mockEncoder.convertToResourceSpans(span)).thenReturn(resourceSpans); - when(mockRecord.getData()).thenReturn(span); - records.add(mockRecord); - } - - target.output(records); - - // 250 total / 100 batch size = 3 calls to httpSender - verify(mockSender, times(3)).send(any(byte[].class)); - verify(mockEncoder, times(recordCount)).convertToResourceSpans(any()); - } - - @Test - void testOutput_shouldSendPartialBatchWhenSomeSpansSucceed() throws Exception { - // Good span - Span goodSpan = mock(Span.class); - when(goodSpan.getSpanId()).thenReturn("good-span"); - Record goodRecord = mock(Record.class); - when(goodRecord.getData()).thenReturn(goodSpan); - - // Bad span (encoder throws) - Span badSpan = mock(Span.class); - when(badSpan.getSpanId()).thenReturn("bad-span"); - Record badRecord = mock(Record.class); - when(badRecord.getData()).thenReturn(badSpan); - - // Good span gets encoded properly - ResourceSpans goodResourceSpans = ResourceSpans.getDefaultInstance(); - when(mockEncoder.convertToResourceSpans(goodSpan)).thenReturn(goodResourceSpans); - - // Bad span causes exception during encode - when(mockEncoder.convertToResourceSpans(badSpan)).thenThrow(new RuntimeException("bad span")); - - target.output(List.of(badRecord, goodRecord)); - - // Encoder is called on both - verify(mockEncoder).convertToResourceSpans(badSpan); - verify(mockEncoder).convertToResourceSpans(goodSpan); - - // Sender should still be called with only the good span - verify(mockSender, times(1)).send(any(byte[].class)); - } - - @Test - void testOutput_shouldNotSendEmptyBatch() throws DecoderException, UnsupportedEncodingException { - // Given: a span that fails encoding - final Span span = mock(Span.class); - when(span.getSpanId()).thenReturn("bad-span"); - final Record record = mock(Record.class); - when(record.getData()).thenReturn(span); - - // Simulate encoding failure (e.g., due to missing fields) - when(mockEncoder.convertToResourceSpans(span)).thenReturn(null); + void testInitialize_startsBuffer() { + // Act + target.initialize(); - // When: the output method is called with that bad span - target.output(List.of(record)); - - // Then: it should not send anything to the OTLP endpoint - verify(mockSender, never()).send(any()); + // Assert + verify(mockBuffer).start(); } - @Test - void testUpdateLatencyMetrics_runsWithoutException_givenValidSpanStartTime() { - Span mockSpan = mock(Span.class); - String startTime = Instant.now().minusSeconds(5).toString(); - when(mockSpan.getStartTime()).thenReturn(startTime); - - Record mockRecord = mock(Record.class); - when(mockRecord.getData()).thenReturn(mockSpan); - - Collection> records = List.of(mockRecord); - - target.updateLatencyMetrics(records); - - // verify that no error counter incremented - verify(mockPluginMetrics.counter("errorsCount"), never()).increment(); + void testOutput_addsEveryRecordToBuffer() { + // Arrange + @SuppressWarnings("unchecked") final Record r1 = mock(Record.class); + @SuppressWarnings("unchecked") final Record r2 = mock(Record.class); + + // Act + target.output(List.of(r1, r2)); + + // Assert + verify(mockBuffer).add(r1); + verify(mockBuffer).add(r2); + verifyNoMoreInteractions(mockBuffer); } @Test - void testUpdateLatencyMetrics_shouldHandleInvalidStartTime() { - final Span badSpan = mock(Span.class); - when(badSpan.getStartTime()).thenReturn("invalid-timestamp"); - - final Record badRecord = mock(Record.class); - when(badRecord.getData()).thenReturn(badSpan); - - final List> records = List.of(badRecord); - - target.updateLatencyMetrics(records); - - // Ensure it attempted to parse and failed gracefully - verify(mockPluginMetrics.summary("deliveryLatency"), never()).record(anyDouble()); - verify(mockPluginMetrics.counter("errorsCount")).increment(1); + void testIsReady_delegatesToBuffer() { + // true case + when(mockBuffer.isRunning()).thenReturn(true); + assertTrue(target.isReady()); + + // false case + when(mockBuffer.isRunning()).thenReturn(false); + assertFalse(target.isReady()); } @Test - void testConstructor_withOnlyConfig_shouldInitializeWithoutException() { - OtlpSink sink = new OtlpSink(mockConfig, mockPluginMetrics, mockPluginSetting); - - sink.initialize(); - sink.shutdown(); - } + void testShutdown_stopsBuffer() { + // Act + target.shutdown(); - @Test - void testIsReady_returnsTrue() { - assert target.isReady(); + // Assert + verify(mockBuffer).stop(); } @Test - void testShutdown_doesNotThrow() { - target.shutdown(); + void testConstructor_doesNotThrow() { + // Just ensure the three-arg constructor still works + assertDoesNotThrow(() -> new OtlpSink(mockConfig, mockMetrics, mockSetting)); } } diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBufferTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBufferTest.java new file mode 100644 index 0000000000..8b64893a74 --- /dev/null +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBufferTest.java @@ -0,0 +1,276 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.sink.otlp.buffer; + +import io.opentelemetry.proto.trace.v1.ResourceSpans; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.model.trace.Span; +import org.opensearch.dataprepper.plugins.otel.codec.OTelProtoStandardCodec; +import org.opensearch.dataprepper.plugins.sink.otlp.configuration.OtlpSinkConfig; +import org.opensearch.dataprepper.plugins.sink.otlp.http.OtlpHttpSender; +import org.opensearch.dataprepper.plugins.sink.otlp.metrics.OtlpSinkMetrics; +import software.amazon.awssdk.regions.Region; + +import java.io.IOException; +import java.lang.reflect.Field; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.TimeUnit; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +class OtlpSinkBufferTest { + + private OtlpSinkConfig config; + private OtlpSinkMetrics metrics; + private OTelProtoStandardCodec.OTelProtoEncoder encoder; + private OtlpHttpSender sender; + private OtlpSinkBuffer buffer; + + @BeforeEach + void setUp() { + System.setProperty("aws.region", Region.US_WEST_2.id()); + + config = mock(OtlpSinkConfig.class); + when(config.getMaxEvents()).thenReturn(2); + when(config.getMaxBatchSize()).thenReturn(1_000_000L); + when(config.getFlushTimeoutMillis()).thenReturn(10L); + + metrics = mock(OtlpSinkMetrics.class); + encoder = mock(OTelProtoStandardCodec.OTelProtoEncoder.class); + sender = mock(OtlpHttpSender.class); + + buffer = new OtlpSinkBuffer(config, metrics, encoder, sender); + buffer.start(); + } + + @AfterEach + void tearDown() { + System.clearProperty("aws.region"); + buffer.stop(); + } + + @Test + void testIsRunningBeforeStartAndAfterStop() throws Exception { + // create a fresh buffer (not started) + final OtlpSinkBuffer localBuffer = new OtlpSinkBuffer(config, metrics, encoder, sender); + assertFalse(localBuffer.isRunning(), "not running until start() is called"); + + localBuffer.start(); + // give the thread a moment to spin up + TimeUnit.MILLISECONDS.sleep(50); + assertTrue(localBuffer.isRunning(), "should be running immediately after start()"); + + localBuffer.stop(); + // give the thread time to terminate + TimeUnit.MILLISECONDS.sleep(50); + assertFalse(localBuffer.isRunning(), "should stop after stop()"); + } + + @Test + void testAddHandlesInterruptedException() throws Exception { + // create a mock queue that throws when put(...) is called + @SuppressWarnings("unchecked") final BlockingQueue> badQueue = mock(BlockingQueue.class); + doThrow(new InterruptedException()).when(badQueue).put(any()); + + // inject it via reflection + final Field queueField = OtlpSinkBuffer.class.getDeclaredField("queue"); + queueField.setAccessible(true); + queueField.set(buffer, badQueue); + + final Record rec = mock(Record.class); + buffer.add(rec); + + // should have hit the InterruptedException branch + verify(metrics).incrementRejectedSpansCount(1); + verify(metrics).incrementErrorsCount(); + } + + @Test + void testFinalFlushOnShutdownWhenNoSizeOrTimeFlush() throws Exception { + // new config: very large batch/time so we never flush by size/time + final OtlpSinkConfig largeConfig = mock(OtlpSinkConfig.class); + when(largeConfig.getMaxEvents()).thenReturn(100); + when(largeConfig.getMaxBatchSize()).thenReturn(Long.MAX_VALUE); + when(largeConfig.getFlushTimeoutMillis()).thenReturn(Long.MAX_VALUE); + + final OtlpSinkBuffer finalBuffer = new OtlpSinkBuffer(largeConfig, metrics, encoder, sender); + finalBuffer.start(); + + // enqueue exactly one record + final Record rec = mock(Record.class); + when(rec.getData()).thenReturn(mock(Span.class)); + final ResourceSpans rs = ResourceSpans.getDefaultInstance(); + when(encoder.convertToResourceSpans(any(Span.class))).thenReturn(rs); + finalBuffer.add(rec); + + // now shutdown + finalBuffer.stop(); + TimeUnit.MILLISECONDS.sleep(50); + + // final‐flush should happen exactly once + verify(sender).send(any(byte[].class)); + verify(metrics).incrementRecordsOut(1); + } + + @Test + void testSendIoExceptionIncrementsRejectedAndError() throws Exception { + // Prepare a tiny batch so that send(...) will be invoked + final ResourceSpans rs = ResourceSpans.getDefaultInstance(); + when(encoder.convertToResourceSpans(any(Span.class))).thenReturn(rs); + doThrow(new IOException("uh-oh")).when(sender).send(any(byte[].class)); + + // Enqueue two spans to hit batch-size flush + for (int i = 0; i < 2; i++) { + final Record rec = mock(Record.class); + when(rec.getData()).thenReturn(mock(Span.class)); + buffer.add(rec); + } + + // Give worker thread time to flush by size + TimeUnit.MILLISECONDS.sleep(50); + buffer.stop(); + + // verify send was attempted + verify(sender).send(any(byte[].class)); + // one batch of 2 spans failed: rejected count should be 2 + verify(metrics).incrementRejectedSpansCount(2); + // error count for the IO error (+ possibly one more when interrupted) + verify(metrics, atLeastOnce()).incrementErrorsCount(); + } + + @Test + void testWorkerThreadFlushesBySize() throws Exception { + final ResourceSpans rs = ResourceSpans.getDefaultInstance(); + when(encoder.convertToResourceSpans(any(Span.class))).thenReturn(rs); + + // Enqueue exactly maxEvents (2) spans + for (int i = 0; i < 2; i++) { + final Record rec = mock(Record.class); + when(rec.getData()).thenReturn(mock(Span.class)); + buffer.add(rec); + } + + TimeUnit.MILLISECONDS.sleep(50); + buffer.stop(); + + // at least one send of our 2-item batch + verify(sender, atLeastOnce()).send(any(byte[].class)); + // at least one successful record-out of 2 + verify(metrics, atLeastOnce()).incrementRecordsOut(2); + } + + @Test + void testQueueCapacityRespectsMinimum() throws Exception { + when(config.getMaxEvents()).thenReturn(1); + buffer = new OtlpSinkBuffer(config, metrics, encoder, sender); + + final Field queueField = OtlpSinkBuffer.class.getDeclaredField("queue"); + queueField.setAccessible(true); + @SuppressWarnings("unchecked") final BlockingQueue queueInstance = (BlockingQueue) queueField.get(buffer); + assertEquals(2000, queueInstance.remainingCapacity()); + } + + @Test + void testQueueCapacityBasedOnMaxEvents() throws Exception { + when(config.getMaxEvents()).thenReturn(300); + buffer = new OtlpSinkBuffer(config, metrics, encoder, sender); + + final Field queueField = OtlpSinkBuffer.class.getDeclaredField("queue"); + queueField.setAccessible(true); + @SuppressWarnings("unchecked") final BlockingQueue queueInstance = (BlockingQueue) queueField.get(buffer); + assertEquals(300 * 10, queueInstance.remainingCapacity()); + } + + @Test + void testWorkerThreadFlushesByBatchByteSize() throws Exception { + // Arrange: make batchSize threshold zero so any non-null ResourceSpans triggers flush + when(config.getMaxEvents()).thenReturn(10); + when(config.getMaxBatchSize()).thenReturn(0L); + when(config.getFlushTimeoutMillis()).thenReturn(Long.MAX_VALUE); + + // restart with new config + buffer.stop(); + buffer = new OtlpSinkBuffer(config, metrics, encoder, sender); + buffer.start(); + + when(encoder.convertToResourceSpans(any(Span.class))) + .thenReturn(ResourceSpans.getDefaultInstance()); + + // Act: add one record + final Record rec = mock(Record.class); + when(rec.getData()).thenReturn(mock(Span.class)); + buffer.add(rec); + + // give the worker a moment to do both the immediate flush and then exit + TimeUnit.MILLISECONDS.sleep(50); + buffer.stop(); + + // Assert: we expect at least one send (could be two) + verify(sender, atLeastOnce()).send(any(byte[].class)); + verify(metrics, atLeastOnce()).incrementRecordsOut(1); + } + + @Test + void testWorkerThreadHandlesEncodeException() throws Exception { + // Bad span first + when(encoder.convertToResourceSpans(any(Span.class))) + .thenThrow(new RuntimeException("boom")) + .thenReturn(ResourceSpans.getDefaultInstance()); + + // Enqueue two spans: one bad, one good + for (int i = 0; i < 2; i++) { + final Record rec = mock(Record.class); + when(rec.getData()).thenReturn(mock(Span.class)); + buffer.add(rec); + } + + TimeUnit.MILLISECONDS.sleep(50); + buffer.stop(); + + // should still send at least one batch for the good span + verify(sender, atLeastOnce()).send(any(byte[].class)); + // one rejected from the encode exception + verify(metrics).incrementRejectedSpansCount(1); + // error count for the encode exception (+ maybe one on interrupt) + verify(metrics, atLeastOnce()).incrementErrorsCount(); + // at least one successful record-out of 1 + verify(metrics, atLeastOnce()).incrementRecordsOut(1); + } + + @Test + void testConstructorDefaults() throws Exception { + // use the two-arg constructor + final OtlpSinkBuffer defaultBuf = new OtlpSinkBuffer(config, metrics); + + // reflectively inspect encoder and sender + final Field encField = OtlpSinkBuffer.class.getDeclaredField("encoder"); + final Field sndField = OtlpSinkBuffer.class.getDeclaredField("sender"); + encField.setAccessible(true); + sndField.setAccessible(true); + + final Object enc = encField.get(defaultBuf); + final Object snd = sndField.get(defaultBuf); + + assertNotNull(enc, "default encoder should not be null"); + assertInstanceOf(OTelProtoStandardCodec.OTelProtoEncoder.class, enc); + + assertNotNull(snd, "default sender should not be null"); + assertInstanceOf(OtlpHttpSender.class, snd); + } +} \ No newline at end of file diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfigurationTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfigTest.java similarity index 85% rename from data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfigurationTest.java rename to data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfigTest.java index 94fe1ae98e..ea7f09fc95 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfigurationTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfigTest.java @@ -13,7 +13,7 @@ import static org.hamcrest.Matchers.nullValue; import static org.junit.jupiter.api.Assertions.assertEquals; -class AwsAuthenticationConfigurationTest { +class AwsAuthenticationConfigTest { private final String expectedRoleArn = "arn:aws:iam::123456789012:role/MyRole"; private final String expectedExternalId = "external-id-123"; @@ -28,7 +28,7 @@ void testDeserializationFromYaml() throws Exception { "sts_external_id: " + expectedExternalId ); - AwsAuthenticationConfiguration config = mapper.readValue(yaml, AwsAuthenticationConfiguration.class); + AwsAuthenticationConfig config = mapper.readValue(yaml, AwsAuthenticationConfig.class); assertEquals(expectedRegion, config.getAwsRegion().toString()); assertEquals(expectedRoleArn, config.getAwsStsRoleArn()); @@ -38,7 +38,7 @@ void testDeserializationFromYaml() throws Exception { @Test void testGetRegion_whenAllIsNull_returnNull() throws JsonProcessingException { final String yaml = "{}"; - AwsAuthenticationConfiguration config = mapper.readValue(yaml, AwsAuthenticationConfiguration.class); + AwsAuthenticationConfig config = mapper.readValue(yaml, AwsAuthenticationConfig.class); assertThat(config.getAwsRegion(), nullValue()); } diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfigTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfigTest.java index bacf23487f..9369321fb0 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfigTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfigTest.java @@ -4,51 +4,144 @@ */ package org.opensearch.dataprepper.plugins.sink.otlp.configuration; +import com.fasterxml.jackson.core.JsonParser; +import com.fasterxml.jackson.databind.DeserializationContext; +import com.fasterxml.jackson.databind.JsonDeserializer; import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.module.SimpleModule; import com.fasterxml.jackson.dataformat.yaml.YAMLFactory; +import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.model.types.ByteCount; +import software.amazon.awssdk.regions.Region; + +import java.io.IOException; +import java.time.Duration; -import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.nullValue; +import static org.junit.jupiter.api.Assertions.assertEquals; class OtlpSinkConfigTest { + private static ObjectMapper mapper; + private static final String EXPECTED_ENDPOINT = "https://example.com/otlp"; - private static final int EXPECTED_BATCH_SIZE = 128; - private static final int EXPECTED_MAX_RETRIES = 4; + private static final int DEFAULT_MAX_RETRIES = 5; + + private static final int CUSTOM_MAX_EVENTS = 100; + private static final String CUSTOM_BATCH_SIZE = "2mb"; + private static final long CUSTOM_BATCH_BYTES = ByteCount.parse(CUSTOM_BATCH_SIZE).getBytes(); + private static final long CUSTOM_FLUSH_TIMEOUT = 500L; + + private static final int DEFAULT_MAX_EVENTS = 512; + private static final long DEFAULT_BATCH_BYTES = ByteCount.parse("1mb").getBytes(); + private static final long DEFAULT_FLUSH_TIMEOUT = 200L; + private static final String EXPECTED_REGION = "us-west-2"; private static final String EXPECTED_ROLE_ARN = "arn:aws:iam::123456789012:role/OtlpRole"; private static final String EXPECTED_EXTERNAL_ID = "my-ext-id"; + @BeforeAll + static void setupMapper() { + mapper = new ObjectMapper(new YAMLFactory()) + .findAndRegisterModules(); // for default Duration support + + // custom deserializer for ByteCount strings like "2mb" + final SimpleModule byteCountModule = new SimpleModule(); + byteCountModule.addDeserializer(ByteCount.class, new JsonDeserializer<>() { + @Override + public ByteCount deserialize(final JsonParser p, final DeserializationContext ctxt) throws IOException { + return ByteCount.parse(p.getValueAsString()); + } + }); + mapper.registerModule(byteCountModule); + + // custom deserializer for Duration strings like "500ms" + final SimpleModule durationModule = new SimpleModule(); + durationModule.addDeserializer(Duration.class, new JsonDeserializer<>() { + @Override + public Duration deserialize(final JsonParser p, final DeserializationContext ctxt) throws IOException { + final String text = p.getValueAsString(); + if (text.endsWith("ms")) { + final long ms = Long.parseLong(text.substring(0, text.length() - 2)); + return Duration.ofMillis(ms); + } + return Duration.parse(text); + } + }); + mapper.registerModule(durationModule); + } + + @Test + void testMinimumConfigDefaults() throws Exception { + final String yaml = "endpoint: \"" + EXPECTED_ENDPOINT + "\""; + + final OtlpSinkConfig config = mapper.readValue(yaml, OtlpSinkConfig.class); + + assertEquals(EXPECTED_ENDPOINT, config.getEndpoint()); + assertEquals(DEFAULT_MAX_RETRIES, config.getMaxRetries()); + assertEquals(DEFAULT_MAX_EVENTS, config.getMaxEvents()); + assertEquals(DEFAULT_BATCH_BYTES, config.getMaxBatchSize()); + assertEquals(DEFAULT_FLUSH_TIMEOUT, config.getFlushTimeoutMillis()); + + assertThat(config.getAwsRegion(), nullValue()); + assertThat(config.getStsRoleArn(), nullValue()); + assertThat(config.getStsExternalId(), nullValue()); + } + + @Test + void testCustomThresholdAndRetries() throws Exception { + final String yaml = String.join("\n", + "endpoint: \"" + EXPECTED_ENDPOINT + "\"", + "max_retries: 3", + "threshold:", + " max_events: " + CUSTOM_MAX_EVENTS, + " max_batch_size: \"" + CUSTOM_BATCH_SIZE + "\"", + " flush_timeout: \"" + CUSTOM_FLUSH_TIMEOUT + "ms\"" + ); + + final OtlpSinkConfig config = mapper.readValue(yaml, OtlpSinkConfig.class); + + assertEquals(EXPECTED_ENDPOINT, config.getEndpoint()); + assertEquals(3, config.getMaxRetries()); + assertEquals(CUSTOM_MAX_EVENTS, config.getMaxEvents()); + assertEquals(CUSTOM_BATCH_BYTES, config.getMaxBatchSize()); + assertEquals(CUSTOM_FLUSH_TIMEOUT, config.getFlushTimeoutMillis()); + } + @Test - void testDeserializationFromYaml() throws Exception { - final String yaml = - "endpoint: \"" + EXPECTED_ENDPOINT + "\"\n" + - "batch_size: " + EXPECTED_BATCH_SIZE + "\n" + - "max_retries: " + EXPECTED_MAX_RETRIES + "\n" + - "aws:\n" + - " region: \"" + EXPECTED_REGION + "\"\n" + - " sts_role_arn: \"" + EXPECTED_ROLE_ARN + "\"\n" + - " sts_external_id: \"" + EXPECTED_EXTERNAL_ID + "\"\n"; - - final ObjectMapper objectMapper = new ObjectMapper(new YAMLFactory()); - final OtlpSinkConfig config = objectMapper.readValue(yaml, OtlpSinkConfig.class); - - assertThat(config.getEndpoint(), equalTo(EXPECTED_ENDPOINT)); - assertThat(config.getBatchSize(), equalTo(EXPECTED_BATCH_SIZE)); - assertThat(config.getMaxRetries(), equalTo(EXPECTED_MAX_RETRIES)); - - assertThat(config.getAwsRegion().toString(), equalTo(EXPECTED_REGION)); - assertThat(config.getStsRoleArn(), equalTo(EXPECTED_ROLE_ARN)); - assertThat(config.getStsExternalId(), equalTo(EXPECTED_EXTERNAL_ID)); + void testAwsBlockDeserialization() throws Exception { + final String yaml = String.join("\n", + "endpoint: \"" + EXPECTED_ENDPOINT + "\"", + "max_retries: " + DEFAULT_MAX_RETRIES, + "aws:", + " region: \"" + EXPECTED_REGION + "\"", + " sts_role_arn: \"" + EXPECTED_ROLE_ARN + "\"", + " sts_external_id: \"" + EXPECTED_EXTERNAL_ID + "\"" + ); + + final OtlpSinkConfig config = mapper.readValue(yaml, OtlpSinkConfig.class); + + assertEquals(EXPECTED_ENDPOINT, config.getEndpoint()); + assertEquals(DEFAULT_MAX_RETRIES, config.getMaxRetries()); + + assertEquals(Region.of(EXPECTED_REGION), config.getAwsRegion()); + assertEquals(EXPECTED_ROLE_ARN, config.getStsRoleArn()); + assertEquals(EXPECTED_EXTERNAL_ID, config.getStsExternalId()); } @Test - void testDeserializationFromYaml_withNullAwsValues() throws Exception { - final String yaml = "{}"; - final ObjectMapper objectMapper = new ObjectMapper(new YAMLFactory()); - final OtlpSinkConfig config = objectMapper.readValue(yaml, OtlpSinkConfig.class); + void testAwsSectionMissing_staysNull() throws Exception { + final String yaml = String.join("\n", + "endpoint: \"" + EXPECTED_ENDPOINT + "\"", + "max_retries: " + DEFAULT_MAX_RETRIES + ); + + final OtlpSinkConfig config = mapper.readValue(yaml, OtlpSinkConfig.class); + + assertEquals(EXPECTED_ENDPOINT, config.getEndpoint()); + assertEquals(DEFAULT_MAX_RETRIES, config.getMaxRetries()); assertThat(config.getAwsRegion(), nullValue()); assertThat(config.getStsRoleArn(), nullValue()); diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/ThresholdConfigTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/ThresholdConfigTest.java new file mode 100644 index 0000000000..6e8d3834a6 --- /dev/null +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/ThresholdConfigTest.java @@ -0,0 +1,73 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.dataprepper.plugins.sink.otlp.configuration; + +import com.fasterxml.jackson.core.JsonParser; +import com.fasterxml.jackson.databind.DeserializationContext; +import com.fasterxml.jackson.databind.JsonDeserializer; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.module.SimpleModule; +import com.fasterxml.jackson.dataformat.yaml.YAMLFactory; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.model.types.ByteCount; + +import java.io.IOException; +import java.time.Duration; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +class ThresholdConfigTest { + + private static ObjectMapper mapper; + + private static final int CUSTOM_MAX_EVENTS = 100; + private static final String CUSTOM_MAX_BATCH_SIZE = "2mb"; + private static final Duration CUSTOM_FLUSH_TIMEOUT = Duration.ofMillis(500); + + private static final int DEFAULT_MAX_EVENTS = 512; + private static final String DEFAULT_MAX_BATCH_SIZE = "1mb"; + private static final Duration DEFAULT_FLUSH_TIMEOUT = Duration.ofMillis(200); + + @BeforeAll + static void setupMapper() { + mapper = new ObjectMapper(new YAMLFactory()) + .findAndRegisterModules(); // for Duration + + // Register a simple deserializer for ByteCount.from “2mb”-style strings + final SimpleModule byteCountModule = new SimpleModule(); + byteCountModule.addDeserializer(ByteCount.class, new JsonDeserializer<>() { + @Override + public ByteCount deserialize(final JsonParser p, final DeserializationContext ctxt) throws IOException { + return ByteCount.parse(p.getValueAsString()); + } + }); + mapper.registerModule(byteCountModule); + } + + @Test + void testDeserializationFromYaml() throws Exception { + final String yaml = String.join("\n", + "max_events: " + CUSTOM_MAX_EVENTS, + "max_batch_size: \"" + CUSTOM_MAX_BATCH_SIZE + "\"", + "flush_timeout: \"" + CUSTOM_FLUSH_TIMEOUT.toString() + "\"" + ); + + final ThresholdConfig config = mapper.readValue(yaml, ThresholdConfig.class); + + assertEquals(CUSTOM_MAX_EVENTS, config.getMaxEvents()); + assertEquals(ByteCount.parse(CUSTOM_MAX_BATCH_SIZE), config.getMaxBatchSize()); + assertEquals(CUSTOM_FLUSH_TIMEOUT, config.getFlushTimeout()); + } + + @Test + void testDefaultsWhenYamlEmpty() throws Exception { + final ThresholdConfig config = mapper.readValue("{}", ThresholdConfig.class); + + assertEquals(DEFAULT_MAX_EVENTS, config.getMaxEvents()); + assertEquals(ByteCount.parse(DEFAULT_MAX_BATCH_SIZE), config.getMaxBatchSize()); + assertEquals(DEFAULT_FLUSH_TIMEOUT, config.getFlushTimeout()); + } +} diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressorTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressorTest.java index d05b47c217..d0c5e86b44 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressorTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressorTest.java @@ -54,6 +54,7 @@ void apply_handlesIOException_andIncrementsErrorMetric() throws IOException { assertTrue(result.isEmpty()); verify(sinkMetrics).incrementErrorsCount(); + verify(sinkMetrics).incrementRejectedSpansCount(1); } private byte[] decompress(byte[] compressed) throws IOException { diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java index 74c1a378d0..e567c5cdaa 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java @@ -5,11 +5,11 @@ package org.opensearch.dataprepper.plugins.sink.otlp.http; -import com.google.protobuf.UnknownFieldSet; -import com.google.protobuf.UnknownFieldSet.Field; import io.opentelemetry.proto.collector.trace.v1.ExportTracePartialSuccess; import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceResponse; import okhttp3.Call; +import okhttp3.ConnectionPool; +import okhttp3.Dispatcher; import okhttp3.HttpUrl; import okhttp3.MediaType; import okhttp3.OkHttpClient; @@ -31,17 +31,20 @@ import java.util.Collections; import java.util.Map; import java.util.Optional; +import java.util.concurrent.ExecutorService; import java.util.function.Consumer; import java.util.function.Function; import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.fail; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.reset; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; @@ -49,7 +52,6 @@ import static org.opensearch.dataprepper.plugins.sink.otlp.http.OtlpHttpSender.NON_RETRYABLE_STATUS_CODES; class OtlpHttpSenderTest { - private static final byte[] PAYLOAD = "test-otlp-payload".getBytes(StandardCharsets.UTF_8); private static final String ERROR_BODY = "{\"error\": \"Something went wrong\"}"; @@ -76,12 +78,14 @@ void setUp() { mockSinkMetrics = mock(OtlpSinkMetrics.class); mockGzipCompressor = mock(GzipCompressor.class); - when(mockGzipCompressor.apply(any())).thenAnswer(invocation -> { - byte[] input = invocation.getArgument(0); - return input == null ? Optional.empty() : Optional.of(input); - }); - - target = new OtlpHttpSender(mockConfig, mockSinkMetrics, mockGzipCompressor, mockSigner, mockHttpClient, mockSleeper); + when(mockGzipCompressor.apply(any())) + .thenAnswer(invocation -> Optional.of((byte[]) invocation.getArgument(0))); + + target = new OtlpHttpSender( + mockConfig, mockSinkMetrics, + mockGzipCompressor, mockSigner, + mockHttpClient, mockSleeper + ); } @AfterEach @@ -93,338 +97,289 @@ void cleanUp() { @Test void testSend_successfulResponse() throws IOException { - SdkHttpFullRequest mockSignedRequest = mock(SdkHttpFullRequest.class); - when(mockSignedRequest.getUri()).thenReturn(HttpUrl.get("https://xray.us-west-2.amazonaws.com/v1/traces").uri()); - when(mockSignedRequest.headers()).thenReturn(Map.of("Authorization", Collections.singletonList("signed-header"))); - - when(mockSigner.signRequest(PAYLOAD)).thenReturn(mockSignedRequest); - - Call mockCall = mock(Call.class); - Response mockResponse = new Response.Builder() - .request(new Request.Builder().url("https://xray.us-west-2.amazonaws.com/v1/traces").build()) + // Arrange + final SdkHttpFullRequest signed = mock(SdkHttpFullRequest.class); + when(signed.getUri()).thenReturn( + HttpUrl.get("https://xray.us-west-2.amazonaws.com/v1/traces").uri()); + when(signed.headers()).thenReturn( + Map.of("Authorization", Collections.singletonList("signed-header"))); + when(mockSigner.signRequest(PAYLOAD)).thenReturn(signed); + + final Call call = mock(Call.class); + final Response resp = new Response.Builder() + .request(new Request.Builder().url(signed.getUri().toString()).build()) .protocol(Protocol.HTTP_1_1) .code(200) .message("OK") - .body(ResponseBody.create("", MediaType.get("application/x-protobuf"))) + .body(ResponseBody.create( + new byte[0], + MediaType.get("application/x-protobuf"))) .build(); - when(mockHttpClient.newCall(any())).thenReturn(mockCall); - when(mockCall.execute()).thenReturn(mockResponse); + when(mockHttpClient.newCall(any())).thenReturn(call); + when(call.execute()).thenReturn(resp); + // Act & Assert assertDoesNotThrow(() -> target.send(PAYLOAD)); } @Test void testSend_doesNotRetryOnNonRetryable4xxResponses() throws IOException { - final SdkHttpFullRequest mockSignedRequest = mock(SdkHttpFullRequest.class); - when(mockSignedRequest.getUri()).thenReturn(HttpUrl.get("https://xray.us-west-2.amazonaws.com/v1/traces").uri()); - when(mockSignedRequest.headers()).thenReturn(Map.of()); - when(mockSigner.signRequest(PAYLOAD)).thenReturn(mockSignedRequest); - - final okhttp3.Request okHttpRequest = new Request.Builder() - .url("https://xray.us-west-2.amazonaws.com/v1/traces") + // Arrange + final SdkHttpFullRequest signed = mock(SdkHttpFullRequest.class); + when(signed.getUri()).thenReturn( + HttpUrl.get("https://xray.us-west-2.amazonaws.com/v1/traces").uri()); + when(signed.headers()).thenReturn(Map.of()); + when(mockSigner.signRequest(PAYLOAD)).thenReturn(signed); + + final Request okReq = new Request.Builder() + .url(signed.getUri().toString()) .build(); - for (int statusCode : NON_RETRYABLE_STATUS_CODES) { - final String responseBodyText = "Non-retryable error from server"; - - final Response mockResponse = new Response.Builder() - .request(okHttpRequest) + for (final int status : NON_RETRYABLE_STATUS_CODES) { + final Response resp = new Response.Builder() + .request(okReq) .protocol(Protocol.HTTP_1_1) - .code(statusCode) + .code(status) .message("Client Error") - .body(ResponseBody.create(responseBodyText, MediaType.get("application/json"))) + .body(ResponseBody.create( + ERROR_BODY.getBytes(StandardCharsets.UTF_8), + MediaType.get("application/json"))) .build(); + final Call call = mock(Call.class); + when(mockHttpClient.newCall(any())).thenReturn(call); + when(call.execute()).thenReturn(resp); - final Call mockCall = mock(Call.class); - when(mockCall.execute()).thenReturn(mockResponse); - when(mockHttpClient.newCall(any())).thenReturn(mockCall); - - assertDoesNotThrow(() -> target.send(PAYLOAD), "Should not throw on non-retryable status " + statusCode); + assertDoesNotThrow(() -> target.send(PAYLOAD)); verify(mockHttpClient, times(1)).newCall(any()); - - reset(mockHttpClient); // reset between iterations + reset(mockHttpClient); } } @Test void testSend_retryOnFailure_thenSuccess() throws IOException { - SdkHttpFullRequest mockSignedRequest = mock(SdkHttpFullRequest.class); - when(mockSignedRequest.getUri()).thenReturn(HttpUrl.get("https://xray.us-west-2.amazonaws.com/v1/traces").uri()); - when(mockSignedRequest.headers()).thenReturn(Map.of()); - - when(mockSigner.signRequest(PAYLOAD)).thenReturn(mockSignedRequest); - - Call mockCall1 = mock(Call.class); - Call mockCall2 = mock(Call.class); - - when(mockHttpClient.newCall(any())) - .thenReturn(mockCall1) - .thenReturn(mockCall2); - - when(mockCall1.execute()).thenThrow(new IOException("first attempt failed")); - Response successResponse = new Response.Builder() - .request(new Request.Builder().url("https://xray.us-west-2.amazonaws.com/v1/traces").build()) + // Arrange + final SdkHttpFullRequest signed = mock(SdkHttpFullRequest.class); + when(signed.getUri()).thenReturn( + HttpUrl.get("https://xray.us-west-2.amazonaws.com/v1/traces").uri()); + when(signed.headers()).thenReturn(Map.of()); + when(mockSigner.signRequest(PAYLOAD)).thenReturn(signed); + + final Call first = mock(Call.class); + final Call second = mock(Call.class); + when(mockHttpClient.newCall(any())).thenReturn(first, second); + when(first.execute()).thenThrow(new IOException("first attempt failed")); + final Response success = new Response.Builder() + .request(new Request.Builder().url(signed.getUri().toString()).build()) .protocol(Protocol.HTTP_1_1) .code(200) .message("OK") - .body(ResponseBody.create("", MediaType.get("application/x-protobuf"))) + .body(ResponseBody.create( + new byte[0], + MediaType.get("application/x-protobuf"))) .build(); - when(mockCall2.execute()).thenReturn(successResponse); + when(second.execute()).thenReturn(success); + // Act & Assert assertDoesNotThrow(() -> target.send(PAYLOAD)); } @Test - void testSend_doesNotThrowException_failsAfterAllRetries() throws IOException { - SdkHttpFullRequest mockSignedRequest = mock(SdkHttpFullRequest.class); - when(mockSignedRequest.getUri()).thenReturn(HttpUrl.get("https://xray.us-west-2.amazonaws.com/v1/traces").uri()); - when(mockSignedRequest.headers()).thenReturn(Map.of()); - - when(mockSigner.signRequest(PAYLOAD)).thenReturn(mockSignedRequest); - - Call mockCall = mock(Call.class); - when(mockHttpClient.newCall(any())).thenReturn(mockCall); - when(mockCall.execute()).thenThrow(new IOException("always fail")); - - assertDoesNotThrow(() -> target.send(PAYLOAD)); + void testSend_throwsIOException_whenFailsAfterAllRetries() throws IOException { + // Arrange + final SdkHttpFullRequest signed = mock(SdkHttpFullRequest.class); + when(signed.getUri()).thenReturn( + HttpUrl.get("https://xray.us-west-2.amazonaws.com/v1/traces").uri()); + when(signed.headers()).thenReturn(Map.of()); + when(mockSigner.signRequest(PAYLOAD)).thenReturn(signed); + + final Call alwaysFail = mock(Call.class); + when(mockHttpClient.newCall(any())).thenReturn(alwaysFail); + when(alwaysFail.execute()).thenThrow(new IOException("always fail")); + + // Act & Assert + final IOException ex = assertThrows(IOException.class, () -> target.send(PAYLOAD)); + assertEquals("always fail", ex.getMessage()); } @Test - void testSend_doesNotThrowsException_on500ResponseWithBody() throws IOException { - // Mock signed request - SdkHttpFullRequest sdkRequest = SdkHttpFullRequest.builder() + void testSend_throwsIOException_on500ResponseWithBody() throws IOException { + final SdkHttpFullRequest signed = SdkHttpFullRequest.builder() .method(software.amazon.awssdk.http.SdkHttpMethod.POST) - .uri(URI.create("https://xray.us-west-2.amazonaws.com/v1/traces")) - .putHeader("Content-Type", "application/x-protobuf") + .uri(URI.create("https://example.com")) + .putHeader("Content-Type", "application/json") .build(); + when(mockSigner.signRequest(PAYLOAD)).thenReturn(signed); - when(mockSigner.signRequest(PAYLOAD)).thenReturn(sdkRequest); + final Request okReq = new Request.Builder() + .url(signed.getUri().toString()).build(); + final Call call = mock(Call.class); + when(mockHttpClient.newCall(any())).thenReturn(call); - // Build actual OkHttp request (we need this to inject into the mocked response) - okhttp3.Request okHttpRequest = new Request.Builder() - .url(sdkRequest.getUri().toString()) - .build(); - - // Mock 500 response with body - Response mockResponse = new Response.Builder() + final Response resp500 = new Response.Builder() + .request(okReq) + .protocol(Protocol.HTTP_1_1) .code(500) .message("Internal Server Error") - .request(okHttpRequest) - .protocol(Protocol.HTTP_1_1) - .body(ResponseBody.create(ERROR_BODY, MediaType.get("application/json"))) + .body(ResponseBody.create( + ERROR_BODY.getBytes(StandardCharsets.UTF_8), + MediaType.get("application/json"))) .build(); + when(call.execute()).thenReturn(resp500); - var call = mock(okhttp3.Call.class); - when(mockHttpClient.newCall(any(Request.class))).thenReturn(call); - when(call.execute()).thenReturn(mockResponse); - - // Run test - assertDoesNotThrow(() -> target.send(PAYLOAD)); + assertThrows(IOException.class, () -> target.send(PAYLOAD)); } @Test - void testSend_doesNotThrowsException_onInterruptedExceptionDuringRetry() throws IOException { - final SdkHttpFullRequest signedRequest = mock(SdkHttpFullRequest.class); - - when(mockSigner.signRequest(any())).thenReturn(signedRequest); - when(signedRequest.getUri()).thenReturn(URI.create("https://example.com")); - when(signedRequest.headers()).thenReturn(Map.of()); - - final Call mockCall = mock(Call.class); - when(mockCall.execute()).thenThrow(new IOException("boom")); - when(mockHttpClient.newCall(any())).thenReturn(mockCall); - - doThrow(new RuntimeException("")).when(mockSleeper).accept(anyInt()); - - target = new OtlpHttpSender(mockConfig, mockSinkMetrics, mockGzipCompressor, mockSigner, mockHttpClient, mockSleeper); - - assertDoesNotThrow(() -> target.send(PAYLOAD)); - assertTrue(Thread.currentThread().isInterrupted()); + void testSend_wrapsInterruptedExceptionDuringRetryAsIOException() throws IOException { + // Arrange: first attempt throws IOException, sleeper throws at retry + final SdkHttpFullRequest signed = mock(SdkHttpFullRequest.class); + when(mockSigner.signRequest(any())).thenReturn(signed); + when(signed.getUri()).thenReturn(URI.create("https://example.com")); + when(signed.headers()).thenReturn(Map.of()); + + final Call call = mock(Call.class); + when(mockHttpClient.newCall(any())).thenReturn(call); + when(call.execute()).thenThrow(new IOException("boom")); + + doThrow(new RuntimeException("sleep failed")) + .when(mockSleeper).accept(anyInt()); + + target = new OtlpHttpSender( + mockConfig, mockSinkMetrics, + mockGzipCompressor, mockSigner, + mockHttpClient, mockSleeper + ); + + // Act & Assert + final IOException ex = assertThrows(IOException.class, () -> target.send(PAYLOAD)); + assertEquals("Sender failed to sleep before retrying.", ex.getMessage()); + assertEquals("sleep failed", ex.getCause().getMessage()); } @Test void testSend_partialSuccessResponse_logsWarning() throws IOException { - final ExportTraceServiceResponse responseProto = ExportTraceServiceResponse.newBuilder() + // Arrange + final ExportTraceServiceResponse proto = ExportTraceServiceResponse.newBuilder() .setPartialSuccess(ExportTracePartialSuccess.newBuilder() .setRejectedSpans(5) .setErrorMessage("Some spans were rejected due to invalid format") .build()) .build(); + final byte[] bytes = proto.toByteArray(); - final byte[] responseBytes = responseProto.toByteArray(); - - final SdkHttpFullRequest sdkRequest = SdkHttpFullRequest.builder() + final SdkHttpFullRequest signed = SdkHttpFullRequest.builder() .method(software.amazon.awssdk.http.SdkHttpMethod.POST) .uri(URI.create("https://xray.us-west-2.amazonaws.com/v1/traces")) .putHeader("Content-Type", "application/x-protobuf") .build(); + when(mockSigner.signRequest(PAYLOAD)).thenReturn(signed); - when(mockSigner.signRequest(PAYLOAD)).thenReturn(sdkRequest); - - final okhttp3.Request okHttpRequest = new Request.Builder() - .url(sdkRequest.getUri().toString()) + final Request okReq = new Request.Builder() + .url(signed.getUri().toString()) .build(); - - final Response mockResponse = new Response.Builder() - .request(okHttpRequest) + final Call call = mock(Call.class); + when(mockHttpClient.newCall(any())).thenReturn(call); + when(call.execute()).thenReturn(new Response.Builder() + .request(okReq) .protocol(Protocol.HTTP_1_1) .code(200) .message("OK") - .body(ResponseBody.create(responseBytes, MediaType.get("application/x-protobuf"))) - .build(); - - final Call mockCall = mock(Call.class); - when(mockHttpClient.newCall(any())).thenReturn(mockCall); - when(mockCall.execute()).thenReturn(mockResponse); + .body(ResponseBody.create(bytes, MediaType.get("application/x-protobuf"))) + .build()); + // Act & Assert assertDoesNotThrow(() -> target.send(PAYLOAD)); } @Test - void testSend_partialSuccessEmpty_logsInfo() throws IOException { - final ExportTraceServiceResponse responseProto = ExportTraceServiceResponse.newBuilder() - .setPartialSuccess(ExportTracePartialSuccess.newBuilder() - .setRejectedSpans(0) - .setErrorMessage("") // Empty = no warning - .build()) - .build(); - - final byte[] responseBytes = responseProto.toByteArray(); - - final SdkHttpFullRequest sdkRequest = SdkHttpFullRequest.builder() - .method(software.amazon.awssdk.http.SdkHttpMethod.POST) - .uri(URI.create("https://xray.us-west-2.amazonaws.com/v1/traces")) - .putHeader("Content-Type", "application/x-protobuf") - .build(); - - when(mockSigner.signRequest(PAYLOAD)).thenReturn(sdkRequest); - - final okhttp3.Request okHttpRequest = new Request.Builder() - .url(sdkRequest.getUri().toString()) - .build(); - - final Response mockResponse = new Response.Builder() - .request(okHttpRequest) - .protocol(Protocol.HTTP_1_1) - .code(200) - .message("OK") - .body(ResponseBody.create(responseBytes, MediaType.get("application/x-protobuf"))) - .build(); - - final Call mockCall = mock(Call.class); - when(mockHttpClient.newCall(any())).thenReturn(mockCall); - when(mockCall.execute()).thenReturn(mockResponse); - + void testSend_skipsSend_whenGzipCompressionFails() { + // Arrange: compressor returns empty → send() should return silently + final Function> skipCompressor = p -> Optional.empty(); + target = new OtlpHttpSender( + mockConfig, mockSinkMetrics, + skipCompressor, mockSigner, + mockHttpClient, mockSleeper + ); + + // Act & Assert: no exception, nothing signed or sent assertDoesNotThrow(() -> target.send(PAYLOAD)); + verify(mockSigner, never()).signRequest(any()); + verify(mockHttpClient, never()).newCall(any()); } @Test - void testSend_successResponseWithNonEmptyBodyNoPartialSuccess_logsInfo() throws IOException { - // Manually add an unknown field to make the serialized body non-empty - final UnknownFieldSet unknownFields = UnknownFieldSet.newBuilder() - .addField(123, Field.newBuilder().addVarint(42).build()) // dummy varint - .build(); - - final ExportTraceServiceResponse responseProto = ExportTraceServiceResponse.newBuilder() - .mergeUnknownFields(unknownFields) - .build(); - - final byte[] responseBytes = responseProto.toByteArray(); - assertTrue(responseBytes.length > 0); // sanity check - - final SdkHttpFullRequest sdkRequest = SdkHttpFullRequest.builder() - .method(software.amazon.awssdk.http.SdkHttpMethod.POST) - .uri(URI.create("https://xray.us-west-2.amazonaws.com/v1/traces")) - .putHeader("Content-Type", "application/x-protobuf") - .build(); - - when(mockSigner.signRequest(PAYLOAD)).thenReturn(sdkRequest); - - final okhttp3.Request okHttpRequest = new Request.Builder() - .url(sdkRequest.getUri().toString()) - .build(); - - final Response mockResponse = new Response.Builder() - .request(okHttpRequest) + void testHandleSuccessfulResponseParseErrorIncrementsError() throws IOException { + // arrange: build a sender with a pass-through compressor lambda + target = new OtlpHttpSender( + mockConfig, + mockSinkMetrics, + payload -> Optional.of(payload), + mockSigner, + mockHttpClient, + mockSleeper + ); + + // stub: signer + final SdkHttpFullRequest signed = mock(SdkHttpFullRequest.class); + when(signed.getUri()).thenReturn(URI.create("https://example.com")); + when(signed.headers()).thenReturn(Map.of()); + when(mockSigner.signRequest(PAYLOAD)).thenReturn(signed); + + // stub: http client returns 200 OK but with invalid protobuf bytes + final Call call = mock(Call.class); + when(mockHttpClient.newCall(any())).thenReturn(call); + final byte[] bad = "not-a-proto".getBytes(StandardCharsets.UTF_8); + final Response resp = new Response.Builder() + .request(new Request.Builder().url(signed.getUri().toString()).build()) .protocol(Protocol.HTTP_1_1) .code(200) .message("OK") - .body(ResponseBody.create(responseBytes, MediaType.get("application/x-protobuf"))) + .body(ResponseBody.create(bad, MediaType.get("application/x-protobuf"))) .build(); + when(call.execute()).thenReturn(resp); - final Call mockCall = mock(Call.class); - when(mockHttpClient.newCall(any())).thenReturn(mockCall); - when(mockCall.execute()).thenReturn(mockResponse); - + // act assertDoesNotThrow(() -> target.send(PAYLOAD)); - } - @Test - void testSend_invalidProtoResponse_logsError() throws IOException { - // Build a response with invalid OTLP proto data (random bytes) - byte[] invalidProtoBytes = "this-is-not-valid-proto".getBytes(StandardCharsets.UTF_8); - - // Mock the signed request - final SdkHttpFullRequest sdkRequest = SdkHttpFullRequest.builder() - .method(software.amazon.awssdk.http.SdkHttpMethod.POST) - .uri(URI.create("https://xray.us-west-2.amazonaws.com/v1/traces")) - .putHeader("Content-Type", "application/x-protobuf") - .build(); - - when(mockSigner.signRequest(PAYLOAD)).thenReturn(sdkRequest); - - // Build OkHttp request to satisfy response builder - final okhttp3.Request okHttpRequest = new Request.Builder() - .url(sdkRequest.getUri().toString()) - .build(); - - // Create mock response with invalid proto payload - final Response mockResponse = new Response.Builder() - .request(okHttpRequest) - .protocol(Protocol.HTTP_1_1) - .code(200) - .message("OK") - .body(ResponseBody.create(invalidProtoBytes, MediaType.get("application/x-protobuf"))) - .build(); - - final Call mockCall = mock(Call.class); - when(mockHttpClient.newCall(any())).thenReturn(mockCall); - when(mockCall.execute()).thenReturn(mockResponse); - - // Should not throw, just logs error - assertDoesNotThrow(() -> target.send(PAYLOAD)); + // assert: parse-failure should have incremented errors + verify(mockSinkMetrics).incrementErrorsCount(); } @Test - void testSend_skipsSend_whenGzipCompressionFails() { - when(mockGzipCompressor.apply(any())).thenReturn(Optional.empty()); - - assertDoesNotThrow(() -> target.send(PAYLOAD)); - - verify(mockSigner, times(0)).signRequest(any()); - verify(mockHttpClient, times(0)).newCall(any()); + void testCloseEvictsAndShutdownsOkHttpResources() { + // arrange: stub out connectionPool and dispatcher on our mockHttpClient + final ConnectionPool pool = mock(ConnectionPool.class); + final Dispatcher dispatcher = mock(Dispatcher.class); + final ExecutorService exec = mock(ExecutorService.class); + when(mockHttpClient.connectionPool()).thenReturn(pool); + when(mockHttpClient.dispatcher()).thenReturn(dispatcher); + when(dispatcher.executorService()).thenReturn(exec); + + // act + target.close(); + + // assert + verify(pool).evictAll(); + verify(exec).shutdown(); } @Test void testDefaultConstructorInitializesDefaults() { target = new OtlpHttpSender(mockConfig, mockSinkMetrics); - - // Reflection to assert internal fields (not great, but useful for unit validation) - assertNotNull(getPrivateField(target, "signer")); - assertNotNull(getPrivateField(target, "httpClient")); - assertNotNull(getPrivateField(target, "sleeper")); + assertNotNull(getField(target, "signer")); + assertNotNull(getField(target, "httpClient")); + assertNotNull(getField(target, "sleeper")); } - private Object getPrivateField(Object instance, String fieldName) { + private Object getField(final Object obj, final String name) { try { - var field = instance.getClass().getDeclaredField(fieldName); - field.setAccessible(true); - return field.get(instance); + final var f = obj.getClass().getDeclaredField(name); + f.setAccessible(true); + return f.get(obj); } catch (Exception e) { - fail("Could not access field: " + fieldName); + fail("Could not access " + name); return null; } } - } diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java index 6a53fb48e1..700823ca9a 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java @@ -10,11 +10,22 @@ import io.micrometer.core.instrument.Timer; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.configuration.PluginSetting; -import static org.mockito.Mockito.anyString; +import java.lang.reflect.Field; +import java.time.Duration; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.BlockingQueue; +import java.util.function.ToDoubleFunction; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -31,86 +42,139 @@ class OtlpSinkMetricsTest { void setUp() { pluginMetrics = mock(PluginMetrics.class); pluginSetting = mock(PluginSetting.class); + + // basic pluginSetting behavior + when(pluginSetting.getPipelineName()).thenReturn("testPipeline"); + when(pluginSetting.getName()).thenReturn("otlp"); + + // stub out all counter() calls to return the same Counter mock counterMock = mock(Counter.class); + when(pluginMetrics.counter(anyString())).thenReturn(counterMock); + summaryMock = mock(DistributionSummary.class); timerMock = mock(Timer.class); - when(pluginMetrics.counter(anyString())).thenReturn(counterMock); - when(pluginSetting.getPipelineName()).thenReturn("otlp_pipeline"); - when(pluginSetting.getName()).thenReturn("otlp"); sinkMetrics = new OtlpSinkMetrics(pluginMetrics, pluginSetting); } - @Test - void testIncrementRecordsIn() { - sinkMetrics.incrementRecordsIn(3); - verify(counterMock).increment(3); - } - @Test void testIncrementRecordsOut() { - sinkMetrics.incrementRecordsOut(2); - verify(counterMock).increment(2); + sinkMetrics.incrementRecordsOut(7); + verify(pluginMetrics).counter("recordsOut"); + verify(counterMock).increment(7.0); } @Test void testIncrementErrorsCount() { sinkMetrics.incrementErrorsCount(); - verify(counterMock).increment(1); + verify(pluginMetrics).counter("errorsCount"); + verify(counterMock).increment(1.0); } @Test - void testIncrementPayloadSize() { - sinkMetrics.incrementPayloadSize(1024); - // Cannot verify summaryMock as DistributionSummary is built statically inside constructor + void testIncrementRetriesCount() { + sinkMetrics.incrementRetriesCount(); + verify(pluginMetrics).counter("retriesCount"); + verify(counterMock).increment(1.0); } @Test - void testIncrementPayloadGzipSize() { - sinkMetrics.incrementPayloadGzipSize(2048); - // Cannot verify summaryMock without injecting mock + void testIncrementRejectedSpansCount() { + sinkMetrics.incrementRejectedSpansCount(5); + verify(pluginMetrics).counter("rejectedSpansCount"); + verify(counterMock).increment(5.0); } @Test - void testRecordDeliveryLatency() { - sinkMetrics.recordDeliveryLatency(150); - // Would require refactoring to inject mock Timer for verification + void testRecordResponseCodeCounters() { + // 5xx + sinkMetrics.recordResponseCode(503); + verify(pluginMetrics).counter("http_5xx_responses"); + verify(counterMock).increment(); + + // 4xx + sinkMetrics.recordResponseCode(404); + verify(pluginMetrics).counter("http_4xx_responses"); + // total increments called twice so far + verify(counterMock, times(2)).increment(); + + // 2xx + sinkMetrics.recordResponseCode(200); + verify(pluginMetrics).counter("http_2xx_responses"); + verify(counterMock, times(3)).increment(); } @Test - void testRecordHttpLatency() { - sinkMetrics.recordHttpLatency(100); - // Would require refactoring to inject mock Timer for verification + void testRegisterQueueGauges() { + final BlockingQueue queue = new ArrayBlockingQueue<>(10); + sinkMetrics.registerQueueGauges(queue); + + // we expect two gauges registered: queueSize and queueCapacity + verify(pluginMetrics).gauge(eq("queueSize"), eq(queue), any()); + verify(pluginMetrics).gauge(eq("queueCapacity"), eq(queue), any()); } @Test - void testIncrementRetriesCount() { - sinkMetrics.incrementRetriesCount(); - verify(counterMock).increment(1); + @SuppressWarnings("unchecked") + void testQueueCapacityGaugeFunction() { + final ArrayBlockingQueue queue = new ArrayBlockingQueue<>(10); + // put 3 elements so size=3, remainingCapacity=7 + queue.add("one"); + queue.add("two"); + queue.add("three"); + + // register the gauges + sinkMetrics.registerQueueGauges(queue); + + // capture the ToDoubleFunction passed to gauge(...) + @SuppressWarnings("rawtypes") final ArgumentCaptor captor = ArgumentCaptor.forClass(ToDoubleFunction.class); + verify(pluginMetrics).gauge(eq("queueCapacity"), eq(queue), captor.capture()); + + // apply it: remainingCapacity + size == 7 + 3 == 10 + final ToDoubleFunction> func = captor.getValue(); + assertEquals(10.0, func.applyAsDouble(queue), 0.0); } @Test - void testIncrementRejectedSpansCount() { - sinkMetrics.incrementRejectedSpansCount(5); - verify(counterMock).increment(5); + void testIncrementPayloadSize_delegatesToSummary() throws Exception { + injectField("payloadSize", summaryMock); + + sinkMetrics.incrementPayloadSize(123L); + + verify(summaryMock).record(123L); } @Test - void testRecordResponseCode_5xx() { - sinkMetrics.recordResponseCode(503); - verify(pluginMetrics).counter("http_5xx_responses"); + void testIncrementPayloadGzipSize_delegatesToSummary() throws Exception { + injectField("payloadGzipSize", summaryMock); + + sinkMetrics.incrementPayloadGzipSize(77L); + + verify(summaryMock).record(77L); } @Test - void testRecordResponseCode_4xx() { - sinkMetrics.recordResponseCode(404); - verify(pluginMetrics).counter("http_4xx_responses"); + void testRecordDeliveryLatency_delegatesToTimer() throws Exception { + injectField("deliveryLatency", timerMock); + + sinkMetrics.recordDeliveryLatency(50L); + + verify(timerMock).record(Duration.ofMillis(50L)); } @Test - void testRecordResponseCode_2xx() { - sinkMetrics.recordResponseCode(200); - verify(pluginMetrics).counter("http_2xx_responses"); + void testRecordHttpLatency_delegatesToTimer() throws Exception { + injectField("httpLatency", timerMock); + + sinkMetrics.recordHttpLatency(250L); + + verify(timerMock).record(Duration.ofMillis(250L)); + } + + private void injectField(final String fieldName, final Object mock) throws Exception { + final Field f = OtlpSinkMetrics.class.getDeclaredField(fieldName); + f.setAccessible(true); + f.set(sinkMetrics, mock); } -} \ No newline at end of file +} diff --git a/data-prepper-plugins/otlp-sink/src/test/resources/data-prepper-config.yaml b/data-prepper-plugins/otlp-sink/src/test/resources/data-prepper-config.yaml deleted file mode 100644 index 544c112df1..0000000000 --- a/data-prepper-plugins/otlp-sink/src/test/resources/data-prepper-config.yaml +++ /dev/null @@ -1,3 +0,0 @@ -ssl: false -metricRegistries: - - CloudWatch diff --git a/data-prepper-plugins/otlp-sink/src/test/resources/pipelines.yaml b/data-prepper-plugins/otlp-sink/src/test/resources/pipelines.yaml deleted file mode 100644 index 449c21032e..0000000000 --- a/data-prepper-plugins/otlp-sink/src/test/resources/pipelines.yaml +++ /dev/null @@ -1,13 +0,0 @@ -otlp_pipeline: - source: - otel_trace_source: - ssl: false - port: 21890 - - buffer: - bounded_blocking: - buffer_size: 1000000 - batch_size: 125000 - - sink: - - otlp: \ No newline at end of file diff --git a/data-prepper-plugins/otlp-sink/src/test/resources/sample-trace.json b/data-prepper-plugins/otlp-sink/src/test/resources/sample-trace.json deleted file mode 100644 index e03778ed6b..0000000000 --- a/data-prepper-plugins/otlp-sink/src/test/resources/sample-trace.json +++ /dev/null @@ -1,28 +0,0 @@ -{ - "resourceSpans": [ - { - "resource": { - "attributes": [ - { - "key": "service.name", - "value": { "stringValue": "test-service" } - } - ] - }, - "scopeSpans": [ - { - "spans": [ - { - "traceId": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", - "spanId": "bbbbbbbbbbbbbbbb", - "name": "test-span", - "kind": "SPAN_KIND_INTERNAL", - "startTimeUnixNano": "1697051838000000000", - "endTimeUnixNano": "1697051839000000000" - } - ] - } - ] - } - ] -} \ No newline at end of file diff --git a/data-prepper-plugins/otlp-sink/src/test/resources/test-span-event.json b/data-prepper-plugins/otlp-sink/src/test/resources/test-span-event.json deleted file mode 100644 index eae0ae4460..0000000000 --- a/data-prepper-plugins/otlp-sink/src/test/resources/test-span-event.json +++ /dev/null @@ -1,23 +0,0 @@ -{ - "traceId": "ffa576d321173ac6cef3601c8f4bde75", - "spanId": "085ac082ffcfbf8d", - "traceState": "TRACE_STATE_1", - "kind": "SPAN_KIND_CLIENT", - "traceGroupFields": { - "endTime": "2020-08-20T05:40:43.217170200Z", - "durationInNanos": 49160000, - "statusCode": 1 - }, - "name": "TRACE_1_ROOT_SPAN", - "traceGroup": "TRACE_1_ROOT_SPAN", - "startTime": "2020-08-20T05:40:43.168010200Z", - "durationInNanos": 49160000, - "endTime": "2020-08-20T05:40:43.217170200Z", - "parentSpanId": "", - "attributes": {}, - "droppedAttributesCount": 0, - "links": [], - "droppedLinksCount": 0, - "events": [], - "droppedEventsCount": 0 -} \ No newline at end of file From 28bd81122d9e8f9f1cfddf06d56d815721f1e546 Mon Sep 17 00:00:00 2001 From: huy pham Date: Tue, 29 Apr 2025 18:59:30 -0700 Subject: [PATCH 12/23] refactor(aws-region): get aws region from endpoint value - Changed logic to get the aws region value from the provided endpoint. Signed-off-by: huy pham --- .../AwsAuthenticationConfig.java | 18 ------- .../otlp/configuration/OtlpSinkConfig.java | 31 ++++++++++-- .../plugins/sink/otlp/http/SigV4Signer.java | 8 +--- .../plugins/sink/otlp/OtlpSinkTest.java | 11 ++--- .../sink/otlp/buffer/OtlpSinkBufferTest.java | 8 ++-- .../AwsAuthenticationConfigTest.java | 14 ------ .../configuration/OtlpSinkConfigTest.java | 48 +++++++++++++++++-- .../sink/otlp/http/OtlpHttpSenderTest.java | 3 +- .../sink/otlp/http/SigV4SignerTest.java | 26 +--------- 9 files changed, 80 insertions(+), 87 deletions(-) diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfig.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfig.java index fc94f5aca6..ec4604cae2 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfig.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfig.java @@ -8,7 +8,6 @@ import jakarta.validation.constraints.Size; import lombok.Getter; import lombok.NoArgsConstructor; -import software.amazon.awssdk.regions.Region; /** * Configuration class for AWS authentication settings. @@ -17,13 +16,6 @@ @Getter @NoArgsConstructor class AwsAuthenticationConfig { - /** - * AWS region. - * Must be a valid AWS region identifier (e.g., us-east-1, us-west-2). - */ - @JsonProperty("region") - @Size(min = 1, message = "Region cannot be empty string") - private String awsRegion; /** * AWS STS Role ARN for assuming role-based access. @@ -42,14 +34,4 @@ class AwsAuthenticationConfig { @JsonProperty("sts_external_id") @Size(min = 2, max = 1224, message = "awsStsExternalId length should be between 2 and 1224 characters") private String awsStsExternalId; - - /** - * Gets the AWS Region object corresponding to the configured region string. - * - * @return Region object if awsRegion is set, otherwise returns null. - * Note: Default region fallback is handled externally by the caller. - */ - Region getAwsRegion() { - return awsRegion != null ? Region.of(awsRegion) : null; - } } diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfig.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfig.java index f0b6e68cfd..2339091534 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfig.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfig.java @@ -12,6 +12,11 @@ import lombok.NoArgsConstructor; import software.amazon.awssdk.regions.Region; +import java.net.URI; +import java.util.Arrays; +import java.util.Set; +import java.util.stream.Collectors; + /** * Configuration class for the OTLP sink plugin. * This class defines the configuration options available when setting up @@ -65,12 +70,30 @@ public long getFlushTimeoutMillis() { @Valid private AwsAuthenticationConfig awsAuthenticationConfig; + /** + * Get AWS region from the provided endpoint. + * + * @return the AWS region + */ public Region getAwsRegion() { - if (awsAuthenticationConfig == null) { - return null; - } + try { + final String host = URI.create(this.endpoint).getHost(); + if (host == null) { + throw new IllegalArgumentException(); + } - return awsAuthenticationConfig.getAwsRegion(); + final Set knownRegions = Region.regions().stream() + .map(Region::id) + .collect(Collectors.toSet()); + + return Arrays.stream(host.split("\\.")) + .filter(knownRegions::contains) + .map(Region::of) + .findFirst() + .orElseThrow(() -> new IllegalArgumentException("No valid AWS region found in endpoint: " + endpoint)); + } catch (final Exception e) { + throw new IllegalArgumentException("Failed to parse AWS region from endpoint: " + endpoint, e); + } } public String getStsRoleArn() { diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4Signer.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4Signer.java index 264aa1c62e..324eb9412e 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4Signer.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4Signer.java @@ -15,7 +15,6 @@ import software.amazon.awssdk.http.SdkHttpFullRequest; import software.amazon.awssdk.http.SdkHttpMethod; import software.amazon.awssdk.regions.Region; -import software.amazon.awssdk.regions.providers.DefaultAwsRegionProviderChain; import software.amazon.awssdk.services.sts.StsClient; import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider; @@ -51,16 +50,11 @@ class SigV4Signer { */ @VisibleForTesting SigV4Signer(@Nonnull final OtlpSinkConfig config, final StsClient stsClient) { - this.region = resolveRegion(config.getAwsRegion()); + this.region = config.getAwsRegion(); this.credentialsProvider = initCredentialsProvider(region, config.getStsRoleArn(), config.getStsExternalId(), stsClient); this.endpointUri = config.getEndpoint() != null ? URI.create(config.getEndpoint()) : URI.create(String.format("https://xray.%s.amazonaws.com%s", region.id(), OTLP_PATH)); - - } - - private static Region resolveRegion(final Region region) { - return region != null ? region : DefaultAwsRegionProviderChain.builder().build().getRegion(); } private static AwsCredentialsProvider initCredentialsProvider( diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java index 5c758650bb..82fbdf34ec 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java @@ -4,7 +4,6 @@ */ package org.opensearch.dataprepper.plugins.sink.otlp; -import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.opensearch.dataprepper.metrics.PluginMetrics; @@ -35,11 +34,12 @@ class OtlpSinkTest { @BeforeEach void setUp() throws Exception { - System.setProperty("aws.region", Region.US_WEST_2.id()); - // Arrange: stub out config, metrics, setting mockConfig = mock(OtlpSinkConfig.class); + when(mockConfig.getAwsRegion()).thenReturn(Region.of("us-west-2")); + mockMetrics = mock(PluginMetrics.class); + mockSetting = mock(PluginSetting.class); when(mockSetting.getPipelineName()).thenReturn("pipeline"); when(mockSetting.getName()).thenReturn("otlp"); @@ -54,11 +54,6 @@ void setUp() throws Exception { bufferField.set(target, mockBuffer); } - @AfterEach - void tearDown() { - System.clearProperty("aws.region"); - } - @Test void testInitialize_startsBuffer() { // Act diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBufferTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBufferTest.java index 8b64893a74..4cdb7605a8 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBufferTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBufferTest.java @@ -44,12 +44,11 @@ class OtlpSinkBufferTest { @BeforeEach void setUp() { - System.setProperty("aws.region", Region.US_WEST_2.id()); - config = mock(OtlpSinkConfig.class); when(config.getMaxEvents()).thenReturn(2); when(config.getMaxBatchSize()).thenReturn(1_000_000L); when(config.getFlushTimeoutMillis()).thenReturn(10L); + when(config.getAwsRegion()).thenReturn(Region.of("us-west-2")); metrics = mock(OtlpSinkMetrics.class); encoder = mock(OTelProtoStandardCodec.OTelProtoEncoder.class); @@ -61,7 +60,6 @@ void setUp() { @AfterEach void tearDown() { - System.clearProperty("aws.region"); buffer.stop(); } @@ -182,7 +180,7 @@ void testQueueCapacityRespectsMinimum() throws Exception { final Field queueField = OtlpSinkBuffer.class.getDeclaredField("queue"); queueField.setAccessible(true); - @SuppressWarnings("unchecked") final BlockingQueue queueInstance = (BlockingQueue) queueField.get(buffer); + final BlockingQueue queueInstance = (BlockingQueue) queueField.get(buffer); assertEquals(2000, queueInstance.remainingCapacity()); } @@ -193,7 +191,7 @@ void testQueueCapacityBasedOnMaxEvents() throws Exception { final Field queueField = OtlpSinkBuffer.class.getDeclaredField("queue"); queueField.setAccessible(true); - @SuppressWarnings("unchecked") final BlockingQueue queueInstance = (BlockingQueue) queueField.get(buffer); + final BlockingQueue queueInstance = (BlockingQueue) queueField.get(buffer); assertEquals(300 * 10, queueInstance.remainingCapacity()); } diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfigTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfigTest.java index ea7f09fc95..1074803c96 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfigTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfigTest.java @@ -4,42 +4,28 @@ */ package org.opensearch.dataprepper.plugins.sink.otlp.configuration; -import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.dataformat.yaml.YAMLFactory; import org.junit.jupiter.api.Test; -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.nullValue; import static org.junit.jupiter.api.Assertions.assertEquals; class AwsAuthenticationConfigTest { private final String expectedRoleArn = "arn:aws:iam::123456789012:role/MyRole"; private final String expectedExternalId = "external-id-123"; - private final String expectedRegion = "us-west-2"; private final ObjectMapper mapper = new ObjectMapper(new YAMLFactory()); @Test void testDeserializationFromYaml() throws Exception { final String yaml = String.join("\n", - "region: " + expectedRegion, "sts_role_arn: " + expectedRoleArn, "sts_external_id: " + expectedExternalId ); AwsAuthenticationConfig config = mapper.readValue(yaml, AwsAuthenticationConfig.class); - assertEquals(expectedRegion, config.getAwsRegion().toString()); assertEquals(expectedRoleArn, config.getAwsStsRoleArn()); assertEquals(expectedExternalId, config.getAwsStsExternalId()); } - - @Test - void testGetRegion_whenAllIsNull_returnNull() throws JsonProcessingException { - final String yaml = "{}"; - AwsAuthenticationConfig config = mapper.readValue(yaml, AwsAuthenticationConfig.class); - - assertThat(config.getAwsRegion(), nullValue()); - } } diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfigTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfigTest.java index 9369321fb0..c33a4b5a45 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfigTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfigTest.java @@ -19,8 +19,10 @@ import java.time.Duration; import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.nullValue; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; class OtlpSinkConfigTest { @@ -38,7 +40,6 @@ class OtlpSinkConfigTest { private static final long DEFAULT_BATCH_BYTES = ByteCount.parse("1mb").getBytes(); private static final long DEFAULT_FLUSH_TIMEOUT = 200L; - private static final String EXPECTED_REGION = "us-west-2"; private static final String EXPECTED_ROLE_ARN = "arn:aws:iam::123456789012:role/OtlpRole"; private static final String EXPECTED_EXTERNAL_ID = "my-ext-id"; @@ -85,7 +86,6 @@ void testMinimumConfigDefaults() throws Exception { assertEquals(DEFAULT_BATCH_BYTES, config.getMaxBatchSize()); assertEquals(DEFAULT_FLUSH_TIMEOUT, config.getFlushTimeoutMillis()); - assertThat(config.getAwsRegion(), nullValue()); assertThat(config.getStsRoleArn(), nullValue()); assertThat(config.getStsExternalId(), nullValue()); } @@ -116,7 +116,6 @@ void testAwsBlockDeserialization() throws Exception { "endpoint: \"" + EXPECTED_ENDPOINT + "\"", "max_retries: " + DEFAULT_MAX_RETRIES, "aws:", - " region: \"" + EXPECTED_REGION + "\"", " sts_role_arn: \"" + EXPECTED_ROLE_ARN + "\"", " sts_external_id: \"" + EXPECTED_EXTERNAL_ID + "\"" ); @@ -126,7 +125,6 @@ void testAwsBlockDeserialization() throws Exception { assertEquals(EXPECTED_ENDPOINT, config.getEndpoint()); assertEquals(DEFAULT_MAX_RETRIES, config.getMaxRetries()); - assertEquals(Region.of(EXPECTED_REGION), config.getAwsRegion()); assertEquals(EXPECTED_ROLE_ARN, config.getStsRoleArn()); assertEquals(EXPECTED_EXTERNAL_ID, config.getStsExternalId()); } @@ -143,8 +141,48 @@ void testAwsSectionMissing_staysNull() throws Exception { assertEquals(EXPECTED_ENDPOINT, config.getEndpoint()); assertEquals(DEFAULT_MAX_RETRIES, config.getMaxRetries()); - assertThat(config.getAwsRegion(), nullValue()); assertThat(config.getStsRoleArn(), nullValue()); assertThat(config.getStsExternalId(), nullValue()); } + + @Test + void testAwsRegion_parsedFromStandardXrayEndpoint() throws Exception { + final String yaml = String.join("\n", + "endpoint: \"https://xray.us-east-1.amazonaws.com\"", + "max_retries: 5" + ); + + final OtlpSinkConfig config = mapper.readValue(yaml, OtlpSinkConfig.class); + + assertEquals("https://xray.us-east-1.amazonaws.com", config.getEndpoint()); + assertEquals(5, config.getMaxRetries()); + + assertThat(config.getAwsRegion(), equalTo(Region.US_EAST_1)); + } + + @Test + void testAwsRegion_invalidEndpoint_throwsException() { + final String yaml = String.join("\n", + "endpoint: \"https://example.invalid-endpoint\"", + "max_retries: 5" + ); + + assertThrows(IllegalArgumentException.class, () -> { + final OtlpSinkConfig config = mapper.readValue(yaml, OtlpSinkConfig.class); + config.getAwsRegion(); // triggers parsing + }); + } + + @Test + void testAwsRegion_throwsException_onInvalidEndpoint() { + final String yaml = String.join("\n", + "endpoint: \"invalid-endpoint\"", + "max_retries: 5" + ); + + assertThrows(IllegalArgumentException.class, () -> { + final OtlpSinkConfig config = mapper.readValue(yaml, OtlpSinkConfig.class); + config.getAwsRegion(); // must trigger parsing logic + }); + } } \ No newline at end of file diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java index e567c5cdaa..b5162b2762 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java @@ -92,7 +92,6 @@ void setUp() { void cleanUp() { System.clearProperty("aws.accessKeyId"); System.clearProperty("aws.secretAccessKey"); - System.clearProperty("aws.region"); } @Test @@ -314,7 +313,7 @@ void testHandleSuccessfulResponseParseErrorIncrementsError() throws IOException target = new OtlpHttpSender( mockConfig, mockSinkMetrics, - payload -> Optional.of(payload), + Optional::of, mockSigner, mockHttpClient, mockSleeper diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4SignerTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4SignerTest.java index 8fb69e754e..647e094cff 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4SignerTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4SignerTest.java @@ -12,7 +12,6 @@ import org.opensearch.dataprepper.plugins.sink.otlp.configuration.OtlpSinkConfig; import software.amazon.awssdk.http.SdkHttpFullRequest; import software.amazon.awssdk.regions.Region; -import software.amazon.awssdk.regions.providers.DefaultAwsRegionProviderChain; import software.amazon.awssdk.services.sts.StsClient; import software.amazon.awssdk.services.sts.StsClientBuilder; import software.amazon.awssdk.services.sts.model.AssumeRoleRequest; @@ -35,6 +34,7 @@ class SigV4SignerTest { private static final byte[] PAYLOAD = "test-payload".getBytes(StandardCharsets.UTF_8); + public static final Region REGION = Region.of("us-west-2"); private OtlpSinkConfig mockXrayConfig; private SigV4Signer target; @@ -50,15 +50,13 @@ void setUp(){ void cleanUp() { System.clearProperty("aws.accessKeyId"); System.clearProperty("aws.secretAccessKey"); - System.clearProperty("aws.region"); } @Test void testSignRequest_withInputEndpoint_whenEndpointIsSet() { // setup final String endpoint = "https://performance.us-west-2.xray.cloudwatch.aws.dev/v1/traces"; - System.setProperty("aws.region", Region.US_WEST_2.toString()); - when(mockXrayConfig.getAwsRegion()).thenReturn(null); + when(mockXrayConfig.getAwsRegion()).thenReturn(REGION); when(mockXrayConfig.getEndpoint()).thenReturn(endpoint); target = new SigV4Signer(mockXrayConfig); @@ -73,26 +71,6 @@ void testSignRequest_withInputEndpoint_whenEndpointIsSet() { assertTrue(signedRequest.getUri().toString().contains(endpoint)); } - @Test - void testSignRequest_withFallbackRegion_whenRegionNotSet() { - // setup - System.setProperty("aws.region", Region.US_WEST_2.toString()); - when(mockXrayConfig.getAwsRegion()).thenReturn(null); - target = new SigV4Signer(mockXrayConfig); - - // run - final SdkHttpFullRequest signedRequest = target.signRequest(PAYLOAD); - - // assert - assertNotNull(signedRequest); - assertEquals("POST", signedRequest.method().name()); - assertTrue(signedRequest.headers().containsKey("Authorization")); - assertEquals("application/x-protobuf", signedRequest.firstMatchingHeader("Content-Type").orElse(null)); - - final String expectedRegion = DefaultAwsRegionProviderChain.builder().build().getRegion().id(); - assertTrue(signedRequest.getUri().toString().contains(String.format("https://xray.%s.amazonaws.com/v1/traces", expectedRegion))); - } - @Test void testSignRequest_withFallbackStsRole_whenStsRoleNotSet() { // setup From e85cd1ce78dfbeb6ac158a84d4e48bbf51b49445 Mon Sep 17 00:00:00 2001 From: huy pham Date: Tue, 29 Apr 2025 19:28:20 -0700 Subject: [PATCH 13/23] refactor(http-retry-code): follow retryable code - Retries on 429, 502, 503, and 504 per OTEL spec. Logs other errors without retry. - See: https://opentelemetry.io/docs/specs/otlp/exporter/#retrying-on-failure Signed-off-by: huy pham --- .../sink/otlp/http/OtlpHttpSender.java | 19 ++++----- .../sink/otlp/http/OtlpHttpSenderTest.java | 41 +++++++++---------- 2 files changed, 28 insertions(+), 32 deletions(-) diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java index 67d1d47f03..5f709bdc1e 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java @@ -33,8 +33,7 @@ * Responsible for sending signed OTLP Protobuf requests to OTLP endpoint using OkHttp. */ public class OtlpHttpSender implements AutoCloseable { - @VisibleForTesting - static final Set NON_RETRYABLE_STATUS_CODES = Set.of(400, 401, 403, 422); + private static final Set RETRYABLE_STATUS_CODES = Set.of(429, 502, 503, 504); private static final int BASE_RETRY_DELAY_MS = 100; private static final Logger LOG = LoggerFactory.getLogger(OtlpHttpSender.class); @@ -165,11 +164,12 @@ public void send(@Nonnull final byte[] payload) throws IOException { } /** - * Handles the response from the OTLP endpoint. - * Logs the response status and body, and throws an exception for retryable errors. + * Handles the OTLP export response. + * Retries on 429, 502, 503, and 504 per OTEL spec. Logs other errors without retry. + * See: https://opentelemetry.io/docs/specs/otlp/exporter/#retrying-on-failure * - * @param response The HTTP response from the OTLP endpoint. - * @throws IOException If the response status is not successful and retryable. + * @param response The HTTP response + * @throws IOException For retryable errors */ private void handleResponse(@Nonnull final Response response) throws IOException { final int status = response.code(); @@ -185,12 +185,11 @@ private void handleResponse(@Nonnull final Response response) throws IOException } final String responseBody = responseBytes != null ? new String(responseBytes, StandardCharsets.UTF_8) : ""; - if (NON_RETRYABLE_STATUS_CODES.contains(status)) { - LOG.error("Non-retryable error. Status: {}, Response: {}", status, responseBody); - return; + if (RETRYABLE_STATUS_CODES.contains(status)) { + throw new IOException(String.format("Retryable error. Status: %d, Response: %s", status, responseBody)); } - throw new IOException(String.format("Failed to send OTLP data. Status: %d, Response: %s", status, responseBody)); + LOG.error("Non-retryable error. Status: {}, Response: {}", status, responseBody); } private void handleSuccessfulResponse(final byte[] responseBytes) { diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java index b5162b2762..9b21ca1c31 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java @@ -49,7 +49,6 @@ import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -import static org.opensearch.dataprepper.plugins.sink.otlp.http.OtlpHttpSender.NON_RETRYABLE_STATUS_CODES; class OtlpHttpSenderTest { private static final byte[] PAYLOAD = "test-otlp-payload".getBytes(StandardCharsets.UTF_8); @@ -135,24 +134,22 @@ void testSend_doesNotRetryOnNonRetryable4xxResponses() throws IOException { .url(signed.getUri().toString()) .build(); - for (final int status : NON_RETRYABLE_STATUS_CODES) { - final Response resp = new Response.Builder() - .request(okReq) - .protocol(Protocol.HTTP_1_1) - .code(status) - .message("Client Error") - .body(ResponseBody.create( - ERROR_BODY.getBytes(StandardCharsets.UTF_8), - MediaType.get("application/json"))) - .build(); - final Call call = mock(Call.class); - when(mockHttpClient.newCall(any())).thenReturn(call); - when(call.execute()).thenReturn(resp); - - assertDoesNotThrow(() -> target.send(PAYLOAD)); - verify(mockHttpClient, times(1)).newCall(any()); - reset(mockHttpClient); - } + final Response resp = new Response.Builder() + .request(okReq) + .protocol(Protocol.HTTP_1_1) + .code(400) + .message("Client Error") + .body(ResponseBody.create( + ERROR_BODY.getBytes(StandardCharsets.UTF_8), + MediaType.get("application/json"))) + .build(); + final Call call = mock(Call.class); + when(mockHttpClient.newCall(any())).thenReturn(call); + when(call.execute()).thenReturn(resp); + + assertDoesNotThrow(() -> target.send(PAYLOAD)); + verify(mockHttpClient, times(1)).newCall(any()); + reset(mockHttpClient); } @Test @@ -202,7 +199,7 @@ void testSend_throwsIOException_whenFailsAfterAllRetries() throws IOException { } @Test - void testSend_throwsIOException_on500ResponseWithBody() throws IOException { + void testSend_throwsIOException_on502ResponseWithBody() throws IOException { final SdkHttpFullRequest signed = SdkHttpFullRequest.builder() .method(software.amazon.awssdk.http.SdkHttpMethod.POST) .uri(URI.create("https://example.com")) @@ -218,8 +215,8 @@ void testSend_throwsIOException_on500ResponseWithBody() throws IOException { final Response resp500 = new Response.Builder() .request(okReq) .protocol(Protocol.HTTP_1_1) - .code(500) - .message("Internal Server Error") + .code(502) + .message("Bad Gateway") .body(ResponseBody.create( ERROR_BODY.getBytes(StandardCharsets.UTF_8), MediaType.get("application/json"))) From 347793280e12305c61779bed867766e5873866ba Mon Sep 17 00:00:00 2001 From: huy pham Date: Tue, 29 Apr 2025 19:36:39 -0700 Subject: [PATCH 14/23] refactor(log/exception): only use log.error - removed log.info statements - Add description to exception Signed-off-by: huy pham --- .../plugins/sink/otlp/http/GzipCompressor.java | 1 - .../plugins/sink/otlp/http/OtlpHttpSender.java | 8 ++------ .../dataprepper/plugins/sink/otlp/http/SigV4Signer.java | 3 +-- .../dataprepper/plugins/sink/otlp/http/ThreadSleeper.java | 5 ++--- .../otlp/configuration/AwsAuthenticationConfigTest.java | 2 +- .../plugins/sink/otlp/http/OtlpHttpSenderTest.java | 4 ++-- .../plugins/sink/otlp/http/SigV4SignerTest.java | 6 +++--- 7 files changed, 11 insertions(+), 18 deletions(-) diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressor.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressor.java index d548e2d7db..cb8b9c915c 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressor.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressor.java @@ -34,7 +34,6 @@ class GzipCompressor implements Function> { /** * Compresses the provided payload using GZIP compression. - * Logs an error if compression fails. * * @param payload The payload to be compressed. * @return Optional containing the compressed payload, or empty if compression failed. diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java index 5f709bdc1e..005269a28e 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java @@ -34,7 +34,6 @@ */ public class OtlpHttpSender implements AutoCloseable { private static final Set RETRYABLE_STATUS_CODES = Set.of(429, 502, 503, 504); - private static final int BASE_RETRY_DELAY_MS = 100; private static final Logger LOG = LoggerFactory.getLogger(OtlpHttpSender.class); private static final MediaType PROTOBUF = MediaType.get("application/x-protobuf"); @@ -150,14 +149,12 @@ public void send(@Nonnull final byte[] payload) throws IOException { final int delay = retryDelaysMs.get(retryIndex) + jitter; try { sleeper.accept(delay); - - LOG.info("Retrying after failure in attempt {}. Sleeping {}ms.", attempt + 1, delay, ioException); sinkMetrics.incrementRetriesCount(); } catch (final RuntimeException runtimeException) { throw new IOException("Sender failed to sleep before retrying.", runtimeException); } } else { - throw ioException; + throw new IOException("Max retries reached", ioException); } } } @@ -166,7 +163,7 @@ public void send(@Nonnull final byte[] payload) throws IOException { /** * Handles the OTLP export response. * Retries on 429, 502, 503, and 504 per OTEL spec. Logs other errors without retry. - * See: https://opentelemetry.io/docs/specs/otlp/exporter/#retrying-on-failure + * See: OTLP/HTTP Response * * @param response The HTTP response * @throws IOException For retryable errors @@ -194,7 +191,6 @@ private void handleResponse(@Nonnull final Response response) throws IOException private void handleSuccessfulResponse(final byte[] responseBytes) { if (responseBytes == null || responseBytes.length == 0) { - LOG.info("OTLP export successful. No response body."); return; } diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4Signer.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4Signer.java index 324eb9412e..ab2e0c7010 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4Signer.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4Signer.java @@ -36,8 +36,7 @@ class SigV4Signer { private final URI endpointUri; /** - * Constructs a SigV4 signer helper based on the AWS authentication configuration. - * Supports optional STS role assumption. + * Constructs a SigV4 signer helper. * * @param config Configuration for region and optional STS role */ diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeper.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeper.java index bd9727919f..6b395fd298 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeper.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeper.java @@ -16,8 +16,7 @@ class ThreadSleeper implements Consumer { /** * Sleeps for the specified duration in milliseconds. - * If the thread is interrupted while sleeping, the interrupted status is cleared - * and a {@link RuntimeException} is thrown to signal the failure. + * Wraps and rethrows {@link InterruptedException} as a runtime exception. * * @param millis the number of milliseconds to sleep */ @@ -27,7 +26,7 @@ public void accept(@Nonnull final Integer millis) { Thread.sleep(millis); } catch (final InterruptedException e) { Thread.currentThread().interrupt(); - throw new RuntimeException(e); + throw new RuntimeException("Sleep interrupted", e); } } } diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfigTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfigTest.java index 1074803c96..be73ba8893 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfigTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfigTest.java @@ -23,7 +23,7 @@ void testDeserializationFromYaml() throws Exception { "sts_external_id: " + expectedExternalId ); - AwsAuthenticationConfig config = mapper.readValue(yaml, AwsAuthenticationConfig.class); + final AwsAuthenticationConfig config = mapper.readValue(yaml, AwsAuthenticationConfig.class); assertEquals(expectedRoleArn, config.getAwsStsRoleArn()); assertEquals(expectedExternalId, config.getAwsStsExternalId()); diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java index 9b21ca1c31..2c8b4c9f19 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java @@ -195,7 +195,7 @@ void testSend_throwsIOException_whenFailsAfterAllRetries() throws IOException { // Act & Assert final IOException ex = assertThrows(IOException.class, () -> target.send(PAYLOAD)); - assertEquals("always fail", ex.getMessage()); + assertEquals("Max retries reached", ex.getMessage()); } @Test @@ -373,7 +373,7 @@ private Object getField(final Object obj, final String name) { final var f = obj.getClass().getDeclaredField(name); f.setAccessible(true); return f.get(obj); - } catch (Exception e) { + } catch (final Exception e) { fail("Could not access " + name); return null; } diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4SignerTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4SignerTest.java index 647e094cff..0767bdf2fa 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4SignerTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4SignerTest.java @@ -94,7 +94,7 @@ void testSignRequest_withCustomCredentials_usingDefaultStsClientFallback() { when(mockXrayConfig.getStsRoleArn()).thenReturn("arn:aws:iam::123456789012:role/test-role"); // Use mocked static builder to simulate StsClient.builder() - try (MockedStatic mockedStsClientStatic = mockStatic(StsClient.class)) { + try (final MockedStatic mockedStsClientStatic = mockStatic(StsClient.class)) { final StsClient mockStsClient = mock(StsClient.class); // Setup fake STS response @@ -118,7 +118,7 @@ void testSignRequest_withCustomCredentials_usingDefaultStsClientFallback() { // run target = new SigV4Signer(mockXrayConfig, null); - SdkHttpFullRequest signedRequest = target.signRequest(PAYLOAD); + final SdkHttpFullRequest signedRequest = target.signRequest(PAYLOAD); // assert assertNotNull(signedRequest); @@ -160,7 +160,7 @@ void testSignRequest_withCustomCredentials_usingMockedSts() { assertEquals("application/x-protobuf", signedRequest.firstMatchingHeader("Content-Type").orElse(null)); assertTrue(signedRequest.getUri().toString().contains("https://xray.us-west-2.amazonaws.com/v1/traces")); - ArgumentMatcher matcher = request -> + final ArgumentMatcher matcher = request -> expectedRoleArn.equals(request.roleArn()) && expectedExternalId.equals(request.externalId()); From 65066273d2cfd9c0175c945e522adc43459557e4 Mon Sep 17 00:00:00 2001 From: huy pham Date: Tue, 29 Apr 2025 21:31:19 -0700 Subject: [PATCH 15/23] fix(otlp-sink-buffer): Ensure final flush drains queue on shutdown - Updated worker loop to retry on interrupt until queue is empty - Avoid logging error on expected shutdown interrupt - test: Added unit tests to verify final flush Signed-off-by: huy pham --- data-prepper-plugins/otlp-sink/build.gradle | 1 + .../sink/otlp/buffer/OtlpSinkBuffer.java | 27 +++++++------- .../sink/otlp/buffer/OtlpSinkBufferTest.java | 36 +++++++++++++++++-- 3 files changed, 50 insertions(+), 14 deletions(-) diff --git a/data-prepper-plugins/otlp-sink/build.gradle b/data-prepper-plugins/otlp-sink/build.gradle index 2e41ea1961..6d536eaab1 100644 --- a/data-prepper-plugins/otlp-sink/build.gradle +++ b/data-prepper-plugins/otlp-sink/build.gradle @@ -66,6 +66,7 @@ dependencies { testRuntimeOnly 'org.junit.jupiter:junit-jupiter-engine:5.9.1' testImplementation 'org.slf4j:slf4j-simple:2.0.7' testImplementation 'org.mockito:mockito-core:5.10.0' + testImplementation 'org.awaitility:awaitility:4.2.0' // Integration Testing integrationTestImplementation project(':data-prepper-test-common') diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBuffer.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBuffer.java index 2a2a435a31..71d1206fd0 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBuffer.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBuffer.java @@ -135,24 +135,21 @@ private void run() { long batchSize = 0; long lastFlush = System.currentTimeMillis(); - while (running || !queue.isEmpty()) { + while (true) { try { final long now = System.currentTimeMillis(); final Record record = queue.poll(100, TimeUnit.MILLISECONDS); if (record != null) { - final ResourceSpans resourceSpans; try { - resourceSpans = encoder.convertToResourceSpans(record.getData()); + final ResourceSpans resourceSpans = encoder.convertToResourceSpans(record.getData()); + batch.add(resourceSpans); + batchSize += resourceSpans.getSerializedSize(); } catch (final Exception e) { LOG.error("Failed to encode span, skipping", e); sinkMetrics.incrementRejectedSpansCount(1); sinkMetrics.incrementErrorsCount(); - continue; } - - batch.add(resourceSpans); - batchSize += resourceSpans.getSerializedSize(); } final boolean flushBySize = batch.size() >= maxEvents || batchSize >= maxBatchBytes; @@ -160,19 +157,25 @@ private void run() { if (flushBySize || flushByTime) { send(batch); - batchSize = 0; lastFlush = now; } + if (!running && queue.isEmpty()) { + break; + } + } catch (final InterruptedException e) { - Thread.currentThread().interrupt(); - LOG.error("Interrupted while polling span", e); - sinkMetrics.incrementErrorsCount(); + if (running) { + LOG.debug("Worker interrupted while polling, continuing..."); + sinkMetrics.incrementErrorsCount(); + } + // Clear interrupt flag to allow queue.poll() again + // Don't break; allow draining } } - // Final flush on shutdown + // Final flush if (!batch.isEmpty()) { send(batch); } diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBufferTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBufferTest.java index 4cdb7605a8..ba8cd81459 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBufferTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBufferTest.java @@ -22,6 +22,8 @@ import java.util.concurrent.BlockingQueue; import java.util.concurrent.TimeUnit; +import static java.util.concurrent.TimeUnit.SECONDS; +import static org.awaitility.Awaitility.await; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertInstanceOf; @@ -122,7 +124,9 @@ void testFinalFlushOnShutdownWhenNoSizeOrTimeFlush() throws Exception { TimeUnit.MILLISECONDS.sleep(50); // final‐flush should happen exactly once - verify(sender).send(any(byte[].class)); + await().atMost(1, SECONDS).untilAsserted(() -> + verify(sender).send(any(byte[].class)) + ); verify(metrics).incrementRecordsOut(1); } @@ -242,7 +246,9 @@ void testWorkerThreadHandlesEncodeException() throws Exception { buffer.stop(); // should still send at least one batch for the good span - verify(sender, atLeastOnce()).send(any(byte[].class)); + await().atMost(1, SECONDS).untilAsserted(() -> + verify(sender).send(any(byte[].class)) + ); // one rejected from the encode exception verify(metrics).incrementRejectedSpansCount(1); // error count for the encode exception (+ maybe one on interrupt) @@ -271,4 +277,30 @@ void testConstructorDefaults() throws Exception { assertNotNull(snd, "default sender should not be null"); assertInstanceOf(OtlpHttpSender.class, snd); } + + @Test + void testWorkerThreadFlushesByTimeoutOnly() throws Exception { + when(config.getMaxEvents()).thenReturn(100); // high enough not to flush by count + when(config.getMaxBatchSize()).thenReturn(Long.MAX_VALUE); // don't flush by size + when(config.getFlushTimeoutMillis()).thenReturn(50L); // very short flush window + + buffer.stop(); + buffer = new OtlpSinkBuffer(config, metrics, encoder, sender); + buffer.start(); + + final Record rec = mock(Record.class); + when(rec.getData()).thenReturn(mock(Span.class)); + when(encoder.convertToResourceSpans(any(Span.class))).thenReturn(ResourceSpans.getDefaultInstance()); + + buffer.add(rec); + + // wait long enough to trigger timeout-based flush + TimeUnit.MILLISECONDS.sleep(100); + buffer.stop(); + + await().atMost(1, SECONDS).untilAsserted(() -> + verify(sender).send(any(byte[].class)) + ); + verify(metrics, atLeastOnce()).incrementRecordsOut(1); + } } \ No newline at end of file From 242ca15f75320123bebf939df52c0067bdb6a049 Mon Sep 17 00:00:00 2001 From: huy pham Date: Tue, 29 Apr 2025 22:18:50 -0700 Subject: [PATCH 16/23] Update README - Removed unused metrics Signed-off-by: huy pham --- data-prepper-plugins/otlp-sink/README.md | 146 ++++++++++++++---- .../sink/otlp/metrics/OtlpSinkMetrics.java | 6 - .../otlp/metrics/OtlpSinkMetricsTest.java | 9 -- 3 files changed, 114 insertions(+), 47 deletions(-) diff --git a/data-prepper-plugins/otlp-sink/README.md b/data-prepper-plugins/otlp-sink/README.md index 63e26e35df..d15e19eaa1 100644 --- a/data-prepper-plugins/otlp-sink/README.md +++ b/data-prepper-plugins/otlp-sink/README.md @@ -1,13 +1,20 @@ # OTLP Sink Plugin -In the first release, the otlp sink plugin sends span data to AWS X-Ray using the OTLP (OpenTelemetry Protocol) format. -Future releases will enhance the sink to send spans, metrics, and traces to any OTLP Protobuf endpoints. +The OTLP sink plugin sends span data using the OpenTelemetry Protocol (OTLP) format. +The initial release supports exporting spans to AWS X-Ray. Future releases will support sending spans, metrics, and logs +to any OTLP Protobuf-compatible endpoint. --- -## Usage +## Known Limitations -### Sample Pipeline Configuration +- Currently, supports only trace data (spans). Support for metrics and logs will be added in future releases. +- No support for DQL-based loss-less delivery in this release. +- Only AWS X-Ray-compatible OTLP endpoints are supported. + +--- + +## Sample Pipeline Configuration ```yaml otlp_pipeline: @@ -25,34 +32,118 @@ otlp_pipeline: sink: - otlp: - endpoint: "https://performance.us-west-2.xray.cloudwatch.aws.dev/v1/traces" - max_retries: 5 # Optional, default: 5 + endpoint: "https://xray.us-west-2.amazonaws.com/v1/traces" + max_retries: 5 threshold: - max_events: 512 # Optional, default: 512 - max_batch_size: 1mb # Optional, default: 1mb - flush_timeout: 200ms # Optional, default: 200ms + max_events: 512 + max_batch_size: 1mb + flush_timeout: 200ms aws: - sts_role_arn: arn:aws:iam::123456789012:role/MyRole # Optional STS Role ARN - sts_external_id: external-id-value # Optional external ID for STS + sts_role_arn: arn:aws:iam::123456789012:role/MyRole + sts_external_id: external-id-value ``` --- ## Configuration Options -| Property | Type | Required | Default | Description | -|----------------------------|----------|----------|---------|---------------------------------------------------------------------------------------| -| `endpoint` | `String` | Yes | — | OTLP gRPC or HTTP endpoint where spans will be sent. | -| `max_retries` | `int` | No | `5` | Maximum number of retry attempts on HTTP send failures. | -| **threshold** | `Object` | No | — | Controls batching behavior. See below for sub-properties. | -| `threshold.max_events` | `int` | No | `512` | Maximum number of spans per batch. | -| `threshold.max_batch_size` | `String` | No | `1mb` | Maximum total payload bytes per batch. Supports human-readable suffixes (`kb`, `mb`). | -| `threshold.flush_timeout` | `String` | No | `200ms` | Time to wait (in milliseconds) before flushing a non-empty batch. | -| **aws** | `Object` | No | — | AWS authentication settings. See below. | -| `aws.sts_role_arn` | `String` | No | — | Amazon Resource Name of the IAM role to assume. | -| `aws.sts_external_id` | `String` | No | — | Optional external ID for assuming IAM roles with STS. | +| Property | Type | Required | Default | Description | +|----------------------------|----------|----------|-----------------------|----------------------------------------------------------------------------------------------------------| +| `endpoint` | `String` | Yes | — | AWS X-Ray OTLP endpoint where spans will be sent. | +| `max_retries` | `int` | No | `5` | Maximum number of retry attempts on HTTP send failures. | +| **threshold** | `Object` | No | — | Controls batching behavior. See below for sub-properties. | +| `threshold.max_events` | `int` | No | `512` (recommended) | Maximum number of spans per batch. Minimum: 1. | +| `threshold.max_batch_size` | `String` | No | `1mb` (recommended) | Maximum total payload bytes per batch. Supports human-readable suffixes (`kb`, `mb`). | +| `threshold.flush_timeout` | `String` | No | `200ms` (recommended) | Maximum time to wait before flushing a non-empty batch. Minimum: 1ms (e.g., `200ms`, `1s`) | +| **aws** | `Object` | No | — | AWS authentication settings. See below. | +| `aws.sts_role_arn` | `String` | No | — | IAM Role ARN that Data Prepper (or OSI) assumes to send spans to X-Ray on behalf of a customer account. | +| `aws.sts_external_id` | `String` | No | — | External ID to use when assuming the role. Required only if the target IAM role enforces sts:ExternalId. | + +**Additional Notes:** + +- Only AWS X-Ray-compatible OTLP endpoints are currently supported (`https://xray..amazonaws.com/v1/traces`). +- `aws.region` is automatically derived from the endpoint. +- When `max_retries` is exceeded: + - Retryable errors (429, 502, 503, 504): exception is logged, plugin error metric is incremented. + - Non-retryable errors (400, 403): error is logged immediately without retry. +- Support for durable delivery (via DQL) will be added in a future release. + +--- + +## Performance Benchmark + +The OTLP Sink plugin is optimized for high-throughput, low-latency trace delivery. + +### Summary + +* Sustains ~3.5K TPS with ≤150ms p99 latency on t4g.large. +* Uses only ~8% CPU, ~100MB heap. +* 0 errors, retries, or drops during a 3-hour soak test. + +### Tuning Recommendations + +| Setting | Recommended | Reason | +|------------------|-------------|-------------------------------------------------------------------------------------------------| +| `max_retries` | `5` | Matches AWS SDK default. Gives ~8s of exponential backoff to tolerate transient 503/5xx errors. | +| `max_events` | `512` | Supports up to 3.5K TPS with 2 workers. Keeps p99 latency around 130ms. | +| `max_batch_size` | `1mb` | Aligns with OTEL + AWS X-Ray guidance. Larger batches get split, increasing latency/load. | +| `flush_timeout` | `200ms` | Short enough to avoid delay, long enough to fill batches and keep CPU/GCs low. | + +### Additional Tuning tips + +* Lower `max_events` to **200–400** to reduce latency below 100 ms +* Decrease `flush_timeout` to **100 ms** for faster flushes (with higher CPU/network cost) +* Increase `max_batch_size` to **≥ 8 MB** only if p99 span > 9 KB +* Add pipeline workers if queue saturates at >4K TPS + +### Queue Sizing Rule + +> Queue capacity = max_events * 10 (minimum 2000) +> +> To keep memory usage under ~50MB: +> max_events ≤ 50_000_000 ÷ (10 × p99_span_size_bytes) +> +> Example: With p99 span size of 1 KB, max_events should be ≤ 5000 + + +--- + +## Protocol Details + +* Protocol: OTLP over HTTP +* Content-Type: `application/x-protobuf` +* Compression: `gzip` (enabled by default) + All outgoing HTTP requests use gzip compression to reduce payload size and bandwidth usage. + +--- + +## Delivery Semantics + +Currently, the sink provides at-most-once delivery. Once retries are exhausted, span batches are dropped. +Future releases will support durable queueing via DQL for loss-less guarantees. + +--- + +## Retry Behavior -**Note:** `aws.region` will be derived from the provided AWS endpoint. +- The sink uses an exponential backoff strategy for retryable HTTP status codes (e.g., 429, 502, 503, 504). +- Maximum number of attempts is controlled by `max_retries`. Once exceeded: + - The span batch is dropped. + - The plugin logs the exception and increments the error metric. +- Non-retryable errors (e.g., 400, 403) are logged and counted immediately without retry. + +--- + +## Logging & Metrics + +* Exceptions are logged with full stack traces. No customer data is logged. +* Metrics are emitted via Micrometer and include: + * recordsIn, recordsOut + * httpLatency, HTTP codes + * errorCount, rejectedSpansCount, retriesCount + * queueSize, queueCapacity + * payloadSize, payloadGzipSize + * JVM stats if configured (e.g., heap usage, GC pauses) --- @@ -60,17 +151,8 @@ otlp_pipeline: See the [CONTRIBUTING](https://github.com/opensearch-project/data-prepper/blob/main/CONTRIBUTING.md) guide for general information on contributions. -The integration tests for this plugin do **not** run as part of the main Data Prepper build and will be added in future -releases. - ### Run unit tests locally ```bash ./gradlew :data-prepper-plugins:otlp-sink:test ``` - -### Run integration tests locally - -```bash -./gradlew :data-prepper-plugins:otlp-sink:integrationTest -``` diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetrics.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetrics.java index 81d2b38c70..89e8d3feac 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetrics.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetrics.java @@ -26,7 +26,6 @@ public class OtlpSinkMetrics { private final PluginMetrics pluginMetrics; private final Timer httpLatency; - private final Timer deliveryLatency; private final DistributionSummary payloadSize; private final DistributionSummary payloadGzipSize; @@ -43,7 +42,6 @@ public OtlpSinkMetrics(@Nonnull final PluginMetrics pluginMetrics, @Nonnull fina final String pluginName = pluginSetting.getName(); httpLatency = buildLatencyTimer(pipelineName, pluginName, "httpLatency"); - deliveryLatency = buildLatencyTimer(pipelineName, pluginName, "deliveryLatency"); payloadSize = buildDistributionSummary(pipelineName, pluginName, "payloadSize"); payloadGzipSize = buildDistributionSummary(pipelineName, pluginName, "payloadGzipSize"); @@ -100,10 +98,6 @@ public void incrementPayloadGzipSize(final long bytes) { payloadGzipSize.record(bytes); } - public void recordDeliveryLatency(final long durationMillis) { - deliveryLatency.record(Duration.ofMillis(durationMillis)); - } - public void recordHttpLatency(final long durationMillis) { httpLatency.record(Duration.ofMillis(durationMillis)); } diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java index 700823ca9a..2b7f61affd 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java @@ -154,15 +154,6 @@ void testIncrementPayloadGzipSize_delegatesToSummary() throws Exception { verify(summaryMock).record(77L); } - @Test - void testRecordDeliveryLatency_delegatesToTimer() throws Exception { - injectField("deliveryLatency", timerMock); - - sinkMetrics.recordDeliveryLatency(50L); - - verify(timerMock).record(Duration.ofMillis(50L)); - } - @Test void testRecordHttpLatency_delegatesToTimer() throws Exception { injectField("httpLatency", timerMock); From 390f951964198f9d77e1a4fd0ea45dd49e626735 Mon Sep 17 00:00:00 2001 From: huy pham Date: Wed, 30 Apr 2025 08:25:37 -0700 Subject: [PATCH 17/23] Removed unused Output codec Signed-off-by: huy pham --- .../otel-proto-common/build.gradle | 2 - .../otel/codec/OtlpTraceOutputCodec.java | 107 --------------- .../otel/codec/OtlpTraceOutputCodecTest.java | 127 ------------------ .../sink/otlp/http/OtlpHttpSender.java | 2 +- 4 files changed, 1 insertion(+), 237 deletions(-) delete mode 100644 data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OtlpTraceOutputCodec.java delete mode 100644 data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OtlpTraceOutputCodecTest.java diff --git a/data-prepper-plugins/otel-proto-common/build.gradle b/data-prepper-plugins/otel-proto-common/build.gradle index ab5ea220e1..40636b049c 100644 --- a/data-prepper-plugins/otel-proto-common/build.gradle +++ b/data-prepper-plugins/otel-proto-common/build.gradle @@ -12,8 +12,6 @@ test { } dependencies { - compileOnly 'org.projectlombok:lombok:1.18.30' - annotationProcessor 'org.projectlombok:lombok:1.18.30' implementation project(':data-prepper-api') implementation libs.opentelemetry.proto implementation libs.protobuf.util diff --git a/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OtlpTraceOutputCodec.java b/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OtlpTraceOutputCodec.java deleted file mode 100644 index b6d26fd373..0000000000 --- a/data-prepper-plugins/otel-proto-common/src/main/java/org/opensearch/dataprepper/plugins/otel/codec/OtlpTraceOutputCodec.java +++ /dev/null @@ -1,107 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.otel.codec; - -import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceRequest; -import io.opentelemetry.proto.trace.v1.ResourceSpans; -import lombok.NonNull; -import org.apache.commons.codec.DecoderException; -import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; -import org.opensearch.dataprepper.model.codec.OutputCodec; -import org.opensearch.dataprepper.model.event.Event; -import org.opensearch.dataprepper.model.sink.OutputCodecContext; -import org.opensearch.dataprepper.model.trace.Span; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import java.io.OutputStream; - -/** - * An implementation of {@link OutputCodec} that converts {@link Span} events - * into OpenTelemetry Protocol (OTLP) binary format using protobuf. - *

- * This codec is primarily intended for use with trace data that will be - * forwarded to OTLP endpoint. - */ -@DataPrepperPlugin(name = "otlp_trace", pluginType = OutputCodec.class) -public class OtlpTraceOutputCodec implements OutputCodec { - - private static final Logger LOG = LoggerFactory.getLogger(OtlpTraceOutputCodec.class); - private static final String OTLP_EXTENSION = "pb"; - - private final OTelProtoStandardCodec.OTelProtoEncoder encoder = new OTelProtoStandardCodec.OTelProtoEncoder(); - - /** - * Initializes the output stream. No-op for OTLP format. - * - * @param outputStream The output stream to write to. - * @param event The event (not used). - * @param context The codec context (not used). - */ - @Override - public void start(final OutputStream outputStream, final Event event, final OutputCodecContext context) { - // No-op for OTLP format - } - - /** - * Writes a single {@link Span} event to the output stream in OTLP binary format. - * - *

This method throws a {@link RuntimeException} (e.g., wrapping {@link DecoderException}) - * if the span is malformed or cannot be encoded. This allows upstream components (such as Sink plugins) - * to track and report failed span encoding attempts via plugin metrics or perform custom error handling. - * - *

Failing fast here instead of silently logging ensures invalid spans are not silently dropped and - * gives pipeline developers better visibility into pipeline health and data loss. - * - * @param event The event to encode. Must be of type {@link Span}. - * @param outputStream The stream to which the encoded bytes will be written. - * @throws IllegalArgumentException If the event is not a {@link Span}. - * @throws RuntimeException If encoding the span fails. - */ - @Override - public void writeEvent(@NonNull final Event event, @NonNull final OutputStream outputStream) { - if (!(event instanceof Span)) { - throw new IllegalArgumentException("OtlpTraceOutputCodec only supports Span events"); - } - - final Span span = (Span) event; - try { - final ResourceSpans resourceSpans = encoder.convertToResourceSpans(span); - final ExportTraceServiceRequest request = ExportTraceServiceRequest.newBuilder() - .addResourceSpans(resourceSpans) - .build(); - - outputStream.write(request.toByteArray()); - } catch (final DecoderException e) { - LOG.warn("Skipping invalid span with ID [{}] due to decoding error.", span.getSpanId(), e); - throw new RuntimeException(e); - } catch (final Exception e) { - LOG.error("Unexpected error while writing span with ID [{}] to OTLP output.", span.getSpanId(), e); - throw new RuntimeException(e); - } - } - - /** - * Finalizes the output stream. No-op for OTLP format. - * - * @param outputStream The output stream to finalize. - */ - @Override - public void complete(final OutputStream outputStream) { - // No-op for OTLP format - } - - /** - * Returns the file extension used by this codec. - * In this case, "pb" for protobuf binary format. - * - * @return The string "pb". - */ - @Override - public String getExtension() { - return OTLP_EXTENSION; - } -} diff --git a/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OtlpTraceOutputCodecTest.java b/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OtlpTraceOutputCodecTest.java deleted file mode 100644 index d6ec2e94d5..0000000000 --- a/data-prepper-plugins/otel-proto-common/src/test/java/org/opensearch/dataprepper/plugins/otel/codec/OtlpTraceOutputCodecTest.java +++ /dev/null @@ -1,127 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.otel.codec; - -import com.fasterxml.jackson.core.type.TypeReference; -import com.fasterxml.jackson.databind.ObjectMapper; -import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceRequest; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.opensearch.dataprepper.model.event.JacksonEvent; -import org.opensearch.dataprepper.model.sink.OutputCodecContext; -import org.opensearch.dataprepper.model.trace.DefaultTraceGroupFields; -import org.opensearch.dataprepper.model.trace.JacksonSpan; -import org.opensearch.dataprepper.model.trace.Span; - -import java.io.ByteArrayOutputStream; -import java.io.InputStream; -import java.util.Map; -import java.util.Objects; - -import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assertions.assertThatThrownBy; - -class OtlpTraceOutputCodecTest { - - private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); - private static final String TEST_SPAN_EVENT_JSON_FILE = "test-span-event.json"; - - private OtlpTraceOutputCodec codec; - - @BeforeEach - void setup() { - codec = new OtlpTraceOutputCodec(); - } - - private Span buildSpanFromTestFile(String fileName, String traceIdOverride) { - try (InputStream inputStream = Objects.requireNonNull( - getClass().getClassLoader().getResourceAsStream(fileName))) { - - final Map spanMap = OBJECT_MAPPER.readValue(inputStream, new TypeReference<>() {}); - final JacksonSpan.Builder builder = JacksonSpan.builder() - .withTraceId(traceIdOverride != null ? traceIdOverride : (String) spanMap.get("traceId")) - .withSpanId((String) spanMap.get("spanId")) - .withParentSpanId((String) spanMap.get("parentSpanId")) - .withTraceState((String) spanMap.get("traceState")) - .withName((String) spanMap.get("name")) - .withKind((String) spanMap.get("kind")) - .withDurationInNanos(((Number) spanMap.get("durationInNanos")).longValue()) - .withStartTime((String) spanMap.get("startTime")) - .withEndTime((String) spanMap.get("endTime")) - .withTraceGroup((String) spanMap.get("traceGroup")); - - final Map traceGroupFieldsMap = (Map) spanMap.get("traceGroupFields"); - if (traceGroupFieldsMap != null) { - builder.withTraceGroupFields(DefaultTraceGroupFields.builder() - .withStatusCode((Integer) traceGroupFieldsMap.getOrDefault("statusCode", 0)) - .withEndTime((String) spanMap.get("endTime")) - .withDurationInNanos(((Number) spanMap.get("durationInNanos")).longValue()) - .build()); - } - - return builder.build(); - } catch (Exception e) { - throw new RuntimeException("Failed to load span from file", e); - } - } - - @Test - void testWriteEvent_withValidSpanFromTestFile_writesSuccessfully() throws Exception { - final Span span = buildSpanFromTestFile(TEST_SPAN_EVENT_JSON_FILE, null); - final ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); - - codec.start(outputStream, span, new OutputCodecContext()); - codec.writeEvent(span, outputStream); - codec.complete(outputStream); - - final byte[] bytes = outputStream.toByteArray(); - assertThat(bytes).isNotEmpty(); - - final ExportTraceServiceRequest request = ExportTraceServiceRequest.parseFrom(bytes); - assertThat(request.getResourceSpansCount()).isGreaterThan(0); - } - - @Test - void testWriteEvent_withBadTraceId_throwsException() throws Exception { - final Span span = buildSpanFromTestFile(TEST_SPAN_EVENT_JSON_FILE,"bad-trace-id" ); - - final ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); - - assertThatThrownBy(() -> codec.writeEvent(span, outputStream)) - .isInstanceOf(RuntimeException.class); - } - - @Test - void testWriteEvent_withNonSpanEvent_throwsException() { - final JacksonEvent nonSpanEvent = JacksonEvent.builder() - .withEventType("fake") - .withData(Map.of("key", "value")) - .build(); - final ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); - - assertThatThrownBy(() -> codec.writeEvent(nonSpanEvent, outputStream)) - .isInstanceOf(IllegalArgumentException.class) - .hasMessage("OtlpTraceOutputCodec only supports Span events"); - } - - @Test - void testWriteEvent_withNullEvent_throwsNullPointerException() { - final ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); - - assertThatThrownBy(() -> codec.writeEvent(null, outputStream)) - .isInstanceOf(NullPointerException.class) - .hasMessage("event is marked non-null but is null"); - } - - @Test - void testWriteEvent_withNullOutputStream_throwsNullPointerException() { - final Span span = buildSpanFromTestFile(TEST_SPAN_EVENT_JSON_FILE, null); - - assertThatThrownBy(() -> codec.writeEvent(span, null)) - .isInstanceOf(NullPointerException.class) - .hasMessage("outputStream is marked non-null but is null"); - } -} diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java index 005269a28e..101987add5 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java @@ -202,7 +202,7 @@ private void handleSuccessfulResponse(final byte[] responseBytes) { final long rejectedSpans = partial.getRejectedSpans(); final String errorMessage = partial.getErrorMessage(); if (rejectedSpans > 0 || !errorMessage.isEmpty()) { - LOG.error("OTLP Partial Success: rejectedSpans={}, message={}", rejectedSpans, errorMessage); + LOG.warn("OTLP Partial Success: rejectedSpans={}, message={}", rejectedSpans, errorMessage); sinkMetrics.incrementRejectedSpansCount(rejectedSpans); sinkMetrics.incrementErrorsCount(); } From 01671b19c45ca117581f93a3967c57689b565cc3 Mon Sep 17 00:00:00 2001 From: huy pham Date: Thu, 15 May 2025 16:49:55 -0700 Subject: [PATCH 18/23] Update OTLP sink with Armeria, metrics, buffering, and config improvements Signed-off-by: huy pham --- data-prepper-plugins/otlp-sink/README.md | 23 +- data-prepper-plugins/otlp-sink/build.gradle | 45 +- .../sink/otlp/buffer/OtlpSinkBuffer.java | 92 ++-- .../AwsAuthenticationConfig.java | 37 -- .../otlp/configuration/OtlpSinkConfig.java | 4 +- .../otlp/configuration/ThresholdConfig.java | 6 +- .../sink/otlp/http/GzipCompressor.java | 10 +- .../sink/otlp/http/OtlpHttpSender.java | 245 +++++----- .../plugins/sink/otlp/http/ThreadSleeper.java | 32 -- .../sink/otlp/metrics/OtlpSinkMetrics.java | 20 +- .../sink/otlp/buffer/OtlpSinkBufferTest.java | 286 +++++------- .../AwsAuthenticationConfigTest.java | 31 -- .../sink/otlp/http/GzipCompressorTest.java | 15 +- .../sink/otlp/http/OtlpHttpSenderTest.java | 432 +++++------------- .../sink/otlp/http/ThreadSleeperTest.java | 52 --- .../otlp/metrics/OtlpSinkMetricsTest.java | 14 +- 16 files changed, 453 insertions(+), 891 deletions(-) delete mode 100644 data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfig.java delete mode 100644 data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeper.java delete mode 100644 data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfigTest.java delete mode 100644 data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeperTest.java diff --git a/data-prepper-plugins/otlp-sink/README.md b/data-prepper-plugins/otlp-sink/README.md index d15e19eaa1..d8c8f40f77 100644 --- a/data-prepper-plugins/otlp-sink/README.md +++ b/data-prepper-plugins/otlp-sink/README.md @@ -10,7 +10,8 @@ to any OTLP Protobuf-compatible endpoint. - Currently, supports only trace data (spans). Support for metrics and logs will be added in future releases. - No support for DQL-based loss-less delivery in this release. -- Only AWS X-Ray-compatible OTLP endpoints are supported. +- Only AWS X-Ray-compatible OTLP endpoints are currently supported (`https://xray..amazonaws.com/v1/traces`). +- Only OTLP over HTTP is supported; gRPC is not yet supported. --- @@ -52,7 +53,7 @@ otlp_pipeline: | `endpoint` | `String` | Yes | — | AWS X-Ray OTLP endpoint where spans will be sent. | | `max_retries` | `int` | No | `5` | Maximum number of retry attempts on HTTP send failures. | | **threshold** | `Object` | No | — | Controls batching behavior. See below for sub-properties. | -| `threshold.max_events` | `int` | No | `512` (recommended) | Maximum number of spans per batch. Minimum: 1. | +| `threshold.max_events` | `int` | No | `512` (recommended) | Maximum number of spans per batch. Use `0` to disable count-based flushing. Must be ≥ 0. | | `threshold.max_batch_size` | `String` | No | `1mb` (recommended) | Maximum total payload bytes per batch. Supports human-readable suffixes (`kb`, `mb`). | | `threshold.flush_timeout` | `String` | No | `200ms` (recommended) | Maximum time to wait before flushing a non-empty batch. Minimum: 1ms (e.g., `200ms`, `1s`) | | **aws** | `Object` | No | — | AWS authentication settings. See below. | @@ -61,19 +62,12 @@ otlp_pipeline: **Additional Notes:** -- Only AWS X-Ray-compatible OTLP endpoints are currently supported (`https://xray..amazonaws.com/v1/traces`). - `aws.region` is automatically derived from the endpoint. -- When `max_retries` is exceeded: - - Retryable errors (429, 502, 503, 504): exception is logged, plugin error metric is incremented. - - Non-retryable errors (400, 403): error is logged immediately without retry. -- Support for durable delivery (via DQL) will be added in a future release. --- ## Performance Benchmark -The OTLP Sink plugin is optimized for high-throughput, low-latency trace delivery. - ### Summary * Sustains ~3.5K TPS with ≤150ms p99 latency on t4g.large. @@ -126,12 +120,17 @@ Future releases will support durable queueing via DQL for loss-less guarantees. ## Retry Behavior -- The sink uses an exponential backoff strategy for retryable HTTP status codes (e.g., 429, 502, 503, 504). +- The sink uses an exponential backoff with jitter strategy for retryable HTTP status codes (e.g., 429, 502, 503, 504). - Maximum number of attempts is controlled by `max_retries`. Once exceeded: - The span batch is dropped. - The plugin logs the exception and increments the error metric. - Non-retryable errors (e.g., 400, 403) are logged and counted immediately without retry. - +- Retry logic follows + the [OTLP/HTTP response specification](https://opentelemetry.io/docs/specs/otlp/#otlphttp-response). +- `Retry-After` header is not used for dynamic backoff because: + - Armeria’s retry rule API only supports boolean conditions or fixed `Backoff` strategies. + - Supporting `Retry-After` would require a custom `Backoff` implementation, adding unnecessary complexity. + - The exponential backoff already handles common retry intervals effectively. --- ## Logging & Metrics @@ -140,7 +139,7 @@ Future releases will support durable queueing via DQL for loss-less guarantees. * Metrics are emitted via Micrometer and include: * recordsIn, recordsOut * httpLatency, HTTP codes - * errorCount, rejectedSpansCount, retriesCount + * errorCount, rejectedSpansCount, failedSpansCount, retriesCount * queueSize, queueCapacity * payloadSize, payloadGzipSize * JVM stats if configured (e.g., heap usage, GC pauses) diff --git a/data-prepper-plugins/otlp-sink/build.gradle b/data-prepper-plugins/otlp-sink/build.gradle index 6d536eaab1..e1846d9fa9 100644 --- a/data-prepper-plugins/otlp-sink/build.gradle +++ b/data-prepper-plugins/otlp-sink/build.gradle @@ -3,15 +3,6 @@ * SPDX-License-Identifier: Apache-2.0 */ -plugins { - id 'java-library' - id 'jacoco' -} - -jacoco { - toolVersion = "0.8.11" -} - sourceSets { integrationTest { java.srcDir file('src/integrationTest/java') @@ -38,9 +29,6 @@ dependencies { // Hibernate implementation 'org.hibernate.validator:hibernate-validator:8.0.1.Final' - // OkHttp - implementation 'com.squareup.okhttp3:okhttp:4.12.0' - // OpenTelemetry Protobuf implementation libs.opentelemetry.proto implementation libs.protobuf.util @@ -58,20 +46,14 @@ dependencies { // Data Prepper Projects implementation project(':data-prepper-api') + implementation project(':data-prepper-plugins:aws-plugin-api') implementation project(':data-prepper-plugins:otel-proto-common') - implementation 'io.micrometer:micrometer-core' - // Unit Testing - testImplementation 'org.junit.jupiter:junit-jupiter-api:5.9.1' - testRuntimeOnly 'org.junit.jupiter:junit-jupiter-engine:5.9.1' - testImplementation 'org.slf4j:slf4j-simple:2.0.7' - testImplementation 'org.mockito:mockito-core:5.10.0' - testImplementation 'org.awaitility:awaitility:4.2.0' + // Armeria + implementation libs.armeria.core - // Integration Testing - integrationTestImplementation project(':data-prepper-test-common') - integrationTestImplementation 'org.junit.jupiter:junit-jupiter-api:5.9.2' - integrationTestRuntimeOnly 'org.junit.jupiter:junit-jupiter-engine:5.9.2' + // Metrics + implementation 'io.micrometer:micrometer-core' } test { @@ -90,14 +72,6 @@ tasks.register('integrationTest', Test) { check.dependsOn integrationTest -jacocoTestReport { - dependsOn test - reports { - xml.required = true - html.required = true - } -} - jacocoTestCoverageVerification { violationRules { rule { @@ -107,15 +81,8 @@ jacocoTestCoverageVerification { limit { counter = 'LINE' value = 'COVEREDRATIO' - minimum = 0.95 + minimum = 1.00 } } } -} - -tasks.named('jacocoTestReport') { - doLast { - def reportPath = layout.buildDirectory.file("reports/jacoco/test/html/index.html").get().asFile.toURI() - println "\nView test coverage report here:\n ${reportPath}\n" - } } \ No newline at end of file diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBuffer.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBuffer.java index 71d1206fd0..3febfdaeb6 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBuffer.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBuffer.java @@ -8,6 +8,7 @@ import com.google.common.annotations.VisibleForTesting; import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceRequest; import io.opentelemetry.proto.trace.v1.ResourceSpans; +import lombok.Getter; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.model.trace.Span; import org.opensearch.dataprepper.plugins.otel.codec.OTelProtoStandardCodec; @@ -18,18 +19,16 @@ import org.slf4j.LoggerFactory; import javax.annotation.Nonnull; -import java.io.IOException; import java.util.ArrayList; import java.util.List; import java.util.concurrent.BlockingQueue; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.TimeUnit; /** - * A lossless, back-pressure aware buffer for OTLP sink. - *

- * Spans submitted via {@link #add(Record)} are enqueued, batched by count, size or time, - * encoded to ResourceSpans, and flushed asynchronously over HTTP. + * A back-pressure buffer for OTLP sink. */ public class OtlpSinkBuffer { private static final Logger LOG = LoggerFactory.getLogger(OtlpSinkBuffer.class); @@ -45,36 +44,33 @@ public class OtlpSinkBuffer { private final long maxBatchBytes; private final long flushTimeoutMillis; - private final Thread workerThread; + private final ExecutorService executor; + + @Getter private volatile boolean running = true; /** - * Creates a new OTLP sink buffer using default encoder and HTTP sender. + * Creates a new OTLP sink buffer. * * @param config the OTLP sink configuration * @param sinkMetrics the metrics collector to use */ public OtlpSinkBuffer(@Nonnull final OtlpSinkConfig config, @Nonnull final OtlpSinkMetrics sinkMetrics) { - this(config, sinkMetrics, null, null); + this(config, sinkMetrics, new OTelProtoStandardCodec.OTelProtoEncoder(), new OtlpHttpSender(config, sinkMetrics)); } /** * Visible for testing only: constructs an OTLP sink buffer with injected encoder and sender. - * - * @param config the OTLP sink configuration - * @param sinkMetrics the metrics collector - * @param encoder custom OTLP encoder (or null to use default) - * @param sender custom HTTP sender (or null to use default) */ @VisibleForTesting OtlpSinkBuffer(@Nonnull final OtlpSinkConfig config, @Nonnull final OtlpSinkMetrics sinkMetrics, - final OTelProtoStandardCodec.OTelProtoEncoder encoder, - final OtlpHttpSender sender) { + @Nonnull final OTelProtoStandardCodec.OTelProtoEncoder encoder, + @Nonnull final OtlpHttpSender sender) { this.sinkMetrics = sinkMetrics; - this.encoder = encoder != null ? encoder : new OTelProtoStandardCodec.OTelProtoEncoder(); - this.sender = sender != null ? sender : new OtlpHttpSender(config, sinkMetrics); + this.encoder = encoder; + this.sender = sender; this.maxEvents = config.getMaxEvents(); this.maxBatchBytes = config.getMaxBatchSize(); @@ -83,43 +79,60 @@ public OtlpSinkBuffer(@Nonnull final OtlpSinkConfig config, @Nonnull final OtlpS this.queue = new LinkedBlockingQueue<>(getQueueCapacity()); sinkMetrics.registerQueueGauges(queue); - this.workerThread = new Thread(this::run, "otlp-sink-buffer-thread"); + this.executor = Executors.newSingleThreadExecutor(r -> { + final Thread t = new Thread(() -> { + try { + r.run(); + } catch (final Throwable t1) { + LOG.error("Worker thread crashed unexpectedly", t1); + sinkMetrics.incrementErrorsCount(); + restartWorker(); + } + }, "otlp-sink-buffer-thread"); + t.setDaemon(false); + return t; + }); } private int getQueueCapacity() { return Math.max(maxEvents * SAFETY_FACTOR, MIN_QUEUE_CAPACITY); } - public boolean isRunning() { - return running && workerThread.isAlive(); - } - public void start() { running = true; - workerThread.start(); + executor.execute(this::run); } public void stop() { running = false; - workerThread.interrupt(); + executor.shutdownNow(); + } + + @VisibleForTesting + void restartWorker() { + if (running && !executor.isShutdown()) { + LOG.info("Restarting OTLP sink buffer worker thread"); + executor.execute(this::run); + } } /** * Enqueues a span record for later batching and sending. *

* This will block if the internal queue is full, guaranteeing - * lossless delivery. On interruption, the span is rejected and + * lossless delivery during normal operations. + * On interruption, the span is still rejected and * error metrics are incremented. * * @param record the span record to enqueue */ public void add(final Record record) { try { - queue.put(record); // block until space available; guaranteeing lossless delivery + queue.put(record); } catch (final InterruptedException e) { Thread.currentThread().interrupt(); LOG.error("Interrupted while enqueuing span", e); - sinkMetrics.incrementRejectedSpansCount(1); + sinkMetrics.incrementFailedSpansCount(1); sinkMetrics.incrementErrorsCount(); } } @@ -147,12 +160,12 @@ private void run() { batchSize += resourceSpans.getSerializedSize(); } catch (final Exception e) { LOG.error("Failed to encode span, skipping", e); - sinkMetrics.incrementRejectedSpansCount(1); + sinkMetrics.incrementFailedSpansCount(1); sinkMetrics.incrementErrorsCount(); } } - final boolean flushBySize = batch.size() >= maxEvents || batchSize >= maxBatchBytes; + final boolean flushBySize = (maxEvents > 0 && batch.size() >= maxEvents) || batchSize >= maxBatchBytes; final boolean flushByTime = !batch.isEmpty() && (now - lastFlush >= flushTimeoutMillis); if (flushBySize || flushByTime) { @@ -170,8 +183,8 @@ private void run() { LOG.debug("Worker interrupted while polling, continuing..."); sinkMetrics.incrementErrorsCount(); } - // Clear interrupt flag to allow queue.poll() again - // Don't break; allow draining + + // Continue to loop if still running } } @@ -182,9 +195,7 @@ private void run() { } /** - * Builds an ExportTraceServiceRequest from the given batch, sends it over HTTP, - * and updates metrics on success or failure. - *

+ * Builds an ExportTraceServiceRequest from the given batch, sends it over HTTP. * The batch is cleared in all cases to prepare for the next batch. * * @param batch the list of ResourceSpans to send @@ -195,15 +206,8 @@ private void send(final List batch) { .build(); final byte[] payload = request.toByteArray(); - try { - sender.send(payload); - sinkMetrics.incrementRecordsOut(batch.size()); - } catch (final IOException e) { - LOG.error("Failed to send payload.", e); - sinkMetrics.incrementRejectedSpansCount(batch.size()); - sinkMetrics.incrementErrorsCount(); - } finally { - batch.clear(); - } + final int spans = batch.size(); + sender.send(payload, spans); + batch.clear(); } } diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfig.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfig.java deleted file mode 100644 index ec4604cae2..0000000000 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfig.java +++ /dev/null @@ -1,37 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ -package org.opensearch.dataprepper.plugins.sink.otlp.configuration; - -import com.fasterxml.jackson.annotation.JsonProperty; -import jakarta.validation.constraints.Size; -import lombok.Getter; -import lombok.NoArgsConstructor; - -/** - * Configuration class for AWS authentication settings. - * This class will be automatically wired by Data-Prepper. - */ -@Getter -@NoArgsConstructor -class AwsAuthenticationConfig { - - /** - * AWS STS Role ARN for assuming role-based access. - * Format: arn:aws:iam::{account}:role/{role-name} - * Length must be between 20 and 2048 characters. - */ - @JsonProperty("sts_role_arn") - @Size(min = 20, max = 2048, message = "awsStsRoleArn length should be between 1 and 2048 characters") - private String awsStsRoleArn; - - /** - * External ID for additional security when assuming an IAM role. - * Required only if the trust policy requires an external ID. - * Length must be between 2 and 1224 characters. - */ - @JsonProperty("sts_external_id") - @Size(min = 2, max = 1224, message = "awsStsExternalId length should be between 2 and 1224 characters") - private String awsStsExternalId; -} diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfig.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfig.java index 2339091534..f9e5021bce 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfig.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfig.java @@ -10,6 +10,7 @@ import jakarta.validation.constraints.NotBlank; import lombok.Getter; import lombok.NoArgsConstructor; +import org.opensearch.dataprepper.aws.api.AwsConfig; import software.amazon.awssdk.regions.Region; import java.net.URI; @@ -63,12 +64,11 @@ public long getFlushTimeoutMillis() { /** * AWS authentication configuration. - * This object contains the AWS region and STS role ARN (if applicable). * This field is kept private and its contents should be accessed via the generated getter methods. */ @JsonProperty("aws") @Valid - private AwsAuthenticationConfig awsAuthenticationConfig; + private AwsConfig awsAuthenticationConfig; /** * Get AWS region from the provided endpoint. diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/ThresholdConfig.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/ThresholdConfig.java index 0375566a9d..d4c87c0c7c 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/ThresholdConfig.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/ThresholdConfig.java @@ -22,8 +22,12 @@ @Getter class ThresholdConfig { + /** + * Max number of spans per batch. + * Use 0 to disable event-count based flushing (unbounded). + */ @JsonProperty("max_events") - @Min(value = 1, message = "max_events must be at least 1") + @Min(value = 0, message = "max_events must be 0 (unbounded) or greater") private int maxEvents = 512; @JsonProperty("max_batch_size") diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressor.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressor.java index cb8b9c915c..44a14134cd 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressor.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressor.java @@ -12,14 +12,13 @@ import java.io.ByteArrayOutputStream; import java.io.IOException; -import java.util.Optional; import java.util.function.Function; import java.util.zip.GZIPOutputStream; /** * Perform GZIP-compression on OTLP byte payloads. */ -class GzipCompressor implements Function> { +class GzipCompressor implements Function { private static final Logger LOG = LoggerFactory.getLogger(GzipCompressor.class); private final OtlpSinkMetrics sinkMetrics; @@ -39,14 +38,13 @@ class GzipCompressor implements Function> { * @return Optional containing the compressed payload, or empty if compression failed. */ @Override - public Optional apply(final byte[] payload) { + public byte[] apply(final byte[] payload) { try { - return Optional.of(compressInternal(payload)); + return compressInternal(payload); } catch (final IOException e) { LOG.error("Failed to compress payload", e); - sinkMetrics.incrementRejectedSpansCount(1); sinkMetrics.incrementErrorsCount(); - return Optional.empty(); + return new byte[0]; } } diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java index 101987add5..7ea371c2ab 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java @@ -5,12 +5,19 @@ package org.opensearch.dataprepper.plugins.sink.otlp.http; import com.google.common.annotations.VisibleForTesting; +import com.linecorp.armeria.client.WebClient; +import com.linecorp.armeria.client.retry.Backoff; +import com.linecorp.armeria.client.retry.RetryRuleWithContent; +import com.linecorp.armeria.client.retry.RetryingClient; +import com.linecorp.armeria.client.retry.RetryingClientBuilder; +import com.linecorp.armeria.common.HttpData; +import com.linecorp.armeria.common.HttpMethod; +import com.linecorp.armeria.common.HttpRequest; +import com.linecorp.armeria.common.HttpResponse; +import com.linecorp.armeria.common.MediaType; +import com.linecorp.armeria.common.RequestHeaders; +import com.linecorp.armeria.common.RequestHeadersBuilder; import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceResponse; -import okhttp3.MediaType; -import okhttp3.OkHttpClient; -import okhttp3.Request; -import okhttp3.RequestBody; -import okhttp3.Response; import org.opensearch.dataprepper.plugins.sink.otlp.configuration.OtlpSinkConfig; import org.opensearch.dataprepper.plugins.sink.otlp.metrics.OtlpSinkMetrics; import org.slf4j.Logger; @@ -18,179 +25,150 @@ import software.amazon.awssdk.http.SdkHttpFullRequest; import javax.annotation.Nonnull; -import java.io.IOException; import java.nio.charset.StandardCharsets; -import java.security.SecureRandom; -import java.util.ArrayList; -import java.util.List; -import java.util.Optional; import java.util.Set; -import java.util.concurrent.TimeUnit; -import java.util.function.Consumer; import java.util.function.Function; /** - * Responsible for sending signed OTLP Protobuf requests to OTLP endpoint using OkHttp. + * Responsible for sending signed OTLP Protobuf requests to OTLP endpoint using an Ameria client. */ -public class OtlpHttpSender implements AutoCloseable { - private static final Set RETRYABLE_STATUS_CODES = Set.of(429, 502, 503, 504); - private static final int BASE_RETRY_DELAY_MS = 100; +public class OtlpHttpSender { private static final Logger LOG = LoggerFactory.getLogger(OtlpHttpSender.class); - private static final MediaType PROTOBUF = MediaType.get("application/x-protobuf"); - private final SecureRandom random = new SecureRandom(); + private static final Set RETRYABLE_STATUS_CODES = Set.of(429, 502, 503, 504); + private static final MediaType PROTOBUF = MediaType.parse("application/x-protobuf"); - private final int maxRetries; private final SigV4Signer signer; - private final OkHttpClient httpClient; - private final Consumer sleeper; - private final List retryDelaysMs; + private final WebClient webClient; private final OtlpSinkMetrics sinkMetrics; - private final Function> gzipCompressor; + private final Function gzipCompressor; /** * Constructor for the OtlpHttpSender. - * Initializes the signer and HTTP client. * * @param config The configuration for the OTLP sink plugin. * @param sinkMetrics The metrics for the OTLP sink plugin. */ public OtlpHttpSender(@Nonnull final OtlpSinkConfig config, @Nonnull final OtlpSinkMetrics sinkMetrics) { - this(config, sinkMetrics, new GzipCompressor(sinkMetrics), null, null, null); + this(sinkMetrics, new GzipCompressor(sinkMetrics), new SigV4Signer(config), buildWebClient(config)); } /** * Constructor for unit testing with injected dependencies. */ @VisibleForTesting - OtlpHttpSender(@Nonnull final OtlpSinkConfig config, @Nonnull final OtlpSinkMetrics sinkMetrics, @Nonnull final Function> gzipCompressor, - final SigV4Signer signer, final OkHttpClient httpClient, final Consumer sleeper) { + OtlpHttpSender(@Nonnull final OtlpSinkMetrics sinkMetrics, @Nonnull final Function gzipCompressor, + final SigV4Signer signer, final WebClient webClient) { this.sinkMetrics = sinkMetrics; this.gzipCompressor = gzipCompressor; - this.signer = signer != null ? signer : new SigV4Signer(config); - this.sleeper = sleeper != null ? sleeper : new ThreadSleeper(); - this.httpClient = httpClient != null ? httpClient : buildOkHttpClient(config.getFlushTimeoutMillis()); - - this.retryDelaysMs = generateExponentialBackoffDelays(config.getMaxRetries()); - this.maxRetries = config.getMaxRetries(); + this.signer = signer; + this.webClient = webClient; } - private static OkHttpClient buildOkHttpClient(final long flushTimeoutMs) { + /** + * Builds a WebClient with retry logic for known OTLP retryable status codes. + *

+ * Retries on 429, 502, 503, and 504 per OTEL spec. + * See: OTLP/HTTP Response + *

+ * We are not using the Retry-After header for dynamic backoff because: + * - Armeria’s retry rule API expects a boolean decision or fixed Backoff. + * - Applying Retry-After semantics would require a custom Backoff implementation, + * adding complexity with minimal benefit for most OTLP endpoints. + * - Our exponential backoff already handles typical retry intervals gracefully. + */ + + private static WebClient buildWebClient(final OtlpSinkConfig config) { + final RetryRuleWithContent retryRule = RetryRuleWithContent.builder() + .onStatus((ctx, status) -> RETRYABLE_STATUS_CODES.contains(status.code())) + .thenBackoff(Backoff.exponential(100, 10_000).withJitter(0.2)); + + final long estimatedContentLimit = Math.max(1, config.getMaxBatchSize()) * (config.getMaxRetries() + 1); + final int safeContentLimit = (int) Math.min(estimatedContentLimit, Integer.MAX_VALUE); + + final RetryingClientBuilder retryingClientBuilder = RetryingClient.builder(retryRule, safeContentLimit) + .maxTotalAttempts(config.getMaxRetries() + 1); + final long httpTimeoutMs = Math.min( - Math.max(flushTimeoutMs * 2, 3_000), - 10_000 + Math.max(config.getFlushTimeoutMillis() * 2, 3_000), 10_000 ); - return new OkHttpClient.Builder() - .callTimeout(httpTimeoutMs, TimeUnit.MILLISECONDS) + return WebClient.builder() + .decorator(retryingClientBuilder.newDecorator()) + .responseTimeoutMillis(httpTimeoutMs) + .maxResponseLength(safeContentLimit) .build(); } /** - * Generates exponential backoff delays with jitter. + * Sends the provided OTLP Protobuf payload to the OTLP endpoint asynchronously. * - * @param retries Number of retries. - * @return List of delay durations in milliseconds. + * @param payload The OTLP Protobuf-encoded data to be sent. + * @param spans The number of spans in the payload. */ - private static List generateExponentialBackoffDelays(final int retries) { - final List delays = new ArrayList<>(); - for (int i = 0; i < retries; i++) { - // Exponential backoff: 100ms, 200ms, 400ms, ... - delays.add(BASE_RETRY_DELAY_MS * (1 << i)); + public void send(@Nonnull final byte[] payload, final int spans) { + final byte[] compressedPayload = gzipCompressor.apply(payload); + if (compressedPayload.length == 0) { + sinkMetrics.incrementFailedSpansCount(spans); + return; } - return delays; - } + final HttpRequest request = buildHttpRequest(compressedPayload); - /** - * Sends the provided OTLP Protobuf payload to the OTLP endpoint. - * Retries with exponential backoff and jitter on failure. - * - * @param payload The OTLP Protobuf-encoded data to be sent. - * @throws IOException when failed to send the payload. - */ - public void send(@Nonnull final byte[] payload) throws IOException { - for (int attempt = 0; attempt <= maxRetries; attempt++) { - try { - final Optional compressedPayload = gzipCompressor.apply(payload); - if (compressedPayload.isEmpty()) { - return; - } - - final SdkHttpFullRequest signedRequest = signer.signRequest(compressedPayload.get()); - final Request.Builder requestBuilder = new Request.Builder() - .url(signedRequest.getUri().toString()) - .post(RequestBody.create(compressedPayload.get(), PROTOBUF)) - .addHeader("Content-Encoding", "gzip"); - - signedRequest.headers().forEach((key, values) -> { - for (final String value : values) { - requestBuilder.addHeader(key, value); - } + final long startTime = System.currentTimeMillis(); + webClient.execute(request) + .aggregate() + .thenAccept(response -> { + final long latency = System.currentTimeMillis() - startTime; + sinkMetrics.recordHttpLatency(latency); + sinkMetrics.incrementPayloadSize(payload.length); + sinkMetrics.incrementPayloadGzipSize(compressedPayload.length); + + final int statusCode = response.status().code(); + final byte[] responseBytes = response.content().array(); + handleResponse(statusCode, responseBytes, spans); + }) + .exceptionally(e -> { + LOG.error("Failed to send {} spans.", spans, e); + sinkMetrics.incrementRejectedSpansCount(spans); + return null; }); + } - final Request request = requestBuilder.build(); - - final long startTime = System.currentTimeMillis(); - try (final Response response = httpClient.newCall(request).execute()) { - final long duration = System.currentTimeMillis() - startTime; - sinkMetrics.recordHttpLatency(duration); - - handleResponse(response); + private HttpRequest buildHttpRequest(final byte[] compressedPayload) { + final SdkHttpFullRequest signedRequest = signer.signRequest(compressedPayload); + final RequestHeadersBuilder headersBuilder = RequestHeaders.builder() + .method(HttpMethod.POST) + .path(signedRequest.getUri().getPath()) + .contentType(PROTOBUF) + .add("Content-Encoding", "gzip"); - sinkMetrics.incrementPayloadSize(payload.length); - sinkMetrics.incrementPayloadGzipSize(compressedPayload.get().length); - return; - } - } catch (final IOException ioException) { - if (attempt < maxRetries) { - final int jitter = random.nextInt(100); - final int retryIndex = Math.min(attempt, retryDelaysMs.size() - 1); - final int delay = retryDelaysMs.get(retryIndex) + jitter; - try { - sleeper.accept(delay); - sinkMetrics.incrementRetriesCount(); - } catch (final RuntimeException runtimeException) { - throw new IOException("Sender failed to sleep before retrying.", runtimeException); - } - } else { - throw new IOException("Max retries reached", ioException); - } - } - } + signedRequest.headers().forEach((k, vList) -> vList.forEach(v -> headersBuilder.add(k, v))); + return HttpRequest.of(headersBuilder.build(), HttpData.wrap(compressedPayload)); } - /** - * Handles the OTLP export response. - * Retries on 429, 502, 503, and 504 per OTEL spec. Logs other errors without retry. - * See: OTLP/HTTP Response - * - * @param response The HTTP response - * @throws IOException For retryable errors - */ - private void handleResponse(@Nonnull final Response response) throws IOException { - final int status = response.code(); - sinkMetrics.recordResponseCode(status); - - final byte[] responseBytes = response.body() != null - ? response.body().bytes() - : null; + private void handleResponse(final int statusCode, final byte[] responseBytes, final int spans) { + sinkMetrics.recordResponseCode(statusCode); - if (status >= 200 && status < 300) { - handleSuccessfulResponse(responseBytes); + if (statusCode >= 200 && statusCode < 300) { + handleSuccessfulResponse(responseBytes, spans); return; } - final String responseBody = responseBytes != null ? new String(responseBytes, StandardCharsets.UTF_8) : ""; - if (RETRYABLE_STATUS_CODES.contains(status)) { - throw new IOException(String.format("Retryable error. Status: %d, Response: %s", status, responseBody)); - } + final String responseBody = responseBytes != null + ? new String(responseBytes, StandardCharsets.UTF_8) + : ""; - LOG.error("Non-retryable error. Status: {}, Response: {}", status, responseBody); + LOG.error("Non-successful OTLP response. Status: {}, Response: {}", statusCode, responseBody); + sinkMetrics.incrementRejectedSpansCount(spans); } - private void handleSuccessfulResponse(final byte[] responseBytes) { - if (responseBytes == null || responseBytes.length == 0) { + /** + * Handles a successful OTLP response with partial success. + */ + private void handleSuccessfulResponse(final byte[] responseBytes, final int spans) { + if (responseBytes == null) { + sinkMetrics.incrementRecordsOut(spans); return; } @@ -201,21 +179,20 @@ private void handleSuccessfulResponse(final byte[] responseBytes) { final var partial = otlpResponse.getPartialSuccess(); final long rejectedSpans = partial.getRejectedSpans(); final String errorMessage = partial.getErrorMessage(); - if (rejectedSpans > 0 || !errorMessage.isEmpty()) { - LOG.warn("OTLP Partial Success: rejectedSpans={}, message={}", rejectedSpans, errorMessage); + if (rejectedSpans > 0) { + LOG.error("OTLP Partial Success: rejectedSpans={}, message={}", rejectedSpans, errorMessage); sinkMetrics.incrementRejectedSpansCount(rejectedSpans); - sinkMetrics.incrementErrorsCount(); } + + final long deliveredSpans = spans - rejectedSpans; + sinkMetrics.incrementRecordsOut(deliveredSpans); + } else { + sinkMetrics.incrementRecordsOut(spans); } } catch (final Exception e) { LOG.error("Could not parse OTLP response as ExportTraceServiceResponse: {}", e.getMessage()); sinkMetrics.incrementErrorsCount(); + sinkMetrics.incrementRecordsOut(spans); } } - - @Override - public void close() { - httpClient.connectionPool().evictAll(); - httpClient.dispatcher().executorService().shutdown(); - } } \ No newline at end of file diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeper.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeper.java deleted file mode 100644 index 6b395fd298..0000000000 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeper.java +++ /dev/null @@ -1,32 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.sink.otlp.http; - -import javax.annotation.Nonnull; -import java.util.function.Consumer; - -/** - * A simple {@link Consumer} that pauses the current thread for a given number - * of milliseconds. - */ -class ThreadSleeper implements Consumer { - - /** - * Sleeps for the specified duration in milliseconds. - * Wraps and rethrows {@link InterruptedException} as a runtime exception. - * - * @param millis the number of milliseconds to sleep - */ - @Override - public void accept(@Nonnull final Integer millis) { - try { - Thread.sleep(millis); - } catch (final InterruptedException e) { - Thread.currentThread().interrupt(); - throw new RuntimeException("Sleep interrupted", e); - } - } -} diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetrics.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetrics.java index 89e8d3feac..0079c8b356 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetrics.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetrics.java @@ -102,17 +102,27 @@ public void recordHttpLatency(final long durationMillis) { httpLatency.record(Duration.ofMillis(durationMillis)); } - public void incrementRetriesCount() { - pluginMetrics.counter("retriesCount").increment(1); + public void registerQueueGauges(final BlockingQueue queue) { + pluginMetrics.gauge("queueSize", queue, BlockingQueue::size); + pluginMetrics.gauge("queueCapacity", queue, q -> q.remainingCapacity() + q.size()); } + /** + * Increments the count of spans that were explicitly rejected by the OTLP endpoint. + * + * @param count The number of spans rejected. + */ public void incrementRejectedSpansCount(final long count) { pluginMetrics.counter("rejectedSpansCount").increment(count); } - public void registerQueueGauges(final BlockingQueue queue) { - pluginMetrics.gauge("queueSize", queue, BlockingQueue::size); - pluginMetrics.gauge("queueCapacity", queue, q -> q.remainingCapacity() + q.size()); + /** + * Increments the count of spans that failed to be processed by the sink. + * + * @param count The number of spans failed. + */ + public void incrementFailedSpansCount(final long count) { + pluginMetrics.counter("failedSpansCount").increment(count); } /** diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBufferTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBufferTest.java index ba8cd81459..2612108d50 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBufferTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBufferTest.java @@ -17,22 +17,22 @@ import org.opensearch.dataprepper.plugins.sink.otlp.metrics.OtlpSinkMetrics; import software.amazon.awssdk.regions.Region; -import java.io.IOException; import java.lang.reflect.Field; import java.util.concurrent.BlockingQueue; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; import static java.util.concurrent.TimeUnit.SECONDS; import static org.awaitility.Awaitility.await; -import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.Mockito.atLeastOnce; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -48,6 +48,7 @@ class OtlpSinkBufferTest { void setUp() { config = mock(OtlpSinkConfig.class); when(config.getMaxEvents()).thenReturn(2); + when(config.getMaxRetries()).thenReturn(2); when(config.getMaxBatchSize()).thenReturn(1_000_000L); when(config.getFlushTimeoutMillis()).thenReturn(10L); when(config.getAwsRegion()).thenReturn(Region.of("us-west-2")); @@ -57,38 +58,25 @@ void setUp() { sender = mock(OtlpHttpSender.class); buffer = new OtlpSinkBuffer(config, metrics, encoder, sender); - buffer.start(); } @AfterEach void tearDown() { buffer.stop(); } - - @Test - void testIsRunningBeforeStartAndAfterStop() throws Exception { - // create a fresh buffer (not started) - final OtlpSinkBuffer localBuffer = new OtlpSinkBuffer(config, metrics, encoder, sender); - assertFalse(localBuffer.isRunning(), "not running until start() is called"); - - localBuffer.start(); - // give the thread a moment to spin up - TimeUnit.MILLISECONDS.sleep(50); - assertTrue(localBuffer.isRunning(), "should be running immediately after start()"); - localBuffer.stop(); - // give the thread time to terminate - TimeUnit.MILLISECONDS.sleep(50); - assertFalse(localBuffer.isRunning(), "should stop after stop()"); + @Test + void testStartAndStopDoesNotThrow() { + buffer.start(); + buffer.stop(); + assertFalse(buffer.isRunning()); } @Test void testAddHandlesInterruptedException() throws Exception { - // create a mock queue that throws when put(...) is called @SuppressWarnings("unchecked") final BlockingQueue> badQueue = mock(BlockingQueue.class); doThrow(new InterruptedException()).when(badQueue).put(any()); - // inject it via reflection final Field queueField = OtlpSinkBuffer.class.getDeclaredField("queue"); queueField.setAccessible(true); queueField.set(buffer, badQueue); @@ -96,211 +84,161 @@ void testAddHandlesInterruptedException() throws Exception { final Record rec = mock(Record.class); buffer.add(rec); - // should have hit the InterruptedException branch - verify(metrics).incrementRejectedSpansCount(1); + verify(metrics).incrementFailedSpansCount(1); verify(metrics).incrementErrorsCount(); } @Test - void testFinalFlushOnShutdownWhenNoSizeOrTimeFlush() throws Exception { - // new config: very large batch/time so we never flush by size/time - final OtlpSinkConfig largeConfig = mock(OtlpSinkConfig.class); - when(largeConfig.getMaxEvents()).thenReturn(100); - when(largeConfig.getMaxBatchSize()).thenReturn(Long.MAX_VALUE); - when(largeConfig.getFlushTimeoutMillis()).thenReturn(Long.MAX_VALUE); - - final OtlpSinkBuffer finalBuffer = new OtlpSinkBuffer(largeConfig, metrics, encoder, sender); - finalBuffer.start(); - - // enqueue exactly one record - final Record rec = mock(Record.class); - when(rec.getData()).thenReturn(mock(Span.class)); - final ResourceSpans rs = ResourceSpans.getDefaultInstance(); - when(encoder.convertToResourceSpans(any(Span.class))).thenReturn(rs); - finalBuffer.add(rec); + void testWorkerThreadHandlesEncodeException() throws Exception { + when(encoder.convertToResourceSpans(any(Span.class))) + .thenThrow(new RuntimeException("boom")) + .thenReturn(ResourceSpans.getDefaultInstance()); - // now shutdown - finalBuffer.stop(); - TimeUnit.MILLISECONDS.sleep(50); + final Record rec1 = mock(Record.class); + final Record rec2 = mock(Record.class); + when(rec1.getData()).thenReturn(mock(Span.class)); + when(rec2.getData()).thenReturn(mock(Span.class)); - // final‐flush should happen exactly once - await().atMost(1, SECONDS).untilAsserted(() -> - verify(sender).send(any(byte[].class)) - ); - verify(metrics).incrementRecordsOut(1); - } + buffer.start(); + buffer.add(rec1); + buffer.add(rec2); - @Test - void testSendIoExceptionIncrementsRejectedAndError() throws Exception { - // Prepare a tiny batch so that send(...) will be invoked - final ResourceSpans rs = ResourceSpans.getDefaultInstance(); - when(encoder.convertToResourceSpans(any(Span.class))).thenReturn(rs); - doThrow(new IOException("uh-oh")).when(sender).send(any(byte[].class)); - - // Enqueue two spans to hit batch-size flush - for (int i = 0; i < 2; i++) { - final Record rec = mock(Record.class); - when(rec.getData()).thenReturn(mock(Span.class)); - buffer.add(rec); - } - - // Give worker thread time to flush by size - TimeUnit.MILLISECONDS.sleep(50); + TimeUnit.MILLISECONDS.sleep(100); buffer.stop(); - // verify send was attempted - verify(sender).send(any(byte[].class)); - // one batch of 2 spans failed: rejected count should be 2 - verify(metrics).incrementRejectedSpansCount(2); - // error count for the IO error (+ possibly one more when interrupted) + await().atMost(1, SECONDS).untilAsserted(() -> verify(sender).send(any(byte[].class), anyInt())); + verify(metrics).incrementFailedSpansCount(1); verify(metrics, atLeastOnce()).incrementErrorsCount(); } @Test - void testWorkerThreadFlushesBySize() throws Exception { - final ResourceSpans rs = ResourceSpans.getDefaultInstance(); - when(encoder.convertToResourceSpans(any(Span.class))).thenReturn(rs); - - // Enqueue exactly maxEvents (2) spans - for (int i = 0; i < 2; i++) { - final Record rec = mock(Record.class); - when(rec.getData()).thenReturn(mock(Span.class)); - buffer.add(rec); - } - - TimeUnit.MILLISECONDS.sleep(50); - buffer.stop(); + void testUncaughtExceptionHandler_logsAndRestarts_actualThread() throws Exception { + final ExecutorService crashingExecutor = Executors.newSingleThreadExecutor(r -> { + final Thread t = new Thread(() -> { + throw new RuntimeException("forced crash"); + }, "otlp-sink-buffer-thread"); + t.setUncaughtExceptionHandler((thread, ex) -> { + metrics.incrementErrorsCount(); + buffer.restartWorker(); + }); + return t; + }); + + final Field executorField = OtlpSinkBuffer.class.getDeclaredField("executor"); + executorField.setAccessible(true); + executorField.set(buffer, crashingExecutor); - // at least one send of our 2-item batch - verify(sender, atLeastOnce()).send(any(byte[].class)); - // at least one successful record-out of 2 - verify(metrics, atLeastOnce()).incrementRecordsOut(2); - } + buffer.start(); - @Test - void testQueueCapacityRespectsMinimum() throws Exception { - when(config.getMaxEvents()).thenReturn(1); - buffer = new OtlpSinkBuffer(config, metrics, encoder, sender); + await().atMost(1, SECONDS).untilAsserted(() -> verify(metrics, atLeastOnce()).incrementErrorsCount()); - final Field queueField = OtlpSinkBuffer.class.getDeclaredField("queue"); - queueField.setAccessible(true); - final BlockingQueue queueInstance = (BlockingQueue) queueField.get(buffer); - assertEquals(2000, queueInstance.remainingCapacity()); + crashingExecutor.shutdownNow(); } @Test - void testQueueCapacityBasedOnMaxEvents() throws Exception { - when(config.getMaxEvents()).thenReturn(300); - buffer = new OtlpSinkBuffer(config, metrics, encoder, sender); + void testRestartWorkerDoesNotSubmitIfShutdown() throws Exception { + final ExecutorService mockExecutor = mock(ExecutorService.class); + when(mockExecutor.isShutdown()).thenReturn(true); - final Field queueField = OtlpSinkBuffer.class.getDeclaredField("queue"); - queueField.setAccessible(true); - final BlockingQueue queueInstance = (BlockingQueue) queueField.get(buffer); - assertEquals(300 * 10, queueInstance.remainingCapacity()); - } + final Field executorField = OtlpSinkBuffer.class.getDeclaredField("executor"); + executorField.setAccessible(true); + executorField.set(buffer, mockExecutor); - @Test - void testWorkerThreadFlushesByBatchByteSize() throws Exception { - // Arrange: make batchSize threshold zero so any non-null ResourceSpans triggers flush - when(config.getMaxEvents()).thenReturn(10); - when(config.getMaxBatchSize()).thenReturn(0L); - when(config.getFlushTimeoutMillis()).thenReturn(Long.MAX_VALUE); + buffer.restartWorker(); - // restart with new config - buffer.stop(); - buffer = new OtlpSinkBuffer(config, metrics, encoder, sender); - buffer.start(); + verify(mockExecutor, never()).execute(any()); + } - when(encoder.convertToResourceSpans(any(Span.class))) - .thenReturn(ResourceSpans.getDefaultInstance()); + @Test + void testRestartWorkerSubmitsRunnableIfRunning() throws Exception { + final ExecutorService mockExecutor = mock(ExecutorService.class); + when(mockExecutor.isShutdown()).thenReturn(false); - // Act: add one record - final Record rec = mock(Record.class); - when(rec.getData()).thenReturn(mock(Span.class)); - buffer.add(rec); + final Field executorField = OtlpSinkBuffer.class.getDeclaredField("executor"); + executorField.setAccessible(true); + executorField.set(buffer, mockExecutor); - // give the worker a moment to do both the immediate flush and then exit - TimeUnit.MILLISECONDS.sleep(50); - buffer.stop(); + buffer.restartWorker(); - // Assert: we expect at least one send (could be two) - verify(sender, atLeastOnce()).send(any(byte[].class)); - verify(metrics, atLeastOnce()).incrementRecordsOut(1); + verify(mockExecutor, atLeastOnce()).execute(any()); } @Test - void testWorkerThreadHandlesEncodeException() throws Exception { - // Bad span first - when(encoder.convertToResourceSpans(any(Span.class))) - .thenThrow(new RuntimeException("boom")) - .thenReturn(ResourceSpans.getDefaultInstance()); + void testRunWithInterruptedException() throws Exception { + final OtlpSinkBuffer localBuffer = new OtlpSinkBuffer(config, metrics, encoder, sender); - // Enqueue two spans: one bad, one good - for (int i = 0; i < 2; i++) { - final Record rec = mock(Record.class); - when(rec.getData()).thenReturn(mock(Span.class)); - buffer.add(rec); - } + // Override the internal queue to throw InterruptedException + final BlockingQueue> interruptingQueue = mock(BlockingQueue.class); + when(interruptingQueue.poll(any(Long.class), any(TimeUnit.class))).thenThrow(new InterruptedException()); - TimeUnit.MILLISECONDS.sleep(50); - buffer.stop(); + final Field queueField = OtlpSinkBuffer.class.getDeclaredField("queue"); + queueField.setAccessible(true); + queueField.set(localBuffer, interruptingQueue); + + localBuffer.start(); + + TimeUnit.MILLISECONDS.sleep(100); // Let the thread hit the poll + + localBuffer.stop(); - // should still send at least one batch for the good span await().atMost(1, SECONDS).untilAsserted(() -> - verify(sender).send(any(byte[].class)) + verify(metrics, atLeastOnce()).incrementErrorsCount() ); - // one rejected from the encode exception - verify(metrics).incrementRejectedSpansCount(1); - // error count for the encode exception (+ maybe one on interrupt) - verify(metrics, atLeastOnce()).incrementErrorsCount(); - // at least one successful record-out of 1 - verify(metrics, atLeastOnce()).incrementRecordsOut(1); } @Test - void testConstructorDefaults() throws Exception { - // use the two-arg constructor - final OtlpSinkBuffer defaultBuf = new OtlpSinkBuffer(config, metrics); + void testWorkerThreadCatchThrowable_andRestart_fromEncoder() throws Exception { + final OTelProtoStandardCodec.OTelProtoEncoder crashingEncoder = mock(OTelProtoStandardCodec.OTelProtoEncoder.class); + when(crashingEncoder.convertToResourceSpans(any())).thenAnswer(invocation -> { + throw new AssertionError("simulated fatal crash"); + }); - // reflectively inspect encoder and sender - final Field encField = OtlpSinkBuffer.class.getDeclaredField("encoder"); - final Field sndField = OtlpSinkBuffer.class.getDeclaredField("sender"); - encField.setAccessible(true); - sndField.setAccessible(true); - - final Object enc = encField.get(defaultBuf); - final Object snd = sndField.get(defaultBuf); + buffer = new OtlpSinkBuffer(config, metrics, crashingEncoder, sender); + buffer.start(); - assertNotNull(enc, "default encoder should not be null"); - assertInstanceOf(OTelProtoStandardCodec.OTelProtoEncoder.class, enc); + final Record record = mock(Record.class); + when(record.getData()).thenReturn(mock(Span.class)); + buffer.add(record); - assertNotNull(snd, "default sender should not be null"); - assertInstanceOf(OtlpHttpSender.class, snd); + await().atMost(2, SECONDS).untilAsserted(() -> { + verify(metrics, atLeastOnce()).incrementErrorsCount(); + }); } @Test - void testWorkerThreadFlushesByTimeoutOnly() throws Exception { - when(config.getMaxEvents()).thenReturn(100); // high enough not to flush by count - when(config.getMaxBatchSize()).thenReturn(Long.MAX_VALUE); // don't flush by size - when(config.getFlushTimeoutMillis()).thenReturn(50L); // very short flush window + void testFinalFlushAfterStopFlushesBatch() throws Exception { + // Configure buffer to *not* flush by size or time + when(config.getMaxEvents()).thenReturn(1000); // very high + when(config.getMaxBatchSize()).thenReturn(Long.MAX_VALUE); // very high + when(config.getFlushTimeoutMillis()).thenReturn(Long.MAX_VALUE); // effectively disables time-based flush - buffer.stop(); buffer = new OtlpSinkBuffer(config, metrics, encoder, sender); buffer.start(); - final Record rec = mock(Record.class); - when(rec.getData()).thenReturn(mock(Span.class)); - when(encoder.convertToResourceSpans(any(Span.class))).thenReturn(ResourceSpans.getDefaultInstance()); + // Prepare one span record + final Record record = mock(Record.class); + when(record.getData()).thenReturn(mock(Span.class)); + when(encoder.convertToResourceSpans(any())).thenReturn(ResourceSpans.getDefaultInstance()); - buffer.add(rec); + buffer.add(record); - // wait long enough to trigger timeout-based flush + // Wait briefly to ensure the record is picked up TimeUnit.MILLISECONDS.sleep(100); + + // Stop should trigger the final flush buffer.stop(); + // Verify final flush triggered send await().atMost(1, SECONDS).untilAsserted(() -> - verify(sender).send(any(byte[].class)) + verify(sender).send(any(byte[].class), anyInt()) ); - verify(metrics, atLeastOnce()).incrementRecordsOut(1); + } + + @Test + void testPublicConstructorInitializesWithDefaults() { + // Only mocks needed are config and metrics + final OtlpSinkBuffer defaultBuffer = new OtlpSinkBuffer(config, metrics); + + assertNotNull(defaultBuffer); // Constructor didn't throw } } \ No newline at end of file diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfigTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfigTest.java deleted file mode 100644 index be73ba8893..0000000000 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/AwsAuthenticationConfigTest.java +++ /dev/null @@ -1,31 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ -package org.opensearch.dataprepper.plugins.sink.otlp.configuration; - -import com.fasterxml.jackson.databind.ObjectMapper; -import com.fasterxml.jackson.dataformat.yaml.YAMLFactory; -import org.junit.jupiter.api.Test; - -import static org.junit.jupiter.api.Assertions.assertEquals; - -class AwsAuthenticationConfigTest { - - private final String expectedRoleArn = "arn:aws:iam::123456789012:role/MyRole"; - private final String expectedExternalId = "external-id-123"; - private final ObjectMapper mapper = new ObjectMapper(new YAMLFactory()); - - @Test - void testDeserializationFromYaml() throws Exception { - final String yaml = String.join("\n", - "sts_role_arn: " + expectedRoleArn, - "sts_external_id: " + expectedExternalId - ); - - final AwsAuthenticationConfig config = mapper.readValue(yaml, AwsAuthenticationConfig.class); - - assertEquals(expectedRoleArn, config.getAwsStsRoleArn()); - assertEquals(expectedExternalId, config.getAwsStsExternalId()); - } -} diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressorTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressorTest.java index d0c5e86b44..36a07327f0 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressorTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/GzipCompressorTest.java @@ -12,11 +12,11 @@ import java.io.ByteArrayInputStream; import java.io.IOException; import java.nio.charset.StandardCharsets; -import java.util.Optional; import java.util.zip.GZIPInputStream; import static org.junit.jupiter.api.Assertions.assertArrayEquals; -import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; @@ -36,12 +36,12 @@ void setUp() { void apply_returnsCompressedPayload() throws IOException { byte[] input = "test-payload".getBytes(); GzipCompressor gzipCompressor = new GzipCompressor(sinkMetrics); - Optional compressed = gzipCompressor.apply(input); - assertTrue(compressed.isPresent(), "Expected compressed payload to be present"); + final byte[] compressed = gzipCompressor.apply(input); // Validate decompression gives original input - byte[] decompressed = decompress(compressed.get()); + assertNotNull(compressed); + final byte[] decompressed = decompress(compressed); assertArrayEquals(input, decompressed); } @@ -50,11 +50,10 @@ void apply_handlesIOException_andIncrementsErrorMetric() throws IOException { GzipCompressor gzipCompressor = spy(new GzipCompressor(sinkMetrics)); doThrow(new IOException("boom")).when(gzipCompressor).compressInternal(any()); - Optional result = gzipCompressor.apply("payload".getBytes(StandardCharsets.UTF_8)); + final byte[] result = gzipCompressor.apply("payload".getBytes(StandardCharsets.UTF_8)); - assertTrue(result.isEmpty()); + assertEquals(0, result.length); verify(sinkMetrics).incrementErrorsCount(); - verify(sinkMetrics).incrementRejectedSpansCount(1); } private byte[] decompress(byte[] compressed) throws IOException { diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java index 2c8b4c9f19..6b7056278f 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java @@ -5,377 +5,195 @@ package org.opensearch.dataprepper.plugins.sink.otlp.http; +import com.linecorp.armeria.client.WebClient; +import com.linecorp.armeria.common.HttpData; +import com.linecorp.armeria.common.HttpRequest; +import com.linecorp.armeria.common.HttpResponse; +import com.linecorp.armeria.common.ResponseHeaders; import io.opentelemetry.proto.collector.trace.v1.ExportTracePartialSuccess; import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceResponse; -import okhttp3.Call; -import okhttp3.ConnectionPool; -import okhttp3.Dispatcher; -import okhttp3.HttpUrl; -import okhttp3.MediaType; -import okhttp3.OkHttpClient; -import okhttp3.Protocol; -import okhttp3.Request; -import okhttp3.Response; -import okhttp3.ResponseBody; -import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.opensearch.dataprepper.plugins.sink.otlp.configuration.OtlpSinkConfig; import org.opensearch.dataprepper.plugins.sink.otlp.metrics.OtlpSinkMetrics; -import software.amazon.awssdk.http.SdkHttpFullRequest; -import software.amazon.awssdk.regions.Region; -import java.io.IOException; +import java.lang.reflect.Method; import java.net.URI; import java.nio.charset.StandardCharsets; -import java.util.Collections; -import java.util.Map; -import java.util.Optional; -import java.util.concurrent.ExecutorService; -import java.util.function.Consumer; import java.util.function.Function; -import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; -import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.awaitility.Awaitility.await; import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.junit.jupiter.api.Assertions.fail; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyInt; -import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.any; +import static org.mockito.Mockito.anyLong; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.reset; -import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; class OtlpHttpSenderTest { - private static final byte[] PAYLOAD = "test-otlp-payload".getBytes(StandardCharsets.UTF_8); - private static final String ERROR_BODY = "{\"error\": \"Something went wrong\"}"; - private OtlpSinkConfig mockConfig; - private SigV4Signer mockSigner; - private OkHttpClient mockHttpClient; - private Consumer mockSleeper; - private Function> mockGzipCompressor; - private OtlpSinkMetrics mockSinkMetrics; - private OtlpHttpSender target; + private static final byte[] PAYLOAD = "test-otlp".getBytes(StandardCharsets.UTF_8); + private static final int SPANS = 3; + + private OtlpSinkMetrics metrics; + private SigV4Signer signer; + private WebClient webClient; + private Function gzipCompressor; + private OtlpHttpSender sender; @BeforeEach - void setUp() { - System.setProperty("aws.accessKeyId", "dummy"); - System.setProperty("aws.secretAccessKey", "dummy"); - - mockConfig = mock(OtlpSinkConfig.class); - when(mockConfig.getAwsRegion()).thenReturn(Region.US_WEST_2); - when(mockConfig.getMaxRetries()).thenReturn(3); - - mockSigner = mock(SigV4Signer.class); - mockHttpClient = mock(OkHttpClient.class); - mockSleeper = mock(ThreadSleeper.class); - mockSinkMetrics = mock(OtlpSinkMetrics.class); - - mockGzipCompressor = mock(GzipCompressor.class); - when(mockGzipCompressor.apply(any())) - .thenAnswer(invocation -> Optional.of((byte[]) invocation.getArgument(0))); - - target = new OtlpHttpSender( - mockConfig, mockSinkMetrics, - mockGzipCompressor, mockSigner, - mockHttpClient, mockSleeper + void setup() { + metrics = mock(OtlpSinkMetrics.class); + signer = mock(SigV4Signer.class); + webClient = mock(WebClient.class); + gzipCompressor = mock(Function.class); + + when(gzipCompressor.apply(any())).thenReturn(PAYLOAD); + when(signer.signRequest(any())).thenReturn( + software.amazon.awssdk.http.SdkHttpFullRequest.builder() + .method(software.amazon.awssdk.http.SdkHttpMethod.POST) + .uri(URI.create("https://localhost/v1/traces")) + .putHeader("Authorization", "sig") + .build() ); - } - @AfterEach - void cleanUp() { - System.clearProperty("aws.accessKeyId"); - System.clearProperty("aws.secretAccessKey"); + sender = new OtlpHttpSender(metrics, gzipCompressor, signer, webClient); } @Test - void testSend_successfulResponse() throws IOException { - // Arrange - final SdkHttpFullRequest signed = mock(SdkHttpFullRequest.class); - when(signed.getUri()).thenReturn( - HttpUrl.get("https://xray.us-west-2.amazonaws.com/v1/traces").uri()); - when(signed.headers()).thenReturn( - Map.of("Authorization", Collections.singletonList("signed-header"))); - when(mockSigner.signRequest(PAYLOAD)).thenReturn(signed); - - final Call call = mock(Call.class); - final Response resp = new Response.Builder() - .request(new Request.Builder().url(signed.getUri().toString()).build()) - .protocol(Protocol.HTTP_1_1) - .code(200) - .message("OK") - .body(ResponseBody.create( - new byte[0], - MediaType.get("application/x-protobuf"))) - .build(); + void testSend_successfulResponse() { + when(webClient.execute(any(HttpRequest.class))).thenReturn( + HttpResponse.of(ResponseHeaders.of(200), HttpData.empty()) + ); - when(mockHttpClient.newCall(any())).thenReturn(call); - when(call.execute()).thenReturn(resp); + sender.send(PAYLOAD, SPANS); - // Act & Assert - assertDoesNotThrow(() -> target.send(PAYLOAD)); + await().untilAsserted(() -> { + verify(metrics).incrementRecordsOut(SPANS); + verify(metrics).incrementPayloadSize(PAYLOAD.length); + verify(metrics).incrementPayloadGzipSize(PAYLOAD.length); + verify(metrics).recordHttpLatency(anyLong()); + }); } @Test - void testSend_doesNotRetryOnNonRetryable4xxResponses() throws IOException { - // Arrange - final SdkHttpFullRequest signed = mock(SdkHttpFullRequest.class); - when(signed.getUri()).thenReturn( - HttpUrl.get("https://xray.us-west-2.amazonaws.com/v1/traces").uri()); - when(signed.headers()).thenReturn(Map.of()); - when(mockSigner.signRequest(PAYLOAD)).thenReturn(signed); - - final Request okReq = new Request.Builder() - .url(signed.getUri().toString()) + void testSend_partialSuccessResponse() { + final ExportTraceServiceResponse proto = ExportTraceServiceResponse.newBuilder() + .setPartialSuccess(ExportTracePartialSuccess.newBuilder() + .setRejectedSpans(2) + .setErrorMessage("invalid span") + .build()) .build(); - final Response resp = new Response.Builder() - .request(okReq) - .protocol(Protocol.HTTP_1_1) - .code(400) - .message("Client Error") - .body(ResponseBody.create( - ERROR_BODY.getBytes(StandardCharsets.UTF_8), - MediaType.get("application/json"))) - .build(); - final Call call = mock(Call.class); - when(mockHttpClient.newCall(any())).thenReturn(call); - when(call.execute()).thenReturn(resp); + when(webClient.execute(any(HttpRequest.class))).thenReturn( + HttpResponse.of(ResponseHeaders.of(200), HttpData.wrap(proto.toByteArray())) + ); + sender.send(PAYLOAD, SPANS); - assertDoesNotThrow(() -> target.send(PAYLOAD)); - verify(mockHttpClient, times(1)).newCall(any()); - reset(mockHttpClient); + await().untilAsserted(() -> { + verify(metrics).incrementRejectedSpansCount(2); + verify(metrics).incrementRecordsOut(SPANS - 2); + }); } @Test - void testSend_retryOnFailure_thenSuccess() throws IOException { - // Arrange - final SdkHttpFullRequest signed = mock(SdkHttpFullRequest.class); - when(signed.getUri()).thenReturn( - HttpUrl.get("https://xray.us-west-2.amazonaws.com/v1/traces").uri()); - when(signed.headers()).thenReturn(Map.of()); - when(mockSigner.signRequest(PAYLOAD)).thenReturn(signed); - - final Call first = mock(Call.class); - final Call second = mock(Call.class); - when(mockHttpClient.newCall(any())).thenReturn(first, second); - when(first.execute()).thenThrow(new IOException("first attempt failed")); - final Response success = new Response.Builder() - .request(new Request.Builder().url(signed.getUri().toString()).build()) - .protocol(Protocol.HTTP_1_1) - .code(200) - .message("OK") - .body(ResponseBody.create( - new byte[0], - MediaType.get("application/x-protobuf"))) - .build(); - when(second.execute()).thenReturn(success); + void testSend_parseErrorOnSuccessResponse() { + when(webClient.execute(any(HttpRequest.class))).thenReturn( + HttpResponse.of(ResponseHeaders.of(200), HttpData.ofUtf8("not-protobuf")) + ); - // Act & Assert - assertDoesNotThrow(() -> target.send(PAYLOAD)); - } + sender.send(PAYLOAD, SPANS); - @Test - void testSend_throwsIOException_whenFailsAfterAllRetries() throws IOException { - // Arrange - final SdkHttpFullRequest signed = mock(SdkHttpFullRequest.class); - when(signed.getUri()).thenReturn( - HttpUrl.get("https://xray.us-west-2.amazonaws.com/v1/traces").uri()); - when(signed.headers()).thenReturn(Map.of()); - when(mockSigner.signRequest(PAYLOAD)).thenReturn(signed); - - final Call alwaysFail = mock(Call.class); - when(mockHttpClient.newCall(any())).thenReturn(alwaysFail); - when(alwaysFail.execute()).thenThrow(new IOException("always fail")); - - // Act & Assert - final IOException ex = assertThrows(IOException.class, () -> target.send(PAYLOAD)); - assertEquals("Max retries reached", ex.getMessage()); + await().untilAsserted(() -> { + verify(metrics).incrementErrorsCount(); + verify(metrics).incrementRecordsOut(SPANS); + }); } @Test - void testSend_throwsIOException_on502ResponseWithBody() throws IOException { - final SdkHttpFullRequest signed = SdkHttpFullRequest.builder() - .method(software.amazon.awssdk.http.SdkHttpMethod.POST) - .uri(URI.create("https://example.com")) - .putHeader("Content-Type", "application/json") - .build(); - when(mockSigner.signRequest(PAYLOAD)).thenReturn(signed); - - final Request okReq = new Request.Builder() - .url(signed.getUri().toString()).build(); - final Call call = mock(Call.class); - when(mockHttpClient.newCall(any())).thenReturn(call); - - final Response resp500 = new Response.Builder() - .request(okReq) - .protocol(Protocol.HTTP_1_1) - .code(502) - .message("Bad Gateway") - .body(ResponseBody.create( - ERROR_BODY.getBytes(StandardCharsets.UTF_8), - MediaType.get("application/json"))) - .build(); - when(call.execute()).thenReturn(resp500); + void testSend_nonSuccessStatus() { + when(webClient.execute(any(HttpRequest.class))).thenReturn( + HttpResponse.of(ResponseHeaders.of(400), HttpData.ofUtf8("{\"error\":\"bad request\"}")) + ); + + sender.send(PAYLOAD, SPANS); - assertThrows(IOException.class, () -> target.send(PAYLOAD)); + await().untilAsserted(() -> { + verify(metrics).recordResponseCode(400); + verify(metrics).incrementRejectedSpansCount(SPANS); + }); } @Test - void testSend_wrapsInterruptedExceptionDuringRetryAsIOException() throws IOException { - // Arrange: first attempt throws IOException, sleeper throws at retry - final SdkHttpFullRequest signed = mock(SdkHttpFullRequest.class); - when(mockSigner.signRequest(any())).thenReturn(signed); - when(signed.getUri()).thenReturn(URI.create("https://example.com")); - when(signed.headers()).thenReturn(Map.of()); - - final Call call = mock(Call.class); - when(mockHttpClient.newCall(any())).thenReturn(call); - when(call.execute()).thenThrow(new IOException("boom")); - - doThrow(new RuntimeException("sleep failed")) - .when(mockSleeper).accept(anyInt()); - - target = new OtlpHttpSender( - mockConfig, mockSinkMetrics, - mockGzipCompressor, mockSigner, - mockHttpClient, mockSleeper - ); + void testSend_skipsSendIfGzipFails() { + sender = new OtlpHttpSender(metrics, ignored -> new byte[0], signer, webClient); + sender.send(PAYLOAD, SPANS); - // Act & Assert - final IOException ex = assertThrows(IOException.class, () -> target.send(PAYLOAD)); - assertEquals("Sender failed to sleep before retrying.", ex.getMessage()); - assertEquals("sleep failed", ex.getCause().getMessage()); + verify(metrics).incrementFailedSpansCount(SPANS); + verifyNoInteractions(webClient); } @Test - void testSend_partialSuccessResponse_logsWarning() throws IOException { - // Arrange - final ExportTraceServiceResponse proto = ExportTraceServiceResponse.newBuilder() - .setPartialSuccess(ExportTracePartialSuccess.newBuilder() - .setRejectedSpans(5) - .setErrorMessage("Some spans were rejected due to invalid format") - .build()) - .build(); - final byte[] bytes = proto.toByteArray(); + void testSend_exceptionDuringSendIncrementsRejected() { + final HttpResponse failingResponse = HttpResponse.ofFailure(new RuntimeException("send failed")); + when(webClient.execute(any(HttpRequest.class))).thenReturn(failingResponse); - final SdkHttpFullRequest signed = SdkHttpFullRequest.builder() - .method(software.amazon.awssdk.http.SdkHttpMethod.POST) - .uri(URI.create("https://xray.us-west-2.amazonaws.com/v1/traces")) - .putHeader("Content-Type", "application/x-protobuf") - .build(); - when(mockSigner.signRequest(PAYLOAD)).thenReturn(signed); + sender.send(PAYLOAD, SPANS); - final Request okReq = new Request.Builder() - .url(signed.getUri().toString()) - .build(); - final Call call = mock(Call.class); - when(mockHttpClient.newCall(any())).thenReturn(call); - when(call.execute()).thenReturn(new Response.Builder() - .request(okReq) - .protocol(Protocol.HTTP_1_1) - .code(200) - .message("OK") - .body(ResponseBody.create(bytes, MediaType.get("application/x-protobuf"))) - .build()); - - // Act & Assert - assertDoesNotThrow(() -> target.send(PAYLOAD)); + await().untilAsserted(() -> verify(metrics).incrementRejectedSpansCount(SPANS)); } @Test - void testSend_skipsSend_whenGzipCompressionFails() { - // Arrange: compressor returns empty → send() should return silently - final Function> skipCompressor = p -> Optional.empty(); - target = new OtlpHttpSender( - mockConfig, mockSinkMetrics, - skipCompressor, mockSigner, - mockHttpClient, mockSleeper - ); + void testConstructor_withDefaultConfig() { + final OtlpSinkConfig config = mock(OtlpSinkConfig.class); + + when(config.getMaxBatchSize()).thenReturn(1_000_000L); + when(config.getMaxRetries()).thenReturn(2); + when(config.getFlushTimeoutMillis()).thenReturn(5000L); + when(config.getAwsRegion()).thenReturn(software.amazon.awssdk.regions.Region.US_WEST_2); - // Act & Assert: no exception, nothing signed or sent - assertDoesNotThrow(() -> target.send(PAYLOAD)); - verify(mockSigner, never()).signRequest(any()); - verify(mockHttpClient, never()).newCall(any()); + final OtlpHttpSender defaultSender = new OtlpHttpSender(config, metrics); + assertNotNull(defaultSender); } @Test - void testHandleSuccessfulResponseParseErrorIncrementsError() throws IOException { - // arrange: build a sender with a pass-through compressor lambda - target = new OtlpHttpSender( - mockConfig, - mockSinkMetrics, - Optional::of, - mockSigner, - mockHttpClient, - mockSleeper - ); - - // stub: signer - final SdkHttpFullRequest signed = mock(SdkHttpFullRequest.class); - when(signed.getUri()).thenReturn(URI.create("https://example.com")); - when(signed.headers()).thenReturn(Map.of()); - when(mockSigner.signRequest(PAYLOAD)).thenReturn(signed); - - // stub: http client returns 200 OK but with invalid protobuf bytes - final Call call = mock(Call.class); - when(mockHttpClient.newCall(any())).thenReturn(call); - final byte[] bad = "not-a-proto".getBytes(StandardCharsets.UTF_8); - final Response resp = new Response.Builder() - .request(new Request.Builder().url(signed.getUri().toString()).build()) - .protocol(Protocol.HTTP_1_1) - .code(200) - .message("OK") - .body(ResponseBody.create(bad, MediaType.get("application/x-protobuf"))) - .build(); - when(call.execute()).thenReturn(resp); - - // act - assertDoesNotThrow(() -> target.send(PAYLOAD)); - - // assert: parse-failure should have incremented errors - verify(mockSinkMetrics).incrementErrorsCount(); + void testConstructor_withMinimumThresholdConfig() { + final OtlpSinkConfig config = mock(OtlpSinkConfig.class); + + // Set all threshold values to minimum valid input + when(config.getMaxBatchSize()).thenReturn(0L); + when(config.getMaxRetries()).thenReturn(0); + when(config.getFlushTimeoutMillis()).thenReturn(1L); + when(config.getAwsRegion()).thenReturn(software.amazon.awssdk.regions.Region.US_WEST_2); + + // Should not throw or crash + final OtlpHttpSender minimalSender = new OtlpHttpSender(config, metrics); + assertNotNull(minimalSender); } @Test - void testCloseEvictsAndShutdownsOkHttpResources() { - // arrange: stub out connectionPool and dispatcher on our mockHttpClient - final ConnectionPool pool = mock(ConnectionPool.class); - final Dispatcher dispatcher = mock(Dispatcher.class); - final ExecutorService exec = mock(ExecutorService.class); - when(mockHttpClient.connectionPool()).thenReturn(pool); - when(mockHttpClient.dispatcher()).thenReturn(dispatcher); - when(dispatcher.executorService()).thenReturn(exec); - - // act - target.close(); - - // assert - verify(pool).evictAll(); - verify(exec).shutdown(); + void testHandleSuccessfulResponse_withNullBody_incrementsRecordsOut() throws Exception { + final Method method = OtlpHttpSender.class.getDeclaredMethod("handleSuccessfulResponse", byte[].class, int.class); + method.setAccessible(true); + method.invoke(sender, null, SPANS); + + verify(metrics).incrementRecordsOut(SPANS); + verifyNoMoreInteractions(metrics); } @Test - void testDefaultConstructorInitializesDefaults() { - target = new OtlpHttpSender(mockConfig, mockSinkMetrics); - assertNotNull(getField(target, "signer")); - assertNotNull(getField(target, "httpClient")); - assertNotNull(getField(target, "sleeper")); - } + void testHandleSuccessfulResponse_withoutPartialSuccess_incrementsRecordsOut() throws Exception { + // Create a response with no partial_success + final ExportTraceServiceResponse response = ExportTraceServiceResponse.newBuilder().build(); + final byte[] responseBytes = response.toByteArray(); + + final Method method = OtlpHttpSender.class.getDeclaredMethod("handleSuccessfulResponse", byte[].class, int.class); + method.setAccessible(true); + method.invoke(sender, responseBytes, SPANS); - private Object getField(final Object obj, final String name) { - try { - final var f = obj.getClass().getDeclaredField(name); - f.setAccessible(true); - return f.get(obj); - } catch (final Exception e) { - fail("Could not access " + name); - return null; - } + verify(metrics).incrementRecordsOut(SPANS); + verifyNoMoreInteractions(metrics); } } diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeperTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeperTest.java deleted file mode 100644 index df368566e4..0000000000 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/ThreadSleeperTest.java +++ /dev/null @@ -1,52 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.dataprepper.plugins.sink.otlp.http; - -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; - -import java.util.function.Consumer; - -import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.junit.jupiter.api.Assertions.fail; - -class ThreadSleeperTest { - private Consumer target; - - @BeforeEach - void setUp() { - target = new ThreadSleeper(); - } - - @Test - void testSleepDoesNotThrowWhenNotInterrupted() { - try { - target.accept(1); - } catch (RuntimeException e) { - fail("Sleep was interrupted unexpectedly"); - } - } - - @Test - void testSleepThrowsInterruptedExceptionIfThreadInterrupted() { - Thread thread = new Thread(() -> { - try { - Thread.currentThread().interrupt(); - target.accept(10); - fail("Expected InterruptedException"); - } catch (RuntimeException e) { - assertTrue(Thread.currentThread().isInterrupted(), "Thread should remain interrupted"); - } - }); - - thread.start(); - try { - thread.join(); - } catch (InterruptedException e) { - fail("Test thread was interrupted"); - } - } -} diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java index 2b7f61affd..10a33c225a 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/metrics/OtlpSinkMetricsTest.java @@ -72,13 +72,6 @@ void testIncrementErrorsCount() { verify(counterMock).increment(1.0); } - @Test - void testIncrementRetriesCount() { - sinkMetrics.incrementRetriesCount(); - verify(pluginMetrics).counter("retriesCount"); - verify(counterMock).increment(1.0); - } - @Test void testIncrementRejectedSpansCount() { sinkMetrics.incrementRejectedSpansCount(5); @@ -86,6 +79,13 @@ void testIncrementRejectedSpansCount() { verify(counterMock).increment(5.0); } + @Test + void testIncrementFailedSpansCount() { + sinkMetrics.incrementFailedSpansCount(5); + verify(pluginMetrics).counter("failedSpansCount"); + verify(counterMock).increment(5.0); + } + @Test void testRecordResponseCodeCounters() { // 5xx From 272859113141da93f78ccbed2156b9e3b656b269 Mon Sep 17 00:00:00 2001 From: huy pham Date: Wed, 21 May 2025 21:40:44 -0700 Subject: [PATCH 19/23] Release EventHandle on both success and failure; switch to AwsCredentialsSupplier. Signed-off-by: huy pham --- data-prepper-plugins/otlp-sink/build.gradle | 12 +- .../plugins/sink/otlp/OtlpSink.java | 8 +- .../sink/otlp/buffer/OtlpSinkBuffer.java | 37 +- .../sink/otlp/http/OtlpHttpSender.java | 68 +++- .../plugins/sink/otlp/http/SigV4Signer.java | 48 +-- .../plugins/sink/otlp/OtlpSinkTest.java | 7 +- .../sink/otlp/buffer/OtlpSinkBufferTest.java | 362 ++++++++++++++++-- .../sink/otlp/http/OtlpHttpSenderTest.java | 272 +++++++++++-- .../sink/otlp/http/SigV4SignerTest.java | 161 ++------ 9 files changed, 711 insertions(+), 264 deletions(-) diff --git a/data-prepper-plugins/otlp-sink/build.gradle b/data-prepper-plugins/otlp-sink/build.gradle index e1846d9fa9..ef3ff0ac7a 100644 --- a/data-prepper-plugins/otlp-sink/build.gradle +++ b/data-prepper-plugins/otlp-sink/build.gradle @@ -19,10 +19,10 @@ configurations { dependencies { // AWS SDK - implementation 'software.amazon.awssdk:core:2.28.23' - implementation 'software.amazon.awssdk:auth:2.28.23' - implementation 'software.amazon.awssdk:sts:2.28.23' - implementation 'software.amazon.awssdk:regions:2.28.23' + implementation 'software.amazon.awssdk:sdk-core' + implementation 'software.amazon.awssdk:auth' + implementation 'software.amazon.awssdk:sts' + implementation 'software.amazon.awssdk:regions' implementation 'software.amazon.awssdk:http-client-spi' implementation 'software.amazon.awssdk:apache-client' @@ -37,8 +37,8 @@ dependencies { // Jackson implementation 'com.fasterxml.jackson.core:jackson-databind' - implementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml:2.15.3' - testImplementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml:2.15.3' + implementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml' + testImplementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml' // Lombok compileOnly 'org.projectlombok:lombok:1.18.30' diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java index f3c756d387..47695cf51e 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java @@ -5,9 +5,11 @@ package org.opensearch.dataprepper.plugins.sink.otlp; +import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin; import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor; +import org.opensearch.dataprepper.model.annotations.Experimental; import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.model.sink.AbstractSink; @@ -23,6 +25,7 @@ /** * OTLP Sink Plugin for Data Prepper. */ +@Experimental @DataPrepperPlugin( name = "otlp", pluginType = Sink.class, @@ -36,16 +39,17 @@ public class OtlpSink extends AbstractSink> { /** * Constructor for the OTLP sink plugin. * + * @param awsCredentialsSupplier the AWS credentials supplier * @param config the configuration for the sink * @param pluginMetrics the plugin metrics to use * @param pluginSetting the plugin setting to use */ @DataPrepperPluginConstructor - public OtlpSink(@Nonnull final OtlpSinkConfig config, @Nonnull final PluginMetrics pluginMetrics, @Nonnull final PluginSetting pluginSetting) { + public OtlpSink(@Nonnull final AwsCredentialsSupplier awsCredentialsSupplier, @Nonnull final OtlpSinkConfig config, @Nonnull final PluginMetrics pluginMetrics, @Nonnull final PluginSetting pluginSetting) { super(pluginSetting); this.sinkMetrics = new OtlpSinkMetrics(pluginMetrics, pluginSetting); - this.buffer = new OtlpSinkBuffer(config, sinkMetrics); + this.buffer = new OtlpSinkBuffer(awsCredentialsSupplier, config, sinkMetrics); } /** diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBuffer.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBuffer.java index 3febfdaeb6..2c1c0e8f04 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBuffer.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBuffer.java @@ -6,9 +6,10 @@ package org.opensearch.dataprepper.plugins.sink.otlp.buffer; import com.google.common.annotations.VisibleForTesting; -import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceRequest; import io.opentelemetry.proto.trace.v1.ResourceSpans; import lombok.Getter; +import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; +import org.opensearch.dataprepper.model.event.EventHandle; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.model.trace.Span; import org.opensearch.dataprepper.plugins.otel.codec.OTelProtoStandardCodec; @@ -17,6 +18,7 @@ import org.opensearch.dataprepper.plugins.sink.otlp.metrics.OtlpSinkMetrics; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import software.amazon.awssdk.utils.Pair; import javax.annotation.Nonnull; import java.util.ArrayList; @@ -52,11 +54,12 @@ public class OtlpSinkBuffer { /** * Creates a new OTLP sink buffer. * + * @param awsCredentialsSupplier the AWS credentials supplier * @param config the OTLP sink configuration * @param sinkMetrics the metrics collector to use */ - public OtlpSinkBuffer(@Nonnull final OtlpSinkConfig config, @Nonnull final OtlpSinkMetrics sinkMetrics) { - this(config, sinkMetrics, new OTelProtoStandardCodec.OTelProtoEncoder(), new OtlpHttpSender(config, sinkMetrics)); + public OtlpSinkBuffer(@Nonnull final AwsCredentialsSupplier awsCredentialsSupplier, @Nonnull final OtlpSinkConfig config, @Nonnull final OtlpSinkMetrics sinkMetrics) { + this(config, sinkMetrics, new OTelProtoStandardCodec.OTelProtoEncoder(), new OtlpHttpSender(awsCredentialsSupplier, config, sinkMetrics)); } /** @@ -144,7 +147,7 @@ public void add(final Record record) { * Handles encoding failures, timeout-based flush, and final flush on shutdown. */ private void run() { - final List batch = new ArrayList<>(); + final List> batch = new ArrayList<>(); long batchSize = 0; long lastFlush = System.currentTimeMillis(); @@ -156,7 +159,8 @@ private void run() { if (record != null) { try { final ResourceSpans resourceSpans = encoder.convertToResourceSpans(record.getData()); - batch.add(resourceSpans); + final EventHandle eventHandle = record.getData().getEventHandle(); + batch.add(Pair.of(resourceSpans, eventHandle)); batchSize += resourceSpans.getSerializedSize(); } catch (final Exception e) { LOG.error("Failed to encode span, skipping", e); @@ -169,7 +173,8 @@ private void run() { final boolean flushByTime = !batch.isEmpty() && (now - lastFlush >= flushTimeoutMillis); if (flushBySize || flushByTime) { - send(batch); + sender.send(batch); + batch.clear(); batchSize = 0; lastFlush = now; } @@ -190,24 +195,8 @@ private void run() { // Final flush if (!batch.isEmpty()) { - send(batch); + sender.send(batch); + batch.clear(); } } - - /** - * Builds an ExportTraceServiceRequest from the given batch, sends it over HTTP. - * The batch is cleared in all cases to prepare for the next batch. - * - * @param batch the list of ResourceSpans to send - */ - private void send(final List batch) { - final ExportTraceServiceRequest request = ExportTraceServiceRequest.newBuilder() - .addAllResourceSpans(batch) - .build(); - final byte[] payload = request.toByteArray(); - - final int spans = batch.size(); - sender.send(payload, spans); - batch.clear(); - } } diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java index 7ea371c2ab..c6ee74e855 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java @@ -17,17 +17,24 @@ import com.linecorp.armeria.common.MediaType; import com.linecorp.armeria.common.RequestHeaders; import com.linecorp.armeria.common.RequestHeadersBuilder; +import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceRequest; import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceResponse; +import io.opentelemetry.proto.trace.v1.ResourceSpans; +import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; +import org.opensearch.dataprepper.model.event.EventHandle; import org.opensearch.dataprepper.plugins.sink.otlp.configuration.OtlpSinkConfig; import org.opensearch.dataprepper.plugins.sink.otlp.metrics.OtlpSinkMetrics; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import software.amazon.awssdk.http.SdkHttpFullRequest; +import software.amazon.awssdk.utils.Pair; import javax.annotation.Nonnull; import java.nio.charset.StandardCharsets; +import java.util.List; import java.util.Set; import java.util.function.Function; +import java.util.stream.Collectors; /** * Responsible for sending signed OTLP Protobuf requests to OTLP endpoint using an Ameria client. @@ -45,11 +52,12 @@ public class OtlpHttpSender { /** * Constructor for the OtlpHttpSender. * + * @param awsCredentialsSupplier the AWS credentials supplier * @param config The configuration for the OTLP sink plugin. * @param sinkMetrics The metrics for the OTLP sink plugin. */ - public OtlpHttpSender(@Nonnull final OtlpSinkConfig config, @Nonnull final OtlpSinkMetrics sinkMetrics) { - this(sinkMetrics, new GzipCompressor(sinkMetrics), new SigV4Signer(config), buildWebClient(config)); + public OtlpHttpSender(@Nonnull final AwsCredentialsSupplier awsCredentialsSupplier, @Nonnull final OtlpSinkConfig config, @Nonnull final OtlpSinkMetrics sinkMetrics) { + this(sinkMetrics, new GzipCompressor(sinkMetrics), new SigV4Signer(awsCredentialsSupplier, config), buildWebClient(config)); } /** @@ -103,38 +111,54 @@ private static WebClient buildWebClient(final OtlpSinkConfig config) { /** * Sends the provided OTLP Protobuf payload to the OTLP endpoint asynchronously. * - * @param payload The OTLP Protobuf-encoded data to be sent. - * @param spans The number of spans in the payload. + * @param batch the batch of spans to send */ - public void send(@Nonnull final byte[] payload, final int spans) { - final byte[] compressedPayload = gzipCompressor.apply(payload); - if (compressedPayload.length == 0) { - sinkMetrics.incrementFailedSpansCount(spans); + public void send(@Nonnull final List> batch) { + if (batch.isEmpty()) { return; } - final HttpRequest request = buildHttpRequest(compressedPayload); + final Pair payloadAndCompressedPayload = getPayloadAndCompressedPayload(batch); + final int spans = batch.size(); + if (payloadAndCompressedPayload.right().length == 0) { + sinkMetrics.incrementFailedSpansCount(spans); + releaseAllEventHandle(batch, false); + return; + } + final HttpRequest request = buildHttpRequest(payloadAndCompressedPayload.right()); final long startTime = System.currentTimeMillis(); + webClient.execute(request) .aggregate() .thenAccept(response -> { final long latency = System.currentTimeMillis() - startTime; sinkMetrics.recordHttpLatency(latency); - sinkMetrics.incrementPayloadSize(payload.length); - sinkMetrics.incrementPayloadGzipSize(compressedPayload.length); + sinkMetrics.incrementPayloadSize(payloadAndCompressedPayload.left().length); + sinkMetrics.incrementPayloadGzipSize(payloadAndCompressedPayload.right().length); final int statusCode = response.status().code(); final byte[] responseBytes = response.content().array(); - handleResponse(statusCode, responseBytes, spans); + handleResponse(statusCode, responseBytes, batch); }) .exceptionally(e -> { LOG.error("Failed to send {} spans.", spans, e); sinkMetrics.incrementRejectedSpansCount(spans); + releaseAllEventHandle(batch, false); return null; }); } + private Pair getPayloadAndCompressedPayload(final List> batch) { + final ExportTraceServiceRequest request = ExportTraceServiceRequest.newBuilder() + .addAllResourceSpans(batch.stream().map(Pair::left).collect(Collectors.toList())) + .build(); + final byte[] payload = request.toByteArray(); + final byte[] compressedPayload = gzipCompressor.apply(payload); + + return Pair.of(payload, compressedPayload); + } + private HttpRequest buildHttpRequest(final byte[] compressedPayload) { final SdkHttpFullRequest signedRequest = signer.signRequest(compressedPayload); final RequestHeadersBuilder headersBuilder = RequestHeaders.builder() @@ -147,11 +171,11 @@ private HttpRequest buildHttpRequest(final byte[] compressedPayload) { return HttpRequest.of(headersBuilder.build(), HttpData.wrap(compressedPayload)); } - private void handleResponse(final int statusCode, final byte[] responseBytes, final int spans) { + private void handleResponse(final int statusCode, final byte[] responseBytes, final List> batch) { sinkMetrics.recordResponseCode(statusCode); if (statusCode >= 200 && statusCode < 300) { - handleSuccessfulResponse(responseBytes, spans); + handleSuccessfulResponse(responseBytes, batch); return; } @@ -160,15 +184,18 @@ private void handleResponse(final int statusCode, final byte[] responseBytes, fi : ""; LOG.error("Non-successful OTLP response. Status: {}, Response: {}", statusCode, responseBody); - sinkMetrics.incrementRejectedSpansCount(spans); + sinkMetrics.incrementRejectedSpansCount(batch.size()); + releaseAllEventHandle(batch, false); } /** * Handles a successful OTLP response with partial success. */ - private void handleSuccessfulResponse(final byte[] responseBytes, final int spans) { + private void handleSuccessfulResponse(final byte[] responseBytes, final List> batch) { + final int spans = batch.size(); if (responseBytes == null) { sinkMetrics.incrementRecordsOut(spans); + releaseAllEventHandle(batch, true); return; } @@ -186,13 +213,22 @@ private void handleSuccessfulResponse(final byte[] responseBytes, final int span final long deliveredSpans = spans - rejectedSpans; sinkMetrics.incrementRecordsOut(deliveredSpans); + + // Optimistically release all as true, no per-span granularity + releaseAllEventHandle(batch, true); } else { sinkMetrics.incrementRecordsOut(spans); + releaseAllEventHandle(batch, true); } } catch (final Exception e) { LOG.error("Could not parse OTLP response as ExportTraceServiceResponse: {}", e.getMessage()); sinkMetrics.incrementErrorsCount(); sinkMetrics.incrementRecordsOut(spans); + releaseAllEventHandle(batch, true); } } + + private void releaseAllEventHandle(@Nonnull final List> batch, final boolean success) { + batch.forEach(pair -> pair.right().release(success)); + } } \ No newline at end of file diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4Signer.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4Signer.java index ab2e0c7010..becfb72960 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4Signer.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4Signer.java @@ -5,18 +5,16 @@ package org.opensearch.dataprepper.plugins.sink.otlp.http; -import com.google.common.annotations.VisibleForTesting; +import org.opensearch.dataprepper.aws.api.AwsCredentialsOptions; +import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; import org.opensearch.dataprepper.plugins.sink.otlp.configuration.OtlpSinkConfig; import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; -import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider; import software.amazon.awssdk.auth.signer.Aws4Signer; import software.amazon.awssdk.auth.signer.params.Aws4SignerParams; import software.amazon.awssdk.core.SdkBytes; import software.amazon.awssdk.http.SdkHttpFullRequest; import software.amazon.awssdk.http.SdkHttpMethod; import software.amazon.awssdk.regions.Region; -import software.amazon.awssdk.services.sts.StsClient; -import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider; import javax.annotation.Nonnull; import java.net.URI; @@ -28,7 +26,6 @@ class SigV4Signer { private static final String SERVICE_NAME = "xray"; private static final String OTLP_PATH = "/v1/traces"; - private static final String OTLP_SINK_SESSION = "otlp-sink-session"; private final Aws4Signer signer = Aws4Signer.create(); private final AwsCredentialsProvider credentialsProvider; @@ -38,46 +35,23 @@ class SigV4Signer { /** * Constructs a SigV4 signer helper. * + * @param awsCredentialsSupplier the AWS credentials supplier * @param config Configuration for region and optional STS role */ - SigV4Signer(@Nonnull final OtlpSinkConfig config) { - this(config, null); - } - - /** - * Package-private constructor for unit testing with mocked STS client. - */ - @VisibleForTesting - SigV4Signer(@Nonnull final OtlpSinkConfig config, final StsClient stsClient) { + SigV4Signer(@Nonnull final AwsCredentialsSupplier awsCredentialsSupplier, @Nonnull final OtlpSinkConfig config) { this.region = config.getAwsRegion(); - this.credentialsProvider = initCredentialsProvider(region, config.getStsRoleArn(), config.getStsExternalId(), stsClient); + + this.credentialsProvider = awsCredentialsSupplier.getProvider(AwsCredentialsOptions.builder() + .withRegion(region) + .withStsRoleArn(config.getStsRoleArn()) + .withStsExternalId(config.getStsExternalId()) + .build()); + this.endpointUri = config.getEndpoint() != null ? URI.create(config.getEndpoint()) : URI.create(String.format("https://xray.%s.amazonaws.com%s", region.id(), OTLP_PATH)); } - private static AwsCredentialsProvider initCredentialsProvider( - @Nonnull final Region region, - final String stsRoleArn, - final String stsExternalId, - final StsClient stsClient - ) { - if (stsRoleArn != null) { - return StsAssumeRoleCredentialsProvider.builder() - .refreshRequest(r -> { - r.roleArn(stsRoleArn); - r.roleSessionName(OTLP_SINK_SESSION); - if (stsExternalId != null) { - r.externalId(stsExternalId); - } - }) - .stsClient(stsClient != null ? stsClient : StsClient.builder().region(region).build()) - .build(); - } - - return DefaultCredentialsProvider.create(); - } - /** * Signs a request payload using AWS SigV4 and returns a fully signed request. * diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java index 82fbdf34ec..bdcfb3032b 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java @@ -6,6 +6,7 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.configuration.PluginSetting; import org.opensearch.dataprepper.model.record.Record; @@ -31,10 +32,12 @@ class OtlpSinkTest { private OtlpSinkConfig mockConfig; private PluginMetrics mockMetrics; private PluginSetting mockSetting; + private AwsCredentialsSupplier mockAwsCredSupplier; @BeforeEach void setUp() throws Exception { // Arrange: stub out config, metrics, setting + mockAwsCredSupplier = mock(AwsCredentialsSupplier.class); mockConfig = mock(OtlpSinkConfig.class); when(mockConfig.getAwsRegion()).thenReturn(Region.of("us-west-2")); @@ -45,7 +48,7 @@ void setUp() throws Exception { when(mockSetting.getName()).thenReturn("otlp"); // Create the real sink - target = new OtlpSink(mockConfig, mockMetrics, mockSetting); + target = new OtlpSink(mockAwsCredSupplier, mockConfig, mockMetrics, mockSetting); // Replace its private buffer with a mock mockBuffer = mock(OtlpSinkBuffer.class); @@ -101,6 +104,6 @@ void testShutdown_stopsBuffer() { @Test void testConstructor_doesNotThrow() { // Just ensure the three-arg constructor still works - assertDoesNotThrow(() -> new OtlpSink(mockConfig, mockMetrics, mockSetting)); + assertDoesNotThrow(() -> new OtlpSink(mockAwsCredSupplier, mockConfig, mockMetrics, mockSetting)); } } diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBufferTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBufferTest.java index 2612108d50..800ace1dc4 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBufferTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBufferTest.java @@ -9,6 +9,8 @@ import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; +import org.opensearch.dataprepper.model.event.EventHandle; import org.opensearch.dataprepper.model.record.Record; import org.opensearch.dataprepper.model.trace.Span; import org.opensearch.dataprepper.plugins.otel.codec.OTelProtoStandardCodec; @@ -27,12 +29,15 @@ import static org.awaitility.Awaitility.await; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyList; +import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.Mockito.atLeastOnce; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -43,6 +48,7 @@ class OtlpSinkBufferTest { private OTelProtoStandardCodec.OTelProtoEncoder encoder; private OtlpHttpSender sender; private OtlpSinkBuffer buffer; + private AwsCredentialsSupplier mockAwsCredSupplier; @BeforeEach void setUp() { @@ -56,22 +62,55 @@ void setUp() { metrics = mock(OtlpSinkMetrics.class); encoder = mock(OTelProtoStandardCodec.OTelProtoEncoder.class); sender = mock(OtlpHttpSender.class); + mockAwsCredSupplier = mock(AwsCredentialsSupplier.class); buffer = new OtlpSinkBuffer(config, metrics, encoder, sender); } @AfterEach void tearDown() { - buffer.stop(); + if (buffer != null) { + buffer.stop(); + } } @Test void testStartAndStopDoesNotThrow() { buffer.start(); + assertTrue(buffer.isRunning()); buffer.stop(); assertFalse(buffer.isRunning()); } + @Test + void testPublicConstructorInitializesWithDefaults() { + final OtlpSinkBuffer defaultBuffer = new OtlpSinkBuffer(mockAwsCredSupplier, config, metrics); + assertNotNull(defaultBuffer); + assertTrue(defaultBuffer.isRunning()); + defaultBuffer.stop(); + } + + @Test + void testQueueCapacityCalculation_withHighMaxEvents() { + when(config.getMaxEvents()).thenReturn(500); + final OtlpSinkBuffer highCapacityBuffer = new OtlpSinkBuffer(config, metrics, encoder, sender); + assertNotNull(highCapacityBuffer); + highCapacityBuffer.stop(); + } + + @Test + void testQueueCapacityCalculation_withLowMaxEvents() { + when(config.getMaxEvents()).thenReturn(1); + final OtlpSinkBuffer lowCapacityBuffer = new OtlpSinkBuffer(config, metrics, encoder, sender); + assertNotNull(lowCapacityBuffer); + lowCapacityBuffer.stop(); + } + + @Test + void testQueueGaugesRegistration() { + verify(metrics).registerQueueGauges(any(BlockingQueue.class)); + } + @Test void testAddHandlesInterruptedException() throws Exception { @SuppressWarnings("unchecked") final BlockingQueue> badQueue = mock(BlockingQueue.class); @@ -90,25 +129,125 @@ void testAddHandlesInterruptedException() throws Exception { @Test void testWorkerThreadHandlesEncodeException() throws Exception { + // First call throws exception, second call succeeds when(encoder.convertToResourceSpans(any(Span.class))) .thenThrow(new RuntimeException("boom")) .thenReturn(ResourceSpans.getDefaultInstance()); - final Record rec1 = mock(Record.class); - final Record rec2 = mock(Record.class); - when(rec1.getData()).thenReturn(mock(Span.class)); - when(rec2.getData()).thenReturn(mock(Span.class)); + final Record rec1 = createMockRecord(); + final Record rec2 = createMockRecord(); buffer.start(); buffer.add(rec1); buffer.add(rec2); + // Wait for processing to complete + await().atMost(2, SECONDS).untilAsserted(() -> { + verify(metrics).incrementFailedSpansCount(1); + verify(metrics, atLeastOnce()).incrementErrorsCount(); + }); + + buffer.stop(); + + // The second record should have been processed and sent + await().atMost(1, SECONDS).untilAsserted(() -> verify(sender).send(anyList())); + } + + @Test + void testFlushByEventCount() throws Exception { + when(config.getMaxEvents()).thenReturn(2); + when(config.getMaxBatchSize()).thenReturn(Long.MAX_VALUE); + when(config.getFlushTimeoutMillis()).thenReturn(Long.MAX_VALUE); + + final ResourceSpans resourceSpans = ResourceSpans.getDefaultInstance(); + when(encoder.convertToResourceSpans(any(Span.class))).thenReturn(resourceSpans); + + buffer = new OtlpSinkBuffer(config, metrics, encoder, sender); + buffer.start(); + + final Record rec1 = createMockRecord(); + final Record rec2 = createMockRecord(); + + buffer.add(rec1); + buffer.add(rec2); + + await().atMost(2, SECONDS).untilAsserted(() -> verify(sender).send(anyList())); + buffer.stop(); + } + + @Test + void testFlushByBatchSize() throws Exception { + when(config.getMaxEvents()).thenReturn(Integer.MAX_VALUE); + when(config.getMaxBatchSize()).thenReturn(100L); + when(config.getFlushTimeoutMillis()).thenReturn(Long.MAX_VALUE); + + final ResourceSpans largeResourceSpans = mock(ResourceSpans.class); + when(largeResourceSpans.getSerializedSize()).thenReturn(150); + when(encoder.convertToResourceSpans(any(Span.class))).thenReturn(largeResourceSpans); + + buffer = new OtlpSinkBuffer(config, metrics, encoder, sender); + buffer.start(); + + final Record rec = createMockRecord(); + buffer.add(rec); + + await().atMost(2, SECONDS).untilAsserted(() -> verify(sender).send(anyList())); + buffer.stop(); + } + + @Test + void testFlushByTimeout() throws Exception { + when(config.getMaxEvents()).thenReturn(Integer.MAX_VALUE); + when(config.getMaxBatchSize()).thenReturn(Long.MAX_VALUE); + when(config.getFlushTimeoutMillis()).thenReturn(50L); + + final ResourceSpans resourceSpans = ResourceSpans.getDefaultInstance(); + when(encoder.convertToResourceSpans(any(Span.class))).thenReturn(resourceSpans); + + buffer = new OtlpSinkBuffer(config, metrics, encoder, sender); + buffer.start(); + + final Record rec = createMockRecord(); + buffer.add(rec); + + await().atMost(2, SECONDS).untilAsserted(() -> verify(sender).send(anyList())); + buffer.stop(); + } + + @Test + void testWorkerLoopWithEmptyQueue() throws Exception { + when(config.getMaxEvents()).thenReturn(Integer.MAX_VALUE); + when(config.getMaxBatchSize()).thenReturn(Long.MAX_VALUE); + when(config.getFlushTimeoutMillis()).thenReturn(Long.MAX_VALUE); + + buffer = new OtlpSinkBuffer(config, metrics, encoder, sender); + buffer.start(); + + TimeUnit.MILLISECONDS.sleep(200); // Let it poll empty queue + + buffer.stop(); + verify(sender, never()).send(anyList()); + } + + @Test + void testEventHandleIncludedInBatch() throws Exception { + final ResourceSpans resourceSpans = ResourceSpans.getDefaultInstance(); + when(encoder.convertToResourceSpans(any(Span.class))).thenReturn(resourceSpans); + + final EventHandle eventHandle = mock(EventHandle.class); + final Record rec = mock(Record.class); + final Span span = mock(Span.class); + when(rec.getData()).thenReturn(span); + when(span.getEventHandle()).thenReturn(eventHandle); + + buffer.start(); + buffer.add(rec); + TimeUnit.MILLISECONDS.sleep(100); buffer.stop(); - await().atMost(1, SECONDS).untilAsserted(() -> verify(sender).send(any(byte[].class), anyInt())); - verify(metrics).incrementFailedSpansCount(1); - verify(metrics, atLeastOnce()).incrementErrorsCount(); + await().atMost(1, SECONDS).untilAsserted(() -> verify(sender).send(anyList())); + verify(span).getEventHandle(); } @Test @@ -163,13 +302,30 @@ void testRestartWorkerSubmitsRunnableIfRunning() throws Exception { verify(mockExecutor, atLeastOnce()).execute(any()); } + @Test + void testRestartWorkerDoesNotSubmitIfNotRunning() throws Exception { + final ExecutorService mockExecutor = mock(ExecutorService.class); + when(mockExecutor.isShutdown()).thenReturn(false); + + final Field executorField = OtlpSinkBuffer.class.getDeclaredField("executor"); + executorField.setAccessible(true); + executorField.set(buffer, mockExecutor); + + buffer.stop(); // Set running to false + + buffer.restartWorker(); + + verify(mockExecutor, never()).execute(any()); + } + @Test void testRunWithInterruptedException() throws Exception { final OtlpSinkBuffer localBuffer = new OtlpSinkBuffer(config, metrics, encoder, sender); // Override the internal queue to throw InterruptedException final BlockingQueue> interruptingQueue = mock(BlockingQueue.class); - when(interruptingQueue.poll(any(Long.class), any(TimeUnit.class))).thenThrow(new InterruptedException()); + when(interruptingQueue.poll(anyLong(), any(TimeUnit.class))).thenThrow(new InterruptedException()); + when(interruptingQueue.isEmpty()).thenReturn(true); final Field queueField = OtlpSinkBuffer.class.getDeclaredField("queue"); queueField.setAccessible(true); @@ -186,6 +342,32 @@ void testRunWithInterruptedException() throws Exception { ); } + @Test + void testRunWithInterruptedException_whileRunning() throws Exception { + final OtlpSinkBuffer localBuffer = new OtlpSinkBuffer(config, metrics, encoder, sender); + + // Override the internal queue to throw InterruptedException first, then return null + final BlockingQueue> interruptingQueue = mock(BlockingQueue.class); + when(interruptingQueue.poll(anyLong(), any(TimeUnit.class))) + .thenThrow(new InterruptedException()) + .thenReturn(null); + when(interruptingQueue.isEmpty()).thenReturn(false, true); + + final Field queueField = OtlpSinkBuffer.class.getDeclaredField("queue"); + queueField.setAccessible(true); + queueField.set(localBuffer, interruptingQueue); + + localBuffer.start(); + + TimeUnit.MILLISECONDS.sleep(200); // Let the thread hit the poll and continue + + localBuffer.stop(); + + await().atMost(2, SECONDS).untilAsserted(() -> + verify(metrics, atLeastOnce()).incrementErrorsCount() + ); + } + @Test void testWorkerThreadCatchThrowable_andRestart_fromEncoder() throws Exception { final OTelProtoStandardCodec.OTelProtoEncoder crashingEncoder = mock(OTelProtoStandardCodec.OTelProtoEncoder.class); @@ -196,8 +378,7 @@ void testWorkerThreadCatchThrowable_andRestart_fromEncoder() throws Exception { buffer = new OtlpSinkBuffer(config, metrics, crashingEncoder, sender); buffer.start(); - final Record record = mock(Record.class); - when(record.getData()).thenReturn(mock(Span.class)); + final Record record = createMockRecord(); buffer.add(record); await().atMost(2, SECONDS).untilAsserted(() -> { @@ -212,14 +393,13 @@ void testFinalFlushAfterStopFlushesBatch() throws Exception { when(config.getMaxBatchSize()).thenReturn(Long.MAX_VALUE); // very high when(config.getFlushTimeoutMillis()).thenReturn(Long.MAX_VALUE); // effectively disables time-based flush + final ResourceSpans resourceSpans = ResourceSpans.getDefaultInstance(); + when(encoder.convertToResourceSpans(any())).thenReturn(resourceSpans); + buffer = new OtlpSinkBuffer(config, metrics, encoder, sender); buffer.start(); - // Prepare one span record - final Record record = mock(Record.class); - when(record.getData()).thenReturn(mock(Span.class)); - when(encoder.convertToResourceSpans(any())).thenReturn(ResourceSpans.getDefaultInstance()); - + final Record record = createMockRecord(); buffer.add(record); // Wait briefly to ensure the record is picked up @@ -230,15 +410,153 @@ void testFinalFlushAfterStopFlushesBatch() throws Exception { // Verify final flush triggered send await().atMost(1, SECONDS).untilAsserted(() -> - verify(sender).send(any(byte[].class), anyInt()) + verify(sender).send(anyList()) ); } @Test - void testPublicConstructorInitializesWithDefaults() { - // Only mocks needed are config and metrics - final OtlpSinkBuffer defaultBuffer = new OtlpSinkBuffer(config, metrics); + void testNoFinalFlushWhenBatchIsEmpty() throws Exception { + when(config.getMaxEvents()).thenReturn(1000); + when(config.getMaxBatchSize()).thenReturn(Long.MAX_VALUE); + when(config.getFlushTimeoutMillis()).thenReturn(Long.MAX_VALUE); + + buffer = new OtlpSinkBuffer(config, metrics, encoder, sender); + buffer.start(); + + // Don't add any records + TimeUnit.MILLISECONDS.sleep(50); + buffer.stop(); + + // Verify no send was called since batch was empty + verify(sender, never()).send(anyList()); + } + + @Test + void testMaxEventsZeroDoesNotTriggerFlush() throws Exception { + when(config.getMaxEvents()).thenReturn(0); + when(config.getMaxBatchSize()).thenReturn(Long.MAX_VALUE); + when(config.getFlushTimeoutMillis()).thenReturn(Long.MAX_VALUE); + + final ResourceSpans resourceSpans = ResourceSpans.getDefaultInstance(); + when(encoder.convertToResourceSpans(any())).thenReturn(resourceSpans); + + buffer = new OtlpSinkBuffer(config, metrics, encoder, sender); + buffer.start(); + + final Record record = createMockRecord(); + buffer.add(record); + + TimeUnit.MILLISECONDS.sleep(100); + buffer.stop(); + + // Should only flush on final flush, not by count + await().atMost(1, SECONDS).untilAsserted(() -> + verify(sender, times(1)).send(anyList()) + ); + } + + @Test + void testDaemonThreadConfiguration() throws Exception { + // This test verifies that the thread is created as non-daemon + // We can't directly test this, but we can verify the thread factory is called + buffer.start(); + assertTrue(buffer.isRunning()); + buffer.stop(); + } + + @Test + void testBatchSizeAccumulation() throws Exception { + when(config.getMaxEvents()).thenReturn(Integer.MAX_VALUE); + when(config.getMaxBatchSize()).thenReturn(50L); + when(config.getFlushTimeoutMillis()).thenReturn(Long.MAX_VALUE); - assertNotNull(defaultBuffer); // Constructor didn't throw + final ResourceSpans resourceSpans1 = mock(ResourceSpans.class); + final ResourceSpans resourceSpans2 = mock(ResourceSpans.class); + when(resourceSpans1.getSerializedSize()).thenReturn(20); + when(resourceSpans2.getSerializedSize()).thenReturn(35); + + when(encoder.convertToResourceSpans(any(Span.class))) + .thenReturn(resourceSpans1) + .thenReturn(resourceSpans2); + + buffer = new OtlpSinkBuffer(config, metrics, encoder, sender); + buffer.start(); + + final Record rec1 = createMockRecord(); + final Record rec2 = createMockRecord(); + + buffer.add(rec1); + buffer.add(rec2); + + // Total size: 20 + 35 = 55, which exceeds 50 + await().atMost(1, SECONDS).untilAsserted(() -> verify(sender).send(anyList())); + buffer.stop(); + } + + @Test + void testMultipleFlushCycles() throws Exception { + when(config.getMaxEvents()).thenReturn(1); + when(config.getMaxBatchSize()).thenReturn(Long.MAX_VALUE); + when(config.getFlushTimeoutMillis()).thenReturn(Long.MAX_VALUE); + + final ResourceSpans resourceSpans = ResourceSpans.getDefaultInstance(); + when(encoder.convertToResourceSpans(any())).thenReturn(resourceSpans); + + buffer = new OtlpSinkBuffer(config, metrics, encoder, sender); + buffer.start(); + + final Record rec1 = createMockRecord(); + final Record rec2 = createMockRecord(); + final Record rec3 = createMockRecord(); + + buffer.add(rec1); + buffer.add(rec2); + buffer.add(rec3); + + await().atMost(2, SECONDS).untilAsserted(() -> verify(sender, times(3)).send(anyList())); + buffer.stop(); + } + + @Test + void testEncodingFailureDoesNotStopProcessing() throws Exception { + // Configure to process multiple records + when(config.getMaxEvents()).thenReturn(2); + when(config.getMaxBatchSize()).thenReturn(Long.MAX_VALUE); + when(config.getFlushTimeoutMillis()).thenReturn(Long.MAX_VALUE); + + // First and third succeed, second fails + when(encoder.convertToResourceSpans(any(Span.class))) + .thenReturn(ResourceSpans.getDefaultInstance()) + .thenThrow(new RuntimeException("encoding error")) + .thenReturn(ResourceSpans.getDefaultInstance()); + + buffer = new OtlpSinkBuffer(config, metrics, encoder, sender); + buffer.start(); + + final Record rec1 = createMockRecord(); + final Record rec2 = createMockRecord(); + final Record rec3 = createMockRecord(); + + buffer.add(rec1); + buffer.add(rec2); + buffer.add(rec3); + + // Should still send the batch with the 2 successful records + await().atMost(2, SECONDS).untilAsserted(() -> { + verify(sender).send(anyList()); + verify(metrics).incrementFailedSpansCount(1); + verify(metrics, atLeastOnce()).incrementErrorsCount(); + }); + + buffer.stop(); + } + + private Record createMockRecord() { + final Record record = mock(Record.class); + final Span span = mock(Span.class); + final EventHandle eventHandle = mock(EventHandle.class); + when(record.getData()).thenReturn(span); + when(span.getEventHandle()).thenReturn(eventHandle); + return record; } } \ No newline at end of file diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java index 6b7056278f..06460c2953 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java @@ -12,14 +12,21 @@ import com.linecorp.armeria.common.ResponseHeaders; import io.opentelemetry.proto.collector.trace.v1.ExportTracePartialSuccess; import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceResponse; +import io.opentelemetry.proto.trace.v1.ResourceSpans; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; +import org.opensearch.dataprepper.model.event.EventHandle; import org.opensearch.dataprepper.plugins.sink.otlp.configuration.OtlpSinkConfig; import org.opensearch.dataprepper.plugins.sink.otlp.metrics.OtlpSinkMetrics; +import software.amazon.awssdk.utils.Pair; import java.lang.reflect.Method; import java.net.URI; import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; import java.util.function.Function; import static org.awaitility.Awaitility.await; @@ -27,6 +34,7 @@ import static org.mockito.Mockito.any; import static org.mockito.Mockito.anyLong; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoInteractions; import static org.mockito.Mockito.verifyNoMoreInteractions; @@ -35,13 +43,18 @@ class OtlpHttpSenderTest { private static final byte[] PAYLOAD = "test-otlp".getBytes(StandardCharsets.UTF_8); - private static final int SPANS = 3; + private static final int SPANS_COUNT = 3; private OtlpSinkMetrics metrics; private SigV4Signer signer; private WebClient webClient; private Function gzipCompressor; private OtlpHttpSender sender; + private AwsCredentialsSupplier mockAwsCredSupplier; + private List> testBatch; + private EventHandle mockEventHandle1; + private EventHandle mockEventHandle2; + private EventHandle mockEventHandle3; @BeforeEach void setup() { @@ -49,6 +62,18 @@ void setup() { signer = mock(SigV4Signer.class); webClient = mock(WebClient.class); gzipCompressor = mock(Function.class); + mockAwsCredSupplier = mock(AwsCredentialsSupplier.class); + + mockEventHandle1 = mock(EventHandle.class); + mockEventHandle2 = mock(EventHandle.class); + mockEventHandle3 = mock(EventHandle.class); + + // Create test batch with ResourceSpans and EventHandles + testBatch = Arrays.asList( + Pair.of(ResourceSpans.newBuilder().build(), mockEventHandle1), + Pair.of(ResourceSpans.newBuilder().build(), mockEventHandle2), + Pair.of(ResourceSpans.newBuilder().build(), mockEventHandle3) + ); when(gzipCompressor.apply(any())).thenReturn(PAYLOAD); when(signer.signRequest(any())).thenReturn( @@ -62,19 +87,37 @@ void setup() { sender = new OtlpHttpSender(metrics, gzipCompressor, signer, webClient); } + @Test + void testSend_emptyBatch_returnsEarly() { + final List> emptyBatch = Collections.emptyList(); + + // Prevent accidental NPE if send logic changes + when(webClient.execute(any(HttpRequest.class))).thenThrow(new AssertionError("Should not call execute")); + + sender.send(emptyBatch); + + verifyNoInteractions(metrics); + } + @Test void testSend_successfulResponse() { when(webClient.execute(any(HttpRequest.class))).thenReturn( HttpResponse.of(ResponseHeaders.of(200), HttpData.empty()) ); - sender.send(PAYLOAD, SPANS); + sender.send(testBatch); await().untilAsserted(() -> { - verify(metrics).incrementRecordsOut(SPANS); - verify(metrics).incrementPayloadSize(PAYLOAD.length); + verify(metrics).incrementRecordsOut(SPANS_COUNT); + verify(metrics).incrementPayloadSize(anyLong()); verify(metrics).incrementPayloadGzipSize(PAYLOAD.length); verify(metrics).recordHttpLatency(anyLong()); + verify(metrics).recordResponseCode(200); + + // Verify all event handles are released with success=true + verify(mockEventHandle1).release(true); + verify(mockEventHandle2).release(true); + verify(mockEventHandle3).release(true); }); } @@ -90,11 +133,45 @@ void testSend_partialSuccessResponse() { when(webClient.execute(any(HttpRequest.class))).thenReturn( HttpResponse.of(ResponseHeaders.of(200), HttpData.wrap(proto.toByteArray())) ); - sender.send(PAYLOAD, SPANS); + + sender.send(testBatch); await().untilAsserted(() -> { verify(metrics).incrementRejectedSpansCount(2); - verify(metrics).incrementRecordsOut(SPANS - 2); + verify(metrics).incrementRecordsOut(SPANS_COUNT - 2); + verify(metrics).recordResponseCode(200); + + // All handles should still be released as true (optimistic) + verify(mockEventHandle1).release(true); + verify(mockEventHandle2).release(true); + verify(mockEventHandle3).release(true); + }); + } + + @Test + void testSend_partialSuccessWithZeroRejected() { + final ExportTraceServiceResponse proto = ExportTraceServiceResponse.newBuilder() + .setPartialSuccess(ExportTracePartialSuccess.newBuilder() + .setRejectedSpans(0) + .setErrorMessage("") + .build()) + .build(); + + when(webClient.execute(any(HttpRequest.class))).thenReturn( + HttpResponse.of(ResponseHeaders.of(200), HttpData.wrap(proto.toByteArray())) + ); + + sender.send(testBatch); + + await().untilAsserted(() -> { + verify(metrics, never()).incrementRejectedSpansCount(anyLong()); + verify(metrics).incrementRecordsOut(SPANS_COUNT); + verify(metrics).recordResponseCode(200); + + // All handles should be released as true + verify(mockEventHandle1).release(true); + verify(mockEventHandle2).release(true); + verify(mockEventHandle3).release(true); }); } @@ -104,11 +181,17 @@ void testSend_parseErrorOnSuccessResponse() { HttpResponse.of(ResponseHeaders.of(200), HttpData.ofUtf8("not-protobuf")) ); - sender.send(PAYLOAD, SPANS); + sender.send(testBatch); await().untilAsserted(() -> { verify(metrics).incrementErrorsCount(); - verify(metrics).incrementRecordsOut(SPANS); + verify(metrics).incrementRecordsOut(SPANS_COUNT); + verify(metrics).recordResponseCode(200); + + // Handles should still be released as true despite parse error + verify(mockEventHandle1).release(true); + verify(mockEventHandle2).release(true); + verify(mockEventHandle3).release(true); }); } @@ -118,21 +201,51 @@ void testSend_nonSuccessStatus() { HttpResponse.of(ResponseHeaders.of(400), HttpData.ofUtf8("{\"error\":\"bad request\"}")) ); - sender.send(PAYLOAD, SPANS); + sender.send(testBatch); await().untilAsserted(() -> { verify(metrics).recordResponseCode(400); - verify(metrics).incrementRejectedSpansCount(SPANS); + verify(metrics).incrementRejectedSpansCount(SPANS_COUNT); + + // All handles should be released with success=false + verify(mockEventHandle1).release(false); + verify(mockEventHandle2).release(false); + verify(mockEventHandle3).release(false); + }); + } + + @Test + void testSend_nonSuccessStatusWithNullResponseBody() { + when(webClient.execute(any(HttpRequest.class))).thenReturn( + HttpResponse.of(ResponseHeaders.of(500), HttpData.empty()) + ); + + sender.send(testBatch); + + await().untilAsserted(() -> { + verify(metrics).recordResponseCode(500); + verify(metrics).incrementRejectedSpansCount(SPANS_COUNT); + + // All handles should be released with success=false + verify(mockEventHandle1).release(false); + verify(mockEventHandle2).release(false); + verify(mockEventHandle3).release(false); }); } @Test void testSend_skipsSendIfGzipFails() { sender = new OtlpHttpSender(metrics, ignored -> new byte[0], signer, webClient); - sender.send(PAYLOAD, SPANS); - verify(metrics).incrementFailedSpansCount(SPANS); + sender.send(testBatch); + + verify(metrics).incrementFailedSpansCount(SPANS_COUNT); verifyNoInteractions(webClient); + + // All handles should be released with success=false + verify(mockEventHandle1).release(false); + verify(mockEventHandle2).release(false); + verify(mockEventHandle3).release(false); } @Test @@ -140,9 +253,16 @@ void testSend_exceptionDuringSendIncrementsRejected() { final HttpResponse failingResponse = HttpResponse.ofFailure(new RuntimeException("send failed")); when(webClient.execute(any(HttpRequest.class))).thenReturn(failingResponse); - sender.send(PAYLOAD, SPANS); + sender.send(testBatch); + + await().untilAsserted(() -> { + verify(metrics).incrementRejectedSpansCount(SPANS_COUNT); - await().untilAsserted(() -> verify(metrics).incrementRejectedSpansCount(SPANS)); + // All handles should be released with success=false + verify(mockEventHandle1).release(false); + verify(mockEventHandle2).release(false); + verify(mockEventHandle3).release(false); + }); } @Test @@ -154,7 +274,7 @@ void testConstructor_withDefaultConfig() { when(config.getFlushTimeoutMillis()).thenReturn(5000L); when(config.getAwsRegion()).thenReturn(software.amazon.awssdk.regions.Region.US_WEST_2); - final OtlpHttpSender defaultSender = new OtlpHttpSender(config, metrics); + final OtlpHttpSender defaultSender = new OtlpHttpSender(mockAwsCredSupplier, config, metrics); assertNotNull(defaultSender); } @@ -169,31 +289,133 @@ void testConstructor_withMinimumThresholdConfig() { when(config.getAwsRegion()).thenReturn(software.amazon.awssdk.regions.Region.US_WEST_2); // Should not throw or crash - final OtlpHttpSender minimalSender = new OtlpHttpSender(config, metrics); + final OtlpHttpSender minimalSender = new OtlpHttpSender(mockAwsCredSupplier, config, metrics); assertNotNull(minimalSender); } @Test - void testHandleSuccessfulResponse_withNullBody_incrementsRecordsOut() throws Exception { - final Method method = OtlpHttpSender.class.getDeclaredMethod("handleSuccessfulResponse", byte[].class, int.class); + void testGetPayloadAndCompressedPayload_privateMethod() throws Exception { + // Test the private method via reflection to ensure 100% coverage + final Method method = OtlpHttpSender.class.getDeclaredMethod("getPayloadAndCompressedPayload", List.class); method.setAccessible(true); - method.invoke(sender, null, SPANS); - verify(metrics).incrementRecordsOut(SPANS); + when(gzipCompressor.apply(any())).thenReturn("compressed".getBytes()); + + @SuppressWarnings("unchecked") final Pair result = (Pair) method.invoke(sender, testBatch); + + assertNotNull(result); + assertNotNull(result.left()); // payload + assertNotNull(result.right()); // compressed payload + verify(gzipCompressor).apply(any()); + } + + @Test + void testBuildHttpRequest_privateMethod() throws Exception { + // Test the private method via reflection + final Method method = OtlpHttpSender.class.getDeclaredMethod("buildHttpRequest", byte[].class); + method.setAccessible(true); + + final HttpRequest result = (HttpRequest) method.invoke(sender, PAYLOAD); + + assertNotNull(result); + verify(signer).signRequest(PAYLOAD); + } + + @Test + void testHandleResponse_privateMethod_success() throws Exception { + final Method method = OtlpHttpSender.class.getDeclaredMethod("handleResponse", int.class, byte[].class, List.class); + method.setAccessible(true); + + method.invoke(sender, 200, null, testBatch); + + verify(metrics).recordResponseCode(200); + verify(metrics).incrementRecordsOut(SPANS_COUNT); + verify(mockEventHandle1).release(true); + verify(mockEventHandle2).release(true); + verify(mockEventHandle3).release(true); + } + + @Test + void testHandleResponse_privateMethod_failure() throws Exception { + final Method method = OtlpHttpSender.class.getDeclaredMethod("handleResponse", int.class, byte[].class, List.class); + method.setAccessible(true); + + method.invoke(sender, 400, "error".getBytes(), testBatch); + + verify(metrics).recordResponseCode(400); + verify(metrics).incrementRejectedSpansCount(SPANS_COUNT); + verify(mockEventHandle1).release(false); + verify(mockEventHandle2).release(false); + verify(mockEventHandle3).release(false); + } + + @Test + void testHandleSuccessfulResponse_privateMethod_withNullBody() throws Exception { + final Method method = OtlpHttpSender.class.getDeclaredMethod("handleSuccessfulResponse", byte[].class, List.class); + method.setAccessible(true); + + method.invoke(sender, null, testBatch); + + verify(metrics).incrementRecordsOut(SPANS_COUNT); + verify(mockEventHandle1).release(true); + verify(mockEventHandle2).release(true); + verify(mockEventHandle3).release(true); verifyNoMoreInteractions(metrics); } @Test - void testHandleSuccessfulResponse_withoutPartialSuccess_incrementsRecordsOut() throws Exception { + void testHandleSuccessfulResponse_privateMethod_withoutPartialSuccess() throws Exception { // Create a response with no partial_success final ExportTraceServiceResponse response = ExportTraceServiceResponse.newBuilder().build(); final byte[] responseBytes = response.toByteArray(); - final Method method = OtlpHttpSender.class.getDeclaredMethod("handleSuccessfulResponse", byte[].class, int.class); + final Method method = OtlpHttpSender.class.getDeclaredMethod("handleSuccessfulResponse", byte[].class, List.class); method.setAccessible(true); - method.invoke(sender, responseBytes, SPANS); - verify(metrics).incrementRecordsOut(SPANS); + method.invoke(sender, responseBytes, testBatch); + + verify(metrics).incrementRecordsOut(SPANS_COUNT); + verify(mockEventHandle1).release(true); + verify(mockEventHandle2).release(true); + verify(mockEventHandle3).release(true); verifyNoMoreInteractions(metrics); } -} + + @Test + void testHandleSuccessfulResponse_privateMethod_withPartialSuccess() throws Exception { + final ExportTraceServiceResponse response = ExportTraceServiceResponse.newBuilder() + .setPartialSuccess(ExportTracePartialSuccess.newBuilder() + .setRejectedSpans(1) + .setErrorMessage("test error") + .build()) + .build(); + final byte[] responseBytes = response.toByteArray(); + + final Method method = OtlpHttpSender.class.getDeclaredMethod("handleSuccessfulResponse", byte[].class, List.class); + method.setAccessible(true); + + method.invoke(sender, responseBytes, testBatch); + + verify(metrics).incrementRejectedSpansCount(1); + verify(metrics).incrementRecordsOut(SPANS_COUNT - 1); + verify(mockEventHandle1).release(true); + verify(mockEventHandle2).release(true); + verify(mockEventHandle3).release(true); + } + + @Test + void testHandleSuccessfulResponse_privateMethod_parseException() throws Exception { + final byte[] invalidBytes = "invalid protobuf".getBytes(); + + final Method method = OtlpHttpSender.class.getDeclaredMethod("handleSuccessfulResponse", byte[].class, List.class); + method.setAccessible(true); + + method.invoke(sender, invalidBytes, testBatch); + + verify(metrics).incrementErrorsCount(); + verify(metrics).incrementRecordsOut(SPANS_COUNT); + verify(mockEventHandle1).release(true); + verify(mockEventHandle2).release(true); + verify(mockEventHandle3).release(true); + } +} \ No newline at end of file diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4SignerTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4SignerTest.java index 0767bdf2fa..3218574377 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4SignerTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4SignerTest.java @@ -4,166 +4,67 @@ */ package org.opensearch.dataprepper.plugins.sink.otlp.http; -import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; -import org.mockito.ArgumentMatcher; -import org.mockito.MockedStatic; +import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; import org.opensearch.dataprepper.plugins.sink.otlp.configuration.OtlpSinkConfig; +import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider; import software.amazon.awssdk.http.SdkHttpFullRequest; import software.amazon.awssdk.regions.Region; -import software.amazon.awssdk.services.sts.StsClient; -import software.amazon.awssdk.services.sts.StsClientBuilder; -import software.amazon.awssdk.services.sts.model.AssumeRoleRequest; -import software.amazon.awssdk.services.sts.model.AssumeRoleResponse; -import software.amazon.awssdk.services.sts.model.Credentials; import java.nio.charset.StandardCharsets; -import java.time.Instant; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.argThat; +import static org.mockito.Mockito.any; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.mockStatic; -import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; class SigV4SignerTest { private static final byte[] PAYLOAD = "test-payload".getBytes(StandardCharsets.UTF_8); - public static final Region REGION = Region.of("us-west-2"); - private OtlpSinkConfig mockXrayConfig; + private static final Region REGION = Region.US_WEST_2; + + private OtlpSinkConfig mockConfig; + private AwsCredentialsSupplier mockSupplier; private SigV4Signer target; @BeforeEach - void setUp(){ - System.setProperty("aws.accessKeyId", "dummy"); - System.setProperty("aws.secretAccessKey", "dummy"); - - mockXrayConfig = mock(OtlpSinkConfig.class); - } + void setup() { + mockConfig = mock(OtlpSinkConfig.class); + mockSupplier = mock(AwsCredentialsSupplier.class); - @AfterEach - void cleanUp() { - System.clearProperty("aws.accessKeyId"); - System.clearProperty("aws.secretAccessKey"); + when(mockConfig.getAwsRegion()).thenReturn(REGION); + when(mockSupplier.getProvider(any())).thenReturn(DefaultCredentialsProvider.create()); } @Test - void testSignRequest_withInputEndpoint_whenEndpointIsSet() { - // setup + void testSignRequest_withExplicitEndpoint() { final String endpoint = "https://performance.us-west-2.xray.cloudwatch.aws.dev/v1/traces"; - when(mockXrayConfig.getAwsRegion()).thenReturn(REGION); - when(mockXrayConfig.getEndpoint()).thenReturn(endpoint); - target = new SigV4Signer(mockXrayConfig); + when(mockConfig.getEndpoint()).thenReturn(endpoint); - // run - final SdkHttpFullRequest signedRequest = target.signRequest(PAYLOAD); + target = new SigV4Signer(mockSupplier, mockConfig); + final SdkHttpFullRequest request = target.signRequest(PAYLOAD); - // assert - assertNotNull(signedRequest); - assertEquals("POST", signedRequest.method().name()); - assertTrue(signedRequest.headers().containsKey("Authorization")); - assertEquals("application/x-protobuf", signedRequest.firstMatchingHeader("Content-Type").orElse(null)); - assertTrue(signedRequest.getUri().toString().contains(endpoint)); + assertNotNull(request); + assertEquals("POST", request.method().name()); + assertTrue(request.headers().containsKey("Authorization")); + assertEquals("application/x-protobuf", request.firstMatchingHeader("Content-Type").orElse(null)); + assertEquals(endpoint, request.getUri().toString()); } @Test - void testSignRequest_withFallbackStsRole_whenStsRoleNotSet() { - // setup - when(mockXrayConfig.getAwsRegion()).thenReturn(Region.US_WEST_2); - target = new SigV4Signer(mockXrayConfig); - - // run - final SdkHttpFullRequest signedRequest = target.signRequest(PAYLOAD); - - // assert - assertNotNull(signedRequest); - assertEquals("POST", signedRequest.method().name()); - assertTrue(signedRequest.headers().containsKey("Authorization")); - assertEquals("application/x-protobuf", signedRequest.firstMatchingHeader("Content-Type").orElse(null)); - assertTrue(signedRequest.getUri().toString().contains(String.format("https://xray.%s.amazonaws.com/v1/traces", Region.US_WEST_2.toString()))); - } - - @Test - void testSignRequest_withCustomCredentials_usingDefaultStsClientFallback() { - when(mockXrayConfig.getAwsRegion()).thenReturn(Region.US_WEST_2); - when(mockXrayConfig.getStsRoleArn()).thenReturn("arn:aws:iam::123456789012:role/test-role"); - - // Use mocked static builder to simulate StsClient.builder() - try (final MockedStatic mockedStsClientStatic = mockStatic(StsClient.class)) { - final StsClient mockStsClient = mock(StsClient.class); - - // Setup fake STS response - when(mockStsClient.assumeRole(any(AssumeRoleRequest.class))) - .thenReturn(AssumeRoleResponse.builder() - .credentials(Credentials.builder() - .accessKeyId("fake-access") - .secretAccessKey("fake-secret") - .sessionToken("fake-token") - .expiration(Instant.now().plusSeconds(3600)) - .build()) - .build()); - - // Mock StsClient builder - final StsClientBuilder mockBuilder = mock(StsClientBuilder.class); - when(mockBuilder.region(Region.US_WEST_2)).thenReturn(mockBuilder); - when(mockBuilder.build()).thenReturn(mockStsClient); - - // Mock the static method StsClient.builder() - mockedStsClientStatic.when(StsClient::builder).thenReturn(mockBuilder); - - // run - target = new SigV4Signer(mockXrayConfig, null); - final SdkHttpFullRequest signedRequest = target.signRequest(PAYLOAD); - - // assert - assertNotNull(signedRequest); - assertTrue(signedRequest.headers().containsKey("Authorization")); - assertTrue(signedRequest.getUri().toString().contains(String.format("https://xray.%s.amazonaws.com/v1/traces", Region.US_WEST_2.toString()))); - } - } - - @Test - void testSignRequest_withCustomCredentials_usingMockedSts() { - // setup - final String expectedRoleArn = "arn:aws:iam::123456789012:role/test-role"; - final String expectedExternalId = "external-id-123"; - when(mockXrayConfig.getAwsRegion()).thenReturn(Region.US_WEST_2); - when(mockXrayConfig.getStsRoleArn()).thenReturn(expectedRoleArn); - when(mockXrayConfig.getStsExternalId()).thenReturn(expectedExternalId); - - final StsClient mockStsClient = mock(StsClient.class); - when(mockStsClient.assumeRole(any(AssumeRoleRequest.class))).thenReturn( - AssumeRoleResponse.builder() - .credentials(Credentials.builder() - .accessKeyId("fake-access-key") - .secretAccessKey("fake-secret-key") - .sessionToken("fake-session-token") - .expiration(Instant.now().plusSeconds(3600)) - .build()) - .build() - ); - - target = new SigV4Signer(mockXrayConfig, mockStsClient); - - // run - final SdkHttpFullRequest signedRequest = target.signRequest(PAYLOAD); - - // assert - assertNotNull(signedRequest); - assertEquals("POST", signedRequest.method().name()); - assertTrue(signedRequest.headers().containsKey("Authorization")); - assertEquals("application/x-protobuf", signedRequest.firstMatchingHeader("Content-Type").orElse(null)); - assertTrue(signedRequest.getUri().toString().contains("https://xray.us-west-2.amazonaws.com/v1/traces")); + void testSignRequest_withDefaultEndpoint() { + when(mockConfig.getEndpoint()).thenReturn(null); - final ArgumentMatcher matcher = request -> - expectedRoleArn.equals(request.roleArn()) && - expectedExternalId.equals(request.externalId()); + target = new SigV4Signer(mockSupplier, mockConfig); + final SdkHttpFullRequest request = target.signRequest(PAYLOAD); - verify(mockStsClient).assumeRole(argThat(matcher)); + assertNotNull(request); + assertEquals("POST", request.method().name()); + assertTrue(request.headers().containsKey("Authorization")); + assertEquals("application/x-protobuf", request.firstMatchingHeader("Content-Type").orElse(null)); + assertEquals("https://xray.us-west-2.amazonaws.com/v1/traces", request.getUri().toString()); } -} +} \ No newline at end of file From 977b524b5e05a44aebbc70d99123eb36f5df506e Mon Sep 17 00:00:00 2001 From: huy pham Date: Fri, 23 May 2025 14:01:55 -0700 Subject: [PATCH 20/23] Avoid breaking change by making the aws config block required. - Updated README - Fixed init bug - Performed e2e and perf tests Signed-off-by: huy pham --- data-prepper-plugins/otlp-sink/README.md | 21 ++++- .../plugins/sink/otlp/OtlpSink.java | 6 +- .../otlp/configuration/OtlpSinkConfig.java | 22 ++++- .../sink/otlp/http/OtlpHttpSender.java | 24 ++--- .../plugins/sink/otlp/http/SigV4Signer.java | 1 + .../plugins/sink/otlp/OtlpSinkTest.java | 32 ++++++- .../sink/otlp/buffer/OtlpSinkBufferTest.java | 3 +- .../configuration/OtlpSinkConfigTest.java | 94 +++++++++++++++++-- .../sink/otlp/http/OtlpHttpSenderTest.java | 8 +- 9 files changed, 180 insertions(+), 31 deletions(-) diff --git a/data-prepper-plugins/otlp-sink/README.md b/data-prepper-plugins/otlp-sink/README.md index d8c8f40f77..09c70d2a9f 100644 --- a/data-prepper-plugins/otlp-sink/README.md +++ b/data-prepper-plugins/otlp-sink/README.md @@ -17,6 +17,25 @@ to any OTLP Protobuf-compatible endpoint. ## Sample Pipeline Configuration +### Minimal Configuration (No STS) + +Use this when Data Prepper has permission to write to AWS X-Ray directly. + +```yaml +otlp_pipeline: + source: + otel_trace_source: + + sink: + - otlp: + endpoint: "https://xray.us-west-2.amazonaws.com/v1/traces" + aws: { } +``` + +### Full Configuration with STS + +Use this when assuming a cross-account role is required. + ```yaml otlp_pipeline: workers: 2 @@ -56,7 +75,7 @@ otlp_pipeline: | `threshold.max_events` | `int` | No | `512` (recommended) | Maximum number of spans per batch. Use `0` to disable count-based flushing. Must be ≥ 0. | | `threshold.max_batch_size` | `String` | No | `1mb` (recommended) | Maximum total payload bytes per batch. Supports human-readable suffixes (`kb`, `mb`). | | `threshold.flush_timeout` | `String` | No | `200ms` (recommended) | Maximum time to wait before flushing a non-empty batch. Minimum: 1ms (e.g., `200ms`, `1s`) | -| **aws** | `Object` | No | — | AWS authentication settings. See below. | +| **aws** | `Object` | Yes | — | AWS authentication settings. Use `{}` if no STS role is needed. See below. | | `aws.sts_role_arn` | `String` | No | — | IAM Role ARN that Data Prepper (or OSI) assumes to send spans to X-Ray on behalf of a customer account. | | `aws.sts_external_id` | `String` | No | — | External ID to use when assuming the role. Required only if the target IAM role enforces sts:ExternalId. | diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java index 47695cf51e..fc7293c4b4 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java @@ -32,6 +32,7 @@ pluginConfigurationType = OtlpSinkConfig.class ) public class OtlpSink extends AbstractSink> { + private volatile boolean initialized = false; private final OtlpSinkBuffer buffer; private final OtlpSinkMetrics sinkMetrics; @@ -48,6 +49,8 @@ public class OtlpSink extends AbstractSink> { public OtlpSink(@Nonnull final AwsCredentialsSupplier awsCredentialsSupplier, @Nonnull final OtlpSinkConfig config, @Nonnull final PluginMetrics pluginMetrics, @Nonnull final PluginSetting pluginSetting) { super(pluginSetting); + config.validate(); + this.sinkMetrics = new OtlpSinkMetrics(pluginMetrics, pluginSetting); this.buffer = new OtlpSinkBuffer(awsCredentialsSupplier, config, sinkMetrics); } @@ -58,6 +61,7 @@ public OtlpSink(@Nonnull final AwsCredentialsSupplier awsCredentialsSupplier, @N @Override public void doInitialize() { buffer.start(); + initialized = true; } /** @@ -79,7 +83,7 @@ public void doOutput(@Nonnull final Collection> records) { */ @Override public boolean isReady() { - return buffer.isRunning(); + return initialized && buffer.isRunning(); } /** diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfig.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfig.java index f9e5021bce..7b2e513725 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfig.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfig.java @@ -68,7 +68,7 @@ public long getFlushTimeoutMillis() { */ @JsonProperty("aws") @Valid - private AwsConfig awsAuthenticationConfig; + private AwsConfig awsConfig; /** * Get AWS region from the provided endpoint. @@ -97,18 +97,30 @@ public Region getAwsRegion() { } public String getStsRoleArn() { - if (awsAuthenticationConfig == null) { + if (awsConfig == null || awsConfig.getAwsStsRoleArn() == null) { return null; } - return awsAuthenticationConfig.getAwsStsRoleArn(); + return awsConfig.getAwsStsRoleArn(); } public String getStsExternalId() { - if (awsAuthenticationConfig == null) { + if (awsConfig == null || awsConfig.getAwsStsExternalId() == null) { return null; } - return awsAuthenticationConfig.getAwsStsExternalId(); + return awsConfig.getAwsStsExternalId(); + } + + /** + * Validate the AWS configuration. + * This method ensures breaking change in future release where non-AWS OTLP endpoints are supported. + * + * @throws IllegalArgumentException if the AWS configuration is invalid + */ + public void validate() { + if (awsConfig == null) { + throw new IllegalArgumentException("aws configuration is required"); + } } } diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java index c6ee74e855..013fe6730f 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSender.java @@ -14,7 +14,6 @@ import com.linecorp.armeria.common.HttpMethod; import com.linecorp.armeria.common.HttpRequest; import com.linecorp.armeria.common.HttpResponse; -import com.linecorp.armeria.common.MediaType; import com.linecorp.armeria.common.RequestHeaders; import com.linecorp.armeria.common.RequestHeadersBuilder; import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceRequest; @@ -42,7 +41,6 @@ public class OtlpHttpSender { private static final Logger LOG = LoggerFactory.getLogger(OtlpHttpSender.class); private static final Set RETRYABLE_STATUS_CODES = Set.of(429, 502, 503, 504); - private static final MediaType PROTOBUF = MediaType.parse("application/x-protobuf"); private final SigV4Signer signer; private final WebClient webClient; @@ -85,7 +83,6 @@ public OtlpHttpSender(@Nonnull final AwsCredentialsSupplier awsCredentialsSuppli * adding complexity with minimal benefit for most OTLP endpoints. * - Our exponential backoff already handles typical retry intervals gracefully. */ - private static WebClient buildWebClient(final OtlpSinkConfig config) { final RetryRuleWithContent retryRule = RetryRuleWithContent.builder() .onStatus((ctx, status) -> RETRYABLE_STATUS_CODES.contains(status.code())) @@ -118,11 +115,14 @@ public void send(@Nonnull final List> batch) { return; } - final Pair payloadAndCompressedPayload = getPayloadAndCompressedPayload(batch); - final int spans = batch.size(); + // Defensive copy to avoid ConcurrentModificationException + final List> immutableBatch = List.copyOf(batch); + + final Pair payloadAndCompressedPayload = getPayloadAndCompressedPayload(immutableBatch); + final int spans = immutableBatch.size(); if (payloadAndCompressedPayload.right().length == 0) { sinkMetrics.incrementFailedSpansCount(spans); - releaseAllEventHandle(batch, false); + releaseAllEventHandle(immutableBatch, false); return; } @@ -139,12 +139,12 @@ public void send(@Nonnull final List> batch) { final int statusCode = response.status().code(); final byte[] responseBytes = response.content().array(); - handleResponse(statusCode, responseBytes, batch); + handleResponse(statusCode, responseBytes, immutableBatch); }) .exceptionally(e -> { LOG.error("Failed to send {} spans.", spans, e); sinkMetrics.incrementRejectedSpansCount(spans); - releaseAllEventHandle(batch, false); + releaseAllEventHandle(immutableBatch, false); return null; }); } @@ -161,12 +161,14 @@ private Pair getPayloadAndCompressedPayload(final List vList.forEach(v -> headersBuilder.add(k, v))); return HttpRequest.of(headersBuilder.build(), HttpData.wrap(compressedPayload)); } diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4Signer.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4Signer.java index becfb72960..9b1b0becab 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4Signer.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4Signer.java @@ -63,6 +63,7 @@ SdkHttpFullRequest signRequest(@Nonnull final byte[] payload) { .method(SdkHttpMethod.POST) .uri(endpointUri) .putHeader("Content-Type", "application/x-protobuf") + .putHeader("Content-Encoding", "gzip") .contentStreamProvider(() -> SdkBytes.fromByteArray(payload).asInputStream()) .build(); diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java index bdcfb3032b..299589d42c 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java @@ -4,8 +4,10 @@ */ package org.opensearch.dataprepper.plugins.sink.otlp; +import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.function.Executable; import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.configuration.PluginSetting; @@ -21,6 +23,7 @@ import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoMoreInteractions; @@ -40,6 +43,7 @@ void setUp() throws Exception { mockAwsCredSupplier = mock(AwsCredentialsSupplier.class); mockConfig = mock(OtlpSinkConfig.class); when(mockConfig.getAwsRegion()).thenReturn(Region.of("us-west-2")); + when(mockConfig.getEndpoint()).thenReturn("https://localhost/v1/traces"); mockMetrics = mock(PluginMetrics.class); @@ -66,6 +70,23 @@ void testInitialize_startsBuffer() { verify(mockBuffer).start(); } + @Test + void testConstructor_throwsWhenAwsConfigIsMissing() { + doThrow(new IllegalArgumentException("aws configuration is required")) + .when(mockConfig).validate(); + + // Act & Assert + final Executable constructorCall = () -> + new OtlpSink(mockAwsCredSupplier, mockConfig, mockMetrics, mockSetting); + + final IllegalArgumentException thrown = Assertions.assertThrows( + IllegalArgumentException.class, + constructorCall + ); + + Assertions.assertEquals("aws configuration is required", thrown.getMessage()); + } + @Test void testOutput_addsEveryRecordToBuffer() { // Arrange @@ -82,12 +103,17 @@ void testOutput_addsEveryRecordToBuffer() { } @Test - void testIsReady_delegatesToBuffer() { - // true case + void testIsReady_returnsTrueOnlyAfterInitialization() { when(mockBuffer.isRunning()).thenReturn(true); + + // Not initialized yet + assertFalse(target.isReady()); + + // Initialize, which sets 'initialized = true' and starts the buffer + target.initialize(); assertTrue(target.isReady()); - // false case + // Now simulate buffer being not running when(mockBuffer.isRunning()).thenReturn(false); assertFalse(target.isReady()); } diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBufferTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBufferTest.java index 800ace1dc4..ec20a6473d 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBufferTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/buffer/OtlpSinkBufferTest.java @@ -53,6 +53,7 @@ class OtlpSinkBufferTest { @BeforeEach void setUp() { config = mock(OtlpSinkConfig.class); + when(config.getEndpoint()).thenReturn("https://localhost/v1/traces"); when(config.getMaxEvents()).thenReturn(2); when(config.getMaxRetries()).thenReturn(2); when(config.getMaxBatchSize()).thenReturn(1_000_000L); @@ -456,7 +457,7 @@ void testMaxEventsZeroDoesNotTriggerFlush() throws Exception { } @Test - void testDaemonThreadConfiguration() throws Exception { + void testDaemonThreadConfiguration() { // This test verifies that the thread is created as non-daemon // We can't directly test this, but we can verify the thread factory is called buffer.start(); diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfigTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfigTest.java index c33a4b5a45..100115c170 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfigTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfigTest.java @@ -21,8 +21,11 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.nullValue; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; class OtlpSinkConfigTest { @@ -74,9 +77,63 @@ public Duration deserialize(final JsonParser p, final DeserializationContext ctx mapper.registerModule(durationModule); } + @Test + void testGetStsRoleArnReturnsNullWhenNotSet() throws Exception { + final String yaml = String.join("\n", + "endpoint: \"" + EXPECTED_ENDPOINT + "\"", + "aws: {}", // sts_role_arn not set + "max_retries: 5" + ); + + final OtlpSinkConfig config = mapper.readValue(yaml, OtlpSinkConfig.class); + + assertNull(config.getStsRoleArn()); + } + + @Test + void testGetStsExternalIdReturnsNullWhenNotSet() throws Exception { + final String yaml = String.join("\n", + "endpoint: \"" + EXPECTED_ENDPOINT + "\"", + "aws: {}", // sts_external_id not set + "max_retries: 5" + ); + + final OtlpSinkConfig config = mapper.readValue(yaml, OtlpSinkConfig.class); + + assertNull(config.getStsExternalId()); + } + + @Test + void testAwsRegion_throwsWhenHostIsNull() throws Exception { + final String yaml = String.join("\n", + "endpoint: \"mailto:user@example.com\"", // No host component + "aws: {}" + ); + + final OtlpSinkConfig config = mapper.readValue(yaml, OtlpSinkConfig.class); + + assertThrows(IllegalArgumentException.class, config::getAwsRegion); + } + + @Test + void testAwsRegion_throwsWhenUriIsMalformed() throws Exception { + final String yaml = String.join("\n", + "endpoint: \"https://xray .us-west-2.amazonaws.com\"", // Invalid URI with space + "aws: {}" + ); + + final OtlpSinkConfig config = mapper.readValue(yaml, OtlpSinkConfig.class); + + final IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class, config::getAwsRegion); + assertTrue(thrown.getMessage().contains("Failed to parse AWS region from endpoint")); + } + @Test void testMinimumConfigDefaults() throws Exception { - final String yaml = "endpoint: \"" + EXPECTED_ENDPOINT + "\""; + final String yaml = String.join("\n", + "endpoint: \"" + EXPECTED_ENDPOINT + "\"", + "aws: {}" + ); final OtlpSinkConfig config = mapper.readValue(yaml, OtlpSinkConfig.class); @@ -94,6 +151,7 @@ void testMinimumConfigDefaults() throws Exception { void testCustomThresholdAndRetries() throws Exception { final String yaml = String.join("\n", "endpoint: \"" + EXPECTED_ENDPOINT + "\"", + "aws: {}", "max_retries: 3", "threshold:", " max_events: " + CUSTOM_MAX_EVENTS, @@ -110,6 +168,26 @@ void testCustomThresholdAndRetries() throws Exception { assertEquals(CUSTOM_FLUSH_TIMEOUT, config.getFlushTimeoutMillis()); } + @Test + void testValidateAwsConfig_doesNotThrowWhenPresent() throws Exception { + final String yaml = String.join("\n", + "endpoint: \"" + EXPECTED_ENDPOINT + "\"", + "aws: {}" + ); + final OtlpSinkConfig config = mapper.readValue(yaml, OtlpSinkConfig.class); + assertDoesNotThrow(config::validate); + } + + @Test + void testValidateAwsConfig_throwsException() throws Exception { + final String yaml = String.join("\n", + "endpoint: \"" + EXPECTED_ENDPOINT + "\"" + ); + final OtlpSinkConfig config = mapper.readValue(yaml, OtlpSinkConfig.class); + + assertThrows(IllegalArgumentException.class, config::validate); + } + @Test void testAwsBlockDeserialization() throws Exception { final String yaml = String.join("\n", @@ -130,25 +208,25 @@ void testAwsBlockDeserialization() throws Exception { } @Test - void testAwsSectionMissing_staysNull() throws Exception { + void testAwsSectionPresentButEmpty_doesNotThrow() throws Exception { final String yaml = String.join("\n", "endpoint: \"" + EXPECTED_ENDPOINT + "\"", + "aws: {}", "max_retries: " + DEFAULT_MAX_RETRIES ); final OtlpSinkConfig config = mapper.readValue(yaml, OtlpSinkConfig.class); - assertEquals(EXPECTED_ENDPOINT, config.getEndpoint()); - assertEquals(DEFAULT_MAX_RETRIES, config.getMaxRetries()); - - assertThat(config.getStsRoleArn(), nullValue()); - assertThat(config.getStsExternalId(), nullValue()); + // aws block exists, so validateAwsConfig() will not throw + assertNull(config.getStsRoleArn()); + assertNull(config.getStsExternalId()); } @Test void testAwsRegion_parsedFromStandardXrayEndpoint() throws Exception { final String yaml = String.join("\n", "endpoint: \"https://xray.us-east-1.amazonaws.com\"", + "aws: {}", "max_retries: 5" ); @@ -164,6 +242,7 @@ void testAwsRegion_parsedFromStandardXrayEndpoint() throws Exception { void testAwsRegion_invalidEndpoint_throwsException() { final String yaml = String.join("\n", "endpoint: \"https://example.invalid-endpoint\"", + "aws: {}", "max_retries: 5" ); @@ -177,6 +256,7 @@ void testAwsRegion_invalidEndpoint_throwsException() { void testAwsRegion_throwsException_onInvalidEndpoint() { final String yaml = String.join("\n", "endpoint: \"invalid-endpoint\"", + "aws: {}", "max_retries: 5" ); diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java index 06460c2953..506802b134 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java @@ -19,6 +19,8 @@ import org.opensearch.dataprepper.model.event.EventHandle; import org.opensearch.dataprepper.plugins.sink.otlp.configuration.OtlpSinkConfig; import org.opensearch.dataprepper.plugins.sink.otlp.metrics.OtlpSinkMetrics; +import software.amazon.awssdk.http.SdkHttpFullRequest; +import software.amazon.awssdk.http.SdkHttpMethod; import software.amazon.awssdk.utils.Pair; import java.lang.reflect.Method; @@ -77,8 +79,8 @@ void setup() { when(gzipCompressor.apply(any())).thenReturn(PAYLOAD); when(signer.signRequest(any())).thenReturn( - software.amazon.awssdk.http.SdkHttpFullRequest.builder() - .method(software.amazon.awssdk.http.SdkHttpMethod.POST) + SdkHttpFullRequest.builder() + .method(SdkHttpMethod.POST) .uri(URI.create("https://localhost/v1/traces")) .putHeader("Authorization", "sig") .build() @@ -269,6 +271,7 @@ void testSend_exceptionDuringSendIncrementsRejected() { void testConstructor_withDefaultConfig() { final OtlpSinkConfig config = mock(OtlpSinkConfig.class); + when(config.getEndpoint()).thenReturn("https://localhost/v1/traces"); when(config.getMaxBatchSize()).thenReturn(1_000_000L); when(config.getMaxRetries()).thenReturn(2); when(config.getFlushTimeoutMillis()).thenReturn(5000L); @@ -283,6 +286,7 @@ void testConstructor_withMinimumThresholdConfig() { final OtlpSinkConfig config = mock(OtlpSinkConfig.class); // Set all threshold values to minimum valid input + when(config.getEndpoint()).thenReturn("https://localhost/v1/traces"); when(config.getMaxBatchSize()).thenReturn(0L); when(config.getMaxRetries()).thenReturn(0); when(config.getFlushTimeoutMillis()).thenReturn(1L); From bb4ffb0e0f997728b4193182dfe44a2e992a4b93 Mon Sep 17 00:00:00 2001 From: huy pham Date: Fri, 23 May 2025 16:23:27 -0700 Subject: [PATCH 21/23] Validate the aws config via the @AssertTrue pattern Signed-off-by: huy pham --- .../plugins/sink/otlp/OtlpSink.java | 2 -- .../otlp/configuration/OtlpSinkConfig.java | 10 ++++------ .../plugins/sink/otlp/OtlpSinkTest.java | 20 ------------------- .../configuration/OtlpSinkConfigTest.java | 11 +++++----- 4 files changed, 9 insertions(+), 34 deletions(-) diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java index fc7293c4b4..7045e882c3 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSink.java @@ -49,8 +49,6 @@ public class OtlpSink extends AbstractSink> { public OtlpSink(@Nonnull final AwsCredentialsSupplier awsCredentialsSupplier, @Nonnull final OtlpSinkConfig config, @Nonnull final PluginMetrics pluginMetrics, @Nonnull final PluginSetting pluginSetting) { super(pluginSetting); - config.validate(); - this.sinkMetrics = new OtlpSinkMetrics(pluginMetrics, pluginSetting); this.buffer = new OtlpSinkBuffer(awsCredentialsSupplier, config, sinkMetrics); } diff --git a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfig.java b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfig.java index 7b2e513725..3a7e698aa1 100644 --- a/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfig.java +++ b/data-prepper-plugins/otlp-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfig.java @@ -6,6 +6,7 @@ import com.fasterxml.jackson.annotation.JsonProperty; import jakarta.validation.Valid; +import jakarta.validation.constraints.AssertTrue; import jakarta.validation.constraints.Min; import jakarta.validation.constraints.NotBlank; import lombok.Getter; @@ -115,12 +116,9 @@ public String getStsExternalId() { /** * Validate the AWS configuration. * This method ensures breaking change in future release where non-AWS OTLP endpoints are supported. - * - * @throws IllegalArgumentException if the AWS configuration is invalid */ - public void validate() { - if (awsConfig == null) { - throw new IllegalArgumentException("aws configuration is required"); - } + @AssertTrue + boolean isAwsConfigValid() { + return awsConfig != null; } } diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java index 299589d42c..d6ed2da99a 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/OtlpSinkTest.java @@ -4,10 +4,8 @@ */ package org.opensearch.dataprepper.plugins.sink.otlp; -import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.function.Executable; import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.configuration.PluginSetting; @@ -23,7 +21,6 @@ import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoMoreInteractions; @@ -70,23 +67,6 @@ void testInitialize_startsBuffer() { verify(mockBuffer).start(); } - @Test - void testConstructor_throwsWhenAwsConfigIsMissing() { - doThrow(new IllegalArgumentException("aws configuration is required")) - .when(mockConfig).validate(); - - // Act & Assert - final Executable constructorCall = () -> - new OtlpSink(mockAwsCredSupplier, mockConfig, mockMetrics, mockSetting); - - final IllegalArgumentException thrown = Assertions.assertThrows( - IllegalArgumentException.class, - constructorCall - ); - - Assertions.assertEquals("aws configuration is required", thrown.getMessage()); - } - @Test void testOutput_addsEveryRecordToBuffer() { // Arrange diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfigTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfigTest.java index 100115c170..85363165b7 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfigTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/configuration/OtlpSinkConfigTest.java @@ -21,8 +21,8 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.nullValue; -import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -169,23 +169,22 @@ void testCustomThresholdAndRetries() throws Exception { } @Test - void testValidateAwsConfig_doesNotThrowWhenPresent() throws Exception { + void testIsAwsConfigValid_returnsTrue_whenPresent() throws Exception { final String yaml = String.join("\n", "endpoint: \"" + EXPECTED_ENDPOINT + "\"", "aws: {}" ); final OtlpSinkConfig config = mapper.readValue(yaml, OtlpSinkConfig.class); - assertDoesNotThrow(config::validate); + assertTrue(config.isAwsConfigValid()); } @Test - void testValidateAwsConfig_throwsException() throws Exception { + void testIsAwsConfigValid_returnsFalse_whenMissing() throws Exception { final String yaml = String.join("\n", "endpoint: \"" + EXPECTED_ENDPOINT + "\"" ); final OtlpSinkConfig config = mapper.readValue(yaml, OtlpSinkConfig.class); - - assertThrows(IllegalArgumentException.class, config::validate); + assertFalse(config.isAwsConfigValid()); } @Test From 01851ea9104803b3af72f2b9d097d9ddfd47af5f Mon Sep 17 00:00:00 2001 From: huy pham Date: Fri, 30 May 2025 07:50:45 -0700 Subject: [PATCH 22/23] Add missing unit test to reach 100% Signed-off-by: huy pham --- .../sink/otlp/http/OtlpHttpSenderTest.java | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java index 506802b134..f50a7c5f82 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/OtlpHttpSenderTest.java @@ -422,4 +422,20 @@ void testHandleSuccessfulResponse_privateMethod_parseException() throws Exceptio verify(mockEventHandle2).release(true); verify(mockEventHandle3).release(true); } + + @Test + void testHandleResponse_withNullResponseBody_logsNoBody() throws Exception { + final Method method = OtlpHttpSender.class.getDeclaredMethod("handleResponse", int.class, byte[].class, List.class); + method.setAccessible(true); + + method.invoke(sender, 500, null, testBatch); + + verify(metrics).recordResponseCode(500); + verify(metrics).incrementRejectedSpansCount(SPANS_COUNT); + + // Verifying event handles released with success=false + verify(mockEventHandle1).release(false); + verify(mockEventHandle2).release(false); + verify(mockEventHandle3).release(false); + } } \ No newline at end of file From 8527a3858e8dd4d0d59dae1183b3561b6251e1e5 Mon Sep 17 00:00:00 2001 From: huy pham Date: Fri, 30 May 2025 11:05:22 -0700 Subject: [PATCH 23/23] Mock AWS credentials provider to allow tests to run without real credentials Signed-off-by: huy pham --- .../plugins/sink/otlp/http/SigV4SignerTest.java | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4SignerTest.java b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4SignerTest.java index 3218574377..45db7131fd 100644 --- a/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4SignerTest.java +++ b/data-prepper-plugins/otlp-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/otlp/http/SigV4SignerTest.java @@ -8,7 +8,8 @@ import org.junit.jupiter.api.Test; import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier; import org.opensearch.dataprepper.plugins.sink.otlp.configuration.OtlpSinkConfig; -import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider; +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; +import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; import software.amazon.awssdk.http.SdkHttpFullRequest; import software.amazon.awssdk.regions.Region; @@ -36,7 +37,10 @@ void setup() { mockSupplier = mock(AwsCredentialsSupplier.class); when(mockConfig.getAwsRegion()).thenReturn(REGION); - when(mockSupplier.getProvider(any())).thenReturn(DefaultCredentialsProvider.create()); + + final AwsBasicCredentials mockCredentials = AwsBasicCredentials.create("mockAccessKey", "mockSecretKey"); + final StaticCredentialsProvider mockCredentialsProvider = StaticCredentialsProvider.create(mockCredentials); + when(mockSupplier.getProvider(any())).thenReturn(mockCredentialsProvider); } @Test