Update nocode instrumentation

bootstrap/src/main/java/com/splunk/opentelemetry/javaagent/bootstrap/nocode/NocodeExpression.java

@@ -0,0 +1,24 @@
+/*
+ * Copyright Splunk Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.splunk.opentelemetry.javaagent.bootstrap.nocode;
+
+public interface NocodeExpression {
+
+  Object evaluate(Object thiz, Object[] params);
+
+  Object evaluateAtEnd(Object thiz, Object[] params, Object returnValue, Throwable error);
+}

bootstrap/src/main/java/com/splunk/opentelemetry/javaagent/bootstrap/nocode/NocodeRules.java

@@ -16,29 +16,34 @@
 
 package com.splunk.opentelemetry.javaagent.bootstrap.nocode;
 
+import io.opentelemetry.api.trace.SpanKind;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
 
 public final class NocodeRules {
 
   public static final class Rule {
-    public final String className;
-    public final String methodName;
-    public final String spanName; // may be null - use default of "class.method"
-    public final String spanKind; // matches the SpanKind enum, null means default to INTERNAL
-    public final String spanStatus; // may be null, should return string from StatusCodes
+    private static final AtomicInteger counter = new AtomicInteger();
 
-    public final Map<String, String> attributes; // key name to jexl expression
+    private final int id = counter.incrementAndGet();
+    private final String className;
+    private final String methodName;
+    private final NocodeExpression spanName; // may be null - use default of "class.method"
+    private final SpanKind spanKind; // may be null
+    private final NocodeExpression spanStatus; // may be null, should return string from StatusCodes
+
+    private final Map<String, NocodeExpression> attributes; // key name to jexl expression
 
     public Rule(
         String className,
         String methodName,
-        String spanName,
-        String spanKind,
-        String spanStatus,
-        Map<String, String> attributes) {
+        NocodeExpression spanName,
+        SpanKind spanKind,
+        NocodeExpression spanStatus,
+        Map<String, NocodeExpression> attributes) {
       this.className = className;
       this.methodName = methodName;
       this.spanName = spanName;
@@ -61,25 +66,52 @@ public String toString() {
           + ",attrs="
           + attributes;
     }
+
+    public int getId() {
+      return id;
+    }
+
+    public String getClassName() {
+      return className;
+    }
+
+    public String getMethodName() {
+      return methodName;
+    }
+
+    public NocodeExpression getSpanName() {
+      return spanName;
+    }
+
+    public SpanKind getSpanKind() {
+      return spanKind;
+    }
+
+    public NocodeExpression getSpanStatus() {
+      return spanStatus;
+    }
+
+    public Map<String, NocodeExpression> getAttributes() {
+      return attributes;
+    }
   }
 
   private NocodeRules() {}
 
-  // Using className.methodName as the key
-  private static final HashMap<String, Rule> name2Rule = new HashMap<>();
+  private static final HashMap<Integer, Rule> ruleMap = new HashMap<>();
 
   // Called by the NocodeInitializer
   public static void setGlobalRules(List<Rule> rules) {
     for (Rule r : rules) {
-      name2Rule.put(r.className + "." + r.methodName, r);
+      ruleMap.put(r.id, r);
     }
   }
 
   public static Iterable<Rule> getGlobalRules() {
-    return name2Rule.values();
+    return ruleMap.values();
   }
 
-  public static Rule find(String className, String methodName) {
-    return name2Rule.get(className + "." + methodName);
+  public static Rule find(int id) {
+    return ruleMap.get(id);
   }
 }

instrumentation/nocode/build.gradle.kts

@@ -6,6 +6,7 @@ plugins {
 dependencies {
   compileOnly(project(":custom"))
   compileOnly("io.opentelemetry.javaagent:opentelemetry-javaagent-tooling")
+  compileOnly("io.opentelemetry.instrumentation:opentelemetry-instrumentation-annotations-support")
   compileOnly("org.snakeyaml:snakeyaml-engine:2.8")
 
   implementation("org.apache.commons:commons-jexl3:3.4.0") {

instrumentation/nocode/src/main/java/com/splunk/opentelemetry/instrumentation/nocode/JexlEvaluator.java

@@ -16,21 +16,21 @@
 
 package com.splunk.opentelemetry.instrumentation.nocode;
 
-import com.splunk.opentelemetry.javaagent.bootstrap.nocode.NocodeEvaluation;
 import java.util.logging.Logger;
 import org.apache.commons.jexl3.JexlBuilder;
 import org.apache.commons.jexl3.JexlContext;
 import org.apache.commons.jexl3.JexlEngine;
+import org.apache.commons.jexl3.JexlExpression;
 import org.apache.commons.jexl3.JexlFeatures;
 import org.apache.commons.jexl3.MapContext;
 import org.apache.commons.jexl3.introspection.JexlPermissions;
 
-public class JexlEvaluator implements NocodeEvaluation.Evaluator {
+class JexlEvaluator {
   private static final Logger logger = Logger.getLogger(JexlEvaluator.class.getName());
 
   private final JexlEngine jexl;
 
-  public JexlEvaluator() {
+  JexlEvaluator() {
     JexlFeatures features =
         new JexlFeatures()
             .register(false) // don't support #register syntax
@@ -59,31 +59,41 @@ private static void setBeginningVariables(JexlContext context, Object thiz, Obje
     }
   }
 
-  private Object evaluateExpression(String expression, JexlContext context) {
+  private static Object evaluateExpression(JexlExpression expression, JexlContext context) {
     try {
-      // could cache the Expression in the Rule if desired
-      return jexl.createExpression(expression).evaluate(context);
+      return expression.evaluate(context);
     } catch (Throwable t) {
       logger.warning("Can't evaluate {" + expression + "}: " + t);
       return null;
     }
   }
 
+  JexlExpression createExpression(String expression) {
+    try {
+      return jexl.createExpression(expression);
+    } catch (Throwable t) {
+      logger.warning("Invalid expression {" + expression + "}: " + t);
+      return null;
+    }
+  }
+
   private void setEndingVariables(JexlContext context, Object returnValue, Throwable error) {
     context.set("returnValue", returnValue);
     context.set("error", error);
   }
 
-  @Override
-  public Object evaluate(String expression, Object thiz, Object[] params) {
+  Object evaluate(JexlExpression expression, Object thiz, Object[] params) {
     JexlContext context = new MapContext();
     setBeginningVariables(context, thiz, params);
     return evaluateExpression(expression, context);
   }
 
-  @Override
-  public Object evaluateAtEnd(
-      String expression, Object thiz, Object[] params, Object returnValue, Throwable error) {
+  Object evaluateAtEnd(
+      JexlExpression expression,
+      Object thiz,
+      Object[] params,
+      Object returnValue,
+      Throwable error) {
     JexlContext context = new MapContext();
     setBeginningVariables(context, thiz, params);
     setEndingVariables(context, returnValue, error);

instrumentation/nocode/src/main/java/com/splunk/opentelemetry/instrumentation/nocode/NocodeAttributesExtractor.java

@@ -16,7 +16,7 @@
 
 package com.splunk.opentelemetry.instrumentation.nocode;
 
-import com.splunk.opentelemetry.javaagent.bootstrap.nocode.NocodeEvaluation;
+import com.splunk.opentelemetry.javaagent.bootstrap.nocode.NocodeExpression;
 import io.opentelemetry.api.common.AttributesBuilder;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.incubator.semconv.code.CodeAttributesExtractor;
@@ -25,8 +25,7 @@
 import java.util.Map;
 import javax.annotation.Nullable;
 
-public final class NocodeAttributesExtractor
-    implements AttributesExtractor<NocodeMethodInvocation, Object> {
+class NocodeAttributesExtractor implements AttributesExtractor<NocodeMethodInvocation, Object> {
   private final AttributesExtractor<ClassAndMethod, Object> codeExtractor;
 
   public NocodeAttributesExtractor() {
@@ -38,10 +37,10 @@ public void onStart(
       AttributesBuilder attributesBuilder, Context context, NocodeMethodInvocation mi) {
     codeExtractor.onStart(attributesBuilder, context, mi.getClassAndMethod());
 
-    Map<String, String> attributes = mi.getRuleAttributes();
+    Map<String, NocodeExpression> attributes = mi.getRuleAttributes();
     for (String key : attributes.keySet()) {
-      String expression = attributes.get(key);
-      Object value = NocodeEvaluation.evaluate(expression, mi.getThiz(), mi.getParameters());
+      NocodeExpression expression = attributes.get(key);
+      Object value = mi.evaluate(expression);
       if (value instanceof Long
           || value instanceof Integer
           || value instanceof Short

instrumentation/nocode/src/main/java/com/splunk/opentelemetry/instrumentation/nocode/NocodeInitializer.java

@@ -19,7 +19,6 @@
 import static io.opentelemetry.sdk.autoconfigure.AutoConfigureUtil.getConfig;
 
 import com.google.auto.service.AutoService;
-import com.splunk.opentelemetry.javaagent.bootstrap.nocode.NocodeEvaluation;
 import com.splunk.opentelemetry.javaagent.bootstrap.nocode.NocodeRules;
 import io.opentelemetry.javaagent.tooling.BeforeAgentListener;
 import io.opentelemetry.sdk.autoconfigure.AutoConfiguredOpenTelemetrySdk;
@@ -33,6 +32,5 @@ public void beforeAgent(AutoConfiguredOpenTelemetrySdk autoConfiguredOpenTelemet
     ConfigProperties config = getConfig(autoConfiguredOpenTelemetrySdk);
     YamlParser yp = new YamlParser(config);
     NocodeRules.setGlobalRules(yp.getInstrumentationRules());
-    NocodeEvaluation.internalSetEvaluator(new JexlEvaluator());
   }
 }

instrumentation/nocode/src/main/java/com/splunk/opentelemetry/instrumentation/nocode/NocodeInstrumentation.java

@@ -24,6 +24,7 @@
 import com.splunk.opentelemetry.javaagent.bootstrap.nocode.NocodeRules;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
+import io.opentelemetry.instrumentation.api.annotation.support.async.AsyncOperationEndSupport;
 import io.opentelemetry.instrumentation.api.incubator.semconv.util.ClassAndMethod;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
@@ -32,6 +33,7 @@
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.implementation.bytecode.assign.Assigner;
 import net.bytebuddy.matcher.ElementMatcher;
+import net.bytebuddy.utility.JavaConstant;
 
 public final class NocodeInstrumentation implements TypeInstrumentation {
   private final NocodeRules.Rule rule;
@@ -42,32 +44,45 @@ public NocodeInstrumentation(NocodeRules.Rule rule) {
 
   @Override
   public ElementMatcher<TypeDescription> typeMatcher() {
-    // names have to match exactly for now to enable rule lookup
-    // at advice time.  In the future, we could support
-    // more complex rules here if we dynamically generated advice classes for
-    // each rule.
-    return rule != null ? named(rule.className) : none();
+    // null rule is used when no rules are configured, this ensures that muzzle can collect helper
+    // classes
+    return rule != null ? named(rule.getClassName()) : none();
   }
 
   @Override
   public void transform(TypeTransformer transformer) {
     transformer.applyAdviceToMethod(
-        rule != null ? named(rule.methodName) : none(),
+        rule != null ? named(rule.getMethodName()) : none(),
+        mapping ->
+            mapping
+                .bind(RuleId.class, JavaConstant.Simple.ofLoaded(rule != null ? rule.getId() : -1))
+                .bind(
+                    MethodReturnType.class,
+                    (instrumentedType, instrumentedMethod, assigner, argumentHandler, sort) ->
+                        Advice.OffsetMapping.Target.ForStackManipulation.of(
+                            instrumentedMethod.getReturnType().asErasure())),
         this.getClass().getName() + "$NocodeAdvice");
   }
 
+  // custom annotation that allows looking up the rule that triggered instrumenting the method
+  @interface RuleId {}
+
+  // custom annotation that represents the return type of the method
+  @interface MethodReturnType {}
+
   @SuppressWarnings("unused")
   public static class NocodeAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void onEnter(
+        @RuleId int ruleId,
         @Advice.Origin("#t") Class<?> declaringClass,
         @Advice.Origin("#m") String methodName,
         @Advice.Local("otelInvocation") NocodeMethodInvocation otelInvocation,
         @Advice.Local("otelContext") Context context,
         @Advice.Local("otelScope") Scope scope,
         @Advice.This Object thiz,
         @Advice.AllArguments Object[] methodParams) {
-      NocodeRules.Rule rule = NocodeRules.find(declaringClass.getName(), methodName);
+      NocodeRules.Rule rule = NocodeRules.find(ruleId);
       otelInvocation =
           new NocodeMethodInvocation(
               rule, ClassAndMethod.create(declaringClass, methodName), thiz, methodParams);
@@ -86,16 +101,16 @@ public static void stopSpan(
         @Advice.Local("otelInvocation") NocodeMethodInvocation otelInvocation,
         @Advice.Local("otelContext") Context context,
         @Advice.Local("otelScope") Scope scope,
-        @Advice.Return(typing = Assigner.Typing.DYNAMIC) Object returnValue,
+        @Advice.Return(readOnly = false, typing = Assigner.Typing.DYNAMIC) Object returnValue,
         @Advice.Thrown Throwable error) {
       if (scope == null) {
         return;
       }
       scope.close();
-      // This is heavily based on the "methods" instrumentation from upstream, but
-      // for now we're not supporting modifying return types/async result codes, etc.
-      // This could be expanded in the future.
-      instrumenter().end(context, otelInvocation, returnValue, error);
+
+      returnValue =
+          AsyncOperationEndSupport.create(instrumenter(), Object.class, method.getReturnType())
+              .asyncEnd(context, otelInvocation, returnValue, error);
     }
   }
 }