Added: event_sample_rate tag to MongoDB integration.

lib/ddtrace/contrib/mongodb/subscribers.rb

@@ -1,3 +1,4 @@
+require 'ddtrace/contrib/sampling'
 require 'ddtrace/contrib/mongodb/ext'
 require 'ddtrace/contrib/mongodb/parsers'
 
@@ -23,6 +24,9 @@ def started(event)
           query = MongoDB.query_builder(event.command_name, event.database_name, event.command)
           serialized_query = query.to_s
 
+          # Set event sample rate
+          Contrib::Sampling.set_event_sample_rate(span, event_sample_rate)
+
           # add operation tags; the full query is stored and used as a resource,
           # since it has been quantized and reduced
           span.set_tag(Ext::TAG_DB, query['database'])
@@ -84,6 +88,14 @@ def clear_span(event)
           return if Thread.current[:datadog_mongo_span].nil?
           Thread.current[:datadog_mongo_span].delete(event.request_id)
         end
+
+        def event_sample_rate
+          datadog_configuration[:event_sample_rate]
+        end
+
+        def datadog_configuration
+          Datadog.configuration[:mongo]
+        end
       end
     end
   end

spec/ddtrace/contrib/mongodb/client_spec.rb

@@ -1,10 +1,12 @@
 require 'spec_helper'
+require 'ddtrace/contrib/sampling_examples'
 
 require 'ddtrace'
 require 'mongo'
 
 RSpec.describe 'Mongo::Client instrumentation' do
   let(:tracer) { get_test_tracer }
+  let(:configuration_options) { { tracer: tracer } }
 
   let(:client) { Mongo::Client.new(["#{host}:#{port}"], client_options) }
   let(:client_options) { { database: database } }
@@ -13,7 +15,6 @@
   let(:database) { 'test' }
   let(:collection) { :artists }
 
-  let(:pin) { Datadog::Pin.get_from(client) }
   let(:spans) { tracer.writer.spans(:keep) }
   let(:span) { spans.first }
 
@@ -26,35 +27,25 @@ def discard_spans!
     Mongo::Logger.logger.level = ::Logger::WARN
 
     Datadog.configure do |c|
-      c.use :mongo
+      c.use :mongo, configuration_options
     end
-
-    # Have to manually update this because its still
-    # using global pin instead of configuration.
-    # Remove this when we remove the pin.
-    pin.tracer = tracer
   end
 
   # Clear data between tests
   let(:drop_database?) { true }
   after(:each) do
     client.database.drop if drop_database?
+    Datadog.registry[:mongo].reset_configuration!
   end
 
   it 'evaluates the block given to the constructor' do
     expect { |b| Mongo::Client.new(["#{host}:#{port}"], client_options, &b) }.to yield_control
   end
 
-  context 'pin' do
-    it 'has the correct attributes' do
-      expect(pin.service).to eq('mongodb')
-      expect(pin.app).to eq('mongodb')
-      expect(pin.app_type).to eq('db')
-    end
-
-    context 'when the service is changed' do
+  context 'when the client is configured' do
+    context 'with a different service name' do
       let(:service) { 'mongodb-primary' }
-      before(:each) { pin.service = service }
+      before(:each) { Datadog.configure(client, service_name: service) }
 
       it 'produces spans with the correct service' do
         client[collection].insert_one(name: 'FKA Twigs')
@@ -63,8 +54,8 @@ def discard_spans!
       end
     end
 
-    context 'when the tracer is disabled' do
-      before(:each) { pin.tracer.enabled = false }
+    context 'to disable the tracer' do
+      before(:each) { tracer.enabled = false }
 
       it 'produces spans with the correct service' do
         client[collection].insert_one(name: 'FKA Twigs')
@@ -78,13 +69,15 @@ def discard_spans!
     shared_examples_for 'a MongoDB trace' do
       it 'has basic properties' do
         expect(spans).to have(1).items
-        expect(span.service).to eq(pin.service)
+        expect(span.service).to eq('mongodb')
         expect(span.span_type).to eq('mongodb')
         expect(span.get_tag('mongodb.db')).to eq(database)
         expect(span.get_tag('mongodb.collection')).to eq(collection.to_s)
         expect(span.get_tag('out.host')).to eq(host)
         expect(span.get_tag('out.port')).to eq(port.to_s)
       end
+
+      it_behaves_like 'event sample rate'
     end
 
     describe '#insert_one operation' do