From cc7af45f0a0c5305793919cdf1cb01e308bae708 Mon Sep 17 00:00:00 2001 From: Prashant Srivastava Date: Thu, 7 Nov 2024 13:57:30 -0800 Subject: [PATCH 1/8] adding span processor for exporting only unsampled spans --- .../javaagent/providers/AwsAttributeKeys.java | 3 + .../AwsUnsampledOnlySpanProcessor.java | 82 ++++++++++++ .../AwsUnsampledOnlySpanProcessorBuilder.java | 51 ++++++++ .../AwsUnsampledOnlySpanProcessorTest.java | 118 ++++++++++++++++++ 4 files changed, 254 insertions(+) create mode 100644 awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessor.java create mode 100644 awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorBuilder.java create mode 100644 awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorTest.java diff --git a/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsAttributeKeys.java b/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsAttributeKeys.java index b73794b9db..f9791a31ee 100644 --- a/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsAttributeKeys.java +++ b/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsAttributeKeys.java @@ -70,6 +70,9 @@ private AwsAttributeKeys() {} static final AttributeKey AWS_LAMBDA_RESOURCE_ID = AttributeKey.stringKey("aws.lambda.resource_mapping.id"); + static final AttributeKey AWS_TRACE_FLAG_SAMPLED = + AttributeKey.booleanKey("aws.trace.flag.sampled"); + // use the same AWS Resource attribute name defined by OTel java auto-instr for aws_sdk_v_1_1 // TODO: all AWS specific attributes should be defined in semconv package and reused cross all // otel packages. Related sim - diff --git a/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessor.java b/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessor.java new file mode 100644 index 0000000000..fa6b3daca6 --- /dev/null +++ b/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessor.java @@ -0,0 +1,82 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.opentelemetry.javaagent.providers; + +import io.opentelemetry.context.Context; +import io.opentelemetry.sdk.common.CompletableResultCode; +import io.opentelemetry.sdk.trace.ReadWriteSpan; +import io.opentelemetry.sdk.trace.ReadableSpan; +import io.opentelemetry.sdk.trace.SpanProcessor; + +final class AwsUnsampledOnlySpanProcessor implements SpanProcessor { + + private final SpanProcessor delegate; + + AwsUnsampledOnlySpanProcessor(SpanProcessor delegate) { + this.delegate = delegate; + } + + public static AwsUnsampledOnlySpanProcessorBuilder builder() { + return new AwsUnsampledOnlySpanProcessorBuilder(); + } + + @Override + public void onStart(Context parentContext, ReadWriteSpan span) { + if (span.getSpanContext().isSampled()) { + span.setAttribute(AwsAttributeKeys.AWS_TRACE_FLAG_SAMPLED, true); + } else { + span.setAttribute(AwsAttributeKeys.AWS_TRACE_FLAG_SAMPLED, false); + } + delegate.onStart(parentContext, span); + } + + @Override + public void onEnd(ReadableSpan span) { + if (!span.getSpanContext().isSampled()) { + delegate.onEnd(span); + } + } + + @Override + public boolean isStartRequired() { + return delegate.isStartRequired(); + } + + @Override + public boolean isEndRequired() { + return delegate.isEndRequired(); + } + + @Override + public CompletableResultCode shutdown() { + return delegate.shutdown(); + } + + @Override + public CompletableResultCode forceFlush() { + return delegate.forceFlush(); + } + + @Override + public void close() { + delegate.close(); + } + + // Visible for testing + SpanProcessor getDelegate() { + return delegate; + } +} diff --git a/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorBuilder.java b/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorBuilder.java new file mode 100644 index 0000000000..854c7b4824 --- /dev/null +++ b/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorBuilder.java @@ -0,0 +1,51 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.opentelemetry.javaagent.providers; + +import io.opentelemetry.sdk.trace.export.BatchSpanProcessor; +import io.opentelemetry.sdk.trace.export.BatchSpanProcessorBuilder; +import io.opentelemetry.sdk.trace.export.SpanExporter; + +import static java.util.Objects.requireNonNull; + +final class AwsUnsampledOnlySpanProcessorBuilder { + + // Default exporter is OtlpUdpSpanExporter with unsampled payload prefix + private SpanExporter exporter = new OtlpUdpSpanExporterBuilder() + .setPayloadSampleDecision(TracePayloadSampleDecision.UNSAMPLED) + .build(); + + public AwsUnsampledOnlySpanProcessorBuilder setSpanExporter(SpanExporter exporter) { + requireNonNull(exporter, "exporter cannot be null"); + this.exporter = exporter; + return this; + } + + public AwsUnsampledOnlySpanProcessorBuilder setMaxQueueSize(int maxQueueSize) { + + return this; + } + + public AwsUnsampledOnlySpanProcessor build() { + BatchSpanProcessor bsp = + BatchSpanProcessor.builder(exporter).setExportUnsampledSpans(true).build(); + return new AwsUnsampledOnlySpanProcessor(bsp); + } + + SpanExporter getSpanExporter() { + return exporter; + } +} diff --git a/awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorTest.java b/awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorTest.java new file mode 100644 index 0000000000..d34bde5e46 --- /dev/null +++ b/awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorTest.java @@ -0,0 +1,118 @@ +package software.amazon.opentelemetry.javaagent.providers; + +import io.opentelemetry.api.trace.SpanContext; +import io.opentelemetry.context.Context; +import io.opentelemetry.sdk.common.CompletableResultCode; +import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; +import io.opentelemetry.sdk.trace.ReadWriteSpan; +import io.opentelemetry.sdk.trace.ReadableSpan; +import io.opentelemetry.sdk.trace.SpanProcessor; +import io.opentelemetry.sdk.trace.data.SpanData; +import io.opentelemetry.sdk.trace.export.BatchSpanProcessor; +import io.opentelemetry.sdk.trace.export.SpanExporter; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; + +import java.util.Collection; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.*; + +public class AwsUnsampledOnlySpanProcessorTest { + + @Test + public void testDefaultSpanProcessor() { + AwsUnsampledOnlySpanProcessorBuilder builder = AwsUnsampledOnlySpanProcessor.builder(); + AwsUnsampledOnlySpanProcessor unsampledSP = builder.build(); + + assertThat(builder.getSpanExporter()).isInstanceOf(OtlpUdpSpanExporter.class); + SpanProcessor delegate = unsampledSP.getDelegate(); + assertThat(delegate).isInstanceOf(BatchSpanProcessor.class); + BatchSpanProcessor delegateBsp = (BatchSpanProcessor) delegate; + String delegateBspString = delegateBsp.toString(); + assertThat(delegateBspString).contains("spanExporter=software.amazon.opentelemetry.javaagent.providers.OtlpUdpSpanExporter"); + assertThat(delegateBspString).contains("exportUnsampledSpans=true"); + } + + @Test + public void testSpanProcessorWithExporter() { + AwsUnsampledOnlySpanProcessorBuilder builder = AwsUnsampledOnlySpanProcessor + .builder() + .setSpanExporter(InMemorySpanExporter.create()); + AwsUnsampledOnlySpanProcessor unsampledSP = builder.build(); + + assertThat(builder.getSpanExporter()).isInstanceOf(InMemorySpanExporter.class); + SpanProcessor delegate = unsampledSP.getDelegate(); + assertThat(delegate).isInstanceOf(BatchSpanProcessor.class); + BatchSpanProcessor delegateBsp = (BatchSpanProcessor) delegate; + String delegateBspString = delegateBsp.toString(); + assertThat(delegateBspString).contains("spanExporter=io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter"); + assertThat(delegateBspString).contains("exportUnsampledSpans=true"); + } + + @Test + public void testStartAddsAttributeToSampledSpan() { + SpanContext mockSpanContext = mock(SpanContext.class); + when(mockSpanContext.isSampled()).thenReturn(true); + Context parentContextMock = mock(Context.class); + ReadWriteSpan spanMock = mock(ReadWriteSpan.class); + when(spanMock.getSpanContext()).thenReturn(mockSpanContext); + + AwsUnsampledOnlySpanProcessor processor = AwsUnsampledOnlySpanProcessor.builder().build(); + processor.onStart(parentContextMock, spanMock); + + //verify setAttribute was called with the correct arguments + verify(spanMock, times(1)).setAttribute(AwsAttributeKeys.AWS_TRACE_FLAG_SAMPLED, true); + } + + @Test + public void testStartAddsAttributeToUnsampledSpan() { + SpanContext mockSpanContext = mock(SpanContext.class); + when(mockSpanContext.isSampled()).thenReturn(false); + Context parentContextMock = mock(Context.class); + ReadWriteSpan spanMock = mock(ReadWriteSpan.class); + when(spanMock.getSpanContext()).thenReturn(mockSpanContext); + + AwsUnsampledOnlySpanProcessor processor = AwsUnsampledOnlySpanProcessor.builder().build(); + processor.onStart(parentContextMock, spanMock); + + //verify setAttribute was called with the correct arguments + verify(spanMock, times(1)).setAttribute(AwsAttributeKeys.AWS_TRACE_FLAG_SAMPLED, false); + } + + @Test + public void testExportsOnlyUnsampledSpans() { + SpanExporter mockExporter = mock(SpanExporter.class); + when(mockExporter.export(anyCollection())).thenReturn(CompletableResultCode.ofSuccess()); + + BatchSpanProcessor delegate = BatchSpanProcessor.builder(mockExporter) + .setExportUnsampledSpans(true) + .setMaxExportBatchSize(1) + .setMaxQueueSize(1) + .build(); + + AwsUnsampledOnlySpanProcessor processor = new AwsUnsampledOnlySpanProcessor(delegate); + + // unsampled span + SpanContext mockSpanContextUnsampled = mock(SpanContext.class); + when(mockSpanContextUnsampled.isSampled()).thenReturn(false); + ReadableSpan mockSpanUnsampled = mock(ReadableSpan.class); + when(mockSpanUnsampled.getSpanContext()).thenReturn(mockSpanContextUnsampled); + + // sampled span + SpanContext mockSpanContextSampled = mock(SpanContext.class); + when(mockSpanContextSampled.isSampled()).thenReturn(true); + ReadableSpan mockSpanSampled = mock(ReadableSpan.class); + when(mockSpanSampled.getSpanContext()).thenReturn(mockSpanContextSampled); + + // flush the unsampled span and verify export was called once + processor.onEnd(mockSpanUnsampled); + processor.forceFlush(); + verify(mockExporter, times(1)).export(anyCollection()); + + // flush the sampled span and verify export was not called again + processor.onEnd(mockSpanSampled); + processor.forceFlush(); + verify(mockExporter, times(1)).export(anyCollection()); + } +} From 56bf70ce3f408fa17ea27f621be53944295e7fcf Mon Sep 17 00:00:00 2001 From: Prashant Srivastava Date: Thu, 7 Nov 2024 14:13:04 -0800 Subject: [PATCH 2/8] spotless apply --- .../AwsUnsampledOnlySpanProcessorBuilder.java | 10 +- .../AwsUnsampledOnlySpanProcessorTest.java | 215 +++++++++--------- 2 files changed, 118 insertions(+), 107 deletions(-) diff --git a/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorBuilder.java b/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorBuilder.java index 854c7b4824..c9f01e58bb 100644 --- a/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorBuilder.java +++ b/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorBuilder.java @@ -15,16 +15,16 @@ package software.amazon.opentelemetry.javaagent.providers; +import static java.util.Objects.requireNonNull; + import io.opentelemetry.sdk.trace.export.BatchSpanProcessor; -import io.opentelemetry.sdk.trace.export.BatchSpanProcessorBuilder; import io.opentelemetry.sdk.trace.export.SpanExporter; -import static java.util.Objects.requireNonNull; - final class AwsUnsampledOnlySpanProcessorBuilder { // Default exporter is OtlpUdpSpanExporter with unsampled payload prefix - private SpanExporter exporter = new OtlpUdpSpanExporterBuilder() + private SpanExporter exporter = + new OtlpUdpSpanExporterBuilder() .setPayloadSampleDecision(TracePayloadSampleDecision.UNSAMPLED) .build(); @@ -40,7 +40,7 @@ public AwsUnsampledOnlySpanProcessorBuilder setMaxQueueSize(int maxQueueSize) { } public AwsUnsampledOnlySpanProcessor build() { - BatchSpanProcessor bsp = + BatchSpanProcessor bsp = BatchSpanProcessor.builder(exporter).setExportUnsampledSpans(true).build(); return new AwsUnsampledOnlySpanProcessor(bsp); } diff --git a/awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorTest.java b/awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorTest.java index d34bde5e46..60de3f31f7 100644 --- a/awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorTest.java +++ b/awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorTest.java @@ -1,5 +1,23 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + package software.amazon.opentelemetry.javaagent.providers; +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.*; + import io.opentelemetry.api.trace.SpanContext; import io.opentelemetry.context.Context; import io.opentelemetry.sdk.common.CompletableResultCode; @@ -7,112 +25,105 @@ import io.opentelemetry.sdk.trace.ReadWriteSpan; import io.opentelemetry.sdk.trace.ReadableSpan; import io.opentelemetry.sdk.trace.SpanProcessor; -import io.opentelemetry.sdk.trace.data.SpanData; import io.opentelemetry.sdk.trace.export.BatchSpanProcessor; import io.opentelemetry.sdk.trace.export.SpanExporter; import org.junit.jupiter.api.Test; -import org.mockito.ArgumentCaptor; - -import java.util.Collection; - -import static org.assertj.core.api.Assertions.assertThat; -import static org.mockito.Mockito.*; public class AwsUnsampledOnlySpanProcessorTest { - @Test - public void testDefaultSpanProcessor() { - AwsUnsampledOnlySpanProcessorBuilder builder = AwsUnsampledOnlySpanProcessor.builder(); - AwsUnsampledOnlySpanProcessor unsampledSP = builder.build(); - - assertThat(builder.getSpanExporter()).isInstanceOf(OtlpUdpSpanExporter.class); - SpanProcessor delegate = unsampledSP.getDelegate(); - assertThat(delegate).isInstanceOf(BatchSpanProcessor.class); - BatchSpanProcessor delegateBsp = (BatchSpanProcessor) delegate; - String delegateBspString = delegateBsp.toString(); - assertThat(delegateBspString).contains("spanExporter=software.amazon.opentelemetry.javaagent.providers.OtlpUdpSpanExporter"); - assertThat(delegateBspString).contains("exportUnsampledSpans=true"); - } - - @Test - public void testSpanProcessorWithExporter() { - AwsUnsampledOnlySpanProcessorBuilder builder = AwsUnsampledOnlySpanProcessor - .builder() - .setSpanExporter(InMemorySpanExporter.create()); - AwsUnsampledOnlySpanProcessor unsampledSP = builder.build(); - - assertThat(builder.getSpanExporter()).isInstanceOf(InMemorySpanExporter.class); - SpanProcessor delegate = unsampledSP.getDelegate(); - assertThat(delegate).isInstanceOf(BatchSpanProcessor.class); - BatchSpanProcessor delegateBsp = (BatchSpanProcessor) delegate; - String delegateBspString = delegateBsp.toString(); - assertThat(delegateBspString).contains("spanExporter=io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter"); - assertThat(delegateBspString).contains("exportUnsampledSpans=true"); - } - - @Test - public void testStartAddsAttributeToSampledSpan() { - SpanContext mockSpanContext = mock(SpanContext.class); - when(mockSpanContext.isSampled()).thenReturn(true); - Context parentContextMock = mock(Context.class); - ReadWriteSpan spanMock = mock(ReadWriteSpan.class); - when(spanMock.getSpanContext()).thenReturn(mockSpanContext); - - AwsUnsampledOnlySpanProcessor processor = AwsUnsampledOnlySpanProcessor.builder().build(); - processor.onStart(parentContextMock, spanMock); - - //verify setAttribute was called with the correct arguments - verify(spanMock, times(1)).setAttribute(AwsAttributeKeys.AWS_TRACE_FLAG_SAMPLED, true); - } - - @Test - public void testStartAddsAttributeToUnsampledSpan() { - SpanContext mockSpanContext = mock(SpanContext.class); - when(mockSpanContext.isSampled()).thenReturn(false); - Context parentContextMock = mock(Context.class); - ReadWriteSpan spanMock = mock(ReadWriteSpan.class); - when(spanMock.getSpanContext()).thenReturn(mockSpanContext); - - AwsUnsampledOnlySpanProcessor processor = AwsUnsampledOnlySpanProcessor.builder().build(); - processor.onStart(parentContextMock, spanMock); - - //verify setAttribute was called with the correct arguments - verify(spanMock, times(1)).setAttribute(AwsAttributeKeys.AWS_TRACE_FLAG_SAMPLED, false); - } - - @Test - public void testExportsOnlyUnsampledSpans() { - SpanExporter mockExporter = mock(SpanExporter.class); - when(mockExporter.export(anyCollection())).thenReturn(CompletableResultCode.ofSuccess()); - - BatchSpanProcessor delegate = BatchSpanProcessor.builder(mockExporter) - .setExportUnsampledSpans(true) - .setMaxExportBatchSize(1) - .setMaxQueueSize(1) - .build(); - - AwsUnsampledOnlySpanProcessor processor = new AwsUnsampledOnlySpanProcessor(delegate); - - // unsampled span - SpanContext mockSpanContextUnsampled = mock(SpanContext.class); - when(mockSpanContextUnsampled.isSampled()).thenReturn(false); - ReadableSpan mockSpanUnsampled = mock(ReadableSpan.class); - when(mockSpanUnsampled.getSpanContext()).thenReturn(mockSpanContextUnsampled); - - // sampled span - SpanContext mockSpanContextSampled = mock(SpanContext.class); - when(mockSpanContextSampled.isSampled()).thenReturn(true); - ReadableSpan mockSpanSampled = mock(ReadableSpan.class); - when(mockSpanSampled.getSpanContext()).thenReturn(mockSpanContextSampled); - - // flush the unsampled span and verify export was called once - processor.onEnd(mockSpanUnsampled); - processor.forceFlush(); - verify(mockExporter, times(1)).export(anyCollection()); - - // flush the sampled span and verify export was not called again - processor.onEnd(mockSpanSampled); - processor.forceFlush(); - verify(mockExporter, times(1)).export(anyCollection()); - } + @Test + public void testDefaultSpanProcessor() { + AwsUnsampledOnlySpanProcessorBuilder builder = AwsUnsampledOnlySpanProcessor.builder(); + AwsUnsampledOnlySpanProcessor unsampledSP = builder.build(); + + assertThat(builder.getSpanExporter()).isInstanceOf(OtlpUdpSpanExporter.class); + SpanProcessor delegate = unsampledSP.getDelegate(); + assertThat(delegate).isInstanceOf(BatchSpanProcessor.class); + BatchSpanProcessor delegateBsp = (BatchSpanProcessor) delegate; + String delegateBspString = delegateBsp.toString(); + assertThat(delegateBspString) + .contains( + "spanExporter=software.amazon.opentelemetry.javaagent.providers.OtlpUdpSpanExporter"); + assertThat(delegateBspString).contains("exportUnsampledSpans=true"); + } + + @Test + public void testSpanProcessorWithExporter() { + AwsUnsampledOnlySpanProcessorBuilder builder = + AwsUnsampledOnlySpanProcessor.builder().setSpanExporter(InMemorySpanExporter.create()); + AwsUnsampledOnlySpanProcessor unsampledSP = builder.build(); + + assertThat(builder.getSpanExporter()).isInstanceOf(InMemorySpanExporter.class); + SpanProcessor delegate = unsampledSP.getDelegate(); + assertThat(delegate).isInstanceOf(BatchSpanProcessor.class); + BatchSpanProcessor delegateBsp = (BatchSpanProcessor) delegate; + String delegateBspString = delegateBsp.toString(); + assertThat(delegateBspString) + .contains("spanExporter=io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter"); + assertThat(delegateBspString).contains("exportUnsampledSpans=true"); + } + + @Test + public void testStartAddsAttributeToSampledSpan() { + SpanContext mockSpanContext = mock(SpanContext.class); + when(mockSpanContext.isSampled()).thenReturn(true); + Context parentContextMock = mock(Context.class); + ReadWriteSpan spanMock = mock(ReadWriteSpan.class); + when(spanMock.getSpanContext()).thenReturn(mockSpanContext); + + AwsUnsampledOnlySpanProcessor processor = AwsUnsampledOnlySpanProcessor.builder().build(); + processor.onStart(parentContextMock, spanMock); + + // verify setAttribute was called with the correct arguments + verify(spanMock, times(1)).setAttribute(AwsAttributeKeys.AWS_TRACE_FLAG_SAMPLED, true); + } + + @Test + public void testStartAddsAttributeToUnsampledSpan() { + SpanContext mockSpanContext = mock(SpanContext.class); + when(mockSpanContext.isSampled()).thenReturn(false); + Context parentContextMock = mock(Context.class); + ReadWriteSpan spanMock = mock(ReadWriteSpan.class); + when(spanMock.getSpanContext()).thenReturn(mockSpanContext); + + AwsUnsampledOnlySpanProcessor processor = AwsUnsampledOnlySpanProcessor.builder().build(); + processor.onStart(parentContextMock, spanMock); + + // verify setAttribute was called with the correct arguments + verify(spanMock, times(1)).setAttribute(AwsAttributeKeys.AWS_TRACE_FLAG_SAMPLED, false); + } + + @Test + public void testExportsOnlyUnsampledSpans() { + SpanExporter mockExporter = mock(SpanExporter.class); + when(mockExporter.export(anyCollection())).thenReturn(CompletableResultCode.ofSuccess()); + + AwsUnsampledOnlySpanProcessor processor = + AwsUnsampledOnlySpanProcessor.builder().setSpanExporter(mockExporter).build(); + + BatchSpanProcessor delegate = + BatchSpanProcessor.builder(mockExporter).setExportUnsampledSpans(true).build(); + + // unsampled span + SpanContext mockSpanContextUnsampled = mock(SpanContext.class); + when(mockSpanContextUnsampled.isSampled()).thenReturn(false); + ReadableSpan mockSpanUnsampled = mock(ReadableSpan.class); + when(mockSpanUnsampled.getSpanContext()).thenReturn(mockSpanContextUnsampled); + + // sampled span + SpanContext mockSpanContextSampled = mock(SpanContext.class); + when(mockSpanContextSampled.isSampled()).thenReturn(true); + ReadableSpan mockSpanSampled = mock(ReadableSpan.class); + when(mockSpanSampled.getSpanContext()).thenReturn(mockSpanContextSampled); + + // flush the unsampled span and verify export was called once + processor.onEnd(mockSpanUnsampled); + processor.forceFlush(); + verify(mockExporter, times(1)).export(anyCollection()); + + // flush the sampled span and verify export was not called again + processor.onEnd(mockSpanSampled); + processor.forceFlush(); + verify(mockExporter, times(1)).export(anyCollection()); + } } From d82031605980bc6626fe71ed443c69675da9b06d Mon Sep 17 00:00:00 2001 From: Prashant Srivastava Date: Thu, 7 Nov 2024 14:55:46 -0800 Subject: [PATCH 3/8] refactor unit test --- .../AwsUnsampledOnlySpanProcessorTest.java | 46 ++++++++++++++----- 1 file changed, 34 insertions(+), 12 deletions(-) diff --git a/awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorTest.java b/awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorTest.java index 60de3f31f7..d91511ab8a 100644 --- a/awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorTest.java +++ b/awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorTest.java @@ -27,6 +27,8 @@ import io.opentelemetry.sdk.trace.SpanProcessor; import io.opentelemetry.sdk.trace.export.BatchSpanProcessor; import io.opentelemetry.sdk.trace.export.SpanExporter; +import java.util.ArrayList; +import java.util.Collection; import org.junit.jupiter.api.Test; public class AwsUnsampledOnlySpanProcessorTest { @@ -98,11 +100,8 @@ public void testExportsOnlyUnsampledSpans() { SpanExporter mockExporter = mock(SpanExporter.class); when(mockExporter.export(anyCollection())).thenReturn(CompletableResultCode.ofSuccess()); - AwsUnsampledOnlySpanProcessor processor = - AwsUnsampledOnlySpanProcessor.builder().setSpanExporter(mockExporter).build(); - - BatchSpanProcessor delegate = - BatchSpanProcessor.builder(mockExporter).setExportUnsampledSpans(true).build(); + TestDelegateProcessor delegate = new TestDelegateProcessor(); + AwsUnsampledOnlySpanProcessor processor = new AwsUnsampledOnlySpanProcessor(delegate); // unsampled span SpanContext mockSpanContextUnsampled = mock(SpanContext.class); @@ -116,14 +115,37 @@ public void testExportsOnlyUnsampledSpans() { ReadableSpan mockSpanSampled = mock(ReadableSpan.class); when(mockSpanSampled.getSpanContext()).thenReturn(mockSpanContextSampled); - // flush the unsampled span and verify export was called once + processor.onEnd(mockSpanSampled); processor.onEnd(mockSpanUnsampled); - processor.forceFlush(); - verify(mockExporter, times(1)).export(anyCollection()); - // flush the sampled span and verify export was not called again - processor.onEnd(mockSpanSampled); - processor.forceFlush(); - verify(mockExporter, times(1)).export(anyCollection()); + // validate that only the unsampled span was delegated + assertThat(delegate.getEndedSpans()).containsExactly(mockSpanUnsampled); + } + + private static class TestDelegateProcessor implements SpanProcessor { + // keep a queue of Readable spans added when onEnd is called + Collection endedSpans = new ArrayList<>(); + + @Override + public void onStart(Context parentContext, ReadWriteSpan span) {} + + @Override + public boolean isStartRequired() { + return false; + } + + @Override + public void onEnd(ReadableSpan span) { + endedSpans.add(span); + } + + @Override + public boolean isEndRequired() { + return false; + } + + public Collection getEndedSpans() { + return endedSpans; + } } } From 11a454c9a3bb372721847e4d9479cc55b4064608 Mon Sep 17 00:00:00 2001 From: Prashant Srivastava Date: Thu, 7 Nov 2024 15:05:32 -0800 Subject: [PATCH 4/8] javadoc and some cleanup --- .../javaagent/providers/AwsUnsampledOnlySpanProcessor.java | 4 ++++ .../providers/AwsUnsampledOnlySpanProcessorBuilder.java | 5 ----- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessor.java b/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessor.java index fa6b3daca6..d5b4b9fac4 100644 --- a/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessor.java +++ b/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessor.java @@ -21,6 +21,10 @@ import io.opentelemetry.sdk.trace.ReadableSpan; import io.opentelemetry.sdk.trace.SpanProcessor; +/** + * {@link SpanProcessor} that only exports unsampled spans in a batch via a delegated @{link BatchSpanProcessor}. + * The processor also adds an attribute to each processed span to indicate that it was sampled or not. + */ final class AwsUnsampledOnlySpanProcessor implements SpanProcessor { private final SpanProcessor delegate; diff --git a/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorBuilder.java b/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorBuilder.java index c9f01e58bb..89efbcf3b4 100644 --- a/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorBuilder.java +++ b/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorBuilder.java @@ -34,11 +34,6 @@ public AwsUnsampledOnlySpanProcessorBuilder setSpanExporter(SpanExporter exporte return this; } - public AwsUnsampledOnlySpanProcessorBuilder setMaxQueueSize(int maxQueueSize) { - - return this; - } - public AwsUnsampledOnlySpanProcessor build() { BatchSpanProcessor bsp = BatchSpanProcessor.builder(exporter).setExportUnsampledSpans(true).build(); From 17fdd89fe5db49f26db18aac8bffa48c32528351 Mon Sep 17 00:00:00 2001 From: Prashant Srivastava Date: Thu, 7 Nov 2024 15:49:22 -0800 Subject: [PATCH 5/8] spotless apply --- .../javaagent/providers/AwsUnsampledOnlySpanProcessor.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessor.java b/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessor.java index d5b4b9fac4..3c4da1dacb 100644 --- a/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessor.java +++ b/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessor.java @@ -22,8 +22,9 @@ import io.opentelemetry.sdk.trace.SpanProcessor; /** - * {@link SpanProcessor} that only exports unsampled spans in a batch via a delegated @{link BatchSpanProcessor}. - * The processor also adds an attribute to each processed span to indicate that it was sampled or not. + * {@link SpanProcessor} that only exports unsampled spans in a batch via a delegated @{link + * BatchSpanProcessor}. The processor also adds an attribute to each processed span to indicate that + * it was sampled or not. */ final class AwsUnsampledOnlySpanProcessor implements SpanProcessor { From d236e90f55a41923ccc26c13984f1153f26c5fe5 Mon Sep 17 00:00:00 2001 From: Prashant Srivastava Date: Thu, 7 Nov 2024 16:15:19 -0800 Subject: [PATCH 6/8] bug fix --- .../javaagent/providers/AwsUnsampledOnlySpanProcessor.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessor.java b/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessor.java index 3c4da1dacb..e15bda66a1 100644 --- a/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessor.java +++ b/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessor.java @@ -57,12 +57,12 @@ public void onEnd(ReadableSpan span) { @Override public boolean isStartRequired() { - return delegate.isStartRequired(); + return true; } @Override public boolean isEndRequired() { - return delegate.isEndRequired(); + return true; } @Override From 3557f6aa1153c63e868d0bcb77b2d4a937f3946c Mon Sep 17 00:00:00 2001 From: Prashant Srivastava Date: Mon, 11 Nov 2024 08:57:53 -0800 Subject: [PATCH 7/8] more unit tests --- .../providers/AwsUnsampledOnlySpanProcessorTest.java | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorTest.java b/awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorTest.java index d91511ab8a..bdcc4873f2 100644 --- a/awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorTest.java +++ b/awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorTest.java @@ -33,6 +33,18 @@ public class AwsUnsampledOnlySpanProcessorTest { + @Test + public void testIsStartRequired() { + SpanProcessor processor = AwsUnsampledOnlySpanProcessor.builder().build(); + assertThat(processor.isStartRequired()).isTrue(); + } + + @Test + public void testIsEndRequired() { + SpanProcessor processor = AwsUnsampledOnlySpanProcessor.builder().build(); + assertThat(processor.isEndRequired()).isTrue(); + } + @Test public void testDefaultSpanProcessor() { AwsUnsampledOnlySpanProcessorBuilder builder = AwsUnsampledOnlySpanProcessor.builder(); From 43af1e542cdb92614465386dc896c24d59e0a009 Mon Sep 17 00:00:00 2001 From: Prashant Srivastava Date: Mon, 11 Nov 2024 22:58:28 -0800 Subject: [PATCH 8/8] populate aws.trace.flag.sampled attribute only if unsampled --- .../javaagent/providers/AwsUnsampledOnlySpanProcessor.java | 4 +--- .../providers/AwsUnsampledOnlySpanProcessorTest.java | 4 ++-- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessor.java b/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessor.java index e15bda66a1..3848016f3b 100644 --- a/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessor.java +++ b/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessor.java @@ -40,9 +40,7 @@ public static AwsUnsampledOnlySpanProcessorBuilder builder() { @Override public void onStart(Context parentContext, ReadWriteSpan span) { - if (span.getSpanContext().isSampled()) { - span.setAttribute(AwsAttributeKeys.AWS_TRACE_FLAG_SAMPLED, true); - } else { + if (!span.getSpanContext().isSampled()) { span.setAttribute(AwsAttributeKeys.AWS_TRACE_FLAG_SAMPLED, false); } delegate.onStart(parentContext, span); diff --git a/awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorTest.java b/awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorTest.java index bdcc4873f2..ba41740bd6 100644 --- a/awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorTest.java +++ b/awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/AwsUnsampledOnlySpanProcessorTest.java @@ -88,8 +88,8 @@ public void testStartAddsAttributeToSampledSpan() { AwsUnsampledOnlySpanProcessor processor = AwsUnsampledOnlySpanProcessor.builder().build(); processor.onStart(parentContextMock, spanMock); - // verify setAttribute was called with the correct arguments - verify(spanMock, times(1)).setAttribute(AwsAttributeKeys.AWS_TRACE_FLAG_SAMPLED, true); + // verify setAttribute was never called + verify(spanMock, never()).setAttribute(any(), anyBoolean()); } @Test