diff --git a/instrumentation/geode-1.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/geode/GeodeRegionInstrumentation.java b/instrumentation/geode-1.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/geode/GeodeRegionInstrumentation.java index 1078d5f9a841..b52e4203e8ba 100644 --- a/instrumentation/geode-1.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/geode/GeodeRegionInstrumentation.java +++ b/instrumentation/geode-1.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/geode/GeodeRegionInstrumentation.java @@ -5,7 +5,6 @@ package io.opentelemetry.javaagent.instrumentation.geode; -import static io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge.currentContext; import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed; import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface; import static io.opentelemetry.javaagent.instrumentation.geode.GeodeSingletons.instrumenter; @@ -19,6 +18,7 @@ import io.opentelemetry.context.Scope; import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation; import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer; +import javax.annotation.Nullable; import net.bytebuddy.asm.Advice; import net.bytebuddy.description.type.TypeDescription; import net.bytebuddy.matcher.ElementMatcher; @@ -61,76 +61,77 @@ public void transform(TypeTransformer transformer) { this.getClass().getName() + "$QueryAdvice"); } - @SuppressWarnings("unused") - public static class SimpleAdvice { + public static class AdviceScope { + private final GeodeRequest request; + private final Context context; + private final Scope scope; - @Advice.OnMethodEnter(suppress = Throwable.class) - public static void onEnter( - @Advice.This Region region, - @Advice.Origin("#m") String methodName, - @Advice.Local("otelRequest") GeodeRequest request, - @Advice.Local("otelContext") Context context, - @Advice.Local("otelScope") Scope scope) { + public AdviceScope(GeodeRequest request, Context context, Scope scope) { + this.request = request; + this.context = context; + this.scope = scope; + } - Context parentContext = currentContext(); - request = GeodeRequest.create(region, methodName, null); + @Nullable + public static AdviceScope start( + Region region, String methodName, @Nullable String query) { + Context parentContext = Context.current(); + GeodeRequest request = GeodeRequest.create(region, methodName, query); if (!instrumenter().shouldStart(parentContext, request)) { - return; + return null; } - context = instrumenter().start(parentContext, request); - scope = context.makeCurrent(); + Context context = instrumenter().start(parentContext, request); + return new AdviceScope(request, context, context.makeCurrent()); + } + + public void end(@Nullable Throwable throwable) { + if (scope != null) { + scope.close(); + } + instrumenter().end(context, request, null, throwable); + } + } + + @SuppressWarnings("unused") + public static class SimpleAdvice { + + @Nullable + @Advice.OnMethodEnter(suppress = Throwable.class) + public static AdviceScope onEnter( + @Advice.This Region region, @Advice.Origin("#m") String methodName) { + return AdviceScope.start(region, methodName, null); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void stopSpan( - @Advice.Thrown Throwable throwable, - @Advice.Local("otelRequest") GeodeRequest request, - @Advice.Local("otelContext") Context context, - @Advice.Local("otelScope") Scope scope) { - if (scope == null) { - return; + @Advice.Thrown @Nullable Throwable throwable, + @Advice.Enter @Nullable AdviceScope adviceScope) { + if (adviceScope != null) { + adviceScope.end(throwable); } - - scope.close(); - instrumenter().end(context, request, null, throwable); } } @SuppressWarnings("unused") public static class QueryAdvice { + @Nullable @Advice.OnMethodEnter(suppress = Throwable.class) - public static void onEnter( + public static AdviceScope onEnter( @Advice.This Region region, @Advice.Origin("#m") String methodName, - @Advice.Argument(0) String query, - @Advice.Local("otelRequest") GeodeRequest request, - @Advice.Local("otelContext") Context context, - @Advice.Local("otelScope") Scope scope) { - - Context parentContext = currentContext(); - request = GeodeRequest.create(region, methodName, query); - if (!instrumenter().shouldStart(parentContext, request)) { - return; - } - - context = instrumenter().start(parentContext, request); - scope = context.makeCurrent(); + @Advice.Argument(0) String query) { + return AdviceScope.start(region, methodName, query); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void stopSpan( - @Advice.Thrown Throwable throwable, - @Advice.Local("otelRequest") GeodeRequest request, - @Advice.Local("otelContext") Context context, - @Advice.Local("otelScope") Scope scope) { - if (scope == null) { - return; + @Advice.Thrown @Nullable Throwable throwable, + @Advice.Enter @Nullable AdviceScope adviceScope) { + if (adviceScope != null) { + adviceScope.end(throwable); } - - scope.close(); - instrumenter().end(context, request, null, throwable); } } }