From dac1522a3f6fef096f379554e84e99f206229302 Mon Sep 17 00:00:00 2001
From: Trask Stalnaker <trask.stalnaker@gmail.com>
Date: Tue, 23 Nov 2021 15:14:54 -0800
Subject: [PATCH] Convert Netty wrapper cache to VirtualField

---
 .../netty/common/FutureListenerWrappers.java  | 61 ++++++++-----------
 1 file changed, 25 insertions(+), 36 deletions(-)

diff --git a/instrumentation/netty/netty-4-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/netty/common/FutureListenerWrappers.java b/instrumentation/netty/netty-4-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/netty/common/FutureListenerWrappers.java
index f1e3881505e0..ce71b689c015 100644
--- a/instrumentation/netty/netty-4-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/netty/common/FutureListenerWrappers.java
+++ b/instrumentation/netty/netty-4-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/netty/common/FutureListenerWrappers.java
@@ -11,22 +11,19 @@
 import io.netty.util.concurrent.ProgressiveFuture;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
-import io.opentelemetry.instrumentation.api.cache.Cache;
-import java.lang.ref.WeakReference;
+import io.opentelemetry.instrumentation.api.field.VirtualField;
 
 public final class FutureListenerWrappers {
-  // Instead of VirtualField use Cache with weak keys and weak values to store link between original
-  // listener and wrapper. VirtualField works fine when wrapper is stored in a field on original
-  // listener, but when listener class is a lambda instead of field it gets stored in a map with
-  // weak keys where original listener is key and wrapper is value. As wrapper has a strong
-  // reference to original listener this causes a memory leak.
-  // Also note that 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.
-  private static final Cache<
-          GenericFutureListener<? extends Future<?>>,
-          WeakReference<GenericFutureListener<? extends Future<?>>>>
-      wrappers = Cache.weak();
+  // important: if this is ever converted to library instrumentation, this will create a memory leak
+  // because while library implementation of VirtualField maintains a weak reference to its keys, it
+  // maintains a strong reference to its values, and the wrapper has a strong reference to original
+  // listener, which will create a memory leak.
+  // this 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.
+  @SuppressWarnings("rawtypes")
+  private static final VirtualField<GenericFutureListener, GenericFutureListener> wrapperField =
+      VirtualField.find(GenericFutureListener.class, GenericFutureListener.class);
 
   private static final ClassValue<Boolean> shouldWrap =
       new ClassValue<Boolean>() {
@@ -46,32 +43,24 @@ public static boolean shouldWrap(GenericFutureListener<? extends Future<?>> list
   @SuppressWarnings("unchecked")
   public static GenericFutureListener<?> wrap(
       Context context, GenericFutureListener<? extends Future<?>> delegate) {
-    WeakReference<GenericFutureListener<? extends Future<?>>> resultReference =
-        wrappers.computeIfAbsent(
-            delegate,
-            key -> {
-              GenericFutureListener<? extends Future<?>> wrapper;
-              if (delegate instanceof GenericProgressiveFutureListener) {
-                wrapper =
-                    new WrappedProgressiveFutureListener(
-                        context, (GenericProgressiveFutureListener<ProgressiveFuture<?>>) delegate);
-              } else {
-                wrapper =
-                    new WrappedFutureListener(context, (GenericFutureListener<Future<?>>) delegate);
-              }
-              return new WeakReference<>(wrapper);
-            });
-    return resultReference.get();
+    GenericFutureListener<? extends Future<?>> wrapper = wrapperField.get(delegate);
+    if (wrapper == null) {
+      if (delegate instanceof GenericProgressiveFutureListener) {
+        wrapper =
+            new WrappedProgressiveFutureListener(
+                context, (GenericProgressiveFutureListener<ProgressiveFuture<?>>) delegate);
+      } else {
+        wrapper = new WrappedFutureListener(context, (GenericFutureListener<Future<?>>) delegate);
+      }
+      wrapperField.set(delegate, wrapper);
+    }
+    return wrapper;
   }
 
+  @SuppressWarnings("unchecked")
   public static GenericFutureListener<? extends Future<?>> getWrapper(
       GenericFutureListener<? extends Future<?>> delegate) {
-    WeakReference<GenericFutureListener<? extends Future<?>>> wrapperReference =
-        wrappers.get(delegate);
-    if (wrapperReference == null) {
-      return delegate;
-    }
-    GenericFutureListener<? extends Future<?>> wrapper = wrapperReference.get();
+    GenericFutureListener<? extends Future<?>> wrapper = wrapperField.get(delegate);
     return wrapper == null ? delegate : wrapper;
   }