Use VirtualField for associating netty listener with wrapper (#5282)

* Use VirtualField for associating netty listener with wrapper

* Move ignoring lambas for injected classes to LambdaTransformer

Author: laurit

instrumentation/internal/internal-lambda/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/internal/lambda/LambdaTransformer.java

@@ -6,6 +6,7 @@
 package io.opentelemetry.javaagent.instrumentation.internal.lambda;
 
 import io.opentelemetry.javaagent.bootstrap.ClassFileTransformerHolder;
+import io.opentelemetry.javaagent.bootstrap.InjectedHelperClassDetector;
 import java.lang.instrument.ClassFileTransformer;
 
 /** Helper class for transforming lambda class bytes. */
@@ -29,6 +30,11 @@ private static boolean isJava9() {
    */
   @SuppressWarnings("unused")
   public static byte[] transform(byte[] classBytes, String slashClassName, Class<?> targetClass) {
+    // Skip transforming lambdas of injected helper classes.
+    if (InjectedHelperClassDetector.isHelperClass(targetClass)) {
+      return classBytes;
+    }
+
     ClassFileTransformer transformer = ClassFileTransformerHolder.getClassFileTransformer();
     if (transformer != null) {
       try {

instrumentation/netty/netty-4-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/netty/common/FutureListenerWrappers.java

@@ -11,25 +11,25 @@
 import io.netty.util.concurrent.ProgressiveFuture;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
-import io.opentelemetry.instrumentation.api.cache.Cache;
+import io.opentelemetry.instrumentation.api.field.VirtualField;
 import java.lang.ref.WeakReference;
 
 public final class FutureListenerWrappers {
   // note: it's ok if the value is collected prior to the key, since this cache is only used to
   // remove the wrapped listener from the netty future, and if the value is collected prior to the
   // key, that means it's no longer used (referenced) by the netty future anyways.
   //
-  // also note: this is not using VirtualField in case this is ever converted to library
-  // instrumentation, because while the library implementation of VirtualField maintains a weak
-  // reference to its keys, it maintains a strong reference to its values, and in this particular
-  // case the wrapper listener (value) has a strong reference to original listener (key), which will
-  // create a memory leak. which is not a problem in the javaagent's implementation of VirtualField,
-  // since it injects the value directly into the key as a field, and so the value is only retained
-  // strongly by the key, and so they can be collected together.
-  private static final Cache<
+  // also note: this is using WeakReference as VirtualField value in case this is ever converted to
+  // library instrumentation, because while the library implementation of VirtualField maintains a
+  // weak reference to its keys, it maintains a strong reference to its values, and in this
+  // particular case the wrapper listener (value) has a strong reference to original listener (key),
+  // which will create a memory leak. which is not a problem in the javaagent's implementation of
+  // VirtualField, since it injects the value directly into the key as a field, and so the value is
+  // only retained strongly by the key, and so they can be collected together.
+  private static final VirtualField<
           GenericFutureListener<? extends Future<?>>,
           WeakReference<GenericFutureListener<? extends Future<?>>>>
-      wrappers = Cache.weak();
+      wrapperVirtualField = VirtualField.find(GenericFutureListener.class, WeakReference.class);
 
   private static final ClassValue<Boolean> shouldWrap =
       new ClassValue<Boolean>() {
@@ -54,7 +54,7 @@ public static GenericFutureListener<?> wrap(
     // collected before we have a chance to make (and return) a strong reference to the wrapper
 
     WeakReference<GenericFutureListener<? extends Future<?>>> resultReference =
-        wrappers.get(delegate);
+        wrapperVirtualField.get(delegate);
 
     if (resultReference != null) {
       GenericFutureListener<? extends Future<?>> wrapper = resultReference.get();
@@ -75,14 +75,14 @@ public static GenericFutureListener<?> wrap(
     } else {
       wrapper = new WrappedFutureListener(context, (GenericFutureListener<Future<?>>) delegate);
     }
-    wrappers.put(delegate, new WeakReference<>(wrapper));
+    wrapperVirtualField.set(delegate, new WeakReference<>(wrapper));
     return wrapper;
   }
 
   public static GenericFutureListener<? extends Future<?>> getWrapper(
       GenericFutureListener<? extends Future<?>> delegate) {
     WeakReference<GenericFutureListener<? extends Future<?>>> wrapperReference =
-        wrappers.get(delegate);
+        wrapperVirtualField.get(delegate);
     if (wrapperReference == null) {
       return delegate;
     }

javaagent-bootstrap/src/main/java/io/opentelemetry/javaagent/bootstrap/InjectedHelperClassDetector.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.bootstrap;
+
+import java.util.function.Function;
+
+/** Helper class for detecting whether given class is an injected helper class. */
+public final class InjectedHelperClassDetector {
+
+  private InjectedHelperClassDetector() {}
+
+  private static volatile Function<Class<?>, Boolean> helperClassDetector;
+
+  /** Sets the {@link Function} for detecting injected helper classes. */
+  public static void internalSetHelperClassDetector(
+      Function<Class<?>, Boolean> helperClassDetector) {
+    if (InjectedHelperClassDetector.helperClassDetector != null) {
+      // Only possible by misuse of this API, just ignore.
+      return;
+    }
+    InjectedHelperClassDetector.helperClassDetector = helperClassDetector;
+  }
+
+  public static boolean isHelperClass(Class<?> clazz) {
+    if (helperClassDetector == null) {
+      return false;
+    }
+    return helperClassDetector.apply(clazz);
+  }
+}

muzzle/src/main/java/io/opentelemetry/javaagent/tooling/HelperInjector.java

@@ -8,6 +8,7 @@
 import io.opentelemetry.instrumentation.api.cache.Cache;
 import io.opentelemetry.javaagent.bootstrap.DefineClassContext;
 import io.opentelemetry.javaagent.bootstrap.HelperResources;
+import io.opentelemetry.javaagent.bootstrap.InjectedHelperClassDetector;
 import io.opentelemetry.javaagent.tooling.muzzle.HelperResource;
 import java.io.File;
 import java.io.IOException;
@@ -50,6 +51,10 @@ public class HelperInjector implements Transformer {
   private static final TransformSafeLogger logger =
       TransformSafeLogger.getLogger(HelperInjector.class);
 
+  static {
+    InjectedHelperClassDetector.internalSetHelperClassDetector(HelperInjector::isInjectedClass);
+  }
+
   // Need this because we can't put null into the injectedClassLoaders map.
   private static final ClassLoader BOOTSTRAP_CLASSLOADER_PLACEHOLDER =
       new SecureClassLoader(null) {