From 47f67d1ccc022dd3ffa99099597e8d853539ce73 Mon Sep 17 00:00:00 2001
From: Tobias Svensson <tob@tobiassvensson.co.uk>
Date: Mon, 16 Sep 2013 17:27:07 +0100
Subject: [PATCH] Clean up locking.

---
 lib/sidetiq/lock/redis.rb   | 81 ++++++++++++++++++++++++-------------
 lib/sidetiq/web.rb          |  8 +---
 test/test_lock.rb           | 10 +++++
 test/test_lock_meta_data.rb |  2 +-
 4 files changed, 65 insertions(+), 36 deletions(-)

diff --git a/lib/sidetiq/lock/redis.rb b/lib/sidetiq/lock/redis.rb
index d420ff3..e9ed4b4 100644
--- a/lib/sidetiq/lock/redis.rb
+++ b/lib/sidetiq/lock/redis.rb
@@ -5,54 +5,79 @@ module Sidetiq
 
       attr_reader :key, :timeout
 
+      def self.all
+        Sidekiq.redis do |redis|
+          redis.keys("sidetiq:*:lock").map do |key|
+            new(key)
+          end
+        end
+      end
+
       def initialize(key, timeout = Sidetiq.config.lock_expire)
-        @key = key.kind_of?(Class) ? "sidetiq:#{key.name}:lock" : "sidetiq:#{key}:lock"
+        @key = extract_key(key)
         @timeout = timeout
       end
 
       def synchronize
         Sidekiq.redis do |redis|
-          acquired, meta = lock(redis)
+          acquired = lock
 
           if acquired
-            debug "Lock: #{meta}"
+            debug "Lock: #{key}"
 
             begin
               yield redis
             ensure
-              unlock(redis)
+              unlock
               debug "Unlock: #{key}"
             end
           end
         end
       end
 
+      def meta_data
+        Sidekiq.redis do |redis|
+          MetaData.from_json(redis.get(key))
+        end
+      end
+
+      def lock
+        Sidekiq.redis do |redis|
+          acquired = false
+
+          watch(redis, key) do
+            if !redis.exists(key)
+              acquired = !!redis.multi do |multi|
+                meta = MetaData.for_new_lock(key)
+                multi.psetex(key, timeout, meta.to_json)
+              end
+            end
+          end
+
+          acquired
+        end
+      end
+
+      def unlock
+        Sidekiq.redis do |redis|
+          watch(redis, key) do
+            if meta_data.owner == Sidetiq::Lock::MetaData::OWNER
+              redis.multi do |multi|
+                multi.del(key)
+              end
+            end
+          end
+        end
+      end
+
       private
 
-      def lock(redis)
-        acquired, meta = false, nil
-
-        watch(redis, key) do
-          if !redis.exists(key)
-            acquired = !!redis.multi do |multi|
-              meta = MetaData.for_new_lock(key)
-              multi.psetex(key, timeout, meta.to_json)
-            end
-          end
-        end
-
-        [acquired, meta]
-      end
-
-      def unlock(redis)
-        watch(redis, key) do
-          meta = MetaData.from_json(redis.get(key))
-
-          if meta.owner == Sidetiq::Lock::MetaData::OWNER
-            redis.multi do |multi|
-              multi.del(key)
-            end
-          end
+      def extract_key(key)
+        case key
+        when Class
+          "sidetiq:#{key.name}:lock"
+        when String
+          key.match(/sidetiq:(.+):lock/) ? key : "sidetiq:#{key}:lock"
         end
       end
 
diff --git a/lib/sidetiq/web.rb b/lib/sidetiq/web.rb
index cfcb2c6..25169d2 100644
--- a/lib/sidetiq/web.rb
+++ b/lib/sidetiq/web.rb
@@ -12,13 +12,7 @@ module Sidetiq
       end
 
       app.get "/sidetiq/locks" do
-        Sidekiq.redis do |redis|
-          lock_keys = redis.keys('sidetiq:*:lock')
-
-          @locks = (lock_keys.any? ? redis.mget(*lock_keys) : []).map do |lock|
-            Sidetiq::Lock::MetaData.from_json(lock)
-          end
-        end
+        @locks = Sidetiq::Lock::Redis.all.map(&:meta_data)
 
         erb File.read(File.join(VIEWS, 'sidetiq_locks.erb'))
       end
diff --git a/test/test_lock.rb b/test/test_lock.rb
index 0dd44e9..862c3f9 100644
--- a/test/test_lock.rb
+++ b/test/test_lock.rb
@@ -21,6 +21,16 @@ class TestLock < Sidetiq::TestCase
     end
   end
 
+  def test_all
+    Sidekiq.redis do |redis|
+      redis.set("sidetiq:Foobar:lock", 1)
+    end
+
+    locks = Sidetiq::Lock::Redis.all
+
+    assert_equal "sidetiq:Foobar:lock", locks.first.key
+  end
+
   def test_lock_sets_correct_meta_data
     key = SecureRandom.hex(8)
     internal_key = "sidetiq:#{key}:lock"
diff --git a/test/test_lock_meta_data.rb b/test/test_lock_meta_data.rb
index 598ade1..a49e901 100644
--- a/test/test_lock_meta_data.rb
+++ b/test/test_lock_meta_data.rb
@@ -45,7 +45,7 @@ class TestLockMetaData < Sidetiq::TestCase
   def test_pttl
     lock = Sidetiq::Lock::Redis.new("foobar", 1000000)
     md = Sidekiq.redis do |r|
-      lock.send(:lock, r)
+      lock.lock
       Sidetiq::Lock::MetaData.from_json(r.get(lock.key))
     end