MONGOID-5882 isolation state

Author: jamis

lib/mongoid/config.rb

@@ -110,6 +110,19 @@ module Config
     # to `:global_thread_pool`.
     option :global_executor_concurrency, default: nil
 
+    # Defines the isolation level that Mongoid uses to store its internal
+    # state.
+    #
+    # This option may be set to either `:thread` (the default) or `:fiber`.
+    #
+    # If set to `:thread`, Mongoid will use thread-local storage to manage
+    # its internal state.
+    #
+    # If set to `:fiber`, Mongoid will use fiber-local storage instead. This
+    # may be necessary if you are using libraries like Falcon, which use
+    # fibers to manage concurrency.
+    option :isolation_level, default: :thread
+
     # When this flag is false, a document will become read-only only once the
     # #readonly! method is called, and an error will be raised on attempting
     # to save or update such documents, instead of just on delete. When this

lib/mongoid/threaded.rb

@@ -6,38 +6,56 @@ module Mongoid
   # This module contains logic for easy access to objects that have a lifecycle
   # on the current thread.
   module Threaded
-    DATABASE_OVERRIDE_KEY = '[mongoid]:db-override'
+    # The key for the shared thread- and fiber-local storage. It must be a
+    # symbol because keys for fiber-local storage must be symbols.
+    STORAGE_KEY = :'[mongoid]'
 
-    # Constant for the key to store clients.
-    CLIENTS_KEY = '[mongoid]:clients'
+    DATABASE_OVERRIDE_KEY = 'db-override'
 
     # The key to override the client.
-    CLIENT_OVERRIDE_KEY = '[mongoid]:client-override'
+    CLIENT_OVERRIDE_KEY = 'client-override'
 
     # The key for the current thread's scope stack.
-    CURRENT_SCOPE_KEY = '[mongoid]:current-scope'
+    CURRENT_SCOPE_KEY = 'current-scope'
 
-    AUTOSAVES_KEY = '[mongoid]:autosaves'
+    AUTOSAVES_KEY = 'autosaves'
 
-    VALIDATIONS_KEY = '[mongoid]:validations'
+    VALIDATIONS_KEY = 'validations'
 
     STACK_KEYS = Hash.new do |hash, key|
-      hash[key] = "[mongoid]:#{key}-stack"
+      hash[key] = "#{key}-stack"
     end
 
     # The key for the current thread's sessions.
-    SESSIONS_KEY = '[mongoid]:sessions'
+    SESSIONS_KEY = 'sessions'
 
     # The key for storing documents modified inside transactions.
-    MODIFIED_DOCUMENTS_KEY = '[mongoid]:modified-documents'
+    MODIFIED_DOCUMENTS_KEY = 'modified-documents'
 
     # The key storing the default value for whether or not callbacks are
     # executed on documents.
-    EXECUTE_CALLBACKS = '[mongoid]:execute-callbacks'
+    EXECUTE_CALLBACKS = 'execute-callbacks'
 
     extend self
 
-    # Queries the thread-local variable with the given name. If a block is
+    # Resets the current thread- or fiber-local storage to its initial state.
+    # This is useful for making sure the state is clean when starting a new
+    # thread or fiber.
+    #
+    # The value of Mongoid::Config.isolation_level is used to determine
+    # whether to reset the storage for the current thread or fiber.
+    def reset!
+      case Config.isolation_level
+      when :thread
+        Thread.current.thread_variable_set(STORAGE_KEY, nil)
+      when :fiber
+        Fiber[STORAGE_KEY] = nil
+      else
+        raise "Unknown isolation level: #{Config.isolation_level.inspect}"
+      end
+    end
+
+    # Queries the thread- or fiber-local variable with the given name. If a block is
     # given, and the variable does not already exist, the return value of the
     # block will be set as the value of the variable before returning it.
     #
@@ -57,7 +75,7 @@ module Threaded
     # @return [ Object | nil ] the value of the queried variable, or nil if
     #   it is not set and no default was given.
     def get(key, &default)
-      result = Thread.current.thread_variable_get(key)
+      result = storage[key]
 
       if result.nil? && default
         result = yield
@@ -67,43 +85,31 @@ def get(key, &default)
       result
     end
 
-    # Sets a thread-local variable with the given name to the given value.
+    # Sets a variable in local storage with the given name to the given value.
     # See #get for a discussion of why this method is necessary, and why
     # Thread#[]= should be avoided in cascading callbacks on embedded children.
     #
     # @param [ String | Symbol ] key the name of the variable to set.
     # @param [ Object | nil ] value the value of the variable to set (or `nil`
     #   if you wish to unset the variable)
     def set(key, value)
-      Thread.current.thread_variable_set(key, value)
+      storage[key] = value
     end
 
-    # Removes the named variable from thread-local storage.
+    # Removes the named variable from local storage.
     #
     # @param [ String | Symbol ] key the name of the variable to remove.
     def delete(key)
-      set(key, nil)
+      storage.delete(key)
     end
 
-    # Queries the presence of a named variable in thread-local storage.
+    # Queries the presence of a named variable in local storage.
     #
     # @param [ String | Symbol ] key the name of the variable to query.
     #
     # @return [ true | false ] whether the given variable is present or not.
     def has?(key)
-      # Here we have a classic example of JRuby not behaving like MRI. In
-      # MRI, if you set a thread variable to nil, it removes it from the list
-      # and subsequent calls to thread_variable?(key) will return false. Not
-      # so with JRuby. Once set, you cannot unset the thread variable.
-      #
-      # However, because setting a variable to nil is supposed to remove it,
-      # we can assume a nil-valued variable doesn't actually exist.
-
-      # So, instead of this:
-      # Thread.current.thread_variable?(key)
-
-      # We have to do this:
-      !get(key).nil?
+      storage.key?(key)
     end
 
     # Begin entry into a named thread local stack.
@@ -508,5 +514,22 @@ def unset_current_scope(klass)
 
       delete(CURRENT_SCOPE_KEY) if scope.empty?
     end
+
+    # Returns the current thread- or fiber-local storage as a Hash. 
+    def storage
+      case Config.isolation_level
+      when :thread
+        if !Thread.current.thread_variable?(STORAGE_KEY)
+          Thread.current.thread_variable_set(STORAGE_KEY, {})
+        end
+
+        Thread.current.thread_variable_get(STORAGE_KEY)
+      when :fiber
+        Fiber[STORAGE_KEY] ||= {}
+      else
+        raise "Unknown isolation level: #{Config.isolation_level.inspect}"
+      end
+    end
+
   end
 end

spec/integration/isolation_state_spec.rb

@@ -0,0 +1,97 @@
+# frozen_string_literal: true
+# rubocop:todo all
+
+require 'spec_helper'
+
+describe 'Mongoid::Config.isolation_level' do
+  def thread_operation(value)
+    Thread.new do
+      Mongoid::Threaded.stack(:testing) << value
+      yield if block_given?
+      Mongoid::Threaded.stack(:testing)
+    end.join.value
+  end
+
+  def fiber_operation(value)
+    Fiber.new do
+      Mongoid::Threaded.stack(:testing) << value
+      yield if block_given?
+      Mongoid::Threaded.stack(:testing)
+    end.resume
+  end
+
+  context 'when set to :thread' do
+    config_override :isolation_level, :thread
+
+    context 'when not operating inside fibers' do
+      let(:result1) { thread_operation('a') { thread_operation('b') } }
+      let(:result2) { thread_operation('b') { thread_operation('c') } }
+
+      it 'isolates state per thread' do
+        expect(result1).to eq(%w[ a ])
+        expect(result2).to eq(%w[ b ])
+      end
+    end
+
+    context 'when operating inside fibers' do
+      let(:result) { thread_operation('a') { fiber_operation('b') } }
+
+      it 'exposes the thread state within the fiber' do
+        expect(result).to eq(%w[ a b ])
+      end
+    end
+  end
+
+  context 'when set to :fiber' do
+    config_override :isolation_level, :fiber
+
+    context 'when operating inside threads' do
+      let(:result) { fiber_operation('a') { thread_operation('b') } }
+
+      it 'exposes the fiber state within the thread' do
+        expect(result).to eq(%w[ a b ])
+      end
+    end
+
+    context 'when operating in nested fibers' do
+      let(:result) { fiber_operation('a') { fiber_operation('b') } }
+
+      it 'propagates fiber state to nested fibers' do
+        expect(result).to eq(%w[ a b ])
+      end
+    end
+
+    context 'when operating in adjacent fibers' do
+      let(:result1) { fiber_operation('a') { fiber_operation('b') } }
+      let(:result2) { fiber_operation('c') { fiber_operation('d') } }
+
+      it 'maintains isolation between adjacent fibers' do
+        expect(result1).to eq(%w[ a b ])
+        expect(result2).to eq(%w[ c d ])
+      end
+    end
+
+    describe '#reset!' do
+      context 'when operating in nested fibers' do
+        let (:result) do
+          fiber_operation('a') do
+            Mongoid::Threaded.reset!
+
+            # once reset, subsequent nested fibers will each have their own
+            # state; they won't touch the reset state here.
+            fiber_operation('b')
+            fiber_operation('c')
+
+            # If we then add to the stack here, it will be unaffected by
+            # the previous fiber operations.
+            Mongoid::Threaded.stack(:testing) << 'd'
+          end
+        end
+
+        it 'clears the fiber state' do
+          expect(result).to eq(%w[ d ])
+        end
+      end
+    end
+  end
+end