ref counting

Author: trask

instrumentation/opentelemetry-api/opentelemetry-api-1.10/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/v1_10/metrics/ApplicationDoubleCounterBuilder.java

@@ -49,7 +49,7 @@ public application.io.opentelemetry.api.metrics.ObservableDoubleCounter buildWit
                 new ApplicationObservableDoubleMeasurement(agentMeasurement));
     return new ApplicationObservableDoubleCounter(
         CallbackAnchor.anchor(agentBuilder::buildWithCallback, callback),
-        () -> CallbackAnchor.remove(callback));
+        CallbackAnchor.releaseOnClose(callback));
   }
 
   // added in 1.15.0

instrumentation/opentelemetry-api/opentelemetry-api-1.10/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/v1_10/metrics/ApplicationDoubleGaugeBuilder.java

@@ -49,7 +49,7 @@ public application.io.opentelemetry.api.metrics.ObservableDoubleGauge buildWithC
                 new ApplicationObservableDoubleMeasurement(agentMeasurement));
     return new ApplicationObservableDoubleGauge(
         CallbackAnchor.anchor(agentBuilder::buildWithCallback, callback),
-        () -> CallbackAnchor.remove(callback));
+        CallbackAnchor.releaseOnClose(callback));
   }
 
   // added in 1.15.0

instrumentation/opentelemetry-api/opentelemetry-api-1.10/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/v1_10/metrics/ApplicationDoubleUpDownCounterBuilder.java

@@ -49,7 +49,7 @@ public application.io.opentelemetry.api.metrics.ObservableDoubleUpDownCounter bu
                 new ApplicationObservableDoubleMeasurement(agentMeasurement));
     return new ApplicationObservableDoubleUpDownCounter(
         CallbackAnchor.anchor(agentBuilder::buildWithCallback, callback),
-        () -> CallbackAnchor.remove(callback));
+        CallbackAnchor.releaseOnClose(callback));
   }
 
   // added in 1.15.0

instrumentation/opentelemetry-api/opentelemetry-api-1.10/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/v1_10/metrics/ApplicationLongCounterBuilder.java

@@ -53,7 +53,7 @@ public application.io.opentelemetry.api.metrics.ObservableLongCounter buildWithC
             applicationCallback.accept(new ApplicationObservableLongMeasurement(agentMeasurement));
     return new ApplicationObservableLongCounter(
         CallbackAnchor.anchor(agentBuilder::buildWithCallback, callback),
-        () -> CallbackAnchor.remove(callback));
+        CallbackAnchor.releaseOnClose(callback));
   }
 
   // added in 1.15.0

instrumentation/opentelemetry-api/opentelemetry-api-1.10/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/v1_10/metrics/ApplicationLongGaugeBuilder.java

@@ -43,7 +43,7 @@ public application.io.opentelemetry.api.metrics.ObservableLongGauge buildWithCal
             applicationCallback.accept(new ApplicationObservableLongMeasurement(agentMeasurement));
     return new ApplicationObservableLongGauge(
         CallbackAnchor.anchor(agentBuilder::buildWithCallback, callback),
-        () -> CallbackAnchor.remove(callback));
+        CallbackAnchor.releaseOnClose(callback));
   }
 
   // added in 1.15.0

instrumentation/opentelemetry-api/opentelemetry-api-1.10/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/v1_10/metrics/ApplicationLongUpDownCounterBuilder.java

@@ -53,7 +53,7 @@ public application.io.opentelemetry.api.metrics.ObservableLongUpDownCounter buil
             applicationCallback.accept(new ApplicationObservableLongMeasurement(agentMeasurement));
     return new ApplicationObservableLongUpDownCounter(
         CallbackAnchor.anchor(agentBuilder::buildWithCallback, callback),
-        () -> CallbackAnchor.remove(callback));
+        CallbackAnchor.releaseOnClose(callback));
   }
 
   // added in 1.15.0

instrumentation/opentelemetry-api/opentelemetry-api-1.10/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/v1_10/metrics/CallbackAnchor.java

@@ -8,8 +8,9 @@
 import io.opentelemetry.javaagent.bootstrap.WeakRefConsumer;
 import io.opentelemetry.javaagent.bootstrap.WeakRefRunnable;
 import java.lang.ref.WeakReference;
-import java.util.Set;
+import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
 import java.util.function.Function;
 
@@ -40,12 +41,13 @@ public final class CallbackAnchor {
   // Anchors callbacks to this class's lifecycle. Since this class is injected as a helper into each
   // application class loader, callbacks are naturally tied to their class loader's lifecycle.
   // Use identity semantics so unusual equals/hashCode implementations don't collapse distinct
-  // callbacks into the same entry.
-  private static final Set<IdentityKey> callbacks = ConcurrentHashMap.newKeySet();
+  // callbacks into the same entry, and ref counts so the same callback can back multiple
+  // instruments without closing one instrument dropping the shared anchor too early.
+  private static final Map<IdentityKey, Integer> callbacks = new ConcurrentHashMap<>();
 
   public static <T, R extends AutoCloseable> R anchor(
       Function<Consumer<T>, R> buildFn, Consumer<T> callback) {
-    callbacks.add(new IdentityKey(callback));
+    callbacks.merge(new IdentityKey(callback), 1, Integer::sum);
     WeakRefConsumer<T> weak = new WeakRefConsumer<>(new WeakReference<>(callback));
     R instrument = buildFn.apply(weak);
     weak.closeWhenCollected(instrument);
@@ -54,15 +56,29 @@ public static <T, R extends AutoCloseable> R anchor(
 
   public static <R extends AutoCloseable> R anchorBatch(
       Function<Runnable, R> buildFn, Runnable callback) {
-    callbacks.add(new IdentityKey(callback));
+    callbacks.merge(new IdentityKey(callback), 1, Integer::sum);
     WeakRefRunnable weak = new WeakRefRunnable(new WeakReference<>(callback));
     R instrument = buildFn.apply(weak);
     weak.closeWhenCollected(instrument);
     return instrument;
   }
 
-  public static void remove(Object callback) {
-    callbacks.remove(new IdentityKey(callback));
+  // Some instrument wrappers may end up calling close() more than once. Guard the deferred
+  // release so each wrapper decrements the callback ref count at most once.
+  public static Runnable releaseOnClose(Object callback) {
+    AtomicBoolean released = new AtomicBoolean();
+    return () -> {
+      if (released.compareAndSet(false, true)) {
+        release(callback);
+      }
+    };
+  }
+
+  private static void release(Object callback) {
+    // Returning null from computeIfPresent removes the mapping
+    // once the ref count reaches zero.
+    callbacks.computeIfPresent(
+        new IdentityKey(callback), (key, count) -> count == 1 ? null : count - 1);
   }
 
   private static final class IdentityKey {

instrumentation/opentelemetry-api/opentelemetry-api-1.15/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/v1_15/metrics/ApplicationMeter115.java

@@ -31,7 +31,7 @@ public application.io.opentelemetry.api.metrics.BatchCallback batchCallback(
                 agentMeter.batchCallback(
                     weak, unwrap(observableMeasurement), unwrap(additionalMeasurements)),
             callback),
-        () -> CallbackAnchor.remove(callback));
+        CallbackAnchor.releaseOnClose(callback));
   }
 
   private static ObservableMeasurement unwrap(

instrumentation/opentelemetry-api/opentelemetry-api-1.15/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/opentelemetryapi/v1_15/metrics/MeterTest.java

@@ -14,8 +14,11 @@
 import io.opentelemetry.api.metrics.Meter;
 import io.opentelemetry.api.metrics.ObservableDoubleMeasurement;
 import io.opentelemetry.api.metrics.ObservableLongMeasurement;
+import io.opentelemetry.instrumentation.test.utils.GcUtils;
 import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
 import io.opentelemetry.sdk.common.InstrumentationScopeInfo;
+import java.lang.ref.WeakReference;
+import java.time.Duration;
 import org.assertj.core.api.AbstractIterableAssert;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -321,4 +324,55 @@ void batchDoubleGauge() throws InterruptedException {
 
     testing.waitAndAssertMetrics(instrumentationName, "test", AbstractIterableAssert::isEmpty);
   }
+
+  @Test
+  void sharedBatchCallbackSurvivesClosingOneHandleTwice() throws Exception {
+    ObservableLongMeasurement firstMeasurement = meter.counterBuilder("test.first").buildObserver();
+    ObservableLongMeasurement secondMeasurement =
+        meter.counterBuilder("test.second").buildObserver();
+    Runnable callback =
+        new Runnable() {
+          @Override
+          public void run() {}
+        };
+    WeakReference<Runnable> callbackRef = new WeakReference<>(callback);
+
+    BatchCallback firstCallback = meter.batchCallback(callback, firstMeasurement);
+    // Intentionally do not retain the second handle so callback reachability depends on the
+    // remaining anchor, not on another wrapper still holding the callback via its onClose action.
+    meter.batchCallback(callback, secondMeasurement);
+
+    firstCallback.close();
+    firstCallback.close();
+    firstCallback = null;
+    callback = null;
+
+    GcUtils.awaitGc(Duration.ofSeconds(10));
+    assertThat(callbackRef.get()).isNotNull();
+  }
+
+  @Test
+  void sharedBatchCallbackCollectedAfterLastHandleClosed() throws Exception {
+    ObservableLongMeasurement firstMeasurement = meter.counterBuilder("test.first").buildObserver();
+    ObservableLongMeasurement secondMeasurement =
+        meter.counterBuilder("test.second").buildObserver();
+    Runnable callback =
+        new Runnable() {
+          @Override
+          public void run() {}
+        };
+    WeakReference<Runnable> callbackRef = new WeakReference<>(callback);
+
+    BatchCallback firstCallback = meter.batchCallback(callback, firstMeasurement);
+    BatchCallback secondCallback = meter.batchCallback(callback, secondMeasurement);
+
+    firstCallback.close();
+    secondCallback.close();
+    firstCallback = null;
+    secondCallback = null;
+    callback = null;
+
+    GcUtils.awaitGc(callbackRef, Duration.ofSeconds(10));
+    assertThat(callbackRef.get()).isNull();
+  }
 }