lib/ddtrace/contrib/mongodb/subscribers.rb in ddtrace-0.18.2 vs lib/ddtrace/contrib/mongodb/subscribers.rb in ddtrace-0.18.3

@@ -15,12 +15,11 @@
           # is safe since it's a unique id used to link events together. Also only one
           # thread is involved in this execution so thread-local storage should be safe. Reference:
           # https://github.com/mongodb/mongo-ruby-driver/blob/master/lib/mongo/monitoring.rb#L70
           # https://github.com/mongodb/mongo-ruby-driver/blob/master/lib/mongo/monitoring/publishable.rb#L38-L56
           span = pin.tracer.trace(Ext::SPAN_COMMAND, service: pin.service, span_type: Ext::SPAN_TYPE_COMMAND)
-          Thread.current[:datadog_mongo_span] ||= {}
-          Thread.current[:datadog_mongo_span][event.request_id] = span
+          set_span(event, span)
 
           # build a quantized Query using the Parser module
           query = MongoDB.query_builder(event.command_name, event.database_name, event.command)
           serialized_query = query.to_s
 
@@ -36,11 +35,11 @@
           # set the resource with the quantized query
           span.resource = serialized_query
         end
 
         def failed(event)
-          span = Thread.current[:datadog_mongo_span][event.request_id]
+          span = get_span(event)
           return unless span
 
           # the failure is not a real exception because it's handled by
           # the framework itself, so we set only the error and the message
           span.set_error(event)
@@ -48,15 +47,15 @@
           Datadog::Tracer.log.debug("error when handling MongoDB 'failed' event: #{e}")
         ensure
           # whatever happens, the Span must be removed from the local storage and
           # it must be finished to prevent any leak
           span.finish unless span.nil?
-          Thread.current[:datadog_mongo_span].delete(event.request_id)
+          clear_span(event)
         end
 
         def succeeded(event)
-          span = Thread.current[:datadog_mongo_span][event.request_id]
+          span = get_span(event)
           return unless span
 
           # add fields that are available only after executing the query
           rows = event.reply.fetch('n', nil)
           span.set_tag(Ext::TAG_ROWS, rows) unless rows.nil?
@@ -64,9 +63,26 @@
           Datadog::Tracer.log.debug("error when handling MongoDB 'succeeded' event: #{e}")
         ensure
           # whatever happens, the Span must be removed from the local storage and
           # it must be finished to prevent any leak
           span.finish unless span.nil?
+          clear_span(event)
+        end
+
+        private
+
+        def get_span(event)
+          Thread.current[:datadog_mongo_span] \
+            && Thread.current[:datadog_mongo_span][event.request_id]
+        end
+
+        def set_span(event, span)
+          Thread.current[:datadog_mongo_span] ||= {}
+          Thread.current[:datadog_mongo_span][event.request_id] = span
+        end
+
+        def clear_span(event)
+          return if Thread.current[:datadog_mongo_span].nil?
           Thread.current[:datadog_mongo_span].delete(event.request_id)
         end
       end
     end
   end