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..f59010ba700 --- /dev/null +++ b/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); + } +} 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",