From a83cee4eb48221807de0b930522a8f024efc8e5b Mon Sep 17 00:00:00 2001
From: Nikita Tkachenko <nikita.tkachenko@datadoghq.com>
Date: Mon, 20 Jan 2025 13:47:24 +0100
Subject: [PATCH] Set test framework and test framework version tags atomically

---
 .../trace/civisibility/utils/SpanUtils.java        | 14 ++++++++++----
 .../src/main/java/datadog/trace/core/DDSpan.java   |  6 ++++++
 .../bootstrap/instrumentation/api/AgentSpan.java   |  3 +++
 .../bootstrap/instrumentation/api/AgentTracer.java |  5 +++++
 4 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/SpanUtils.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/SpanUtils.java
index c016272353e..04aeb7a5770 100644
--- a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/SpanUtils.java
+++ b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/SpanUtils.java
@@ -7,7 +7,9 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.Iterator;
+import java.util.Map;
 import java.util.TreeSet;
 import java.util.function.BinaryOperator;
 import java.util.function.Consumer;
@@ -83,8 +85,10 @@ private static void setFrameworks(AgentSpan span, Collection<TestFramework> fram
     }
     if (frameworks.size() == 1) {
       TestFramework framework = frameworks.iterator().next();
-      span.setTag(Tags.TEST_FRAMEWORK, framework.getName());
-      span.setTag(Tags.TEST_FRAMEWORK_VERSION, framework.getVersion());
+      Map<String, String> tags = new HashMap<>();
+      tags.put(Tags.TEST_FRAMEWORK, framework.getName());
+      tags.put(Tags.TEST_FRAMEWORK_VERSION, framework.getVersion());
+      span.setAllTags(tags);
       return;
     }
     Collection<String> names = new ArrayList<>(frameworks.size());
@@ -93,8 +97,10 @@ private static void setFrameworks(AgentSpan span, Collection<TestFramework> fram
       names.add(framework.getName());
       versions.add(framework.getVersion());
     }
-    span.setTag(Tags.TEST_FRAMEWORK, names);
-    span.setTag(Tags.TEST_FRAMEWORK_VERSION, versions);
+    Map<String, Collection<String>> tags = new HashMap<>();
+    tags.put(Tags.TEST_FRAMEWORK, names);
+    tags.put(Tags.TEST_FRAMEWORK_VERSION, versions);
+    span.setAllTags(tags);
   }
 
   private static void propagateStatus(AgentSpan parentSpan, AgentSpan childSpan) {
diff --git a/dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java b/dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java
index dcf49e974c9..6a8e6f713c1 100644
--- a/dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java
+++ b/dd-trace-core/src/main/java/datadog/trace/core/DDSpan.java
@@ -485,6 +485,12 @@ public DDSpan setTag(final String tag, final Object value) {
     return this;
   }
 
+  @Override
+  public AgentSpan setAllTags(Map<String, ?> map) {
+    context.setAllTags(map);
+    return this;
+  }
+
   // FIXME [API] this is not on AgentSpan or MutableSpan
   @Override
   public DDSpan removeTag(final String tag) {
diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java
index 0df378b912c..0202096113b 100644
--- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java
+++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java
@@ -10,6 +10,7 @@
 import datadog.trace.api.gateway.IGSpanInfo;
 import datadog.trace.api.gateway.RequestContext;
 import datadog.trace.api.interceptor.MutableSpan;
+import java.util.Map;
 import javax.annotation.Nullable;
 
 public interface AgentSpan
@@ -35,6 +36,8 @@ public interface AgentSpan
 
   AgentSpan setTag(String key, Object value);
 
+  AgentSpan setAllTags(Map<String, ?> map);
+
   @Override
   AgentSpan setTag(String key, Number value);
 
diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java
index 9fb97ae2c74..2cb1947d90c 100644
--- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java
+++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java
@@ -642,6 +642,11 @@ public AgentSpan setTag(final String key, final Object value) {
       return this;
     }
 
+    @Override
+    public AgentSpan setAllTags(Map<String, ?> map) {
+      return this;
+    }
+
     @Override
     public AgentSpan setMetric(final CharSequence key, final int value) {
       return this;