MONGOID-5417 Fix memory leak when call 'with' (#5409)

Author: comandeo-mongo

lib/mongoid/persistence_context.rb

@@ -189,8 +189,7 @@ class << self
       #
       # @return [ Mongoid::PersistenceContext ] The persistence context for the object.
       def set(object, options_or_context)
-        key = "[mongoid][#{object.object_id}]:context"
-        existing_context = Thread.current[key]
+        existing_context = get_context(object)
         existing_options = if existing_context
           existing_context.options
         else
@@ -201,7 +200,7 @@ def set(object, options_or_context)
         end
         new_options = existing_options.merge(options_or_context)
         context = PersistenceContext.new(object, new_options)
-        Thread.current[key] = context
+        store_context(object, context)
       end
 
       # Get the persistence context for a particular class or model instance.
@@ -213,7 +212,7 @@ def set(object, options_or_context)
       #
       # @return [ Mongoid::PersistenceContext ] The persistence context for the object.
       def get(object)
-        Thread.current["[mongoid][#{object.object_id}]:context"]
+        get_context(object)
       end
 
       # Clear the persistence context for a particular class or model instance.
@@ -232,7 +231,44 @@ def clear(object, cluster = nil, original_context = nil)
           end
         end
       ensure
-        Thread.current["[mongoid][#{object.object_id}]:context"] = original_context
+        store_context(object, original_context)
+      end
+
+      private
+
+      # Key to store persistence contexts in the thread local storage.
+      #
+      # @api private
+      PERSISTENCE_CONTEXT_KEY = :"[mongoid]:persistence_context"
+
+      # Get the persistence context for a given object from the thread local
+      #   storage.
+      #
+      # @param [ Object ] object Object to get the persistance context for.
+      #
+      # @return [ Mongoid::PersistenceContext | nil ] The persistence context
+      #   for the object if previously stored, otherwise nil.
+      #
+      # @api private
+      def get_context(object)
+        Thread.current[PERSISTENCE_CONTEXT_KEY] ||= {}
+        Thread.current[PERSISTENCE_CONTEXT_KEY][object.object_id]
+      end
+
+      # Store persistence context for a given object in the thread local
+      #   storage.
+      #
+      # @param [ Object ] object Object to store the persistance context for.
+      # @param [ Mongoid::PersistenceContext ] context Context to store
+      #
+      # @api private
+      def store_context(object, context)
+        if context.nil?
+          Thread.current[PERSISTENCE_CONTEXT_KEY]&.delete(object.object_id)
+        else
+          Thread.current[PERSISTENCE_CONTEXT_KEY] ||= {}
+          Thread.current[PERSISTENCE_CONTEXT_KEY][object.object_id] = context
+        end
       end
     end
   end

spec/mongoid/clients_spec.rb

@@ -1042,6 +1042,31 @@
         end
       end
     end
+
+    context 'when using on different objects' do
+      require_mri
+
+      let(:first_band) do
+        Band.create!(name: "The Beatles")
+      end
+
+      let(:second_band) do
+        Band.create!(name: 'Led Zeppelin')
+      end
+
+      it 'does not create extra symbols symbols' do
+        first_band.with(write: { w: 0, j: false }) do |band|
+          band.set(active: false)
+        end
+        initial_symbols_count = Symbol.all_symbols.size
+        second_band.with(write: { w: 0, j: false }) do |band|
+          band.set(active: false)
+        end
+        expect(Symbol.all_symbols.size).to eq(initial_symbols_count)
+      end
+
+    end
+
   end
 
   context "when overriding the default database" do