Add incubator ComposableAnnotatingSampler (#7804)

Author: anuraaga

sdk-extensions/incubator/src/main/java/io/opentelemetry/sdk/extension/incubator/trace/samplers/ComposableAnnotatingSampler.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.sdk.extension.incubator.trace.samplers;
+
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.sdk.trace.data.LinkData;
+import java.util.List;
+
+final class ComposableAnnotatingSampler implements ComposableSampler {
+  private final ComposableSampler delegate;
+  private final Attributes attributes;
+  private final String description;
+
+  ComposableAnnotatingSampler(ComposableSampler delegate, Attributes attributes) {
+    this.delegate = delegate;
+    this.attributes = attributes;
+
+    this.description =
+        "ComposableAnnotatingSampler{" + delegate.getDescription() + "," + attributes + "}";
+  }
+
+  @Override
+  public SamplingIntent getSamplingIntent(
+      Context parentContext,
+      String traceId,
+      String name,
+      SpanKind spanKind,
+      Attributes attributes,
+      List<LinkData> parentLinks) {
+    SamplingIntent intent =
+        delegate.getSamplingIntent(parentContext, traceId, name, spanKind, attributes, parentLinks);
+    return SamplingIntent.create(
+        intent.getThreshold(),
+        intent.isThresholdReliable(),
+        intent.getAttributes().toBuilder().putAll(this.attributes).build(),
+        intent.getTraceStateUpdater());
+  }
+
+  @Override
+  public String getDescription() {
+    return description;
+  }
+}

sdk-extensions/incubator/src/main/java/io/opentelemetry/sdk/extension/incubator/trace/samplers/ComposableSampler.java

@@ -45,6 +45,14 @@ static ComposableRuleBasedSamplerBuilder ruleBasedBuilder() {
     return new ComposableRuleBasedSamplerBuilder();
   }
 
+  /**
+   * Returns a {@link ComposableSampler} that adds the given {@link Attributes} to all sampled
+   * spans.
+   */
+  static ComposableSampler annotating(ComposableSampler sampler, Attributes attributes) {
+    return new ComposableAnnotatingSampler(sampler, attributes);
+  }
+
   /** Returns the {@link SamplingIntent} to use to make a sampling decision. */
   SamplingIntent getSamplingIntent(
       Context parentContext,

sdk-extensions/incubator/src/test/java/io/opentelemetry/sdk/extension/incubator/trace/samplers/ComposableAnnotatingSamplerTest.java

@@ -0,0 +1,85 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.sdk.extension.incubator.trace.samplers;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.api.trace.TraceId;
+import io.opentelemetry.context.Context;
+import java.util.Collections;
+import org.junit.jupiter.api.Test;
+
+class ComposableAnnotatingSamplerTest {
+
+  private static final Attributes ATTRIBUTES =
+      Attributes.of(
+          AttributeKey.stringKey("http.route"), "/bear", AttributeKey.longKey("size"), 100L);
+
+  @Test
+  void testDescription() {
+    assertThat(
+            ComposableSampler.annotating(ComposableSampler.alwaysOn(), Attributes.empty())
+                .getDescription())
+        .isEqualTo("ComposableAnnotatingSampler{ComposableAlwaysOnSampler,{}}");
+    assertThat(
+            ComposableSampler.annotating(ComposableSampler.alwaysOn(), ATTRIBUTES).getDescription())
+        .isEqualTo(
+            "ComposableAnnotatingSampler{ComposableAlwaysOnSampler,{http.route=\"/bear\", size=100}}");
+  }
+
+  @Test
+  void setsAttributes() {
+    SamplingIntent intent =
+        ComposableSampler.annotating(ComposableSampler.alwaysOn(), ATTRIBUTES)
+            .getSamplingIntent(
+                Context.root(),
+                TraceId.getInvalid(),
+                "span",
+                SpanKind.SERVER,
+                Attributes.empty(),
+                Collections.emptyList());
+    assertThat(intent.getThreshold()).isEqualTo(0);
+    assertThat(intent.getAttributes()).isEqualTo(ATTRIBUTES);
+  }
+
+  @Test
+  void mergesAttributes() {
+    // Easiest way to have a SamplingIntent with attributes is to use annotating sampler itself.
+    // This effectively creates a "test coverage dependency" where we know it's ok since we verify
+    // the base case in setsAttributes().
+    ComposableSampler baseSampler =
+        ComposableSampler.annotating(
+            ComposableSampler.alwaysOn(),
+            Attributes.of(
+                AttributeKey.stringKey("http.route"),
+                "/cat",
+                AttributeKey.stringKey("type"),
+                "mammal"));
+
+    SamplingIntent intent =
+        ComposableSampler.annotating(baseSampler, ATTRIBUTES)
+            .getSamplingIntent(
+                Context.root(),
+                TraceId.getInvalid(),
+                "span",
+                SpanKind.SERVER,
+                Attributes.empty(),
+                Collections.emptyList());
+    assertThat(intent.getThreshold()).isEqualTo(0);
+    assertThat(intent.getAttributes())
+        .isEqualTo(
+            Attributes.of(
+                AttributeKey.stringKey("http.route"),
+                "/bear",
+                AttributeKey.longKey("size"),
+                100L,
+                AttributeKey.stringKey("type"),
+                "mammal"));
+  }
+}