Fix memory leak on netty

instrumentation-api-caching/src/main/java/io/opentelemetry/instrumentation/api/caching/CacheBuilder.java

@@ -14,6 +14,7 @@ public final class CacheBuilder {
  private static final long UNSET = -1;
 
  private boolean weakKeys;
+ private boolean weakValues;
  private long maximumSize = UNSET;
  private Executor executor = null;
 
@@ -32,6 +33,12 @@ public CacheBuilder setWeakKeys() {
  return this;
  }
 
+ /** Sets that values should be referenced weakly. */
+ public CacheBuilder setWeakValues() {
+ this.weakValues = true;
+ return this;
+ }
+
  // Visible for testing
  CacheBuilder setExecutor(Executor executor) {
  this.executor = executor;
@@ -47,6 +54,9 @@ public <K, V> Cache<K, V> build() {
  if (weakKeys) {
  caffeine.weakKeys();
  }
+ if (weakValues) {
+ caffeine.weakValues();
+ }
  if (maximumSize != UNSET) {
  caffeine.maximumSize(maximumSize);
  }

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

@@ -11,85 +11,78 @@
 import io.netty.util.concurrent.ProgressiveFuture;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
-import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
+import io.opentelemetry.instrumentation.api.caching.Cache;
 
 public final class FutureListenerWrappers {
+ private static final Cache<
+ GenericFutureListener<? extends Future<?>>, GenericFutureListener<? extends Future<?>>>
+ wrappers = Cache.newBuilder().setWeakKeys().setWeakValues().build();
+
  @SuppressWarnings("unchecked")
- public static GenericFutureListener<? extends Future<? super Void>> wrap(
- ContextStore<GenericFutureListener, GenericFutureListener> contextStore,
- Context context,
- GenericFutureListener<? extends Future<? super Void>> delegate) {
+ public static GenericFutureListener<?> wrap(
+ Context context, GenericFutureListener<? extends Future<?>> delegate) {
  if (delegate instanceof WrappedFutureListener
  || delegate instanceof WrappedProgressiveFutureListener) {
  return delegate;
  }
- return (GenericFutureListener<? extends Future<? super Void>>)
- contextStore.putIfAbsent(
- delegate,
- () -> {
- if (delegate instanceof GenericProgressiveFutureListener) {
- return new WrappedProgressiveFutureListener(
- context,
- (GenericProgressiveFutureListener<ProgressiveFuture<? super Void>>) delegate);
- } else {
- return new WrappedFutureListener(
- context, (GenericFutureListener<Future<? super Void>>) delegate);
- }
- });
+ return wrappers.computeIfAbsent(
+ delegate,
+ (key) -> {
+ if (delegate instanceof GenericProgressiveFutureListener) {
+ return new WrappedProgressiveFutureListener(
+ context, (GenericProgressiveFutureListener<ProgressiveFuture<?>>) delegate);
+ } else {
+ return new WrappedFutureListener(context, (GenericFutureListener<Future<?>>) delegate);
+ }
+ });
  }
 
- public static GenericFutureListener<? extends Future<? super Void>> getWrapper(
- ContextStore<GenericFutureListener, GenericFutureListener> contextStore,
- GenericFutureListener<? extends Future<? super Void>> delegate) {
- GenericFutureListener<? extends Future<? super Void>> wrapper =
- (GenericFutureListener<? extends Future<? super Void>>) contextStore.get(delegate);
+ public static GenericFutureListener<? extends Future<?>> getWrapper(
+ GenericFutureListener<? extends Future<?>> delegate) {
+ GenericFutureListener<? extends Future<?>> wrapper = wrappers.get(delegate);
  return wrapper == null ? delegate : wrapper;
  }
 
- private static final class WrappedFutureListener
- implements GenericFutureListener<Future<? super Void>> {
+ private static final class WrappedFutureListener implements GenericFutureListener<Future<?>> {
 
  private final Context context;
- private final GenericFutureListener<Future<? super Void>> delegate;
+ private final GenericFutureListener<Future<?>> delegate;
 
- private WrappedFutureListener(
- Context context, GenericFutureListener<Future<? super Void>> delegate) {
+ private WrappedFutureListener(Context context, GenericFutureListener<Future<?>> delegate) {
  this.context = context;
  this.delegate = delegate;
  }
 
  @Override
- public void operationComplete(Future<? super Void> future) throws Exception {
+ public void operationComplete(Future<?> future) throws Exception {
  try (Scope ignored = context.makeCurrent()) {
  delegate.operationComplete(future);
  }
  }
  }
 
  private static final class WrappedProgressiveFutureListener
- implements GenericProgressiveFutureListener<ProgressiveFuture<? super Void>> {
+ implements GenericProgressiveFutureListener<ProgressiveFuture<?>> {
 
  private final Context context;
- private final GenericProgressiveFutureListener<ProgressiveFuture<? super Void>> delegate;
+ private final GenericProgressiveFutureListener<ProgressiveFuture<?>> delegate;
 
  private WrappedProgressiveFutureListener(
- Context context,
- GenericProgressiveFutureListener<ProgressiveFuture<? super Void>> delegate) {
+ Context context, GenericProgressiveFutureListener<ProgressiveFuture<?>> delegate) {
  this.context = context;
  this.delegate = delegate;
  }
 
  @Override
- public void operationProgressed(
- ProgressiveFuture<? super Void> progressiveFuture, long l, long l1) throws Exception {
+ public void operationProgressed(ProgressiveFuture<?> progressiveFuture, long l, long l1)
+ throws Exception {
  try (Scope ignored = context.makeCurrent()) {
  delegate.operationProgressed(progressiveFuture, l, l1);
  }
  }
 
  @Override
- public void operationComplete(ProgressiveFuture<? super Void> progressiveFuture)
- throws Exception {
+ public void operationComplete(ProgressiveFuture<?> progressiveFuture) throws Exception {
  try (Scope ignored = context.makeCurrent()) {
  delegate.operationComplete(progressiveFuture);
  }

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

@@ -17,8 +17,6 @@
 import io.opentelemetry.context.Context;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
-import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
-import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
 import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
@@ -60,28 +58,23 @@ public static class AddListenerAdvice {
  @Advice.OnMethodEnter
  public static void wrapListener(
  @Advice.Argument(value = 0, readOnly = false)
- GenericFutureListener<? extends Future<? super Void>> listener) {
- ContextStore<GenericFutureListener, GenericFutureListener> contextStore =
- InstrumentationContext.get(GenericFutureListener.class, GenericFutureListener.class);
- listener =
- FutureListenerWrappers.wrap(contextStore, Java8BytecodeBridge.currentContext(), listener);
+ GenericFutureListener<? extends Future<?>> listener) {
+ listener = FutureListenerWrappers.wrap(Java8BytecodeBridge.currentContext(), listener);
  }
  }
 
  public static class AddListenersAdvice {
  @Advice.OnMethodEnter
  public static void wrapListener(
  @Advice.Argument(value = 0, readOnly = false)
- GenericFutureListener<? extends Future<? super Void>>[] listeners) {
+ GenericFutureListener<? extends Future<?>>[] listeners) {
 
- ContextStore<GenericFutureListener, GenericFutureListener> contextStore =
- InstrumentationContext.get(GenericFutureListener.class, GenericFutureListener.class);
  Context context = Java8BytecodeBridge.currentContext();
  @SuppressWarnings("unchecked")
- GenericFutureListener<? extends Future<? super Void>>[] wrappedListeners =
+ GenericFutureListener<? extends Future<?>>[] wrappedListeners =
  new GenericFutureListener[listeners.length];
  for (int i = 0; i < listeners.length; ++i) {
- wrappedListeners[i] = FutureListenerWrappers.wrap(contextStore, context, listeners[i]);
+ wrappedListeners[i] = FutureListenerWrappers.wrap(context, listeners[i]);
  }
  listeners = wrappedListeners;
  }
@@ -91,26 +84,22 @@ public static class RemoveListenerAdvice {
  @Advice.OnMethodEnter
  public static void wrapListener(
  @Advice.Argument(value = 0, readOnly = false)
- GenericFutureListener<? extends Future<? super Void>> listener) {
- ContextStore<GenericFutureListener, GenericFutureListener> contextStore =
- InstrumentationContext.get(GenericFutureListener.class, GenericFutureListener.class);
- listener = FutureListenerWrappers.getWrapper(contextStore, listener);
+ GenericFutureListener<? extends Future<?>> listener) {
+ listener = FutureListenerWrappers.getWrapper(listener);
  }
  }
 
  public static class RemoveListenersAdvice {
  @Advice.OnMethodEnter
  public static void wrapListener(
  @Advice.Argument(value = 0, readOnly = false)
- GenericFutureListener<? extends Future<? super Void>>[] listeners) {
+ GenericFutureListener<? extends Future<?>>[] listeners) {
 
- ContextStore<GenericFutureListener, GenericFutureListener> contextStore =
- InstrumentationContext.get(GenericFutureListener.class, GenericFutureListener.class);
  @SuppressWarnings("unchecked")
- GenericFutureListener<? extends Future<? super Void>>[] wrappedListeners =
+ GenericFutureListener<? extends Future<?>>[] wrappedListeners =
  new GenericFutureListener[listeners.length];
  for (int i = 0; i < listeners.length; ++i) {
- wrappedListeners[i] = FutureListenerWrappers.getWrapper(contextStore, listeners[i]);
+ wrappedListeners[i] = FutureListenerWrappers.getWrapper(listeners[i]);
  }
  listeners = wrappedListeners;
  }