diff --git a/async-profiler-context/src/main/java/io/pyroscope/labels/v2/Pyroscope.java b/async-profiler-context/src/main/java/io/pyroscope/labels/v2/Pyroscope.java index 48cfc898..5d852368 100644 --- a/async-profiler-context/src/main/java/io/pyroscope/labels/v2/Pyroscope.java +++ b/async-profiler-context/src/main/java/io/pyroscope/labels/v2/Pyroscope.java @@ -60,6 +60,11 @@ public static void clear() { ScopedContext.CONTEXTS.clear(); } + public static LabelsSet getCurrentLabels() { + ScopedContext context = ScopedContext.getCurrentContext(); + return context != null ? context.labels : new LabelsSet(); + } + public static JfrLabels.LabelsSnapshot dump() { final JfrLabels.LabelsSnapshot.Builder sb = JfrLabels.LabelsSnapshot.newBuilder(); final StringTableBuilder stb = new StringTableBuilder(); diff --git a/async-profiler-context/src/main/java/io/pyroscope/labels/v2/ScopedContext.java b/async-profiler-context/src/main/java/io/pyroscope/labels/v2/ScopedContext.java index b11f0558..6cb9a029 100644 --- a/async-profiler-context/src/main/java/io/pyroscope/labels/v2/ScopedContext.java +++ b/async-profiler-context/src/main/java/io/pyroscope/labels/v2/ScopedContext.java @@ -52,6 +52,7 @@ public final class ScopedContext implements AutoCloseable { public static final AtomicBoolean ENABLED = new AtomicBoolean(false); static final AtomicLong CONTEXT_COUNTER = new AtomicLong(0); + static ThreadLocal CURRENT_CONTEXT_ID = ThreadLocal.withInitial(() -> 0L); static final ConcurrentHashMap CONTEXTS = new ConcurrentHashMap<>(); static final ConcurrentHashMap CONSTANT_CONTEXTS = new ConcurrentHashMap<>(); @@ -70,6 +71,10 @@ static AsyncProfiler getAsyncProfiler() { final long prevContextId; final AtomicBoolean closed = new AtomicBoolean(false); + public static ScopedContext getCurrentContext() { + return CONTEXTS.get(CURRENT_CONTEXT_ID.get()); + } + /** * Creates a new ScopedContext with the given labels. * The previous context ID is set to 0 (root context). @@ -113,6 +118,7 @@ public ScopedContext(@NotNull LabelsSet labels, @NotNull ScopedContext prev) { this.prevContextId = prevContextId; CONTEXTS.put(contextId, this); getAsyncProfiler().setContextId(contextId); + CURRENT_CONTEXT_ID.set(contextId); } else { this.contextId = 0; this.prevContextId = 0; @@ -133,6 +139,7 @@ public void close() { } if (ENABLED.get()) { getAsyncProfiler().setContextId(this.prevContextId); + CURRENT_CONTEXT_ID.set(this.prevContextId); } } diff --git a/async-profiler-context/src/test/java/io/pyroscope/labels/v2/LabelsTest.java b/async-profiler-context/src/test/java/io/pyroscope/labels/v2/LabelsTest.java index 1291981d..0e779ca9 100644 --- a/async-profiler-context/src/test/java/io/pyroscope/labels/v2/LabelsTest.java +++ b/async-profiler-context/src/test/java/io/pyroscope/labels/v2/LabelsTest.java @@ -96,6 +96,14 @@ void testNestedEqualLabelSets() { } } + @Test + void testCaptureContext() { + try (ScopedContext s = new ScopedContext(new LabelsSet("k1", "v1"))) { + assertEquals(s.contextId, ScopedContext.CURRENT_CONTEXT_ID.get()); + } + assertEquals(0, ScopedContext.CURRENT_CONTEXT_ID.get()); + } + @Test void exception() {