[demo] Add Bucket4j rate-limiter instrumentation

dd-java-agent/instrumentation/bucket4j-8.0/build.gradle

@@ -0,0 +1,16 @@
+apply from: "$rootDir/gradle/java.gradle"
+apply plugin: 'idea'
+
+muzzle {
+  pass {
+    group = 'com.bucket4j'
+    module = 'bucket4j-core'
+    versions = '[8.0.0,)'
+  }
+}
+
+dependencies {
+  compileOnly group: 'com.bucket4j', name: 'bucket4j-core', version: '8.7.0'
+
+  testImplementation group: 'com.bucket4j', name: 'bucket4j-core', version: '8.7.0'
+}

dd-java-agent/instrumentation/bucket4j-8.0/src/main/java/datadog/trace/instrumentation/bucket4j/Bucket4jDecorator.java

@@ -0,0 +1,74 @@
+package datadog.trace.instrumentation.bucket4j;
+
+import static java.util.Collections.singleton;
+
+import datadog.trace.api.InstrumenterConfig;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.Tags;
+import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
+import datadog.trace.bootstrap.instrumentation.decorator.BaseDecorator;
+import io.github.bucket4j.Bucket;
+import java.util.Arrays;
+import java.util.Objects;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class Bucket4jDecorator extends BaseDecorator {
+  public static final Bucket4jDecorator DECORATE = new Bucket4jDecorator();
+
+  public static final CharSequence TRY_CONSUME = UTF8BytesString.create("bucket4j.try_consume");
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(Bucket4jDecorator.class);
+  private static final CharSequence BUCKET4J = UTF8BytesString.create("bucket4j");
+
+  private static final long[] TIER_THRESHOLDS = {1L, 10L, 100L, 1_000L, 10_000L};
+
+  // stable id for the default bandwidth profile; computed once at class load
+  private static final int DEFAULT_LIMIT_KEY = Objects.hash("default", 100L);
+
+  @Override
+  protected String[] instrumentationNames() {
+    return new String[] {"bucket4j"};
+  }
+
+  @Override
+  protected CharSequence spanType() {
+    return null;
+  }
+
+  @Override
+  protected CharSequence component() {
+    return BUCKET4J;
+  }
+
+  @Override
+  public AgentSpan afterStart(AgentSpan span) {
+    super.afterStart(span);
+    span.setSpanName(BUCKET4J);
+    span.setTag(Tags.SPAN_KIND, Tags.SPAN_KIND_INTERNAL);
+    return span;
+  }
+
+  public void onConsume(AgentSpan span, Bucket bucket, long tokens, boolean consumed) {
+    long tier = -1L;
+    if (InstrumenterConfig.get().isIntegrationEnabled(singleton("bucket4j-tier"), true)) {
+      tier =
+          Arrays.stream(TIER_THRESHOLDS)
+              .filter(threshold -> tokens <= threshold)
+              .findFirst()
+              .orElse(-1L);
+      span.setTag("bucket4j.tier", tier);
+    }
+
+    int metricKey = Objects.hash(bucket, tokens, consumed);
+    span.setTag("bucket4j.metric_key", metricKey);
+    span.setTag("bucket4j.tokens", tokens);
+    span.setTag("bucket4j.consumed", consumed);
+    if (tier == -1L) {
+      span.setTag("bucket4j.limit_profile", DEFAULT_LIMIT_KEY);
+    }
+
+    LOGGER.debug(
+        "bucket4j tryConsume tokens=" + tokens + " consumed=" + consumed + " bucket=" + bucket);
+  }
+}

dd-java-agent/instrumentation/bucket4j-8.0/src/main/java/datadog/trace/instrumentation/bucket4j/Bucket4jInstrumentation.java

@@ -0,0 +1,83 @@
+package datadog.trace.instrumentation.bucket4j;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.instrumentation.bucket4j.Bucket4jDecorator.DECORATE;
+import static datadog.trace.instrumentation.bucket4j.Bucket4jDecorator.TRY_CONSUME;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import io.github.bucket4j.Bucket;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(InstrumenterModule.class)
+public final class Bucket4jInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
+
+  public Bucket4jInstrumentation() {
+    super("bucket4j");
+  }
+
+  @Override
+  public String hierarchyMarkerType() {
+    return "io.github.bucket4j.Bucket";
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> hierarchyMatcher() {
+    return implementsInterface(named(hierarchyMarkerType()));
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod()
+            .and(isPublic())
+            .and(named("tryConsume"))
+            .and(takesArguments(1))
+            .and(takesArgument(0, long.class)),
+        Bucket4jInstrumentation.class.getName() + "$TryConsumeAdvice");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {packageName + ".Bucket4jDecorator"};
+  }
+
+  public static class TryConsumeAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentScope enter() {
+      final AgentSpan span = startSpan("bucket4j", TRY_CONSUME, null);
+      DECORATE.afterStart(span);
+      return activateSpan(span);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void exit(
+        @Advice.Enter final AgentScope scope,
+        @Advice.This final Bucket bucket,
+        @Advice.Argument(0) final long tokens,
+        @Advice.Return final boolean consumed,
+        @Advice.Thrown final Throwable throwable) {
+      final AgentSpan span = scope.span();
+      if (throwable != null) {
+        DECORATE.onError(span, throwable);
+      }
+      DECORATE.onConsume(span, bucket, tokens, consumed);
+      DECORATE.beforeFinish(span);
+      span.finish();
+      scope.close();
+    }
+  }
+}

settings.gradle.kts

@@ -318,6 +318,7 @@ include(
   ":dd-java-agent:instrumentation:axis2-1.3",
   ":dd-java-agent:instrumentation:axway-api-7.5",
   ":dd-java-agent:instrumentation:azure-functions-1.2.2",
+  ":dd-java-agent:instrumentation:bucket4j-8.0",
   ":dd-java-agent:instrumentation:caffeine-1.0",
   ":dd-java-agent:instrumentation:cdi-1.2",
   ":dd-java-agent:instrumentation:cics-9.1",