locker.rb in mongoid-locker-0.2.0

- old
+ new

@@ -1,5 +1,7 @@
+require File.expand_path(File.join(File.dirname(__FILE__), 'locker', 'wrapper'))
+
 module Mongoid
   module Locker
     module ClassMethods
       # A scope to retrieve all locked documents in the collection.
       #
@@ -46,68 +48,74 @@
     # @return [Boolean] true if locked, false otherwise
     def locked?
       !!(self.locked_until && self.locked_until > Time.now)
     end
 
+    # Returns whether the current instance has the lock or not.
+    #
+    # @return [Boolean] true if locked, false otherwise
+    def has_lock?
+      @has_lock && self.locked?
+    end
+
     # Primary method of plugin: execute the provided code once the document has been successfully locked.
     #
     # @param [Hash] opts for the locking mechanism
     # @option opts [Fixnum] :timeout The number of seconds until the lock is considered "expired" - defaults to the {ClassMethods#lock_timeout}
     # @option opts [Boolean] :wait If the document is currently locked, wait until the lock expires and try again
     # @return [void]
     def with_lock opts={}, &block
-      self.lock opts
+      # don't try to re-lock/unlock on recursive calls
+      had_lock = self.has_lock?
+      self.lock(opts) unless had_lock
+
       begin
         yield
       ensure
-        self.unlock
+        self.unlock unless had_lock
       end
     end
 
 
     protected
 
     def lock opts={}
-      coll = self.class.collection
       time = Time.now
       timeout = opts[:timeout] || self.class.lock_timeout
       expiration = time + timeout
 
       # lock the document atomically in the DB without persisting entire doc
-      record = coll.find_and_modify(
-        :query => {
+      locked = Mongoid::Locker::Wrapper.update(
+        self.class,
+        {
           :_id => self.id,
           '$or' => [
             # not locked
             {:locked_until => nil},
             # expired
             {:locked_until => {'$lte' => time}}
           ]
         },
-        :update => {
+        {
           '$set' => {
             :locked_at => time,
             :locked_until => expiration
           }
         }
       )
 
-      if record
-        # lock successful
+      if locked
+        # document successfully updated, meaning it was locked
         self.locked_at = time
         self.locked_until = expiration
+        @has_lock = true
       else
         # couldn't grab lock
 
-        existing_query = {
-          :_id => self.id,
-          :locked_until => {'$exists' => true}
-        }
-
-        if opts[:wait] && existing = coll.find(existing_query, :limit => 1).first
+        if opts[:wait] && locked_until = Mongoid::Locker::Wrapper.locked_until(self)
           # doc is locked - wait until it expires
-          wait_time = existing.locked_until - Time.now
+          wait_time = locked_until - Time.now
           sleep wait_time if wait_time > 0
 
           # only wait once
           opts.dup
           opts.delete :wait
@@ -120,18 +128,23 @@
       end
     end
 
     def unlock
       # unlock the document in the DB without persisting entire doc
-      self.class.collection.update({:_id => self.id}, {
-        '$set' => {
-          :locked_at => nil,
-          :locked_until => nil,
+      Mongoid::Locker::Wrapper.update(
+        self.class,
+        {:_id => self.id},
+        {
+          '$set' => {
+            :locked_at => nil,
+            :locked_until => nil,
+          }
         }
-      }, {:safe => true})
+      )
 
       self.locked_at = nil
       self.locked_until = nil
+      @has_lock = false
     end
   end
 
   # Error thrown if document could not be successfully locked.
   class LockError < Exception; end