Tag profiles for serverless

DataDog · Jan 24, 2025 · 601469c · 601469c
1 parent 56d4919
commit 601469c
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 4 deletions.
diff --git a/dd-java-agent/agent-profiling/profiling-uploader/build.gradle b/dd-java-agent/agent-profiling/profiling-uploader/build.gradle
@@ -23,6 +23,7 @@ dependencies {
   implementation libs.slf4j
   implementation project(':communication')
   implementation project(':internal-api')
+  implementation project(':utils:container-utils')
   implementation project(':utils:socket-utils')
   implementation project(':utils:version-utils')
 

diff --git a/...ling/profiling-uploader/src/main/java/com/datadog/profiling/uploader/ProfileUploader.java b/...ling/profiling-uploader/src/main/java/com/datadog/profiling/uploader/ProfileUploader.java
@@ -21,6 +21,7 @@
 import com.squareup.moshi.JsonAdapter;
 import com.squareup.moshi.JsonReader;
 import com.squareup.moshi.JsonWriter;
+import datadog.common.container.ServerlessInfo;
 import datadog.common.version.VersionInfo;
 import datadog.communication.http.OkHttpUtils;
 import datadog.trace.api.Config;
@@ -110,6 +111,8 @@ public final class ProfileUploader {
               + V4_ATTACHMENT_FILENAME
               + "\"");
 
+  static final String SERVELESS_TAG = "functionname";
+
   private final Config config;
   private final ConfigProvider configProvider;
 
@@ -173,6 +176,9 @@ public ProfileUploader(final Config config, final ConfigProvider configProvider)
     if (Platform.isNativeImage()) {
       tagsMap.put(DDTags.RUNTIME_VERSION_TAG, tagsMap.get(DDTags.RUNTIME_VERSION_TAG) + "-aot");
     }
+    if (ServerlessInfo.get().isRunningInServerlessEnvironment()) {
+      tagsMap.put(SERVELESS_TAG, ServerlessInfo.get().getFunctionName());
+    }
 
     // Comma separated tags string for V2.4 format
     Pattern quotes = Pattern.compile("\"");

diff --git a/.../profiling-uploader/src/test/java/com/datadog/profiling/uploader/ProfileUploaderTest.java b/.../profiling-uploader/src/test/java/com/datadog/profiling/uploader/ProfileUploaderTest.java
@@ -15,8 +15,7 @@
  */
 package com.datadog.profiling.uploader;
 
-import static com.datadog.profiling.uploader.ProfileUploader.V4_PROFILE_END_PARAM;
-import static com.datadog.profiling.uploader.ProfileUploader.V4_PROFILE_START_PARAM;
+import static com.datadog.profiling.uploader.ProfileUploader.*;
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
@@ -81,6 +80,7 @@
 import org.apache.commons.fileupload.FileItem;
 import org.apache.commons.io.IOUtils;
 import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
@@ -104,6 +104,8 @@ public class ProfileUploaderTest {
 
   private static final Map<String, String> TAGS;
 
+  private static final String FUNCTION_NAME = "my_function";
+
   static {
     // Not using Guava's ImmutableMap because we want to test null value
     final Map<String, String> tags = new HashMap<>();
@@ -120,6 +122,7 @@ public class ProfileUploaderTest {
       new ImmutableMap.Builder<String, String>()
           .put("baz", "123")
           .put("foo", "bar")
+          .put(SERVELESS_TAG, FUNCTION_NAME)
           .put("quoted", "quoted")
           .put(DDTags.PID_TAG, PidHelper.getPid())
           .put(VersionInfo.PROFILER_VERSION_TAG, VersionInfo.VERSION)
@@ -149,6 +152,20 @@ public class ProfileUploaderTest {
 
   private ProfileUploader uploader;
 
+  @BeforeAll
+  static void setUpAll() throws Exception {
+    Map<String, String> env = System.getenv();
+
+    Field field = env.getClass().getDeclaredField("m");
+    field.setAccessible(true);
+
+    @SuppressWarnings("unchecked")
+    Map<String, String> modifiableEnv = (Map<String, String>) field.get(env);
+    // hard-coding the env variable here; we could theoretically make the field from ServerlessInfo
+    // public instead
+    modifiableEnv.put("AWS_LAMBDA_FUNCTION_NAME", FUNCTION_NAME);
+  }
+
   @BeforeEach
   public void setup() throws IOException {
     server.start();

diff --git a/internal-api/src/main/java/datadog/trace/api/Config.java b/internal-api/src/main/java/datadog/trace/api/Config.java
@@ -2,7 +2,6 @@
 
 import static datadog.trace.api.ConfigDefaults.*;
 import static datadog.trace.api.DDTags.*;
-import static datadog.trace.api.DDTags.PROFILING_ENABLED;
 import static datadog.trace.api.config.AppSecConfig.*;
 import static datadog.trace.api.config.CiVisibilityConfig.*;
 import static datadog.trace.api.config.CrashTrackingConfig.*;
@@ -3410,7 +3409,7 @@ public Map<String, Object> getLocalRootSpanTags() {
     result.putAll(runtimeTags);
     result.put(LANGUAGE_TAG_KEY, LANGUAGE_TAG_VALUE);
     result.put(SCHEMA_VERSION_TAG_KEY, SpanNaming.instance().version());
-    result.put(PROFILING_ENABLED, isProfilingEnabled() ? 1 : 0);
+    result.put(DDTags.PROFILING_ENABLED, isProfilingEnabled() ? 1 : 0);
     if (isAppSecStandaloneEnabled()) {
       result.put(APM_ENABLED, 0);
     }