From df276a069c5f5c989481f455e2f245f637fdff0b Mon Sep 17 00:00:00 2001
From: Andrea Marziali <andrea.marziali@datadoghq.com>
Date: Wed, 7 May 2025 14:11:20 +0200
Subject: [PATCH] Collect process tags for profiling upload requests

---
 .../profiling/uploader/ProfileUploader.java   |  6 ++++
 .../uploader/ProfileUploaderTest.java         | 29 +++++++++++++++++++
 .../java/datadog/trace/api/ProcessTags.java   |  7 ++++-
 3 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/dd-java-agent/agent-profiling/profiling-uploader/src/main/java/com/datadog/profiling/uploader/ProfileUploader.java b/dd-java-agent/agent-profiling/profiling-uploader/src/main/java/com/datadog/profiling/uploader/ProfileUploader.java
index c8b239d9326..ab588da6e1a 100644
--- a/dd-java-agent/agent-profiling/profiling-uploader/src/main/java/com/datadog/profiling/uploader/ProfileUploader.java
+++ b/dd-java-agent/agent-profiling/profiling-uploader/src/main/java/com/datadog/profiling/uploader/ProfileUploader.java
@@ -27,6 +27,7 @@
 import datadog.trace.api.Config;
 import datadog.trace.api.DDTags;
 import datadog.trace.api.Platform;
+import datadog.trace.api.ProcessTags;
 import datadog.trace.api.git.GitInfo;
 import datadog.trace.api.git.GitInfoProvider;
 import datadog.trace.api.profiling.RecordingData;
@@ -469,6 +470,7 @@ public void toJson(JsonWriter writer, RecordingData recordingData) throws IOExce
       if (recordingData == null) {
         return;
       }
+      final CharSequence processTags = ProcessTags.getTagsForSerialization();
       writer.beginObject();
       writer.name("attachments");
       writer.beginArray();
@@ -476,6 +478,10 @@ public void toJson(JsonWriter writer, RecordingData recordingData) throws IOExce
       writer.endArray();
       writer.name(V4_PROFILE_TAGS_PARAM);
       writer.value(tags + ",snapshot:" + recordingData.getKind().name().toLowerCase(Locale.ROOT));
+      if (processTags != null) {
+        writer.name("process_tags");
+        writer.value(processTags.toString());
+      }
       writer.name(V4_PROFILE_START_PARAM);
       writer.value(recordingData.getStart().toString());
       writer.name(V4_PROFILE_END_PARAM);
diff --git a/dd-java-agent/agent-profiling/profiling-uploader/src/test/java/com/datadog/profiling/uploader/ProfileUploaderTest.java b/dd-java-agent/agent-profiling/profiling-uploader/src/test/java/com/datadog/profiling/uploader/ProfileUploaderTest.java
index 3edb3764c50..6b5f5ad7a8c 100644
--- a/dd-java-agent/agent-profiling/profiling-uploader/src/test/java/com/datadog/profiling/uploader/ProfileUploaderTest.java
+++ b/dd-java-agent/agent-profiling/profiling-uploader/src/test/java/com/datadog/profiling/uploader/ProfileUploaderTest.java
@@ -43,6 +43,7 @@
 import datadog.trace.api.Config;
 import datadog.trace.api.DDTags;
 import datadog.trace.api.Platform;
+import datadog.trace.api.ProcessTags;
 import datadog.trace.api.profiling.ProfilingSnapshot;
 import datadog.trace.api.profiling.RecordingData;
 import datadog.trace.api.profiling.RecordingInputStream;
@@ -830,6 +831,34 @@ public void testShutdown() throws Exception {
     verify(recording).release();
   }
 
+  @ParameterizedTest(name = "process tags enabled ''{0}''")
+  @ValueSource(booleans = {true, false})
+  public void testRequestWithProcessTags(boolean processTagsEnabled) throws Exception {
+    when(config.isExperimentalPropagateProcessTagsEnabled()).thenReturn(processTagsEnabled);
+    ProcessTags.reset(config);
+    uploader =
+        new ProfileUploader(
+            config, configProvider, ioLogger, (int) TERMINATION_TIMEOUT.getSeconds());
+
+    server.enqueue(new MockResponse().setResponseCode(200));
+    uploadAndWait(RECORDING_TYPE, mockRecordingData());
+
+    final RecordedRequest recordedRequest = server.takeRequest(5, TimeUnit.SECONDS);
+    assertNotNull(recordedRequest);
+    final List<FileItem> multiPartItems =
+        FileUpload.parse(
+            recordedRequest.getBody().readByteArray(), recordedRequest.getHeader("Content-Type"));
+
+    final FileItem rawEvent = multiPartItems.get(0);
+    final Map<String, ?> parsed = new ObjectMapper().readValue(rawEvent.get(), Map.class);
+    if (processTagsEnabled) {
+      assertNotNull(ProcessTags.getTagsForSerialization());
+      assertEquals(ProcessTags.getTagsForSerialization().toString(), parsed.get("process_tags"));
+    } else {
+      assertNull(parsed.get("process_tags"));
+    }
+  }
+
   private RecordingData mockRecordingData() throws IOException {
     return mockRecordingData(false, ProfilingSnapshot.Kind.PERIODIC);
   }
diff --git a/internal-api/src/main/java/datadog/trace/api/ProcessTags.java b/internal-api/src/main/java/datadog/trace/api/ProcessTags.java
index 53bfb8d979a..184930b26c1 100644
--- a/internal-api/src/main/java/datadog/trace/api/ProcessTags.java
+++ b/internal-api/src/main/java/datadog/trace/api/ProcessTags.java
@@ -147,9 +147,14 @@ static void empty() {
 
   /** Visible for testing. */
   static void reset() {
+    reset(Config.get());
+  }
+
+  /** Visible for testing. */
+  public static void reset(Config config) {
     synchronized (Lazy.TAGS) {
       empty();
-      enabled = Config.get().isExperimentalPropagateProcessTagsEnabled();
+      enabled = config.isExperimentalPropagateProcessTagsEnabled();
       Lazy.TAGS.putAll(Lazy.loadTags());
     }
   }