From cc503225b0908f12b791d0e051f56712ce94b673 Mon Sep 17 00:00:00 2001 From: Rafal Wojsznis Date: Thu, 12 Jul 2018 21:24:16 +0200 Subject: [PATCH] Make lock container configurable --- README.md | 41 +++++++++++++++++--- lib/sidekiq/lock.rb | 12 +++++- lib/sidekiq/lock/container.rb | 15 ++++++++ lib/sidekiq/lock/middleware.rb | 4 +- lib/sidekiq/lock/testing/inline.rb | 4 +- lib/sidekiq/lock/worker.rb | 2 +- test/lib/container_test.rb | 23 ++++++++++++ test/lib/middleware_test.rb | 19 +++++----- test/lib/testing/inline_test.rb | 15 +++++--- test/lib/worker_test.rb | 60 ++++++++++++++++++++++++------ test/test_helper.rb | 12 ++++-- 11 files changed, 164 insertions(+), 43 deletions(-) create mode 100644 lib/sidekiq/lock/container.rb create mode 100644 test/lib/container_test.rb diff --git a/README.md b/README.md index ec425ad..12f3991 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,8 @@ Redis-based simple locking mechanism for [sidekiq][2]. Uses [SET command][1] introduced in Redis 2.6.16. -It can be handy if you push a lot of jobs into the queue(s), but you don't want to execute specific jobs at the same time - it provides a `lock` method that you can use in whatever way you want. +It can be handy if you push a lot of jobs into the queue(s), but you don't want to execute specific jobs at the same +time - it provides a `lock` method that you can use in whatever way you want. ## Installation @@ -33,7 +34,8 @@ $ bundle Sidekiq-lock is a middleware/module combination, let me go through my thought process here :). -In your worker class include `Sidekiq::Lock::Worker` module and provide `lock` attribute inside `sidekiq_options`, for example: +In your worker class include `Sidekiq::Lock::Worker` module and provide `lock` attribute inside `sidekiq_options`, +for example: ``` ruby class Worker @@ -51,13 +53,17 @@ end What will happen is: -- middleware will setup a `Sidekiq::Lock::RedisLock` object under `Thread.current[Sidekiq::Lock::THREAD_KEY]` (well, I had no better idea for this) - assuming you provided `lock` options, otherwise it will do nothing, just execute your worker's code +- middleware will setup a `Sidekiq::Lock::RedisLock` object under `Thread.current[Sidekiq::Lock::THREAD_KEY]` +(it should work in most use cases without any problems - but it's configurable, more below) - assuming you provided +`lock` options, otherwise it will do nothing, just execute your worker's code -- `Sidekiq::Lock::Worker` module provides a `lock` method that just simply points to that thread variable, just as a convenience +- `Sidekiq::Lock::Worker` module provides a `lock` method that just simply points to that thread variable, just as +a convenience So now in your worker class you can call (whenever you need): -- `lock.acquire!` - will try to acquire the lock, if returns false on failure (that means some other process / thread took the lock first) +- `lock.acquire!` - will try to acquire the lock, if returns false on failure (that means some other process / thread +took the lock first) - `lock.acquired?` - set to `true` when lock is successfully acquired - `lock.release!` - deletes the lock (only if it's: acquired by current thread and not already expired) @@ -114,9 +120,32 @@ Sidekiq.configure_server do |config| end ``` +### Customizing lock _container_ + +If you would like to change default behavior of storing lock instance in `Thread.current` for whatever reason you +can do that as well via server configuration: + +``` ruby +# Any thread-safe class that implements .fetch and .store methods will do +class CustomStorage + def fetch + # returns stored lock instance + end + + def store(lock_instance) + # store lock + end +end + +Sidekiq.configure_server do |config| + config.lock_container = CustomStorage.new +end +``` + ### Inline testing -As you know middleware is not invoked when testing jobs inline, you can require in your test/spec helper file `sidekiq/lock/testing/inline` to include two methods that will help you setting / clearing up lock manually: +As you know middleware is not invoked when testing jobs inline, you can require in your test/spec helper file +`sidekiq/lock/testing/inline` to include two methods that will help you setting / clearing up lock manually: - `set_sidekiq_lock(worker_class, payload)` - note: payload should be an array of worker arguments - `clear_sidekiq_lock` diff --git a/lib/sidekiq/lock.rb b/lib/sidekiq/lock.rb index c4eb71e..7d94d7b 100644 --- a/lib/sidekiq/lock.rb +++ b/lib/sidekiq/lock.rb @@ -1,11 +1,20 @@ +require 'sidekiq/lock/container' require 'sidekiq/lock/middleware' require 'sidekiq/lock/redis_lock' require 'sidekiq/lock/version' require 'sidekiq/lock/worker' module Sidekiq + def self.lock_container + @lock_container ||= Lock::Container.new + end + def self.lock_method - @lock_method ||= :lock + @lock_method ||= Lock::METHOD_NAME + end + + def self.lock_container=(container) + @lock_container = container end def self.lock_method=(method) @@ -14,6 +23,7 @@ def self.lock_method=(method) module Lock THREAD_KEY = :sidekiq_lock + METHOD_NAME = :lock end end diff --git a/lib/sidekiq/lock/container.rb b/lib/sidekiq/lock/container.rb new file mode 100644 index 0000000..cb085bd --- /dev/null +++ b/lib/sidekiq/lock/container.rb @@ -0,0 +1,15 @@ +module Sidekiq + module Lock + class Container + THREAD_KEY = :sidekiq_lock + + def fetch + Thread.current[THREAD_KEY] + end + + def store(lock) + Thread.current[THREAD_KEY] = lock + end + end + end +end diff --git a/lib/sidekiq/lock/middleware.rb b/lib/sidekiq/lock/middleware.rb index 5f96c74..e8e1d7e 100644 --- a/lib/sidekiq/lock/middleware.rb +++ b/lib/sidekiq/lock/middleware.rb @@ -1,7 +1,7 @@ module Sidekiq module Lock class Middleware - def call(worker, msg, queue) + def call(worker, msg, _queue) options = lock_options(worker) setup_lock(options, msg['args']) unless options.nil? @@ -11,7 +11,7 @@ def call(worker, msg, queue) private def setup_lock(options, payload) - Thread.current[Sidekiq::Lock::THREAD_KEY] = RedisLock.new(options, payload) + Sidekiq.lock_container.store(RedisLock.new(options, payload)) end def lock_options(worker) diff --git a/lib/sidekiq/lock/testing/inline.rb b/lib/sidekiq/lock/testing/inline.rb index 69080db..897f518 100644 --- a/lib/sidekiq/lock/testing/inline.rb +++ b/lib/sidekiq/lock/testing/inline.rb @@ -1,8 +1,8 @@ def set_sidekiq_lock(worker_class, payload) options = worker_class.get_sidekiq_options['lock'] - Thread.current[Sidekiq::Lock::THREAD_KEY] = Sidekiq::Lock::RedisLock.new(options, payload) + Sidekiq.lock_container.store(Sidekiq::Lock::RedisLock.new(options, payload)) end def clear_sidekiq_lock - Thread.current[Sidekiq::Lock::THREAD_KEY] = nil + Sidekiq.lock_container.store(nil) end diff --git a/lib/sidekiq/lock/worker.rb b/lib/sidekiq/lock/worker.rb index 6c3abe3..f941855 100644 --- a/lib/sidekiq/lock/worker.rb +++ b/lib/sidekiq/lock/worker.rb @@ -3,7 +3,7 @@ module Lock module Worker def self.included(base) base.send(:define_method, Sidekiq.lock_method) do - Thread.current[Sidekiq::Lock::THREAD_KEY] + Sidekiq.lock_container.fetch end end end diff --git a/test/lib/container_test.rb b/test/lib/container_test.rb new file mode 100644 index 0000000..4b6a635 --- /dev/null +++ b/test/lib/container_test.rb @@ -0,0 +1,23 @@ +require 'test_helper' +require 'open3' + +module Sidekiq + module Lock + describe Container do + it 'stores and fetches given value under Thread.current' do + begin + container = Sidekiq::Lock::Container.new + thread_key = Sidekiq::Lock::Container::THREAD_KEY + + Thread.current[thread_key] = 'value' + assert_equal 'value', container.fetch + + container.store 'new-value' + assert_equal Thread.current[thread_key], 'new-value' + ensure + Thread.current[thread_key] = nil + end + end + end + end +end diff --git a/test/lib/middleware_test.rb b/test/lib/middleware_test.rb index 1b61597..241ef73 100644 --- a/test/lib/middleware_test.rb +++ b/test/lib/middleware_test.rb @@ -6,36 +6,35 @@ module Lock before do Sidekiq.redis = REDIS Sidekiq.redis { |c| c.flushdb } - set_lock_variable! + reset_lock_variable! end - let(:handler){ Sidekiq::Lock::Middleware.new } + let(:handler) { Sidekiq::Lock::Middleware.new } it 'sets lock variable with provided static lock options' do - handler.call(LockWorker.new, {'class' => LockWorker, 'args' => []}, 'default') do + handler.call(LockWorker.new, { 'class' => LockWorker, 'args' => [] }, 'default') do true end - assert_kind_of RedisLock, lock_thread_variable + assert_kind_of RedisLock, lock_container_variable end it 'sets lock variable with provided dynamic options' do - handler.call(DynamicLockWorker.new, {'class' => DynamicLockWorker, 'args' => [1234, 1000]}, 'default') do + handler.call(DynamicLockWorker.new, { 'class' => DynamicLockWorker, 'args' => [1234, 1000] }, 'default') do true end - assert_equal "lock:1234", lock_thread_variable.name - assert_equal 2000, lock_thread_variable.timeout + assert_equal "lock:1234", lock_container_variable.name + assert_equal 2000, lock_container_variable.timeout end it 'sets nothing for workers without lock options' do - handler.call(RegularWorker.new, {'class' => RegularWorker, 'args' => []}, 'default') do + handler.call(RegularWorker.new, { 'class' => RegularWorker, 'args' => [] }, 'default') do true end - assert_nil lock_thread_variable + assert_nil lock_container_variable end - end end end diff --git a/test/lib/testing/inline_test.rb b/test/lib/testing/inline_test.rb index 3139bb4..1b026fd 100644 --- a/test/lib/testing/inline_test.rb +++ b/test/lib/testing/inline_test.rb @@ -2,20 +2,23 @@ require "sidekiq/lock/testing/inline" describe "inline test helper" do - after { set_lock_variable! } + after { reset_lock_variable! } it "has helper fuction for setting lock" do - Sidekiq::Lock::RedisLock.expects(:new).with({ timeout: 1, name: 'lock-worker' }, 'worker argument').returns('lock set') + Sidekiq::Lock::RedisLock + .expects(:new) + .with({ timeout: 1, name: 'lock-worker' }, 'worker argument') + .returns('lock set') + set_sidekiq_lock(LockWorker, 'worker argument') - assert_equal 'lock set', lock_thread_variable + assert_equal 'lock set', lock_container_variable end it "has helper fuction for clearing lock" do set_lock_variable! "test" - assert_equal "test", lock_thread_variable + assert_equal "test", lock_container_variable clear_sidekiq_lock - assert_nil lock_thread_variable + assert_nil lock_container_variable end - end diff --git a/test/lib/worker_test.rb b/test/lib/worker_test.rb index af86e10..9d1541b 100644 --- a/test/lib/worker_test.rb +++ b/test/lib/worker_test.rb @@ -3,26 +3,64 @@ module Sidekiq module Lock describe Worker do - after { set_lock_variable! } + # after { } - it 'sets lock method that points to thread variable' do - set_lock_variable! "test" - assert_equal "test", LockWorker.new.lock + class CustomContainer + def initialize + @lock = nil + end + + def fetch + @lock + end + + def store(lock) + @lock = lock + end end + # it 'sets lock method that points to thread variable' do + # set_lock_variable! "test" + # assert_equal "test", LockWorker.new.lock + # end + it 'allows method name configuration' do - Sidekiq.lock_method = :custom_lock_name + begin + Sidekiq.lock_method = :custom_lock_name - class WorkerWithCustomLockName - include Sidekiq::Worker - include Sidekiq::Lock::Worker - end + class WorkerWithCustomLockName + include Sidekiq::Worker + include Sidekiq::Lock::Worker + end + + set_lock_variable! "custom_name" - set_lock_variable! "custom_name" + assert_equal "custom_name", WorkerWithCustomLockName.new.custom_lock_name - assert_equal "custom_name", WorkerWithCustomLockName.new.custom_lock_name + reset_lock_variable! + ensure + + Sidekiq.lock_method = Sidekiq::Lock::METHOD_NAME + end end + it 'allows container configuration' do + begin + container = CustomContainer.new + Sidekiq.lock_container = container + + class WorkerWithCustomContainer + include Sidekiq::Worker + include Sidekiq::Lock::Worker + end + + container.store "lock-variable" + + assert_equal "lock-variable", WorkerWithCustomContainer.new.lock + ensure + Sidekiq.lock_container = Sidekiq::Lock::Container.new + end + end end end end diff --git a/test/test_helper.rb b/test/test_helper.rb index f8511a4..c0f4eba 100644 --- a/test/test_helper.rb +++ b/test/test_helper.rb @@ -20,10 +20,14 @@ def redis(command, *args) end end -def set_lock_variable!(value = nil) - Thread.current[Sidekiq::Lock::THREAD_KEY] = value +def set_lock_variable!(value) + Sidekiq.lock_container.store(value) end -def lock_thread_variable - Thread.current[Sidekiq::Lock::THREAD_KEY] +def reset_lock_variable! + set_lock_variable!(nil) +end + +def lock_container_variable + Sidekiq.lock_container.fetch end