Refactor AwsMetricAttributesSpanExporter to an onEnding processor

Author: majanjua-amzn

appsignals-tests/contract-tests/src/test/java/software/amazon/opentelemetry/appsignals/test/httpservers/base/BaseHttpServerTest.java

@@ -52,7 +52,7 @@
 public abstract class BaseHttpServerTest extends ContractTestBase {
 
   /**
-   * Assert span attributes inserted by the AwsMetricAttributesSpanExporter
+   * Assert span attributes inserted by the AwsAttributeGeneratingSpanProcessor
    *
    * @param resourceScopeSpans list of spans that were exported by the application
    * @param method the http method that was used (GET, PUT, DELETE...)

awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsApplicationSignalsCustomizerProvider.java

@@ -76,7 +76,7 @@
  *   <li>Add SpanMetricsProcessor to create RED metrics.
  *   <li>Add AttributePropagatingSpanProcessor to propagate span attributes from parent to child
  *       spans.
- *   <li>Add AwsMetricAttributesSpanExporter to add more attributes to all spans.
+ *   <li>Add AwsAttributeGeneratingSpanProcessor to add more attributes to all spans.
  * </ul>
  *
  * <p>You can control when these customizations are applied using the property
@@ -337,6 +337,9 @@ private SdkTracerProviderBuilder customizeTracerProviderBuilder(
       // Construct and set local and remote attributes span processor
       tracerProviderBuilder.addSpanProcessor(
           AttributePropagatingSpanProcessorBuilder.create().build());
+      // Construct and set AWS Attribute Generating Span Processor
+      tracerProviderBuilder.addSpanProcessor(
+          AwsAttributeGeneratingSpanProcessorBuilder.create(ResourceHolder.getResource()).build());
 
       // If running on Lambda, we just need to export 100% spans and skip generating any Application
       // Signals metrics.
@@ -351,16 +354,9 @@ private SdkTracerProviderBuilder customizeTracerProviderBuilder(
                 .setEndpoint(tracesEndpoint)
                 .build();
 
-        // Wrap the udp exporter with the AwsMetricsAttributesSpanExporter to add Application
-        // Signals attributes to unsampled spans too
-        SpanExporter appSignalsSpanExporter =
-            AwsMetricAttributesSpanExporterBuilder.create(
-                    spanExporter, ResourceHolder.getResource())
-                .build();
-
         tracerProviderBuilder.addSpanProcessor(
             AwsUnsampledOnlySpanProcessorBuilder.create()
-                .setSpanExporter(appSignalsSpanExporter)
+                .setSpanExporter(spanExporter)
                 .setMaxExportBatchSize(LAMBDA_SPAN_EXPORT_BATCH_SIZE)
                 .build());
         return tracerProviderBuilder;
@@ -461,12 +457,6 @@ SpanExporter customizeSpanExporter(SpanExporter spanExporter, ConfigProperties c
       }
     }
 
-    if (isApplicationSignalsEnabled(configProps)) {
-      spanExporter =
-          AwsMetricAttributesSpanExporterBuilder.create(spanExporter, ResourceHolder.getResource())
-              .build();
-    }
-
     if (this.sampler instanceof AwsXrayRemoteSampler) {
       ((AwsXrayRemoteSampler) this.sampler).setSpanExporter(spanExporter);
     }

awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsAttributeGeneratingSpanProcessor.java

@@ -0,0 +1,95 @@
+/*
+ * 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 software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_SPAN_KIND;
+
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.sdk.resources.Resource;
+import io.opentelemetry.sdk.trace.ReadWriteSpan;
+import io.opentelemetry.sdk.trace.ReadableSpan;
+import io.opentelemetry.sdk.trace.data.SpanData;
+import io.opentelemetry.sdk.trace.internal.ExtendedSpanProcessor;
+import java.util.Map;
+
+public class AwsAttributeGeneratingSpanProcessor implements ExtendedSpanProcessor {
+
+  private final MetricAttributeGenerator generator;
+  private final Resource resource;
+
+  public static AwsAttributeGeneratingSpanProcessor create(
+      MetricAttributeGenerator generator, Resource resource) {
+    return new AwsAttributeGeneratingSpanProcessor(generator, resource);
+  }
+
+  private AwsAttributeGeneratingSpanProcessor(
+      MetricAttributeGenerator generator, Resource resource) {
+    this.generator = generator;
+    this.resource = resource;
+  }
+
+  @Override
+  public void onStart(Context parentContext, ReadWriteSpan span) {}
+
+  @Override
+  public boolean isStartRequired() {
+    return false;
+  }
+
+  @Override
+  public void onEnding(ReadWriteSpan span) {
+    // If the map has no items, no modifications are required. If there is one item, it means the
+    // span either produces Service or Dependency metric attributes, and in either case we want to
+    // modify the span with them. If there are two items, the span produces both Service and
+    // Dependency metric attributes indicating the span is a local dependency root. The Service
+    // Attributes must be a subset of the Dependency, with the exception of AWS_SPAN_KIND. The
+    // knowledge that the span is a local root is more important than knowing that it is a
+    // Dependency metric, so we take all the Dependency metrics but replace AWS_SPAN_KIND with
+    // LOCAL_ROOT.
+    SpanData spanData = span.toSpanData();
+    Map<String, Attributes> attributeMap =
+        generator.generateMetricAttributeMapFromSpan(span.toSpanData(), resource);
+
+    boolean generatesServiceMetrics =
+        AwsSpanProcessingUtil.shouldGenerateServiceMetricAttributes(spanData);
+    boolean generatesDependencyMetrics =
+        AwsSpanProcessingUtil.shouldGenerateDependencyMetricAttributes(spanData);
+
+    if (generatesServiceMetrics && generatesDependencyMetrics) {
+      // Order matters: dependency metric attributes include AWS_SPAN_KIND key
+      span.setAllAttributes(attributeMap.get(MetricAttributeGenerator.DEPENDENCY_METRIC));
+      span.setAttribute(AWS_SPAN_KIND, AwsSpanProcessingUtil.LOCAL_ROOT);
+    } else if (generatesServiceMetrics) {
+      span.setAllAttributes(attributeMap.get(MetricAttributeGenerator.SERVICE_METRIC));
+    } else if (generatesDependencyMetrics) {
+      span.setAllAttributes(attributeMap.get(MetricAttributeGenerator.DEPENDENCY_METRIC));
+    }
+  }
+
+  @Override
+  public boolean isOnEndingRequired() {
+    return true;
+  }
+
+  @Override
+  public void onEnd(ReadableSpan span) {}
+
+  @Override
+  public boolean isEndRequired() {
+    return false;
+  }
+}

awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsMetricAttributesSpanExporterBuilder.java renamed to awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsAttributeGeneratingSpanProcessorBuilder.java

@@ -19,43 +19,40 @@
 
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import io.opentelemetry.sdk.resources.Resource;
-import io.opentelemetry.sdk.trace.export.SpanExporter;
 
-public class AwsMetricAttributesSpanExporterBuilder {
+public class AwsAttributeGeneratingSpanProcessorBuilder {
 
   // Defaults
   private static final MetricAttributeGenerator DEFAULT_GENERATOR =
       new AwsMetricAttributeGenerator();
 
   // Required builder elements
-  private final SpanExporter delegate;
   private final Resource resource;
 
   // Optional builder elements
   private MetricAttributeGenerator generator = DEFAULT_GENERATOR;
 
-  public static AwsMetricAttributesSpanExporterBuilder create(
-      SpanExporter delegate, Resource resource) {
-    return new AwsMetricAttributesSpanExporterBuilder(delegate, resource);
+  public static AwsAttributeGeneratingSpanProcessorBuilder create(Resource resource) {
+    return new AwsAttributeGeneratingSpanProcessorBuilder(resource);
   }
 
-  private AwsMetricAttributesSpanExporterBuilder(SpanExporter delegate, Resource resource) {
-    this.delegate = delegate;
+  private AwsAttributeGeneratingSpanProcessorBuilder(Resource resource) {
     this.resource = resource;
   }
 
   /**
-   * Sets the generator used to generate attributes used spancs exported by the exporter. If unset,
+   * Sets the generator used to generate attributes added to spans in the processor. If unset,
    * defaults to {@link #DEFAULT_GENERATOR}. Must not be null.
    */
   @CanIgnoreReturnValue
-  public AwsMetricAttributesSpanExporterBuilder setGenerator(MetricAttributeGenerator generator) {
+  public AwsAttributeGeneratingSpanProcessorBuilder setGenerator(
+      MetricAttributeGenerator generator) {
     requireNonNull(generator, "generator");
     this.generator = generator;
     return this;
   }
 
-  public AwsMetricAttributesSpanExporter build() {
-    return AwsMetricAttributesSpanExporter.create(delegate, generator, resource);
+  public AwsAttributeGeneratingSpanProcessor build() {
+    return AwsAttributeGeneratingSpanProcessor.create(generator, resource);
   }
 }

awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsMetricAttributesSpanExporter.java

@@ -23,7 +23,7 @@
 /**
  * Metric attribute generator defines an interface for classes that can generate specific attributes
  * to be used by an {@link AwsSpanMetricsProcessor} to produce metrics and by {@link
- * AwsMetricAttributesSpanExporter} to wrap the original span.
+ * AwsAttributeGeneratingSpanProcessor} to update the original span.
  */
 public interface MetricAttributeGenerator {
   static final String SERVICE_METRIC = "Service";