lib/elastic_apm/spies/mongo.rb in elastic-apm-4.5.0 vs lib/elastic_apm/spies/mongo.rb in elastic-apm-4.5.1

@@ -35,20 +35,14 @@
         SUBTYPE = 'mongodb'
         ACTION = 'query'
 
         EVENT_KEY = :__elastic_instrumenter_mongo_events_key
 
-        class Collection
-          def events
-            Thread.current[EVENT_KEY] ||= {}
-          end
+        def events
+          Thread.current[EVENT_KEY] ||= []
         end
 
-        def initialize
-          @collection = Collection.new
-        end
-
         def started(event)
           push_event(event)
         end
 
         def failed(event)
@@ -76,11 +70,11 @@
           # key is the integer 1.
           # For getMore commands, the value at `command_name` is the cursor id
           # and the collection name is at the key `collection`
           collection =
             if event.command[event.command_name] == 1 ||
-               event.command[event.command_name].is_a?(BSON::Int64)
+              event.command[event.command_name].is_a?(BSON::Int64)
               event.command[:collection]
             else
               event.command[event.command_name]
             end
 
@@ -95,17 +89,16 @@
               subtype: SUBTYPE,
               action: ACTION,
               context: build_context(event)
             )
 
-          @collection.events[event.operation_id] = span
+          events << span
         end
 
         def pop_event(event)
-          span = @collection.events.delete(event.operation_id)
           return unless (curr = ElasticAPM.current_span)
 
-          curr == span && ElasticAPM.end_span
+          curr == events[-1] && ElasticAPM.end_span(events.pop)
         end
 
         def build_context(event)
           Span::Context.new(
             db: {