From 29990721157a17b28bee1933b327f9187bc62afd Mon Sep 17 00:00:00 2001
From: Aleksandar Petrov <8142643+aleks-p@users.noreply.github.com>
Date: Fri, 20 Dec 2024 12:11:18 -0400
Subject: [PATCH] Add ProfilerApi for consumption by other libraries

---
 .../io/pyroscope/javaagent/ProfilerSdk.java   | 28 +++++++++++++++++++
 .../pyroscope/javaagent/api/ProfilerApi.java  |  9 ++++++
 .../javaagent/api/ProfilerScopedContext.java  |  8 ++++++
 .../impl/ProfilerScopedContextWrapper.java    | 24 ++++++++++++++++
 .../io/pyroscope/labels/ScopedContext.java    |  5 ++--
 5 files changed, 71 insertions(+), 3 deletions(-)
 create mode 100644 agent/src/main/java/io/pyroscope/javaagent/ProfilerSdk.java
 create mode 100644 agent/src/main/java/io/pyroscope/javaagent/api/ProfilerApi.java
 create mode 100644 agent/src/main/java/io/pyroscope/javaagent/api/ProfilerScopedContext.java
 create mode 100644 agent/src/main/java/io/pyroscope/javaagent/impl/ProfilerScopedContextWrapper.java

diff --git a/agent/src/main/java/io/pyroscope/javaagent/ProfilerSdk.java b/agent/src/main/java/io/pyroscope/javaagent/ProfilerSdk.java
new file mode 100644
index 0000000..d4054fb
--- /dev/null
+++ b/agent/src/main/java/io/pyroscope/javaagent/ProfilerSdk.java
@@ -0,0 +1,28 @@
+package io.pyroscope.javaagent;
+
+import io.pyroscope.javaagent.api.ProfilerScopedContext;
+import io.pyroscope.javaagent.api.ProfilerApi;
+import io.pyroscope.javaagent.config.Config;
+import io.pyroscope.javaagent.impl.ProfilerScopedContextWrapper;
+import io.pyroscope.labels.LabelsSet;
+import io.pyroscope.labels.ScopedContext;
+
+import java.util.Map;
+
+public class ProfilerSdk implements ProfilerApi {
+
+    @Override
+    public void startProfiling() {
+        PyroscopeAgent.start(Config.build());
+    }
+
+    @Override
+    public boolean isProfilingStarted() {
+        return PyroscopeAgent.isStarted();
+    }
+
+    @Override
+    public ProfilerScopedContext createScopedContext(Map<String, String> labels) {
+        return new ProfilerScopedContextWrapper(new ScopedContext(new LabelsSet(labels)));
+    }
+}
diff --git a/agent/src/main/java/io/pyroscope/javaagent/api/ProfilerApi.java b/agent/src/main/java/io/pyroscope/javaagent/api/ProfilerApi.java
new file mode 100644
index 0000000..f984b7d
--- /dev/null
+++ b/agent/src/main/java/io/pyroscope/javaagent/api/ProfilerApi.java
@@ -0,0 +1,9 @@
+package io.pyroscope.javaagent.api;
+
+import java.util.Map;
+
+public interface ProfilerApi {
+    void startProfiling();
+    boolean isProfilingStarted();
+    ProfilerScopedContext createScopedContext(Map<String, String> labels);
+}
diff --git a/agent/src/main/java/io/pyroscope/javaagent/api/ProfilerScopedContext.java b/agent/src/main/java/io/pyroscope/javaagent/api/ProfilerScopedContext.java
new file mode 100644
index 0000000..80dccc3
--- /dev/null
+++ b/agent/src/main/java/io/pyroscope/javaagent/api/ProfilerScopedContext.java
@@ -0,0 +1,8 @@
+package io.pyroscope.javaagent.api;
+
+import java.util.function.BiConsumer;
+
+public interface ProfilerScopedContext {
+    void forEachLabel(BiConsumer<String, String> consumer);
+    void close();
+}
diff --git a/agent/src/main/java/io/pyroscope/javaagent/impl/ProfilerScopedContextWrapper.java b/agent/src/main/java/io/pyroscope/javaagent/impl/ProfilerScopedContextWrapper.java
new file mode 100644
index 0000000..e46e2e7
--- /dev/null
+++ b/agent/src/main/java/io/pyroscope/javaagent/impl/ProfilerScopedContextWrapper.java
@@ -0,0 +1,24 @@
+package io.pyroscope.javaagent.impl;
+
+import io.pyroscope.javaagent.api.ProfilerScopedContext;
+import io.pyroscope.labels.ScopedContext;
+
+import java.util.function.BiConsumer;
+
+public class ProfilerScopedContextWrapper implements ProfilerScopedContext {
+    private final ScopedContext ctx;
+
+    public ProfilerScopedContextWrapper(ScopedContext ctx) {
+        this.ctx = ctx;
+    }
+
+    @Override
+    public void forEachLabel(BiConsumer<String, String> consumer) {
+        ctx.forEachLabel(consumer);
+    }
+
+    @Override
+    public void close() {
+        ctx.close();
+    }
+}
diff --git a/async-profiler-context/src/main/java/io/pyroscope/labels/ScopedContext.java b/async-profiler-context/src/main/java/io/pyroscope/labels/ScopedContext.java
index a2b9213..30673b6 100644
--- a/async-profiler-context/src/main/java/io/pyroscope/labels/ScopedContext.java
+++ b/async-profiler-context/src/main/java/io/pyroscope/labels/ScopedContext.java
@@ -61,7 +61,6 @@ public ScopedContext(LabelsSet labels) {
         context.set(current);
     }
 
-
     @Override
     public void close() {
         if (closed) {
@@ -73,9 +72,9 @@ public void close() {
         AsyncProfiler.getInstance().setContextId(previous.id);
     }
 
-    public void forEach(BiConsumer<String, String> consumer) {
+    public void forEachLabel(BiConsumer<String, String> labelConsumer) {
         for (Map.Entry<Ref<String>, Ref<String>> it : current.labels.entrySet()) {
-            consumer.accept(it.getKey().val, it.getValue().val);
+            labelConsumer.accept(it.getKey().val, it.getValue().val);
         }
     }