extract Record#start_datadog_tracer_span with Datadog::Tracing logic

DataDog · Nov 30, 2023 · a97ad49 · a97ad49
1 parent 878a745
commit a97ad49
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 17 deletions.
diff --git a/lib/datadog/ci/recorder.rb b/lib/datadog/ci/recorder.rb
@@ -41,10 +41,7 @@ def start_test_session(service_name: nil, tags: {})
           span_type: Ext::AppTypes::TYPE_TEST_SESSION
         }
 
-        tracer_span = Datadog::Tracing.trace("test.session", **span_options)
-        trace = Datadog::Tracing.active_trace
-
-        set_trace_origin(trace)
+        tracer_span = start_datadog_tracer_span("test.session", span_options)
 
         tags[Ext::Test::TAG_TEST_SESSION_ID] = tracer_span.id
 
@@ -71,10 +68,7 @@ def start_test_module(test_module_name, service_name: nil, tags: {})
           span_type: Ext::AppTypes::TYPE_TEST_MODULE
         }
 
-        tracer_span = Datadog::Tracing.trace(test_module_name, **span_options)
-        trace = Datadog::Tracing.active_trace
-
-        set_trace_origin(trace)
+        tracer_span = start_datadog_tracer_span(test_module_name, span_options)
 
         tags[Ext::Test::TAG_TEST_MODULE_ID] = tracer_span.id
         tags[Ext::Test::TAG_MODULE] = tracer_span.name
@@ -115,20 +109,15 @@ def trace_test(test_name, service_name: nil, operation_name: "test", tags: {}, &
         }
 
         if block
-          Datadog::Tracing.trace(operation_name, **span_options) do |tracer_span, trace|
-            set_trace_origin(trace)
-
+          start_datadog_tracer_span(operation_name, span_options) do |tracer_span|
             test = build_test(tracer_span, tags)
 
             @local_context.activate_test!(test) do
               block.call(test)
             end
           end
         else
-          tracer_span = Datadog::Tracing.trace(operation_name, **span_options)
-          trace = Datadog::Tracing.active_trace
-
-          set_trace_origin(trace)
+          tracer_span = start_datadog_tracer_span(operation_name, span_options)
 
           test = build_test(tracer_span, tags)
           @local_context.activate_test!(test)
@@ -145,11 +134,11 @@ def trace(span_type, span_name, tags: {}, &block)
         }
 
         if block
-          Datadog::Tracing.trace(span_name, **span_options) do |tracer_span, trace|
+          start_datadog_tracer_span(span_name, span_options) do |tracer_span|
             block.call(build_span(tracer_span, tags))
           end
         else
-          tracer_span = Datadog::Tracing.trace(span_name, **span_options)
+          tracer_span = start_datadog_tracer_span(span_name, span_options)
 
           build_span(tracer_span, tags)
         end
@@ -232,6 +221,22 @@ def set_initial_tags(ci_span, tags)
         ci_span.set_tags(environment_tags)
       end
 
+      def start_datadog_tracer_span(span_name, span_options, &block)
+        if block
+          Datadog::Tracing.trace(span_name, **span_options) do |tracer_span, trace|
+            set_trace_origin(trace)
+
+            yield tracer_span
+          end
+        else
+          tracer_span = Datadog::Tracing.trace(span_name, **span_options)
+          trace = Datadog::Tracing.active_trace
+          set_trace_origin(trace)
+
+          tracer_span
+        end
+      end
+
       def null_span
         @null_span ||= NullSpan.new
       end

diff --git a/sig/datadog/ci/recorder.rbs b/sig/datadog/ci/recorder.rbs
@@ -57,6 +57,8 @@ module Datadog
       def null_span: () -> Datadog::CI::Span
 
       def skip_tracing: (?untyped block) -> untyped
+
+      def start_datadog_tracer_span: (String span_name, Hash[untyped, untyped] span_options) ?{ (untyped) -> untyped } -> untyped
     end
   end
 end