From c75a2c2979b474faf485b8e0d299e2d6572fce01 Mon Sep 17 00:00:00 2001
From: Nikita Tkachenko
 <121111529+nikita-tkachenko-datadog@users.noreply.github.com>
Date: Tue, 24 Sep 2024 16:54:15 +0200
Subject: [PATCH] Add CPU count tag to Test Visibility events (#7659)

---
 .../trace/civisibility/decorator/TestDecoratorImpl.java        | 3 +++
 .../trace/civisibility/decorator/TestDecoratorImplTest.groovy  | 1 +
 dd-trace-api/src/main/java/datadog/trace/api/DDTags.java       | 1 +
 3 files changed, 5 insertions(+)

diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/decorator/TestDecoratorImpl.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/decorator/TestDecoratorImpl.java
index d1ba5c9ee56..c75c85b3b9b 100644
--- a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/decorator/TestDecoratorImpl.java
+++ b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/decorator/TestDecoratorImpl.java
@@ -15,6 +15,7 @@ public class TestDecoratorImpl implements TestDecorator {
   private final String component;
   private final String sessionName;
   private final Map<String, String> ciTags;
+  private final int cpuCount;
 
   public TestDecoratorImpl(
       String component, String sessionName, String testCommand, Map<String, String> ciTags) {
@@ -27,6 +28,7 @@ public TestDecoratorImpl(
       this.sessionName =
           Strings.isNotBlank(ciJobName) ? ciJobName + "-" + testCommand : testCommand;
     }
+    cpuCount = Runtime.getRuntime().availableProcessors();
   }
 
   protected String testType() {
@@ -46,6 +48,7 @@ public CharSequence component() {
   public AgentSpan afterStart(final AgentSpan span) {
     span.setSamplingPriority(PrioritySampling.SAMPLER_KEEP);
     span.setTag(DDTags.ORIGIN_KEY, CIAPP_TEST_ORIGIN);
+    span.setTag(DDTags.HOST_VCPU_COUNT, cpuCount);
     span.setTag(Tags.TEST_TYPE, testType());
     span.setTag(Tags.COMPONENT, component());
     span.setTag(Tags.TEST_SESSION_NAME, sessionName);
diff --git a/dd-java-agent/agent-ci-visibility/src/test/groovy/datadog/trace/civisibility/decorator/TestDecoratorImplTest.groovy b/dd-java-agent/agent-ci-visibility/src/test/groovy/datadog/trace/civisibility/decorator/TestDecoratorImplTest.groovy
index c4e531d099c..6cc4bc1ce6b 100644
--- a/dd-java-agent/agent-ci-visibility/src/test/groovy/datadog/trace/civisibility/decorator/TestDecoratorImplTest.groovy
+++ b/dd-java-agent/agent-ci-visibility/src/test/groovy/datadog/trace/civisibility/decorator/TestDecoratorImplTest.groovy
@@ -23,6 +23,7 @@ class TestDecoratorImplTest extends Specification {
     1 * span.setTag(Tags.TEST_TYPE, decorator.testType())
     1 * span.setSamplingPriority(PrioritySampling.SAMPLER_KEEP)
     1 * span.setTag(DDTags.ORIGIN_KEY, decorator.origin())
+    1 * span.setTag(DDTags.HOST_VCPU_COUNT, Runtime.runtime.availableProcessors())
     1 * span.setTag("ci-tag-1", "value")
     1 * span.setTag("ci-tag-2", "another value")
 
diff --git a/dd-trace-api/src/main/java/datadog/trace/api/DDTags.java b/dd-trace-api/src/main/java/datadog/trace/api/DDTags.java
index a9b48550a79..f456eeb02fe 100644
--- a/dd-trace-api/src/main/java/datadog/trace/api/DDTags.java
+++ b/dd-trace-api/src/main/java/datadog/trace/api/DDTags.java
@@ -67,6 +67,7 @@ public class DDTags {
   public static final String PEER_SERVICE_REMAPPED_FROM = "_dd.peer.service.remapped_from";
   public static final String INTERNAL_GIT_REPOSITORY_URL = "_dd.git.repository_url";
   public static final String INTERNAL_GIT_COMMIT_SHA = "_dd.git.commit.sha";
+  public static final String HOST_VCPU_COUNT = "_dd.host.vcpu_count";
 
   public static final String PROFILING_ENABLED = "_dd.profiling.enabled";
   public static final String DSM_ENABLED = "_dd.dsm.enabled";