From ce570e09843f1ff861fc37d9b35289204690ccea Mon Sep 17 00:00:00 2001 From: Jaroslav Bachorik Date: Wed, 14 Oct 2026 16:37:14 +0200 Subject: [PATCH] feat(profiling): Reference-count RecordingData lifecycle RecordingData gains explicit retain/release semantics so multiple consumers (recording streaming and the OTLP uploader) can share a recording without one side closing it prematurely. Implemented across the OpenJDK, OracleJDK and ddprof controller backends. --- .../openjdk/OpenJdkRecordingData.java | 2 +- .../oracle/OracleJdkRecordingData.java | 18 +- .../RecordingDataRefCountingTest.java | 223 ++++++++++++++++++ .../ddprof/DatadogProfilerRecordingData.java | 2 +- .../profiling/agent/CompositeController.java | 2 +- .../agent/ScrubRecordingDataListener.java | 2 +- .../trace/api/profiling/RecordingData.java | 76 +++++- 7 files changed, 309 insertions(+), 16 deletions(-) create mode 100644 dd-java-agent/agent-profiling/profiling-controller/src/test/java/com/datadog/profiling/controller/RecordingDataRefCountingTest.java diff --git a/dd-java-agent/agent-profiling/profiling-controller-openjdk/src/main/java/com/datadog/profiling/controller/openjdk/OpenJdkRecordingData.java b/dd-java-agent/agent-profiling/profiling-controller-openjdk/src/main/java/com/datadog/profiling/controller/openjdk/OpenJdkRecordingData.java index da022c226ed..3c25699ac04 100644 --- a/dd-java-agent/agent-profiling/profiling-controller-openjdk/src/main/java/com/datadog/profiling/controller/openjdk/OpenJdkRecordingData.java +++ b/dd-java-agent/agent-profiling/profiling-controller-openjdk/src/main/java/com/datadog/profiling/controller/openjdk/OpenJdkRecordingData.java @@ -45,7 +45,7 @@ public RecordingInputStream getStream() throws IOException { } @Override - public void release() { + protected void doRelease() { recording.close(); } diff --git a/dd-java-agent/agent-profiling/profiling-controller-oracle/src/main/java/com/datadog/profiling/controller/oracle/OracleJdkRecordingData.java b/dd-java-agent/agent-profiling/profiling-controller-oracle/src/main/java/com/datadog/profiling/controller/oracle/OracleJdkRecordingData.java index 226e1b0f24d..04c98e1ae7a 100644 --- a/dd-java-agent/agent-profiling/profiling-controller-oracle/src/main/java/com/datadog/profiling/controller/oracle/OracleJdkRecordingData.java +++ b/dd-java-agent/agent-profiling/profiling-controller-oracle/src/main/java/com/datadog/profiling/controller/oracle/OracleJdkRecordingData.java @@ -23,9 +23,13 @@ import java.util.Date; import javax.annotation.Nonnull; import javax.management.ObjectName; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** Implementation for profiling recordings. */ public class OracleJdkRecordingData extends RecordingData { + private static final Logger log = LoggerFactory.getLogger(OracleJdkRecordingData.class); + private final ObjectName recordingId; private final String name; @@ -51,8 +55,15 @@ public RecordingInputStream getStream() throws IOException { } @Override - public void release() { - // noop + protected void doRelease() { + // the recording is shared by every listener's stream — the OTLP listener may close its + // stream while the classic uploader still needs to open its own, so the recording itself + // can only be closed when the last reference is released + try { + helper.closeRecording(recordingId); + } catch (Exception e) { + log.warn("Failed to close Oracle JDK recording {}", recordingId, e); + } } @Override @@ -116,7 +127,8 @@ public void close() throws IOException { if (streamId != -1) { helper.closeStream(streamId); } - helper.closeRecording(recordingId); + // the recording itself is closed in doRelease() — closing it here would prevent any + // other listener from opening its own stream on the same recording } catch (Exception e) { throw new IOException(e); } diff --git a/dd-java-agent/agent-profiling/profiling-controller/src/test/java/com/datadog/profiling/controller/RecordingDataRefCountingTest.java b/dd-java-agent/agent-profiling/profiling-controller/src/test/java/com/datadog/profiling/controller/RecordingDataRefCountingTest.java new file mode 100644 index 00000000000..99358a6b7ba --- /dev/null +++ b/dd-java-agent/agent-profiling/profiling-controller/src/test/java/com/datadog/profiling/controller/RecordingDataRefCountingTest.java @@ -0,0 +1,223 @@ +/* + * Copyright 2026 Datadog + * + * 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.datadog.profiling.controller; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.trace.api.profiling.ProfilingSnapshot; +import datadog.trace.api.profiling.RecordingData; +import datadog.trace.api.profiling.RecordingInputStream; +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.nio.file.Path; +import java.time.Instant; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import org.junit.jupiter.api.Test; + +/** + * Tests for RecordingData reference counting with multiple handlers. The reference count starts at + * 1 (the base reference). The primary listener releases the base reference; each additional + * listener calls retain() then release(). + */ +public class RecordingDataRefCountingTest { + + /** Test RecordingData implementation that tracks release calls. */ + private static class TestRecordingData extends RecordingData { + private final AtomicInteger releaseCount = new AtomicInteger(0); + private final CountDownLatch releaseLatch = new CountDownLatch(1); + + public TestRecordingData() { + super(Instant.now(), Instant.now(), ProfilingSnapshot.Kind.PERIODIC); + } + + @Nonnull + @Override + public RecordingInputStream getStream() throws IOException { + return new RecordingInputStream(new ByteArrayInputStream(new byte[0])); + } + + @Override + protected void doRelease() { + releaseCount.incrementAndGet(); + releaseLatch.countDown(); + } + + @Nullable + @Override + public Path getPath() { + return null; + } + + @Override + public String getName() { + return "test-recording"; + } + + public int getReleaseCount() { + return releaseCount.get(); + } + + public boolean awaitRelease(long timeout, TimeUnit unit) throws InterruptedException { + return releaseLatch.await(timeout, unit); + } + } + + @Test + public void testSingleHandler() throws InterruptedException { + TestRecordingData data = new TestRecordingData(); + + // Single handler: just release the base reference (no retain needed) + data.release(); + + assertTrue(data.awaitRelease(1, TimeUnit.SECONDS), "Release should be called"); + assertEquals(1, data.getReleaseCount(), "doRelease() should be called exactly once"); + } + + @Test + public void testTwoHandlers() throws InterruptedException { + TestRecordingData data = new TestRecordingData(); + + // Two handlers: retain once for the additional handler, release twice + data.retain(); // Additional handler (e.g., OTLP) + assertEquals(0, data.getReleaseCount(), "Should not be released yet"); + + // First handler releases (base reference) + data.release(); + assertEquals(0, data.getReleaseCount(), "Should not be released after first release"); + + // Second handler releases + data.release(); + + assertTrue(data.awaitRelease(1, TimeUnit.SECONDS), "Release should be called"); + assertEquals(1, data.getReleaseCount(), "doRelease() should be called exactly once"); + } + + @Test + public void testThreeHandlers() throws InterruptedException { + TestRecordingData data = new TestRecordingData(); + + // Three handlers: retain twice, release three times + data.retain(); // Additional handler 1 + data.retain(); // Additional handler 2 + assertEquals(0, data.getReleaseCount(), "Should not be released yet"); + + // First two handlers release + data.release(); + data.release(); + assertEquals(0, data.getReleaseCount(), "Should not be released after two releases"); + + // Third handler releases + data.release(); + + assertTrue(data.awaitRelease(1, TimeUnit.SECONDS), "Release should be called"); + assertEquals(1, data.getReleaseCount(), "doRelease() should be called exactly once"); + } + + @Test + public void testOverReleaseIsSilentNoOp() throws InterruptedException { + TestRecordingData data = new TestRecordingData(); + + // Release the base reference + data.release(); + assertTrue(data.awaitRelease(1, TimeUnit.SECONDS), "Release should be called"); + + // Additional release after full release is a silent no-op — unlike retain() after full release + data.release(); + assertEquals(1, data.getReleaseCount(), "doRelease() should still be called exactly once"); + } + + @Test + public void testRetainAfterFullRelease() throws InterruptedException { + TestRecordingData data = new TestRecordingData(); + + data.release(); + assertTrue(data.awaitRelease(1, TimeUnit.SECONDS), "Release should be called"); + + // Cannot retain after full release + assertThrows( + IllegalStateException.class, data::retain, "Should throw when retaining after release"); + } + + @Test + public void testMultipleReleaseIdempotent() throws InterruptedException { + TestRecordingData data = new TestRecordingData(); + + data.release(); + assertTrue(data.awaitRelease(1, TimeUnit.SECONDS), "Release should be called"); + + // Additional release calls should be silent no-ops (released flag guards them) + data.release(); + data.release(); + + assertEquals(1, data.getReleaseCount(), "doRelease() should still be called exactly once"); + } + + @Test + public void testConcurrentHandlers() throws InterruptedException { + TestRecordingData data = new TestRecordingData(); + int numAdditionalHandlers = 9; // total 10 with the base reference + + // Retain for all additional handlers + for (int i = 0; i < numAdditionalHandlers; i++) { + data.retain(); + } + + int totalHandlers = numAdditionalHandlers + 1; + + // Simulate concurrent release from multiple threads + CountDownLatch startLatch = new CountDownLatch(1); + CountDownLatch doneLatch = new CountDownLatch(totalHandlers); + + for (int i = 0; i < totalHandlers; i++) { + new Thread( + () -> { + try { + startLatch.await(); + data.release(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + // count down even on interrupt, otherwise it masks as a liveness timeout below + doneLatch.countDown(); + } + }) + .start(); + } + + // Start all threads + startLatch.countDown(); + + // Wait for all threads to complete + assertTrue(doneLatch.await(5, TimeUnit.SECONDS), "All threads should complete"); + assertTrue(data.awaitRelease(1, TimeUnit.SECONDS), "Release should be called"); + assertEquals(1, data.getReleaseCount(), "doRelease() should be called exactly once"); + } + + @Test + public void testRetainChaining() { + TestRecordingData data = new TestRecordingData(); + + // retain() should return this for chaining + RecordingData result = data.retain(); + assertEquals(data, result, "retain() should return the same instance"); + } +} diff --git a/dd-java-agent/agent-profiling/profiling-ddprof/src/main/java/com/datadog/profiling/ddprof/DatadogProfilerRecordingData.java b/dd-java-agent/agent-profiling/profiling-ddprof/src/main/java/com/datadog/profiling/ddprof/DatadogProfilerRecordingData.java index b1a6c7eea41..fd142bcde43 100644 --- a/dd-java-agent/agent-profiling/profiling-ddprof/src/main/java/com/datadog/profiling/ddprof/DatadogProfilerRecordingData.java +++ b/dd-java-agent/agent-profiling/profiling-ddprof/src/main/java/com/datadog/profiling/ddprof/DatadogProfilerRecordingData.java @@ -24,7 +24,7 @@ public RecordingInputStream getStream() throws IOException { } @Override - public void release() { + protected void doRelease() { try { Files.deleteIfExists(recordingFile); } catch (IOException e) { diff --git a/dd-java-agent/agent-profiling/src/main/java/com/datadog/profiling/agent/CompositeController.java b/dd-java-agent/agent-profiling/src/main/java/com/datadog/profiling/agent/CompositeController.java index 236bf52b1ae..581d099ff51 100644 --- a/dd-java-agent/agent-profiling/src/main/java/com/datadog/profiling/agent/CompositeController.java +++ b/dd-java-agent/agent-profiling/src/main/java/com/datadog/profiling/agent/CompositeController.java @@ -118,7 +118,7 @@ public RecordingInputStream getStream() throws IOException { } @Override - public void release() { + protected void doRelease() { for (RecordingData data : data) { data.release(); } diff --git a/dd-java-agent/agent-profiling/src/main/java/com/datadog/profiling/agent/ScrubRecordingDataListener.java b/dd-java-agent/agent-profiling/src/main/java/com/datadog/profiling/agent/ScrubRecordingDataListener.java index 211d0e23c64..f5931bfa68c 100644 --- a/dd-java-agent/agent-profiling/src/main/java/com/datadog/profiling/agent/ScrubRecordingDataListener.java +++ b/dd-java-agent/agent-profiling/src/main/java/com/datadog/profiling/agent/ScrubRecordingDataListener.java @@ -139,7 +139,7 @@ public RecordingInputStream getStream() throws IOException { } @Override - public void release() { + protected void doRelease() { try { Files.deleteIfExists(scrubbedFile); } catch (IOException e) { diff --git a/internal-api/src/main/java/datadog/trace/api/profiling/RecordingData.java b/internal-api/src/main/java/datadog/trace/api/profiling/RecordingData.java index d62f1541c3e..69b21fc0e88 100644 --- a/internal-api/src/main/java/datadog/trace/api/profiling/RecordingData.java +++ b/internal-api/src/main/java/datadog/trace/api/profiling/RecordingData.java @@ -21,12 +21,29 @@ import javax.annotation.Nonnull; import javax.annotation.Nullable; -/** Platform-agnostic API for operations required when retrieving data using the ProfilingSystem. */ +/** + * Platform-agnostic API for operations required when retrieving data using the ProfilingSystem. + * + *

Subclassing note: {@link #release()} is final and reference-counted; subclasses implement + * resource cleanup by overriding {@link #doRelease()}, invoked exactly once when the reference + * count reaches zero. + */ public abstract class RecordingData implements ProfilingSnapshot { protected final Instant start; protected final Instant end; protected final Kind kind; + // Reference counting for multiple listeners. Starts at 1 (the base reference owned by the + // first listener in the chain). Additional listeners call retain() before use and release() + // when done. doRelease() fires when the count reaches 0. + // Plain fields suffice: both are only ever accessed under the private lock below. + private int refCount = 1; + private boolean released = false; + + // Private lock so the intrinsic lock of this publicly exposed class is not used for + // synchronization (SpotBugs USO_UNSAFE_METHOD_SYNCHRONIZATION / CERT LCK00-J). + private final Object lock = new Object(); + public RecordingData(final Instant start, final Instant end, Kind kind) { this.start = start; this.end = end; @@ -41,18 +58,59 @@ public RecordingData(final Instant start, final Instant end, Kind kind) { public abstract RecordingInputStream getStream() throws IOException; /** - * Releases the resources associated with the recording, for example the underlying file. + * Must be called once for each handler beyond the first, before it processes the recording. The + * counting scheme is documented on the {@code refCount} field. + * + * @return this instance for chaining + * @throws IllegalStateException if the recording has already been released + */ + @Nonnull + public final RecordingData retain() { + synchronized (lock) { + if (released) { + throw new IllegalStateException("Cannot retain released RecordingData"); + } + refCount++; + return this; + } + } + + /** + * Releases one reference to the recording. The resources are released when the reference count + * reaches zero (see the {@code refCount} field for the counting scheme). * - *

Forgetting to releasing this when done streaming, will lead to one or more of the following: + *

Forgetting to release when done streaming leaks the recording file and memory. * - *

+ *

Releasing an already fully-released recording is a silent no-op: the {@code released} flag + * guards against double release. Retaining after full release, in contrast, throws {@link + * IllegalStateException}. + */ + public final void release() { + boolean shouldRelease = false; + synchronized (lock) { + if (released) { + // retain() refuses to hand out new references once released, so the count can never + // reach zero a second time + return; + } + int remaining = --refCount; + if (remaining == 0) { + released = true; + shouldRelease = true; + } + } + if (shouldRelease) { + doRelease(); + } + } + + /** + * Actual resource cleanup implementation. Subclasses must override this method instead of {@link + * #release()}. * - *

Please don't forget to call release when done streaming... + *

This method is called exactly once when the reference count reaches zero. */ - public abstract void release(); + protected abstract void doRelease(); /** * Returns the name of the recording from which the data is originating.