Make advisory lock more robust

fatkodima · Jan 16, 2025 · a1aaf37 · a1aaf37
1 parent d598591
commit a1aaf37
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/lib/online_migrations/advisory_lock.rb b/lib/online_migrations/advisory_lock.rb
@@ -23,13 +23,10 @@ def unlock
 
     # Runs the given block if an advisory lock is able to be acquired.
     def try_with_lock
-      if try_lock
-        begin
-          yield
-        ensure
-          unlock
-        end
-      end
+      locked = try_lock
+      yield if locked
+    ensure
+      unlock if locked
     end
 
     def active?

diff --git a/test/advisory_lock_test.rb b/test/advisory_lock_test.rb
@@ -27,9 +27,12 @@ def test_unlock
   def test_try_with_lock
     assert_not @lock.active?
 
+    called = false
     @lock.try_with_lock do
+      called = true
       assert @lock.active?
     end
+    assert called
 
     assert_not @lock.active?
   end