Refactor to optimize memory and delay #to_digest

DataDog · Dec 6, 2022 · 6ce4135 · 6ce4135
1 parent ea27c9b
commit 6ce4135
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 27 deletions.
diff --git a/lib/datadog/tracing/contrib/mysql2/instrumentation.rb b/lib/datadog/tracing/contrib/mysql2/instrumentation.rb
@@ -46,7 +46,7 @@ def query(sql, options = {})
                 propagation_mode = Contrib::Propagation::SqlComment::Mode.new(comment_propagation)
 
                 Contrib::Propagation::SqlComment.annotate!(span, propagation_mode)
-                sql = Contrib::Propagation::SqlComment.prepend_comment(sql, trace_op.to_digest, propagation_mode)
+                sql = Contrib::Propagation::SqlComment.prepend_comment(sql, span, trace_op, propagation_mode)
 
                 super(sql, options)
               end

diff --git a/lib/datadog/tracing/contrib/pg/instrumentation.rb b/lib/datadog/tracing/contrib/pg/instrumentation.rb
@@ -95,7 +95,8 @@ def trace(name, sql: nil, statement_name: nil)
                   Contrib::Propagation::SqlComment.annotate!(span, propagation_mode)
                   propagated_sql_statement = Contrib::Propagation::SqlComment.prepend_comment(
                     sql,
-                    trace_op.to_digest,
+                    span,
+                    trace_op,
                     propagation_mode
                   )
                 end

diff --git a/lib/datadog/tracing/contrib/propagation/sql_comment.rb b/lib/datadog/tracing/contrib/propagation/sql_comment.rb
@@ -17,19 +17,21 @@ def self.annotate!(span_op, mode)
             span_op.set_tag(Ext::TAG_DBM_TRACE_INJECTED, true) if mode.full?
           end
 
-          def self.prepend_comment(sql, digest, mode)
+          # Inject span_op and trace_op instead of TraceDigest to improve memory usage
+          # for `disabled` and `service` mode
+          def self.prepend_comment(sql, span_op, trace_op, mode)
             return sql unless mode.enabled?
 
             tags = {
-              Ext::KEY_DATABASE_SERVICE => digest.span_service,
+              Ext::KEY_DATABASE_SERVICE => span_op.service,
               Ext::KEY_ENVIRONMENT => datadog_configuration.env,
               Ext::KEY_PARENT_SERVICE => datadog_configuration.service,
               Ext::KEY_VERSION => datadog_configuration.version
             }
 
             if mode.full?
               tags[Ext::KEY_TRACEPARENT] =
-                Tracing::Distributed::TraceContext.new(fetcher: nil).send(:build_traceparent, digest)
+                Tracing::Distributed::TraceContext.new(fetcher: nil).send(:build_traceparent, trace_op.to_digest)
             end
 
             "#{Comment.new(tags)} #{sql}"

diff --git a/spec/datadog/tracing/contrib/propagation/sql_comment_spec.rb b/spec/datadog/tracing/contrib/propagation/sql_comment_spec.rb
@@ -56,16 +56,19 @@
     end
 
     let(:sql_statement) { 'SELECT 1' }
-    let(:digest) do
-      Datadog::Tracing::TraceDigest.new(
-        span_service: 'database_service',
-        trace_id: 0xC0FFEE,
-        span_id: 0xBEE,
-        trace_flags: 0xFE
+
+    let(:span_op) { double(service: 'database_service') }
+    let(:trace_op) do
+      double(
+        to_digest: Datadog::Tracing::TraceDigest.new(
+          trace_id: 0xC0FFEE,
+          span_id: 0xBEE,
+          trace_flags: 0xFE
+        )
       )
     end
 
-    subject { described_class.prepend_comment(sql_statement, digest, propagation_mode) }
+    subject { described_class.prepend_comment(sql_statement, span_op, trace_op, propagation_mode) }
 
     context 'when `disabled` mode' do
       let(:mode) { 'disabled' }

diff --git a/spec/datadog/tracing/contrib/sql_comment_propagation_examples.rb b/spec/datadog/tracing/contrib/sql_comment_propagation_examples.rb
@@ -34,15 +34,6 @@
 end
 
 RSpec.shared_examples_for 'propagates with sql comment' do |mode:, span_op_name:, error: nil|
-  RSpec::Matchers.define :a_trace_digest_with do |expected|
-    match do |actual|
-      actual.instance_of?(Datadog::Tracing::TraceDigest) &&
-        expected.all? do |key, value|
-          actual.__send__(key) == value
-        end
-    end
-  end
-
   it "propagates with mode: #{mode}" do
     expect(Datadog::Tracing::Contrib::Propagation::SqlComment::Mode)
       .to receive(:new).with(mode).and_return(propagation_mode)
@@ -77,12 +68,8 @@
 
     expect(Datadog::Tracing::Contrib::Propagation::SqlComment).to have_received(:prepend_comment).with(
       sql_statement,
-      a_trace_digest_with(
-        span_service: service_name,
-        trace_id: trace.id,
-        span_id: span.id,
-        trace_sampling_priority: 1
-      ),
+      a_span_operation_with(service: service_name),
+      duck_type(:to_digest),
       propagation_mode
     )
   end