Use built-in wait and timeout functionality if the database engine supports it

lib/with_advisory_lock/base.rb

@@ -84,27 +84,16 @@ def stable_hashcode(input)
     end
 
     def yield_with_lock_and_timeout(&block)
-      give_up_at = Time.now + @timeout_seconds if @timeout_seconds
-      while @timeout_seconds.nil? || Time.now < give_up_at
-        r = yield_with_lock(&block)
-        return r if r.lock_was_acquired?
-
-        # Randomizing sleep time may help reduce contention.
-        sleep(rand(0.05..0.15))
+      if lock
+        yield_with_acquired_lock(&block)
+      else
+        FAILED_TO_LOCK
       end
-      FAILED_TO_LOCK
     end
 
-    def yield_with_lock
+    def yield_with_lock(&block)
       if try_lock
-        begin
-          lock_stack.push(lock_stack_item)
-          result = block_given? ? yield : nil
-          Result.new(true, result)
-        ensure
-          lock_stack.pop
-          release_lock
-        end
+        yield_with_acquired_lock(&block)
       else
         FAILED_TO_LOCK
       end
@@ -114,5 +103,30 @@ def yield_with_lock
     def unique_column_name
       "t#{SecureRandom.hex}"
     end
+
+    private
+
+    def yield_with_acquired_lock
+      begin
+        lock_stack.push(lock_stack_item)
+        result = block_given? ? yield : nil
+        Result.new(true, result)
+      ensure
+        lock_stack.pop
+        release_lock
+      end
+    end
+
+    def lock_via_sleep_loop
+      give_up_at = Time.now + timeout_seconds if timeout_seconds
+      loop do
+        return true if try_lock
+
+        # Randomizing sleep time may help reduce contention.
+        sleep(rand(0.05..0.15))
+
+        return false if timeout_seconds && Time.now > give_up_at
+      end
+    end
   end
 end

lib/with_advisory_lock/flock.rb

@@ -26,6 +26,10 @@ def try_lock
       0 == file_io.flock((shared ? File::LOCK_SH : File::LOCK_EX) | File::LOCK_NB)
     end
 
+    def lock
+      lock_via_sleep_loop
+    end
+
     def release_lock
       0 == file_io.flock(File::LOCK_UN)
     end

lib/with_advisory_lock/mysql.rb

@@ -2,14 +2,22 @@
 
 module WithAdvisoryLock
   class MySQL < Base
-    # See https://dev.mysql.com/doc/refman/5.7/en/miscellaneous-functions.html#function_get-lock
+    # See https://dev.mysql.com/doc/refman/en/locking-functions.html
     def try_lock
       raise ArgumentError, 'shared locks are not supported on MySQL' if shared
       raise ArgumentError, 'transaction level locks are not supported on MySQL' if transaction
 
       execute_successful?("GET_LOCK(#{quoted_lock_str}, 0)")
     end
 
+    # See https://dev.mysql.com/doc/refman/en/locking-functions.html
+    def lock
+      raise ArgumentError, 'shared locks are not supported on MySQL' if shared
+      raise ArgumentError, 'transaction level locks are not supported on MySQL' if transaction
+
+      execute_successful?("GET_LOCK(#{quoted_lock_str}, #{timeout.nil? ? -1 : timeout})")
+    end
+
     def release_lock
       execute_successful?("RELEASE_LOCK(#{quoted_lock_str})")
     end

lib/with_advisory_lock/postgresql.rb

@@ -8,6 +8,16 @@ def try_lock
       execute_successful?(pg_function)
     end
 
+    # See http://www.postgresql.org/docs/9.1/static/functions-admin.html#FUNCTIONS-ADVISORY-LOCKS
+    def lock
+      if timeout.nil?
+        pg_function = "pg_advisory#{transaction ? '_xact' : ''}_lock#{shared ? '_shared' : ''}"
+        execute_successful?(pg_function)
+      else
+        lock_via_sleep_loop
+      end
+    end
+
     def release_lock
       return if transaction
 

test/concern_test.rb

@@ -21,7 +21,7 @@ class WithAdvisoryLockConcernTest < GemTestCase
 end
 
 class ActiveRecordQueryCacheTest < GemTestCase
-  test 'does not disable quary cache by default' do
+  test 'does not disable query cache by default' do
     ActiveRecord::Base.expects(:uncached).never
     Tag.with_advisory_lock('lock') { Tag.first }
   end

test/thread_test.rb

@@ -7,13 +7,14 @@ class SeparateThreadTest < GemTestCase
     @lock_name = 'testing 1,2,3' # OMG COMMAS
     @mutex = Mutex.new
     @t1_acquired_lock = false
+    @t1_locking = true
     @t1_return_value = nil
 
     @t1 = Thread.new do
       ActiveRecord::Base.connection_pool.with_connection do
         @t1_return_value = Label.with_advisory_lock(@lock_name) do
           @mutex.synchronize { @t1_acquired_lock = true }
-          sleep
+          sleep(0.1) while @t1_locking
           't1 finished'
         end
       end
@@ -25,14 +26,46 @@ class SeparateThreadTest < GemTestCase
   end
 
   teardown do
+    @t1_locking = false
     @t1.wakeup if @t1.status == 'sleep'
     @t1.join
   end
 
-  test '#with_advisory_lock with a 0 timeout returns false immediately' do
+  test '#with_advisory_lock with no timeout waits until lock can be acquired, yields to the provided block, and then returns true' do
+    yielded_to = false
+    t2 = Thread.new { sleep(1); @t1_locking = false }
+    start_time = Time.now
+
+    response = Label.with_advisory_lock(@lock_name) do
+      yielded_to = true
+    end
+
+    t2.join
+
+    assert_in_delta(Time.now - start_time, 1, 0.5, "Expected with_advisory_lock to wait 1 second")
+    assert(yielded_to, "Expected with_advisory_lock to yield to the block")
+    assert(response, "Expect with_advisory_lock to return true")
+  end
+
+  test '#with_advisory_lock with a 0 timeout returns false immediately and does not yield to the provided block' do
+    start_time = Time.now
+
     response = Label.with_advisory_lock(@lock_name, 0) do
       raise 'should not be yielded to'
     end
+
+    assert_in_delta(Time.now - start_time, 0, 0.5, "Expected with_advisory_lock to return immediately")
+    assert_not(response)
+  end
+
+  test '#with_advisory_lock with a 1 timeout waits 1 second, returns false, and does not yield to the provided block' do
+    start_time = Time.now
+
+    response = Label.with_advisory_lock(@lock_name, 1) do
+      raise 'should not be yielded to'
+    end
+
+    assert_in_delta(Time.now - start_time, 1, 0.5, "Expected with_advisory_lock to wait 1 second")
     assert_not(response)
   end
 
@@ -45,6 +78,7 @@ class SeparateThreadTest < GemTestCase
   end
 
   test 'can re-establish the lock after the other thread releases it' do
+    @t1_locking = false
     @t1.wakeup
     @t1.join
     assert_equal('t1 finished', @t1_return_value)

test/transaction_test.rb

@@ -29,7 +29,7 @@ class PostgresqlTest < TransactionScopingTest
       end
     end
 
-    test 'session locks release after the block executes' do
+    test 'without timeout, the session locks release after the block executes' do
       Tag.transaction do
         assert_equal(0, @pg_lock_count.call)
         Tag.with_advisory_lock 'test' do
@@ -39,7 +39,17 @@ class PostgresqlTest < TransactionScopingTest
       end
     end
 
-    test 'session locks release when transaction fails inside block' do
+    test 'with timeout, the session locks release after the block executes' do
+      Tag.transaction do
+        assert_equal(0, @pg_lock_count.call)
+        Tag.with_advisory_lock 'test', timeout_seconds: 1  do
+          assert_equal(1, @pg_lock_count.call)
+        end
+        assert_equal(0, @pg_lock_count.call)
+      end
+    end
+
+    test 'without timeout, the session locks release when transaction fails inside block' do
       Tag.transaction do
         assert_equal(0, @pg_lock_count.call)
 
@@ -54,7 +64,22 @@ class PostgresqlTest < TransactionScopingTest
       end
     end
 
-    test 'transaction level locks hold until the transaction completes' do
+    test 'with timeout, the session locks release when transaction fails inside block' do
+      Tag.transaction do
+        assert_equal(0, @pg_lock_count.call)
+
+        exception = assert_raises(ActiveRecord::StatementInvalid) do
+          Tag.with_advisory_lock 'test', timeout_seconds: 1 do
+            Tag.connection.execute 'SELECT 1/0;'
+          end
+        end
+
+        assert_match(/#{Regexp.escape('division by zero')}/, exception.message)
+        assert_equal(0, @pg_lock_count.call)
+      end
+    end
+
+    test 'without timeout, the transaction level locks hold until the transaction completes' do
       Tag.transaction do
         assert_equal(0, @pg_lock_count.call)
         Tag.with_advisory_lock 'test', transaction: true do
@@ -64,5 +89,16 @@ class PostgresqlTest < TransactionScopingTest
       end
       assert_equal(0, @pg_lock_count.call)
     end
+
+    test 'with timeout, the transaction level locks hold until the transaction completes' do
+      Tag.transaction do
+        assert_equal(0, @pg_lock_count.call)
+        Tag.with_advisory_lock 'test', timeout_seconds: 1, transaction: true do
+          assert_equal(1, @pg_lock_count.call)
+        end
+        assert_equal(1, @pg_lock_count.call)
+      end
+      assert_equal(0, @pg_lock_count.call)
+    end
   end
 end