From 9c209739d4f26045ec916c4873d4a8b894986b20 Mon Sep 17 00:00:00 2001 From: Tobias Schoknecht Date: Thu, 15 Oct 2020 17:57:39 +0200 Subject: [PATCH] Add support for automatically reconnecting if a readonly host was detected --- Gemfile.lock | 2 +- README.md | 6 +- lib/gcra/redis_store.rb | 26 +++++-- lib/gcra/version.rb | 2 +- spec/lib/gcra/redis_store_spec.rb | 109 +++++++++++++++++++++++++++++- 5 files changed, 137 insertions(+), 8 deletions(-) diff --git a/Gemfile.lock b/Gemfile.lock index dbd52e6..f296a81 100644 --- a/Gemfile.lock +++ b/Gemfile.lock @@ -1,7 +1,7 @@ PATH remote: . specs: - gcra (1.1.0) + gcra (1.2.0) GEM remote: https://rubygems.org/ diff --git a/README.md b/README.md index 659b162..1b0d84f 100644 --- a/README.md +++ b/README.md @@ -24,7 +24,11 @@ require 'gcra/redis_store' redis = Redis.new(host: 'localhost', port: 6379, timeout: 0.1) key_prefix = 'rate-limit-app1:' -store = GCRA::RedisStore.new(redis, key_prefix) +store = GCRA::RedisStore.new( + redis, + key_prefix, + { reconnect_on_readonly: false }, + ) rate_period = 0.5 # Two requests per second max_burst = 10 # Allow 10 additional requests as a burst diff --git a/lib/gcra/redis_store.rb b/lib/gcra/redis_store.rb index 0164694..1df30ac 100644 --- a/lib/gcra/redis_store.rb +++ b/lib/gcra/redis_store.rb @@ -16,11 +16,15 @@ class RedisStore # Digest::SHA1.hexdigest(CAS_SCRIPT) CAS_SHA = "89118e702230c0d65969c5fc557a6e942a2f4d31".freeze CAS_SCRIPT_MISSING_KEY_RESPONSE = 'key does not exist'.freeze - SCRIPT_NOT_IN_CACHE_RESPONSE = 'NOSCRIPT No matching script. Please use EVAL.' + SCRIPT_NOT_IN_CACHE_RESPONSE = 'NOSCRIPT No matching script. Please use EVAL.'.freeze - def initialize(redis, key_prefix) + CONNECTED_TO_READONLY = "READONLY You can't write against a read only slave.".freeze + + def initialize(redis, key_prefix, options = {}) @redis = redis @key_prefix = key_prefix + + @reconnect_on_readonly = options[:reconnect_on_readonly] || false end # Returns the value of the key or nil, if it isn't in the store. @@ -43,8 +47,18 @@ def get_with_time(key) # Also set the key's expiration (ttl, in seconds). def set_if_not_exists_with_ttl(key, value, ttl_nano) full_key = @key_prefix + key - ttl_milli = calculate_ttl_milli(ttl_nano) - @redis.set(full_key, value, nx: true, px: ttl_milli) + retried = false + begin + ttl_milli = calculate_ttl_milli(ttl_nano) + @redis.set(full_key, value, nx: true, px: ttl_milli) + rescue Redis::CommandError => e + if e.message == CONNECTED_TO_READONLY && @reconnect_on_readonly && !retried + @redis.client.reconnect + retried = true + retry + end + raise + end end # Atomically compare the value at key to the old value. If it matches, set it to the new value @@ -63,6 +77,10 @@ def compare_and_set_with_ttl(key, old_value, new_value, ttl_nano) @redis.script('load', CAS_SCRIPT) retried = true retry + elsif e.message == CONNECTED_TO_READONLY && @reconnect_on_readonly && !retried + @redis.client.reconnect + retried = true + retry end raise end diff --git a/lib/gcra/version.rb b/lib/gcra/version.rb index 2575ddf..2e21b60 100644 --- a/lib/gcra/version.rb +++ b/lib/gcra/version.rb @@ -1,3 +1,3 @@ module GCRA - VERSION = '1.1.0'.freeze + VERSION = '1.2.0'.freeze end diff --git a/spec/lib/gcra/redis_store_spec.rb b/spec/lib/gcra/redis_store_spec.rb index bf2dcca..3486859 100644 --- a/spec/lib/gcra/redis_store_spec.rb +++ b/spec/lib/gcra/redis_store_spec.rb @@ -8,7 +8,8 @@ # Needs redis running on localhost:6379 (default port) let(:redis) { Redis.new } let(:key_prefix) { 'gcra-ruby-specs:' } - let(:store) { described_class.new(redis, key_prefix) } + let(:options) { {} } + let(:store) { described_class.new(redis, key_prefix, options) } def cleanup_redis keys = redis.keys("#{key_prefix}*") @@ -67,6 +68,47 @@ def cleanup_redis expect(did_set).to eq(false) end + it 'with a readonly host and no readonly configured' do + exception = Redis::CommandError.new(GCRA::RedisStore::CONNECTED_TO_READONLY) + + expect(redis) + .to receive(:set) + .with('gcra-ruby-specs:foo', 2_000_000_000_000_000_000, nx: true, px: 1) + .and_raise(exception) + + expect do + store.set_if_not_exists_with_ttl('foo', 2_000_000_000_000_000_000, 1) + end.to raise_error(exception) + end + + context 'with reconnect on readonly configured' do + let(:options) do + { + reconnect_on_readonly: true, + } + end + + it 'with a readonly host' do + exception = Redis::CommandError.new(GCRA::RedisStore::CONNECTED_TO_READONLY) + + expect(redis.client) + .to receive(:reconnect) + .and_call_original + + expect(redis) + .to receive(:set) + .with('gcra-ruby-specs:foo', 2_000_000_000_000_000_000, nx: true, px: 1) + .and_raise(exception) + + expect(redis) + .to receive(:set) + .with('gcra-ruby-specs:foo', 2_000_000_000_000_000_000, nx: true, px: 1) + .and_call_original + + store.set_if_not_exists_with_ttl('foo', 2_000_000_000_000_000_000, 1) + end + end + it 'with no existing key, returns true' do did_set = store.set_if_not_exists_with_ttl( 'foo', 3_000_000_000_000_000_000, 10 * 1_000_000_000 @@ -156,6 +198,71 @@ def cleanup_redis expect(redis.ttl('gcra-ruby-specs:foo')).to be > 8 expect(redis.ttl('gcra-ruby-specs:foo')).to be <= 10 end + + context 'with reconnect on readonly not configured' do + it 'raises an error when the request is executed against a readonly host' do + exception = Redis::CommandError.new(GCRA::RedisStore::CONNECTED_TO_READONLY) + + expect(redis) + .to receive(:evalsha) + .with( + GCRA::RedisStore::CAS_SHA, + keys: ['gcra-ruby-specs:foo'], + argv: [3000000000000000000, 4000000000000000000, 10000], + ) + .and_raise(exception) + + expect do + store.compare_and_set_with_ttl( + 'foo', + 3_000_000_000_000_000_000, + 4_000_000_000_000_000_000, + 10 * 1_000_000_000, + ) + end.to raise_error(exception) + end + end + + context 'with reconnect on readonly configured' do + let(:options) do + { + reconnect_on_readonly: true, + } + end + + it 'attempts a reconnect once and then executes evalsha again' do + exception = Redis::CommandError.new(GCRA::RedisStore::CONNECTED_TO_READONLY) + + expect(redis.client) + .to receive(:reconnect) + .and_call_original + + expect(redis) + .to receive(:evalsha) + .with( + GCRA::RedisStore::CAS_SHA, + keys: ['gcra-ruby-specs:foo'], + argv: [3000000000000000000, 4000000000000000000, 10000], + ) + .and_raise(exception) + + expect(redis) + .to receive(:evalsha) + .with( + GCRA::RedisStore::CAS_SHA, + keys: ['gcra-ruby-specs:foo'], + argv: [3000000000000000000, 4000000000000000000, 10000], + ) + + store.compare_and_set_with_ttl( + 'foo', + 3_000_000_000_000_000_000, + 4_000_000_000_000_000_000, + 10 * 1_000_000_000, + ) + end + end + end context 'functional test with RateLimiter' do