From 576529eeaa2fe181264523a9db344a0978a14754 Mon Sep 17 00:00:00 2001
From: Andrey <andrey.marchenko@datadoghq.com>
Date: Fri, 20 Sep 2024 12:34:00 +0200
Subject: [PATCH] set _dd.host.vcpu_count metric for all CI events

---
 lib/datadog/ci/ext/test.rb                                 | 3 +++
 lib/datadog/ci/span.rb                                     | 7 +++++++
 lib/datadog/ci/test_visibility/context.rb                  | 4 ++++
 sig/datadog/ci/ext/test.rbs                                | 2 ++
 sig/datadog/ci/span.rbs                                    | 2 ++
 spec/datadog/ci/test_visibility/component_spec.rb          | 1 +
 .../datadog/ci/test_visibility/serializers/test_v1_spec.rb | 2 +-
 .../datadog/ci/test_visibility/serializers/test_v2_spec.rb | 2 +-
 8 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/lib/datadog/ci/ext/test.rb b/lib/datadog/ci/ext/test.rb
index bf7493f4..3bebeeac 100644
--- a/lib/datadog/ci/ext/test.rb
+++ b/lib/datadog/ci/ext/test.rb
@@ -71,6 +71,9 @@ module Test
         # common tags that are serialized directly in msgpack header in metadata field
         METADATA_TAG_TEST_SESSION_NAME = "test_session.name"
 
+        # internal metric with the number of virtual CPUs
+        METRIC_CPU_COUNT = "_dd.host.vcpu_count"
+
         # tags that are common for the whole session and can be inherited from the test session
         INHERITABLE_TAGS = [TAG_FRAMEWORK, TAG_FRAMEWORK_VERSION].freeze
 
diff --git a/lib/datadog/ci/span.rb b/lib/datadog/ci/span.rb
index 430e509a..62e8806b 100644
--- a/lib/datadog/ci/span.rb
+++ b/lib/datadog/ci/span.rb
@@ -114,6 +114,13 @@ def clear_tag(key)
         tracer_span.clear_tag(key)
       end
 
+      # Gets metric value by key.
+      # @param [String] key the key of the metric.
+      # @return [Numeric] value the value of the metric.
+      def get_metric(key)
+        tracer_span.get_metric(key)
+      end
+
       # Sets metric value by key.
       # @param [String] key the key of the metric.
       # @param [Numeric] value the value of the metric.
diff --git a/lib/datadog/ci/test_visibility/context.rb b/lib/datadog/ci/test_visibility/context.rb
index 2405bcf5..37af17e0 100644
--- a/lib/datadog/ci/test_visibility/context.rb
+++ b/lib/datadog/ci/test_visibility/context.rb
@@ -11,6 +11,8 @@
 require_relative "../ext/environment"
 require_relative "../ext/test"
 
+require_relative "../utils/test_run"
+
 require_relative "../span"
 require_relative "../test"
 require_relative "../test_session"
@@ -203,6 +205,8 @@ def set_initial_tags(ci_span, tags)
 
           ci_span.set_tags(tags)
           ci_span.set_tags(@environment_tags)
+
+          ci_span.set_metric(Ext::Test::METRIC_CPU_COUNT, Utils::TestRun.virtual_cpu_count)
         end
 
         # PROPAGATING CONTEXT FROM TOP-LEVEL TO THE LOWER LEVELS
diff --git a/sig/datadog/ci/ext/test.rbs b/sig/datadog/ci/ext/test.rbs
index 62ef6eda..ae3d5f49 100644
--- a/sig/datadog/ci/ext/test.rbs
+++ b/sig/datadog/ci/ext/test.rbs
@@ -98,6 +98,8 @@ module Datadog
 
         METADATA_TAG_TEST_SESSION_NAME: "test_session.name"
 
+        METRIC_CPU_COUNT: "_dd.host.vcpu_count"
+
         module Status
           PASS: "pass"
 
diff --git a/sig/datadog/ci/span.rbs b/sig/datadog/ci/span.rbs
index 53525bd5..5d32ee58 100644
--- a/sig/datadog/ci/span.rbs
+++ b/sig/datadog/ci/span.rbs
@@ -35,6 +35,8 @@ module Datadog
 
       def clear_tag: (String key) -> void
 
+      def get_metric: (String key) -> Numeric?
+
       def set_metric: (String key, untyped value) -> void
 
       def set_tags: (Hash[untyped, untyped] tags) -> void
diff --git a/spec/datadog/ci/test_visibility/component_spec.rb b/spec/datadog/ci/test_visibility/component_spec.rb
index e3dd67f1..ba0badcf 100644
--- a/spec/datadog/ci/test_visibility/component_spec.rb
+++ b/spec/datadog/ci/test_visibility/component_spec.rb
@@ -47,6 +47,7 @@
         expect(span_under_test).to have_test_tag(tag)
       end
       expect(span_under_test).to have_test_tag(:command, test_command)
+      expect(span_under_test.get_metric(Datadog::CI::Ext::Test::METRIC_CPU_COUNT)).to eq(Etc.nprocessors)
     end
   end
 
diff --git a/spec/datadog/ci/test_visibility/serializers/test_v1_spec.rb b/spec/datadog/ci/test_visibility/serializers/test_v1_spec.rb
index 2c7ddeb9..7749f975 100644
--- a/spec/datadog/ci/test_visibility/serializers/test_v1_spec.rb
+++ b/spec/datadog/ci/test_visibility/serializers/test_v1_spec.rb
@@ -39,7 +39,7 @@
           }
         )
         expect(metrics).to eq(
-          {"_dd.top_level" => 1.0, "memory_allocations" => 16}
+          {"_dd.top_level" => 1.0, "memory_allocations" => 16, "_dd.host.vcpu_count" => Etc.nprocessors}
         )
       end
     end
diff --git a/spec/datadog/ci/test_visibility/serializers/test_v2_spec.rb b/spec/datadog/ci/test_visibility/serializers/test_v2_spec.rb
index 5577110b..250224c0 100644
--- a/spec/datadog/ci/test_visibility/serializers/test_v2_spec.rb
+++ b/spec/datadog/ci/test_visibility/serializers/test_v2_spec.rb
@@ -49,7 +49,7 @@
         expect(meta["_test.session_id"]).to be_nil
         expect(meta["_test.module_id"]).to be_nil
 
-        expect(metrics).to eq({"_dd.top_level" => 1, "memory_allocations" => 16})
+        expect(metrics).to eq({"_dd.top_level" => 1, "memory_allocations" => 16, "_dd.host.vcpu_count" => Etc.nprocessors})
       end
     end