diff --git a/lib/online_migrations/advisory_lock.rb b/lib/online_migrations/advisory_lock.rb
index 82fbc82..27998d1 100644
--- 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
index 6ca4172..92e3b93 100644
--- 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