Merge pull request #283 from HoyaBoya/feature/pessimstic_locking_on_ar_persistence

Introduce Pessimistic Locking For AR Persistence

README.md

@@ -781,6 +781,41 @@ end
 
 which then leads to `transaction(:requires_new => false)`, the Rails default.
 
+### Pessimistic Locking
+
+AASM supports [Active Record pessimistic locking via `with_lock`](http://api.rubyonrails.org/classes/ActiveRecord/Locking/Pessimistic.html#method-i-with_lock) for database persistence layers.
+
+| Option | Purpose |
+| ------ | ------- |
+| `false` (default) | No lock is obtained | |
+| `true` | Obtain a blocking pessimistic lock e.g. `FOR UPDATE` |
+| String | Obtain a lock based on the SQL string e.g. `FOR UPDATE NOWAIT` |
+
+
+```ruby
+class Job < ActiveRecord::Base
+  include AASM
+
+  aasm :requires_lock => true do
+    ...
+  end
+
+  ...
+end
+```
+
+```ruby
+class Job < ActiveRecord::Base
+  include AASM
+
+  aasm :requires_lock => 'FOR UPDATE NOWAIT' do
+    ...
+  end
+
+  ...
+end
+```
+
 
 ### Column name & migration
 

lib/aasm/base.rb

@@ -25,6 +25,11 @@ module AASM
       # use requires_new for nested transactions (in ActiveRecord)
       configure :requires_new_transaction, true
 
+      # use pessimistic locking (in ActiveRecord)
+      # true for FOR UPDATE lock
+      # string for a specific lock type i.e. FOR UPDATE NOWAIT
+      configure :requires_lock, false
+
       # set to true to forbid direct assignment of aasm_state column (in ActiveRecord)
       configure :no_direct_assignment, false
 

lib/aasm/configuration.rb

@@ -15,6 +15,9 @@ module AASM
     # for ActiveRecord: use requires_new for nested transactions?
     attr_accessor :requires_new_transaction
 
+    # for ActiveRecord: use pessimistic locking
+    attr_accessor :requires_lock
+
     # forbid direct assignment in aasm_state column (in ActiveRecord)
     attr_accessor :no_direct_assignment
 

lib/aasm/persistence/active_record_persistence.rb

@@ -165,7 +165,15 @@ module AASM
           end
 
           begin
-            success = options[:persist] ? self.class.transaction(:requires_new => requires_new?(state_machine_name)) { super } : super
+            success = if options[:persist]
+              self.class.transaction(:requires_new => requires_new?(state_machine_name)) do
+                lock!(requires_lock?(state_machine_name)) if requires_lock?(state_machine_name)
+                super
+              end
+            else
+              super
+            end
+
             if options[:persist] && success
               event.fire_callbacks(:after_commit, self, *args)
               event.fire_global_callbacks(:after_all_commits, self, *args)
@@ -184,6 +192,10 @@ module AASM
           AASM::StateMachine[self.class][state_machine_name].config.requires_new_transaction
         end
 
+        def requires_lock?(state_machine_name)
+          AASM::StateMachine[self.class][state_machine_name].config.requires_lock
+        end
+
         def aasm_validate_states
           AASM::StateMachine[self.class].keys.each do |state_machine_name|
             unless aasm_skipping_validations(state_machine_name)

spec/database.rb

@@ -17,24 +17,19 @@ ActiveRecord::Migration.suppress_messages do
     t.string "right"
   end
 
-  ActiveRecord::Migration.create_table "validators", :force => true do |t|
-    t.string "name"
-    t.string "status"
-  end
-  ActiveRecord::Migration.create_table "multiple_validators", :force => true do |t|
-    t.string "name"
-    t.string "status"
+  %w(validators multiple_validators).each do |table_name|
+    ActiveRecord::Migration.create_table table_name, :force => true do |t|
+      t.string "name"
+      t.string "status"
+    end
   end
 
-  ActiveRecord::Migration.create_table "transactors", :force => true do |t|
-    t.string "name"
-    t.string "status"
-    t.integer "worker_id"
-  end
-  ActiveRecord::Migration.create_table "multiple_transactors", :force => true do |t|
-    t.string "name"
-    t.string "status"
-    t.integer "worker_id"
+  %w(transactors no_lock_transactors lock_transactors lock_no_wait_transactors multiple_transactors).each do |table_name|
+    ActiveRecord::Migration.create_table table_name, :force => true do |t|
+      t.string "name"
+      t.string "status"
+      t.integer "worker_id"
+    end
   end
 
   ActiveRecord::Migration.create_table "workers", :force => true do |t|

spec/models/transactor.rb

@@ -26,6 +26,54 @@ private
 
 end
 
+class NoLockTransactor < ActiveRecord::Base
+
+  belongs_to :worker
+
+  include AASM
+
+  aasm :column => :status do
+    state :sleeping, :initial => true
+    state :running
+
+    event :run do
+      transitions :to => :running, :from => :sleeping
+    end
+  end
+end
+
+class LockTransactor < ActiveRecord::Base
+
+  belongs_to :worker
+
+  include AASM
+
+  aasm :column => :status, requires_lock: true do
+    state :sleeping, :initial => true
+    state :running
+
+    event :run do
+      transitions :to => :running, :from => :sleeping
+    end
+  end
+end
+
+class LockNoWaitTransactor < ActiveRecord::Base
+
+  belongs_to :worker
+
+  include AASM
+
+  aasm :column => :status, requires_lock: 'FOR UPDATE NOWAIT' do
+    state :sleeping, :initial => true
+    state :running
+
+    event :run do
+      transitions :to => :running, :from => :sleeping
+    end
+  end
+end
+
 class MultipleTransactor < ActiveRecord::Base
 
   belongs_to :worker

spec/unit/persistence/active_record_persistence_spec.rb

@@ -425,6 +425,39 @@ describe 'transitions with persistence' do
     expect(persistor).not_to be_sleeping
   end
 
+  describe 'pessimistic locking' do
+    let(:worker) { Worker.create!(:name => 'worker', :status => 'sleeping') }
+
+    subject { transactor.run! }
+
+    context 'no lock' do
+      let(:transactor) { NoLockTransactor.create!(:name => 'no_lock_transactor', :worker => worker) }
+
+      it 'should not invoke lock!' do
+        expect(transactor).to_not receive(:lock!)
+        subject
+      end
+    end
+
+    context 'a default lock' do
+      let(:transactor) { LockTransactor.create!(:name => 'lock_transactor', :worker => worker) }
+
+      it 'should invoke lock! with true' do
+        expect(transactor).to receive(:lock!).with(true).and_call_original
+        subject
+      end
+    end
+
+    context 'a FOR UPDATE NOWAIT lock' do
+      let(:transactor) { LockNoWaitTransactor.create!(:name => 'lock_no_wait_transactor', :worker => worker) }
+
+      it 'should invoke lock! with FOR UPDATE NOWAIT' do
+        expect(transactor).to receive(:lock!).with('FOR UPDATE NOWAIT').and_call_original
+        subject
+      end
+    end
+  end
+
   describe 'transactions' do
     let(:worker) { Worker.create!(:name => 'worker', :status => 'sleeping') }
     let(:transactor) { Transactor.create!(:name => 'transactor', :worker => worker) }