From ebd60327b06b1ca4fbb3b1ee444e23d3ba95734b Mon Sep 17 00:00:00 2001 From: Douglas Q Hawkins Date: Thu, 9 Jul 2026 14:45:11 -0400 Subject: [PATCH 1/2] Add Bucket4j (tryConsume) instrumentation Traces bucket4j rate-limit checks: one span per Bucket#tryConsume with tokens/consumed tags and a rate tier. Co-Authored-By: Claude Opus 4.8 --- .../instrumentation/bucket4j-8.0/build.gradle | 16 ++++ .../bucket4j/Bucket4jDecorator.java | 79 ++++++++++++++++++ .../bucket4j/Bucket4jInstrumentation.java | 83 +++++++++++++++++++ settings.gradle.kts | 1 + 4 files changed, 179 insertions(+) create mode 100644 dd-java-agent/instrumentation/bucket4j-8.0/build.gradle create mode 100644 dd-java-agent/instrumentation/bucket4j-8.0/src/main/java/datadog/trace/instrumentation/bucket4j/Bucket4jDecorator.java create mode 100644 dd-java-agent/instrumentation/bucket4j-8.0/src/main/java/datadog/trace/instrumentation/bucket4j/Bucket4jInstrumentation.java diff --git a/dd-java-agent/instrumentation/bucket4j-8.0/build.gradle b/dd-java-agent/instrumentation/bucket4j-8.0/build.gradle new file mode 100644 index 00000000000..98a10b6b03a --- /dev/null +++ b/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' +} diff --git a/dd-java-agent/instrumentation/bucket4j-8.0/src/main/java/datadog/trace/instrumentation/bucket4j/Bucket4jDecorator.java b/dd-java-agent/instrumentation/bucket4j-8.0/src/main/java/datadog/trace/instrumentation/bucket4j/Bucket4jDecorator.java new file mode 100644 index 00000000000..ac758c19ed1 --- /dev/null +++ b/dd-java-agent/instrumentation/bucket4j-8.0/src/main/java/datadog/trace/instrumentation/bucket4j/Bucket4jDecorator.java @@ -0,0 +1,79 @@ +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.HashMap; +import java.util.Map; +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}; + + // pre-compute keys for the well-known default bandwidths once at load time + public static final Map KNOWN_LIMIT_KEYS = buildKnownLimitKeys(); + + private static Map buildKnownLimitKeys() { + Map keys = new HashMap<>(); + keys.put("default", Objects.hash("default", 100L)); + keys.put("burst", Objects.hash("burst", 1000L)); + return keys; + } + + @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) { + if (InstrumenterConfig.get().isIntegrationEnabled(singleton("bucket4j-tier"), true)) { + long 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); + + LOGGER.debug( + "bucket4j tryConsume tokens=" + tokens + " consumed=" + consumed + " bucket=" + bucket); + } +} diff --git a/dd-java-agent/instrumentation/bucket4j-8.0/src/main/java/datadog/trace/instrumentation/bucket4j/Bucket4jInstrumentation.java b/dd-java-agent/instrumentation/bucket4j-8.0/src/main/java/datadog/trace/instrumentation/bucket4j/Bucket4jInstrumentation.java new file mode 100644 index 00000000000..53af81ac30c --- /dev/null +++ b/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 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(); + } + } +} diff --git a/settings.gradle.kts b/settings.gradle.kts index dda1f432e6f..ea1d223d076 100644 --- a/settings.gradle.kts +++ b/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", From 62abebcbef6c2f0bb03ac96bf842a1f7a9622e4f Mon Sep 17 00:00:00 2001 From: Douglas Q Hawkins Date: Thu, 9 Jul 2026 15:38:55 -0400 Subject: [PATCH 2/2] Simplify bucket4j limit-key handling Co-Authored-By: Claude Opus 4.8 --- .../bucket4j/Bucket4jDecorator.java | 19 +++++++------------ 1 file changed, 7 insertions(+), 12 deletions(-) diff --git a/dd-java-agent/instrumentation/bucket4j-8.0/src/main/java/datadog/trace/instrumentation/bucket4j/Bucket4jDecorator.java b/dd-java-agent/instrumentation/bucket4j-8.0/src/main/java/datadog/trace/instrumentation/bucket4j/Bucket4jDecorator.java index ac758c19ed1..f59010ba700 100644 --- a/dd-java-agent/instrumentation/bucket4j-8.0/src/main/java/datadog/trace/instrumentation/bucket4j/Bucket4jDecorator.java +++ b/dd-java-agent/instrumentation/bucket4j-8.0/src/main/java/datadog/trace/instrumentation/bucket4j/Bucket4jDecorator.java @@ -9,8 +9,6 @@ import datadog.trace.bootstrap.instrumentation.decorator.BaseDecorator; import io.github.bucket4j.Bucket; import java.util.Arrays; -import java.util.HashMap; -import java.util.Map; import java.util.Objects; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -25,15 +23,8 @@ public class Bucket4jDecorator extends BaseDecorator { private static final long[] TIER_THRESHOLDS = {1L, 10L, 100L, 1_000L, 10_000L}; - // pre-compute keys for the well-known default bandwidths once at load time - public static final Map KNOWN_LIMIT_KEYS = buildKnownLimitKeys(); - - private static Map buildKnownLimitKeys() { - Map keys = new HashMap<>(); - keys.put("default", Objects.hash("default", 100L)); - keys.put("burst", Objects.hash("burst", 1000L)); - return keys; - } + // 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() { @@ -59,8 +50,9 @@ public AgentSpan afterStart(AgentSpan span) { } public void onConsume(AgentSpan span, Bucket bucket, long tokens, boolean consumed) { + long tier = -1L; if (InstrumenterConfig.get().isIntegrationEnabled(singleton("bucket4j-tier"), true)) { - long tier = + tier = Arrays.stream(TIER_THRESHOLDS) .filter(threshold -> tokens <= threshold) .findFirst() @@ -72,6 +64,9 @@ public void onConsume(AgentSpan span, Bucket bucket, long tokens, boolean consum 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);