From 800a70096780a1d290860a0f88bf1b120a1c5998 Mon Sep 17 00:00:00 2001 From: Aryan Kumar Date: Fri, 10 Oct 2025 11:32:46 +0530 Subject: [PATCH 1/7] Key Server code + rspec --- .gitignore | 2 + .rspec | 1 + Gemfile | 10 +++ Gemfile.lock | 48 ++++++++++++ README.md | 19 ++++- app.rb | 54 +++++++++++++ config.ru | 6 ++ key_server.rb | 112 +++++++++++++++++++++++++++ spec/keyserver_spec.rb | 170 +++++++++++++++++++++++++++++++++++++++++ spec/spec_helper.rb | 2 + 10 files changed, 423 insertions(+), 1 deletion(-) create mode 100644 .gitignore create mode 100644 .rspec create mode 100644 Gemfile create mode 100644 Gemfile.lock create mode 100644 app.rb create mode 100644 config.ru create mode 100644 key_server.rb create mode 100644 spec/keyserver_spec.rb create mode 100644 spec/spec_helper.rb diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..1f6ae9e --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +.bundle +vendor \ No newline at end of file diff --git a/.rspec b/.rspec new file mode 100644 index 0000000..c99d2e7 --- /dev/null +++ b/.rspec @@ -0,0 +1 @@ +--require spec_helper diff --git a/Gemfile b/Gemfile new file mode 100644 index 0000000..2eaf5a6 --- /dev/null +++ b/Gemfile @@ -0,0 +1,10 @@ +source "https://rubygems.org" + +gem "json" +gem "securerandom" +gem "sinatra" + +group :development, :test do + gem 'rspec' + gem 'rack-test' + end \ No newline at end of file diff --git a/Gemfile.lock b/Gemfile.lock new file mode 100644 index 0000000..c8432e5 --- /dev/null +++ b/Gemfile.lock @@ -0,0 +1,48 @@ +GEM + remote: https://rubygems.org/ + specs: + base64 (0.3.0) + diff-lcs (1.6.2) + json (2.7.6) + mustermann (3.0.4) + ruby2_keywords (~> 0.0.1) + rack (2.2.19) + rack-protection (3.2.0) + base64 (>= 0.1.0) + rack (~> 2.2, >= 2.2.4) + rack-test (2.2.0) + rack (>= 1.3) + rspec (3.13.1) + rspec-core (~> 3.13.0) + rspec-expectations (~> 3.13.0) + rspec-mocks (~> 3.13.0) + rspec-core (3.13.5) + rspec-support (~> 3.13.0) + rspec-expectations (3.13.5) + diff-lcs (>= 1.2.0, < 2.0) + rspec-support (~> 3.13.0) + rspec-mocks (3.13.5) + diff-lcs (>= 1.2.0, < 2.0) + rspec-support (~> 3.13.0) + rspec-support (3.13.6) + ruby2_keywords (0.0.5) + securerandom (0.3.2) + sinatra (3.2.0) + mustermann (~> 3.0) + rack (~> 2.2, >= 2.2.4) + rack-protection (= 3.2.0) + tilt (~> 2.0) + tilt (2.6.1) + +PLATFORMS + ruby + +DEPENDENCIES + json + rack-test + rspec + securerandom + sinatra + +BUNDLED WITH + 1.17.2 diff --git a/README.md b/README.md index 025f33d..627bfd2 100644 --- a/README.md +++ b/README.md @@ -1 +1,18 @@ -# keyserver-onboarding \ No newline at end of file +Write a server which can generate random api keys, assign them for usage and release them after sometime. Following endpoints should be available on the server to interact with it. +1. There should be one endpoint to generate keys. Once a key is generated it has an expiry of 5 minutes, after which the key is no longer available. +2. There should be an endpoint to get an available key. On hitting this endpoint server should serve a random key which is not already being used. This key should be blocked and should not be served again by E2 while it is in this state. If no eligible key is available then it should serve 404. +3. There should be an endpoint to unblock a key. Unblocked keys can be served via E2 again. Once a key becomes unblocked we can take it as having a new expiry 5 mins from when it was unblocked. +4. There should be an endpoint to delete a key. Deleted keys should be purged. +5. All keys are to be kept alive by clients calling this endpoint every 5 minutes. If a particular key has not received a keep alive in last five minutes then it should be deleted and never used again. + +Apart from these endpoints, following rules should be enforced: +1. All blocked keys should get released automatically within 60 secs if E3 is not called. +2. No endpoint call should result in an iteration of whole set of keys i.e. no endpoint request should be O(n). They should either be O(lg n) or O(1). + +RSpecs (Unit tests) expected for the code. + +You should build this without full fledged Rails framework. + +#### Technical Requirements: +Ruby verison: 2.6.6 +rspec version: 3.10.x \ No newline at end of file diff --git a/app.rb b/app.rb new file mode 100644 index 0000000..5288909 --- /dev/null +++ b/app.rb @@ -0,0 +1,54 @@ +require './key_server.rb' +require 'sinatra' +keyserver = KeyServer.new +Thread.new do + while true + keyserver.free_up + sleep 1 + end +end + +get '/' do + 'Hi, Welcome to my app. Try the following APIs: + /generatekeys + /key + /realeaseKey/:id + delete /keys/:id + /keepalive/:id' +end +#E1 +post '/generateKeys' do + content_type :json + keys=keyserver.generate_keys + {status: "success", keys_gen: keys}.to_json +end +#E2 +get '/key' do + content_type :json + available_key = keyserver.get_key + if available_key == nil + halt 404, "No keys available." + end + {key: available_key}.to_json +end +#E3 +put '/releasekey/:id' do + content_type :json + key = params["id"] + halt 400 if keyserver.unblock_key(key)==false + {status: "Successfully unblocked the key."}.to_json +end +#E4 +delete '/deletekey/:id' do + content_type :json + key = params["id"] + halt 400 if keyserver.delete_key(key)==false + {status: "successfully deleted"}.to_json +end +#E5 +put '/keepalive/:id' do + content_type :json + key = params["id"] + halt 400 if keyserver.keep_alive(key)==false + {status: "Key Refreshed."}.to_json +end diff --git a/config.ru b/config.ru new file mode 100644 index 0000000..19c9bc5 --- /dev/null +++ b/config.ru @@ -0,0 +1,6 @@ +require 'rubygems' +require 'bundler' +Bundler.require +require './app.rb' + +run Sinatra::Application \ No newline at end of file diff --git a/key_server.rb b/key_server.rb new file mode 100644 index 0000000..6bfb033 --- /dev/null +++ b/key_server.rb @@ -0,0 +1,112 @@ +require 'securerandom' + +class KeyServer + attr_reader :keys, :used + + def initialize + @keys = Hash.new + @used = Hash.new + @lock = Mutex.new + end + + def generate_keys + count = 5 + while count>0 do + key = SecureRandom.hex.to_sym + @lock.synchronize { + if @keys[key] == nil + @keys[key] = {created_at: Time.now} + end + } + count-=1 + end + @keys + end + + def get_key + key, val = nil, nil + @lock.synchronize { + key, val = @keys.first + if key!=nil + @used[key.to_sym] = {assigned_at: Time.now} + @keys.delete(key.to_sym) + end + } + return key + end + + def unblock_key(key) + val = nil + @lock.synchronize{ + val = @used[key.to_sym] + } + if val==nil + return false + else + + @lock.synchronize{ + @used.delete(key.to_sym) + @keys[key.to_sym] = {created_at: Time.now} + } + end + return true + end + + def delete_key(key) + val = nil + @lock.synchronize{ + val = @used[key.to_sym] + } + if val==nil + return false + else + @lock.synchronize{ + @used.delete(key.to_sym) + } + end + return true + end + + def keep_alive(key) + val_used = nil + val_free = nil + + @lock.synchronize{ + val_used = @used[key.to_sym] + val_free = @keys[key.to_sym] + } + + if val_free == nil + if val_used==nil + return false + else + @lock.synchronize{ + @used[key.to_sym] = {assigned_at: Time.now} + } + end + else + @lock.synchronize{ + @keys[key.to_sym] = {created_at: Time.now} + } + end + true + end + + def free_up + @used.each do |key, time| + @lock.synchronize{ + if Time.now - @used[key][:assigned_at]>60 + @used.delete(key) + @keys[key] = {created_at: Time.now} + end + } + end + @keys.each do |key, creation_time| + @lock.synchronize{ + if Time.now - (@keys[key][:created_at]) > 300 + @keys.delete(key) + end + } + end + end +end diff --git a/spec/keyserver_spec.rb b/spec/keyserver_spec.rb new file mode 100644 index 0000000..836522f --- /dev/null +++ b/spec/keyserver_spec.rb @@ -0,0 +1,170 @@ +require 'spec_helper' +require 'time' + +describe KeyServer do + let (:keyserver) {KeyServer.new} + + describe '#initialize' do + it 'initializes all the hashes used for storing the keys' do + expect(keyserver.keys).to be_empty + expect(keyserver.used).to be_empty + end + end + + describe '#generate_keys' do + it 'generates 5 unique keys' do + keyserver.generate_keys + expect(keyserver.keys.size).to eq(5) + end + end + + describe '#get_key' do + it 'returns a key which is available and not in use' do + keyserver.generate_keys + key = keyserver.get_key + expect(keyserver.keys).not_to have_key(key) + expect(keyserver.used).to have_key(key) + end + + it 'returns nil if no key is available' do + key = keyserver.get_key + expect(key).to eq(nil) + end + end + + describe '#unblock_key' do + context 'when the key is currently used' do + it 'unblocks the key and returns true' do + keyserver.generate_keys + key = keyserver.get_key + res = keyserver.unblock_key(key) + expect(res).to be true + end + + it 'removes the key from the @used pool' do + keyserver.generate_keys + key = keyserver.get_key + expect { + keyserver.unblock_key(key) + }.to change { keyserver.used.size }.by(-1) + + expect(keyserver.used).not_to have_key(key) + end + + it 'adds the key back to the @keys pool' do + keyserver.generate_keys + key = keyserver.get_key + expect { + keyserver.unblock_key(key) + }.to change { keyserver.keys.size }.by(1) + + expect(keyserver.keys).to have_key(key) + end + end + + context 'when the key is not found' do + it 'returns false if it did not get the key to unblock' do + key = "not_there".to_sym + result = keyserver.unblock_key(key) + expect(result).to be false + end + + it 'does not change the size of the @keys or @used pools' do + key = "not_there".to_sym + expect { + keyserver.unblock_key(key) + }.to change { keyserver.keys.size }.by(0) + + expect { + keyserver.unblock_key(key) + }.to change { keyserver.used.size }.by(0) + end + end + end + + describe '#keep_alive' do + let(:epsilon) { 0.1 } + + it 'updates the assigned_at timestamp for the key in @used' do + keyserver.generate_keys + key = keyserver.get_key + sleep(0.001) + start_time = Time.now + + expect(keyserver.keep_alive(key)).to be true + + expect(keyserver.used[key][:assigned_at]).to be_within(epsilon).of(start_time) + end + + it 'returns false if the key is not in use or free' do + key = "not_there".to_sym + expect(keyserver.keep_alive(key)).to be false + end + end + + describe '#delete_key' do + context 'when the key is currently used' do + it 'returns true and removes the key from @used' do + keyserver.generate_keys + key = keyserver.get_key + expect(keyserver.delete_key(key)).to be true + + expect(keyserver.used).not_to have_key(key) + expect(keyserver.keys).not_to have_key(key) + end + + it 'decreases the @used pool size by 1' do + keyserver.generate_keys + key = keyserver.get_key + expect { + keyserver.delete_key(key) + }.to change { keyserver.used.size }.by(-1) + end + end + + context 'when the key is not found in @used' do + it 'returns false and does not change pool sizes' do + key = "not_there".to_sym + expect(keyserver.delete_key(key)).to be false + expect { + keyserver.delete_key(key) + }.to change { keyserver.used.size }.by(0) + end + end + end + describe '#free_up' do + let(:epsilon) { 0.1 } + + before do + old_time_used = Time.now - 65 + old_time_free = Time.now - 305 + + keyserver.instance_variable_set(:@used, { + :old_used_key => { assigned_at: old_time_used } + }) + + keyserver.instance_variable_set(:@keys, { + :expired_free_key => { created_at: old_time_free } + }) + end + + it 'moves a key older than 60 seconds from @used back to @keys' do + key_to_free = :old_used_key + + expect(keyserver.used).to have_key(key_to_free) + keyserver.free_up + + expect(keyserver.used).not_to have_key(key_to_free) + expect(keyserver.keys).to have_key(key_to_free) + end + + it 'deletes a key older than 5 minutes from the @keys pool' do + key_to_delete = :expired_free_key + + expect(keyserver.keys).to have_key(key_to_delete) + keyserver.free_up + + expect(keyserver.keys).not_to have_key(key_to_delete) + end + end +end \ No newline at end of file diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb new file mode 100644 index 0000000..cb876c5 --- /dev/null +++ b/spec/spec_helper.rb @@ -0,0 +1,2 @@ +require_relative '../app.rb' +require_relative '../key_server.rb' \ No newline at end of file From 1420d36d677f95a5dcd0afd0e77a808c9c97bff6 Mon Sep 17 00:00:00 2001 From: Aryan Kumar Date: Fri, 10 Oct 2025 11:41:26 +0530 Subject: [PATCH 2/7] Added Home page message --- app.rb | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/app.rb b/app.rb index 5288909..df03747 100644 --- a/app.rb +++ b/app.rb @@ -9,12 +9,7 @@ end get '/' do - 'Hi, Welcome to my app. Try the following APIs: - /generatekeys - /key - /realeaseKey/:id - delete /keys/:id - /keepalive/:id' + 'Hi, Welcome to my app.' end #E1 post '/generateKeys' do From 9d2795a87d39299f71d5cbe430cb191f9077607f Mon Sep 17 00:00:00 2001 From: Aryan Kumar Date: Fri, 10 Oct 2025 12:31:18 +0530 Subject: [PATCH 3/7] added run commands in readme --- README.md | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 627bfd2..b3e6641 100644 --- a/README.md +++ b/README.md @@ -15,4 +15,11 @@ You should build this without full fledged Rails framework. #### Technical Requirements: Ruby verison: 2.6.6 -rspec version: 3.10.x \ No newline at end of file +rspec version: 3.10.x + +To run: +bundle install +bundle exec rackup + +To test: +bundle exec rspec \ No newline at end of file From 1f9957e895b9462d8005834215f9eccf7ef00f1b Mon Sep 17 00:00:00 2001 From: Aryan Kumar Date: Mon, 27 Oct 2025 13:05:21 +0530 Subject: [PATCH 4/7] fixed the race conditions and made the cleanup code consistent --- Gemfile | 2 + Gemfile.lock | 8 ++- key_server.rb | 184 +++++++++++++++++++++++--------------------------- 3 files changed, 93 insertions(+), 101 deletions(-) diff --git a/Gemfile b/Gemfile index 2eaf5a6..df42260 100644 --- a/Gemfile +++ b/Gemfile @@ -3,6 +3,8 @@ source "https://rubygems.org" gem "json" gem "securerandom" gem "sinatra" +gem "pqueue" +gem "puma" group :development, :test do gem 'rspec' diff --git a/Gemfile.lock b/Gemfile.lock index c8432e5..63e57d5 100644 --- a/Gemfile.lock +++ b/Gemfile.lock @@ -6,6 +6,10 @@ GEM json (2.7.6) mustermann (3.0.4) ruby2_keywords (~> 0.0.1) + nio4r (2.7.4) + pqueue (2.1.0) + puma (7.1.0) + nio4r (~> 2.0) rack (2.2.19) rack-protection (3.2.0) base64 (>= 0.1.0) @@ -39,10 +43,12 @@ PLATFORMS DEPENDENCIES json + pqueue + puma rack-test rspec securerandom sinatra BUNDLED WITH - 1.17.2 + 2.7.2 diff --git a/key_server.rb b/key_server.rb index 6bfb033..7b5fcc6 100644 --- a/key_server.rb +++ b/key_server.rb @@ -1,112 +1,96 @@ require 'securerandom' +require 'thread' class KeyServer - attr_reader :keys, :used + attr_reader :keys, :used + def initialize + @keys = {} + @used = {} + @lock = Mutex.new + end - def initialize - @keys = Hash.new - @used = Hash.new - @lock = Mutex.new - end + def generate_keys + @lock.synchronize { + count = 5 + while count > 0 do + key = SecureRandom.hex.to_sym + if @keys[key].nil? + @keys[key] = {created_at: Time.now} + count -= 1 + end + end + } + @keys.keys + end - def generate_keys - count = 5 - while count>0 do - key = SecureRandom.hex.to_sym - @lock.synchronize { - if @keys[key] == nil - @keys[key] = {created_at: Time.now} - end - } - count-=1 - end - @keys - end - - def get_key - key, val = nil, nil - @lock.synchronize { - key, val = @keys.first - if key!=nil - @used[key.to_sym] = {assigned_at: Time.now} - @keys.delete(key.to_sym) - end - } - return key - end - - def unblock_key(key) - val = nil - @lock.synchronize{ - val = @used[key.to_sym] - } - if val==nil - return false - else + def get_key + key = nil + @lock.synchronize { + key, = @keys.first + if key + key = key.to_sym + @used[key] = {assigned_at: Time.now} + @keys.delete(key) + end + } + return key + end - @lock.synchronize{ - @used.delete(key.to_sym) - @keys[key.to_sym] = {created_at: Time.now} - } - end - return true - end + def unblock_key(key) + key = key.to_sym + @lock.synchronize { + if @used.key?(key) + @used.delete(key) + @keys[key] = {created_at: Time.now} + return true + end + return false + } + end - def delete_key(key) - val = nil - @lock.synchronize{ - val = @used[key.to_sym] - } - if val==nil - return false - else - @lock.synchronize{ - @used.delete(key.to_sym) - } - end - return true - end + def delete_key(key) + key = key.to_sym + @lock.synchronize { + if @used.key?(key) + @used.delete(key) + return true + end + if @keys.key?(key) + @keys.delete(key) + return true + end - def keep_alive(key) - val_used = nil - val_free = nil - - @lock.synchronize{ - val_used = @used[key.to_sym] - val_free = @keys[key.to_sym] - } - - if val_free == nil - if val_used==nil - return false - else - @lock.synchronize{ - @used[key.to_sym] = {assigned_at: Time.now} - } - end - else - @lock.synchronize{ - @keys[key.to_sym] = {created_at: Time.now} - } - end - true - end + return false + } + end - def free_up + def keep_alive(key) + key = key.to_sym + @lock.synchronize { + if @used.key?(key) + @used[key][:assigned_at] = Time.now + return true + elsif @keys.key?(key) + @keys[key][:created_at] = Time.now + return true + else + return false + end + } + end + def free_up + @lock.synchronize{ @used.each do |key, time| - @lock.synchronize{ - if Time.now - @used[key][:assigned_at]>60 - @used.delete(key) - @keys[key] = {created_at: Time.now} - end - } + if Time.now - @used[key][:assigned_at]>60 + @used.delete(key) + @keys[key] = {created_at: Time.now} + end end - @keys.each do |key, creation_time| - @lock.synchronize{ - if Time.now - (@keys[key][:created_at]) > 300 - @keys.delete(key) - end - } + @keys.each do |key, creation_time| + if Time.now - (@keys[key][:created_at]) > 300 + @keys.delete(key) end - end -end + end + } + end +end \ No newline at end of file From f578f73e44a1aae4e2072fdd5bb870525c165fa6 Mon Sep 17 00:00:00 2001 From: Aryan Kumar Date: Wed, 29 Oct 2025 13:50:38 +0530 Subject: [PATCH 5/7] changed the datastore to a redis based storage ensuring all operations to be less than O(n) --- Gemfile | 2 +- Gemfile.lock | 6 ++ app.rb | 15 ++-- constants.rb | 7 ++ key_server.rb | 169 ++++++++++++++++++++--------------- spec/keyserver_spec.rb | 198 ++++++++++++++--------------------------- spec/spec_helper.rb | 3 +- 7 files changed, 189 insertions(+), 211 deletions(-) create mode 100644 constants.rb diff --git a/Gemfile b/Gemfile index df42260..72310a1 100644 --- a/Gemfile +++ b/Gemfile @@ -5,7 +5,7 @@ gem "securerandom" gem "sinatra" gem "pqueue" gem "puma" - +gem "redis" group :development, :test do gem 'rspec' gem 'rack-test' diff --git a/Gemfile.lock b/Gemfile.lock index 63e57d5..8017b9a 100644 --- a/Gemfile.lock +++ b/Gemfile.lock @@ -2,6 +2,7 @@ GEM remote: https://rubygems.org/ specs: base64 (0.3.0) + connection_pool (2.5.4) diff-lcs (1.6.2) json (2.7.6) mustermann (3.0.4) @@ -16,6 +17,10 @@ GEM rack (~> 2.2, >= 2.2.4) rack-test (2.2.0) rack (>= 1.3) + redis (5.4.1) + redis-client (>= 0.22.0) + redis-client (0.26.1) + connection_pool rspec (3.13.1) rspec-core (~> 3.13.0) rspec-expectations (~> 3.13.0) @@ -46,6 +51,7 @@ DEPENDENCIES pqueue puma rack-test + redis rspec securerandom sinatra diff --git a/app.rb b/app.rb index df03747..f61ccc7 100644 --- a/app.rb +++ b/app.rb @@ -1,10 +1,15 @@ require './key_server.rb' require 'sinatra' -keyserver = KeyServer.new -Thread.new do - while true - keyserver.free_up +require 'json' +require 'redis' + +redis = Redis.new +redis.flushall +keyserver = KeyServer.new(redis) +Thread.new do + loop do sleep 1 + keyserver.free_up end end @@ -14,7 +19,7 @@ #E1 post '/generateKeys' do content_type :json - keys=keyserver.generate_keys + keys=keyserver.generate {status: "success", keys_gen: keys}.to_json end #E2 diff --git a/constants.rb b/constants.rb new file mode 100644 index 0000000..f085b66 --- /dev/null +++ b/constants.rb @@ -0,0 +1,7 @@ +module KeyServerConstants + T_USED = 60 + T_FREE = 300 + UNBLOCKED_SET = 'UNBLOCKED' + BLOCKED_ZSET = 'BLOCKED' # score = expire_at (float) + FREE_ZSET = 'FREE' # score = expire_at (float) +end diff --git a/key_server.rb b/key_server.rb index 7b5fcc6..37756f0 100644 --- a/key_server.rb +++ b/key_server.rb @@ -1,96 +1,117 @@ require 'securerandom' -require 'thread' +require 'redis' +require 'time' +require_relative 'constants' + class KeyServer - attr_reader :keys, :used - def initialize - @keys = {} - @used = {} - @lock = Mutex.new + include KeyServerConstants + + def initialize(redis) + @redis = redis end - def generate_keys - @lock.synchronize { - count = 5 - while count > 0 do - key = SecureRandom.hex.to_sym - if @keys[key].nil? - @keys[key] = {created_at: Time.now} - count -= 1 - end - end - } - @keys.keys + def generate + key = SecureRandom.hex + return "Key Generation failed" if key.nil? + + now = Time.now.to_f + @redis.multi do |r| + r.setex(key, T_FREE, Time.now.to_s) + r.sadd(UNBLOCKED_SET, key) + r.zadd(FREE_ZSET, now + T_FREE, key) + r.setex(key, T_FREE, Time.now.to_s) + end + key end def get_key - key = nil - @lock.synchronize { - key, = @keys.first - if key - key = key.to_sym - @used[key] = {assigned_at: Time.now} - @keys.delete(key) - end - } - return key + key = @redis.spop(UNBLOCKED_SET) + return nil if key.nil? + + now = Time.now.to_f + @redis.multi do |r| + r.setex(key, T_USED, Time.now.to_s) + r.zadd(BLOCKED_ZSET, now + T_USED, key) + r.zrem(FREE_ZSET, key) + r.setex(key, T_USED, Time.now.to_s) + end + key end def unblock_key(key) - key = key.to_sym - @lock.synchronize { - if @used.key?(key) - @used.delete(key) - @keys[key] = {created_at: Time.now} - return true - end - return false - } + if @redis.exists(key)==1 && !@redis.sismember(UNBLOCKED_SET, key) + now = Time.now.to_f + @redis.multi do |r| + r.zrem(BLOCKED_ZSET, key) + r.sadd(UNBLOCKED_SET, key) + r.zadd(FREE_ZSET, now + T_FREE, key) + r.setex(key, T_FREE, Time.now.to_s) + end + return true + else + return false + end + end def delete_key(key) - key = key.to_sym - @lock.synchronize { - if @used.key?(key) - @used.delete(key) - return true - end - if @keys.key?(key) - @keys.delete(key) - return true - end - - return false - } + if @redis.exists(key)==1 + if @redis.sismember(UNBLOCKED_SET, key) + @redis.multi do |r| + r.srem(UNBLOCKED_SET, key) + r.zrem(FREE_ZSET, key) + r.del(key) + end + else + @redis.multi do |r| + r.zrem(BLOCKED_ZSET, key) + r.del(key) + end + end + return true + else + return false + end end def keep_alive(key) - key = key.to_sym - @lock.synchronize { - if @used.key?(key) - @used[key][:assigned_at] = Time.now - return true - elsif @keys.key?(key) - @keys[key][:created_at] = Time.now - return true - else - return false - end - } + return false unless @redis.exists(key)==1 + + now = Time.now.to_f + if @redis.zscore(BLOCKED_ZSET, key) + @redis.multi do |r| + r.zadd(BLOCKED_ZSET, now + T_USED, key) + r.setex(key, T_USED, Time.now.to_s) + end + else + @redis.multi do |r| + r.zadd(FREE_ZSET, now + T_FREE, key) + r.setex(key, T_FREE, Time.now.to_s) + end + end + return true end + def free_up - @lock.synchronize{ - @used.each do |key, time| - if Time.now - @used[key][:assigned_at]>60 - @used.delete(key) - @keys[key] = {created_at: Time.now} - end - end - @keys.each do |key, creation_time| - if Time.now - (@keys[key][:created_at]) > 300 - @keys.delete(key) - end + now = Time.now.to_f + expired_blocked = @redis.zrangebyscore(BLOCKED_ZSET, '-inf', now) + expired_blocked.each do |k| + @redis.multi do |r| + r.zrem(BLOCKED_ZSET, k) + r.sadd(UNBLOCKED_SET, k) + r.zadd(FREE_ZSET, now + T_FREE, k) + r.setex(k, T_FREE, Time.now.to_s) + end + end + expired_free = @redis.zrangebyscore(FREE_ZSET, '-inf', now) + expired_free.each do |k| + @redis.multi do |r| + r.zrem(FREE_ZSET, k) + r.srem(UNBLOCKED_SET, k) + r.del(k) end - } + end + nil end end \ No newline at end of file diff --git a/spec/keyserver_spec.rb b/spec/keyserver_spec.rb index 836522f..fc22e86 100644 --- a/spec/keyserver_spec.rb +++ b/spec/keyserver_spec.rb @@ -1,170 +1,108 @@ require 'spec_helper' -require 'time' +require 'redis' describe KeyServer do - let (:keyserver) {KeyServer.new} + let(:redis) { Redis.new } + let(:keyserver) { KeyServer.new(redis) } - describe '#initialize' do - it 'initializes all the hashes used for storing the keys' do - expect(keyserver.keys).to be_empty - expect(keyserver.used).to be_empty - end + before(:each) do + redis.flushall end - describe '#generate_keys' do - it 'generates 5 unique keys' do - keyserver.generate_keys - expect(keyserver.keys.size).to eq(5) + describe '#generate' do + it 'generates a key and adds it to UNBLOCKED_SET and FREE_ZSET' do + key = keyserver.generate + expect(redis.sismember(KeyServer::UNBLOCKED_SET, key)).to be true + expect(redis.zscore(KeyServer::FREE_ZSET, key)).not_to be_nil + expect(redis.exists(key)).to eq 1 end end describe '#get_key' do - it 'returns a key which is available and not in use' do - keyserver.generate_keys - key = keyserver.get_key - expect(keyserver.keys).not_to have_key(key) - expect(keyserver.used).to have_key(key) + it 'returns a key from UNBLOCKED_SET and moves it to BLOCKED_ZSET' do + key = keyserver.generate + got_key = keyserver.get_key + expect(got_key).to eq(key) + expect(redis.sismember(KeyServer::UNBLOCKED_SET, key)).to be false + expect(redis.zscore(KeyServer::BLOCKED_ZSET, key)).not_to be_nil end it 'returns nil if no key is available' do - key = keyserver.get_key - expect(key).to eq(nil) + expect(keyserver.get_key).to be_nil end end describe '#unblock_key' do - context 'when the key is currently used' do - it 'unblocks the key and returns true' do - keyserver.generate_keys - key = keyserver.get_key - res = keyserver.unblock_key(key) - expect(res).to be true - end + it 'moves a blocked key back to UNBLOCKED_SET and FREE_ZSET' do + key = keyserver.generate + keyserver.get_key + expect(redis.zscore(KeyServer::BLOCKED_ZSET, key)).not_to be_nil + keyserver.unblock_key(key) + expect(redis.sismember(KeyServer::UNBLOCKED_SET, key)).to be true + expect(redis.zscore(KeyServer::FREE_ZSET, key)).not_to be_nil + expect(redis.zscore(KeyServer::BLOCKED_ZSET, key)).to be_nil + end - it 'removes the key from the @used pool' do - keyserver.generate_keys - key = keyserver.get_key - expect { - keyserver.unblock_key(key) - }.to change { keyserver.used.size }.by(-1) - - expect(keyserver.used).not_to have_key(key) - end - - it 'adds the key back to the @keys pool' do - keyserver.generate_keys - key = keyserver.get_key - expect { - keyserver.unblock_key(key) - }.to change { keyserver.keys.size }.by(1) - - expect(keyserver.keys).to have_key(key) - end + it 'returns false if key is not blocked' do + expect(keyserver.unblock_key('not_there')).to be false end + end - context 'when the key is not found' do - it 'returns false if it did not get the key to unblock' do - key = "not_there".to_sym - result = keyserver.unblock_key(key) - expect(result).to be false - end + describe '#delete_key' do + it 'removes key from all sets and deletes it' do + key = keyserver.generate + keyserver.get_key + expect(keyserver.delete_key(key)).to be true + expect(redis.exists(key)).to eq 0 + expect(redis.sismember(KeyServer::UNBLOCKED_SET, key)).to be false + expect(redis.zscore(KeyServer::FREE_ZSET, key)).to be_nil + expect(redis.zscore(KeyServer::BLOCKED_ZSET, key)).to be_nil + end - it 'does not change the size of the @keys or @used pools' do - key = "not_there".to_sym - expect { - keyserver.unblock_key(key) - }.to change { keyserver.keys.size }.by(0) - - expect { - keyserver.unblock_key(key) - }.to change { keyserver.used.size }.by(0) - end + it 'returns false if key does not exist' do + expect(keyserver.delete_key('not_there')).to be false end end describe '#keep_alive' do - let(:epsilon) { 0.1 } - - it 'updates the assigned_at timestamp for the key in @used' do - keyserver.generate_keys - key = keyserver.get_key - sleep(0.001) - start_time = Time.now - + it 'extends TTL for blocked key' do + key = keyserver.generate + keyserver.get_key expect(keyserver.keep_alive(key)).to be true - - expect(keyserver.used[key][:assigned_at]).to be_within(epsilon).of(start_time) - end - - it 'returns false if the key is not in use or free' do - key = "not_there".to_sym - expect(keyserver.keep_alive(key)).to be false + expect(redis.ttl(key)).to be > 0 end - end - describe '#delete_key' do - context 'when the key is currently used' do - it 'returns true and removes the key from @used' do - keyserver.generate_keys - key = keyserver.get_key - expect(keyserver.delete_key(key)).to be true - - expect(keyserver.used).not_to have_key(key) - expect(keyserver.keys).not_to have_key(key) - end - - it 'decreases the @used pool size by 1' do - keyserver.generate_keys - key = keyserver.get_key - expect { - keyserver.delete_key(key) - }.to change { keyserver.used.size }.by(-1) - end + it 'extends TTL for free key' do + key = keyserver.generate + expect(keyserver.keep_alive(key)).to be true + expect(redis.ttl(key)).to be > 0 end - context 'when the key is not found in @used' do - it 'returns false and does not change pool sizes' do - key = "not_there".to_sym - expect(keyserver.delete_key(key)).to be false - expect { - keyserver.delete_key(key) - }.to change { keyserver.used.size }.by(0) - end + it 'returns false if key does not exist' do + expect(keyserver.keep_alive('not_there')).to be false end end - describe '#free_up' do - let(:epsilon) { 0.1 } - before do - old_time_used = Time.now - 65 - old_time_free = Time.now - 305 - - keyserver.instance_variable_set(:@used, { - :old_used_key => { assigned_at: old_time_used } - }) - - keyserver.instance_variable_set(:@keys, { - :expired_free_key => { created_at: old_time_free } - }) - end - - it 'moves a key older than 60 seconds from @used back to @keys' do - key_to_free = :old_used_key - - expect(keyserver.used).to have_key(key_to_free) + describe '#free_up' do + it 'moves expired blocked keys back to UNBLOCKED_SET and FREE_ZSET' do + key = keyserver.generate + keyserver.get_key + # Fast-forward time by manipulating the zset score + redis.zadd(KeyServer::BLOCKED_ZSET, Time.now.to_f - 100, key) keyserver.free_up - - expect(keyserver.used).not_to have_key(key_to_free) - expect(keyserver.keys).to have_key(key_to_free) + expect(redis.sismember(KeyServer::UNBLOCKED_SET, key)).to be true + expect(redis.zscore(KeyServer::BLOCKED_ZSET, key)).to be_nil + expect(redis.zscore(KeyServer::FREE_ZSET, key)).not_to be_nil end - it 'deletes a key older than 5 minutes from the @keys pool' do - key_to_delete = :expired_free_key - - expect(keyserver.keys).to have_key(key_to_delete) + it 'removes expired free keys from all sets' do + key = keyserver.generate + # Fast-forward time by manipulating the zset score + redis.zadd(KeyServer::FREE_ZSET, Time.now.to_f - 100, key) keyserver.free_up - - expect(keyserver.keys).not_to have_key(key_to_delete) + expect(redis.sismember(KeyServer::UNBLOCKED_SET, key)).to be false + expect(redis.zscore(KeyServer::FREE_ZSET, key)).to be_nil + expect(redis.exists(key)).to eq 0 end end end \ No newline at end of file diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb index cb876c5..69c8fb4 100644 --- a/spec/spec_helper.rb +++ b/spec/spec_helper.rb @@ -1,2 +1,3 @@ require_relative '../app.rb' -require_relative '../key_server.rb' \ No newline at end of file +require_relative '../key_server.rb' +require_relative '../constants.rb' \ No newline at end of file From 2f3acbd5ad155e5269b460a4a682a627ca8364f1 Mon Sep 17 00:00:00 2001 From: Aryan Kumar Date: Thu, 30 Oct 2025 19:28:13 +0530 Subject: [PATCH 6/7] remved all loops, using multi for atomicity and removing expired keys using zremrangebyscore --- app.rb | 16 ++--- key_server.rb | 152 ++++++++++++++++------------------------- spec/keyserver_spec.rb | 111 +++++++++++++----------------- 3 files changed, 114 insertions(+), 165 deletions(-) diff --git a/app.rb b/app.rb index f61ccc7..0094118 100644 --- a/app.rb +++ b/app.rb @@ -5,13 +5,13 @@ redis = Redis.new redis.flushall -keyserver = KeyServer.new(redis) -Thread.new do - loop do - sleep 1 - keyserver.free_up - end -end +keyserver = KeyServer.new +# Thread.new do +# loop do +# sleep 1 +# keyserver.free_up +# end +# end get '/' do 'Hi, Welcome to my app.' @@ -19,7 +19,7 @@ #E1 post '/generateKeys' do content_type :json - keys=keyserver.generate + keys=keyserver.generate(5) {status: "success", keys_gen: keys}.to_json end #E2 diff --git a/key_server.rb b/key_server.rb index 37756f0..a37d495 100644 --- a/key_server.rb +++ b/key_server.rb @@ -1,117 +1,79 @@ -require 'securerandom' require 'redis' -require 'time' -require_relative 'constants' - +require 'securerandom' +require 'json' class KeyServer - include KeyServerConstants + AVAILABLE_SET = "available_keys" + KEEPALIVE_TTL = 300 + BLOCK_TTL = 60 - def initialize(redis) - @redis = redis + def initialize + @redis = Redis.new end - def generate - key = SecureRandom.hex - return "Key Generation failed" if key.nil? - - now = Time.now.to_f - @redis.multi do |r| - r.setex(key, T_FREE, Time.now.to_s) - r.sadd(UNBLOCKED_SET, key) - r.zadd(FREE_ZSET, now + T_FREE, key) - r.setex(key, T_FREE, Time.now.to_s) + def generate(n = 1) + created = [] + now = Time.now.to_i + + n.times do + id = SecureRandom.hex + @redis.multi do + expire_at = now + KEEPALIVE_TTL + @redis.setex("key:#{id}", KEEPALIVE_TTL, id) + @redis.zadd(AVAILABLE_SET, expire_at, id) + end + created << id end - key + + { created: created, count: created.size } end + def get_key - key = @redis.spop(UNBLOCKED_SET) - return nil if key.nil? - - now = Time.now.to_f - @redis.multi do |r| - r.setex(key, T_USED, Time.now.to_s) - r.zadd(BLOCKED_ZSET, now + T_USED, key) - r.zrem(FREE_ZSET, key) - r.setex(key, T_USED, Time.now.to_s) + now = Time.now.to_i + @redis.zremrangebyscore(AVAILABLE_SET, 0, now) + id, _score = @redis.zpopmin(AVAILABLE_SET) + return nil unless id + @redis.multi do + @redis.zrem(AVAILABLE_SET, id) + @redis.setex("blocked:#{id}", BLOCK_TTL, "1") end - key + { id: id, assigned_at: Time.now.to_i } end - def unblock_key(key) - if @redis.exists(key)==1 && !@redis.sismember(UNBLOCKED_SET, key) - now = Time.now.to_f - @redis.multi do |r| - r.zrem(BLOCKED_ZSET, key) - r.sadd(UNBLOCKED_SET, key) - r.zadd(FREE_ZSET, now + T_FREE, key) - r.setex(key, T_FREE, Time.now.to_s) - end - return true - else - return false - end - + def unblock_key(id) + return false unless @redis.exists?("blocked:#{id}") + expire_at = Time.now.to_i + KEEPALIVE_TTL + @redis.multi do + @redis.del("blocked:#{id}") + @redis.zadd(AVAILABLE_SET, expire_at, id) + @redis.expire("key:#{id}", KEEPALIVE_TTL) + end + { message: "Key #{id} unblocked" } end - def delete_key(key) - if @redis.exists(key)==1 - if @redis.sismember(UNBLOCKED_SET, key) - @redis.multi do |r| - r.srem(UNBLOCKED_SET, key) - r.zrem(FREE_ZSET, key) - r.del(key) - end - else - @redis.multi do |r| - r.zrem(BLOCKED_ZSET, key) - r.del(key) - end - end - return true - else - return false + def delete_key(id) + return false unless @redis.exists?("key:#{id}") + @redis.multi do + @redis.del("key:#{id}") + @redis.del("blocked:#{id}") + @redis.zrem(AVAILABLE_SET, id) end + { message: "Key #{id} deleted" } end - def keep_alive(key) - return false unless @redis.exists(key)==1 - - now = Time.now.to_f - if @redis.zscore(BLOCKED_ZSET, key) - @redis.multi do |r| - r.zadd(BLOCKED_ZSET, now + T_USED, key) - r.setex(key, T_USED, Time.now.to_s) - end - else - @redis.multi do |r| - r.zadd(FREE_ZSET, now + T_FREE, key) - r.setex(key, T_FREE, Time.now.to_s) - end + def keep_alive(id) + return false unless @redis.exists?("key:#{id}") + @redis.multi do + expire_at = Time.now.to_i + KEEPALIVE_TTL + @redis.zadd(AVAILABLE_SET, expire_at, id) + @redis.expire("key:#{id}", @redis.ttl("key:#{id}")+KEEPALIVE_TTL) end - return true + { message: "Keepalive refreshed for #{id}" } end - def free_up - now = Time.now.to_f - expired_blocked = @redis.zrangebyscore(BLOCKED_ZSET, '-inf', now) - expired_blocked.each do |k| - @redis.multi do |r| - r.zrem(BLOCKED_ZSET, k) - r.sadd(UNBLOCKED_SET, k) - r.zadd(FREE_ZSET, now + T_FREE, k) - r.setex(k, T_FREE, Time.now.to_s) - end - end - expired_free = @redis.zrangebyscore(FREE_ZSET, '-inf', now) - expired_free.each do |k| - @redis.multi do |r| - r.zrem(FREE_ZSET, k) - r.srem(UNBLOCKED_SET, k) - r.del(k) - end - end - nil + def ttl(id) + return nil unless @redis.exists?("key:#{id}") + @redis.ttl("key:#{id}") end -end \ No newline at end of file +end diff --git a/spec/keyserver_spec.rb b/spec/keyserver_spec.rb index fc22e86..41968f6 100644 --- a/spec/keyserver_spec.rb +++ b/spec/keyserver_spec.rb @@ -3,28 +3,32 @@ describe KeyServer do let(:redis) { Redis.new } - let(:keyserver) { KeyServer.new(redis) } + let(:keyserver) { KeyServer.new } before(:each) do redis.flushall end describe '#generate' do - it 'generates a key and adds it to UNBLOCKED_SET and FREE_ZSET' do - key = keyserver.generate - expect(redis.sismember(KeyServer::UNBLOCKED_SET, key)).to be true - expect(redis.zscore(KeyServer::FREE_ZSET, key)).not_to be_nil - expect(redis.exists(key)).to eq 1 + it 'creates keys and adds them to AVAILABLE_SET with TTL' do + result = keyserver.generate(5) + result[:created].each do |key| + expect(redis.exists("key:#{key}")).to eq(1) + expect(redis.zscore(KeyServer::AVAILABLE_SET, key)).not_to be_nil + end + expect(result[:count]).to eq(5) end end describe '#get_key' do - it 'returns a key from UNBLOCKED_SET and moves it to BLOCKED_ZSET' do - key = keyserver.generate - got_key = keyserver.get_key - expect(got_key).to eq(key) - expect(redis.sismember(KeyServer::UNBLOCKED_SET, key)).to be false - expect(redis.zscore(KeyServer::BLOCKED_ZSET, key)).not_to be_nil + it 'returns a key from AVAILABLE_SET and blocks it' do + keyserver.generate(5) + assigned = keyserver.get_key + key = assigned[:id] + + expect(assigned[:id]).not_to be_nil + expect(redis.zscore(KeyServer::AVAILABLE_SET, key)).to be_nil + expect(redis.exists("blocked:#{key}")).to eq(1) end it 'returns nil if no key is available' do @@ -33,30 +37,32 @@ end describe '#unblock_key' do - it 'moves a blocked key back to UNBLOCKED_SET and FREE_ZSET' do - key = keyserver.generate - keyserver.get_key - expect(redis.zscore(KeyServer::BLOCKED_ZSET, key)).not_to be_nil - keyserver.unblock_key(key) - expect(redis.sismember(KeyServer::UNBLOCKED_SET, key)).to be true - expect(redis.zscore(KeyServer::FREE_ZSET, key)).not_to be_nil - expect(redis.zscore(KeyServer::BLOCKED_ZSET, key)).to be_nil + it 'moves a blocked key back to AVAILABLE_SET and refreshes TTL' do + keyserver.generate(5) + key = keyserver.get_key[:id] + + expect(redis.exists("blocked:#{key}")).to eq(1) + + result = keyserver.unblock_key(key) + expect(result[:message]).to match(/unblocked/) + expect(redis.exists("blocked:#{key}")).to eq(0) + expect(redis.zscore(KeyServer::AVAILABLE_SET, key)).not_to be_nil + expect(redis.ttl("key:#{key}")).to be > 0 end - it 'returns false if key is not blocked' do + it 'returns false if key does not exist' do expect(keyserver.unblock_key('not_there')).to be false end end describe '#delete_key' do - it 'removes key from all sets and deletes it' do - key = keyserver.generate + it 'removes key from AVAILABLE_SET and deletes it' do + key = keyserver.generate(5)[:created].first keyserver.get_key - expect(keyserver.delete_key(key)).to be true - expect(redis.exists(key)).to eq 0 - expect(redis.sismember(KeyServer::UNBLOCKED_SET, key)).to be false - expect(redis.zscore(KeyServer::FREE_ZSET, key)).to be_nil - expect(redis.zscore(KeyServer::BLOCKED_ZSET, key)).to be_nil + expect(keyserver.delete_key(key)[:message]).to match(/deleted/) + expect(redis.exists("key:#{key}")).to eq(0) + expect(redis.exists("blocked:#{key}")).to eq(0) + expect(redis.zscore(KeyServer::AVAILABLE_SET, key)).to be_nil end it 'returns false if key does not exist' do @@ -64,45 +70,26 @@ end end - describe '#keep_alive' do - it 'extends TTL for blocked key' do - key = keyserver.generate - keyserver.get_key - expect(keyserver.keep_alive(key)).to be true - expect(redis.ttl(key)).to be > 0 + describe '#keepalive' do + it 'refreshes TTL for available key' do + key = keyserver.generate[:created].first + ttl_before = redis.ttl("key:#{key}") + keyserver.keep_alive(key) + ttl_after = redis.ttl("key:#{key}") + expect(ttl_after).to be > ttl_before end - it 'extends TTL for free key' do - key = keyserver.generate - expect(keyserver.keep_alive(key)).to be true - expect(redis.ttl(key)).to be > 0 + it 'refreshes TTL for blocked key' do + key = keyserver.generate[:created].first + keyserver.get_key + ttl_before = redis.ttl("key:#{key}") + keyserver.keep_alive(key) + ttl_after = redis.ttl("key:#{key}") + expect(ttl_after).to be > ttl_before end it 'returns false if key does not exist' do expect(keyserver.keep_alive('not_there')).to be false end end - - describe '#free_up' do - it 'moves expired blocked keys back to UNBLOCKED_SET and FREE_ZSET' do - key = keyserver.generate - keyserver.get_key - # Fast-forward time by manipulating the zset score - redis.zadd(KeyServer::BLOCKED_ZSET, Time.now.to_f - 100, key) - keyserver.free_up - expect(redis.sismember(KeyServer::UNBLOCKED_SET, key)).to be true - expect(redis.zscore(KeyServer::BLOCKED_ZSET, key)).to be_nil - expect(redis.zscore(KeyServer::FREE_ZSET, key)).not_to be_nil - end - - it 'removes expired free keys from all sets' do - key = keyserver.generate - # Fast-forward time by manipulating the zset score - redis.zadd(KeyServer::FREE_ZSET, Time.now.to_f - 100, key) - keyserver.free_up - expect(redis.sismember(KeyServer::UNBLOCKED_SET, key)).to be false - expect(redis.zscore(KeyServer::FREE_ZSET, key)).to be_nil - expect(redis.exists(key)).to eq 0 - end - end -end \ No newline at end of file +end From 003d4ee66197eecb38b6d2d095f3bb1713ee2943 Mon Sep 17 00:00:00 2001 From: Aryan Kumar Date: Fri, 31 Oct 2025 10:39:21 +0530 Subject: [PATCH 7/7] added LUA scripts to ensure no race condition occurs --- key_server.rb | 142 ++++++++++++++++++++++++++++++++++++-------------- 1 file changed, 102 insertions(+), 40 deletions(-) diff --git a/key_server.rb b/key_server.rb index a37d495..bfe59c9 100644 --- a/key_server.rb +++ b/key_server.rb @@ -5,7 +5,7 @@ class KeyServer AVAILABLE_SET = "available_keys" KEEPALIVE_TTL = 300 - BLOCK_TTL = 60 + BLOCK_TTL = 60 def initialize @redis = Redis.new @@ -13,67 +13,129 @@ def initialize def generate(n = 1) created = [] - now = Time.now.to_i - + n.times do id = SecureRandom.hex - @redis.multi do - expire_at = now + KEEPALIVE_TTL - @redis.setex("key:#{id}", KEEPALIVE_TTL, id) - @redis.zadd(AVAILABLE_SET, expire_at, id) - end + lua_script = <<-LUA + local key_id = KEYS[1] + local ttl = ARGV[1] + local available_set = ARGV[2] + local expire_at = tonumber(redis.call('TIME')[1]) + tonumber(ttl) + + redis.call('SETEX', 'key:' .. key_id, ttl, key_id) + redis.call('ZADD', available_set, expire_at, key_id) + + return 1 + LUA + + @redis.eval(lua_script, keys: [id], argv: [KEEPALIVE_TTL, AVAILABLE_SET]) created << id end - + { created: created, count: created.size } end - def get_key - now = Time.now.to_i - @redis.zremrangebyscore(AVAILABLE_SET, 0, now) - id, _score = @redis.zpopmin(AVAILABLE_SET) - return nil unless id - @redis.multi do - @redis.zrem(AVAILABLE_SET, id) - @redis.setex("blocked:#{id}", BLOCK_TTL, "1") - end - { id: id, assigned_at: Time.now.to_i } + lua_script = <<-LUA + local now = tonumber(ARGV[1]) + local block_ttl = ARGV[2] + local available_set = ARGV[3] + + redis.call('ZREMRANGEBYSCORE', available_set, 0, now) + + local result = redis.call('ZPOPMIN', available_set, 1) + + if result and result[1] then + local id = result[1] + + redis.call('SETEX', 'blocked:' .. id, block_ttl, "1") + + return {id, now} + else + return nil + end + LUA + + result = @redis.eval(lua_script, keys: [], argv: [Time.now.to_i, BLOCK_TTL, AVAILABLE_SET]) + + return nil unless result && result[0] + + { id: result[0], assigned_at: result[1].to_i } end def unblock_key(id) - return false unless @redis.exists?("blocked:#{id}") - expire_at = Time.now.to_i + KEEPALIVE_TTL - @redis.multi do - @redis.del("blocked:#{id}") - @redis.zadd(AVAILABLE_SET, expire_at, id) - @redis.expire("key:#{id}", KEEPALIVE_TTL) - end + lua_script = <<-LUA + local id = KEYS[1] + local ttl = ARGV[1] + local available_set = ARGV[2] + + if redis.call('EXISTS', 'blocked:' .. id) == 0 then + return 0 + end + + local expire_at = tonumber(redis.call('TIME')[1]) + tonumber(ttl) + + redis.call('DEL', 'blocked:' .. id) + redis.call('ZADD', available_set, expire_at, id) + redis.call('EXPIRE', 'key:' .. id, ttl) + + return 1 + LUA + + result = @redis.eval(lua_script, keys: [id], argv: [KEEPALIVE_TTL, AVAILABLE_SET]) + + return false if result == 0 { message: "Key #{id} unblocked" } end def delete_key(id) - return false unless @redis.exists?("key:#{id}") - @redis.multi do - @redis.del("key:#{id}") - @redis.del("blocked:#{id}") - @redis.zrem(AVAILABLE_SET, id) - end + lua_script = <<-LUA + local id = KEYS[1] + local available_set = ARGV[1] + + if redis.call('EXISTS', 'key:' .. id) == 0 then + return 0 + end + + redis.call('DEL', 'key:' .. id) + redis.call('DEL', 'blocked:' .. id) + redis.call('ZREM', available_set, id) + + return 1 + LUA + + result = @redis.eval(lua_script, keys: [id], argv: [AVAILABLE_SET]) + + return false if result == 0 { message: "Key #{id} deleted" } end def keep_alive(id) - return false unless @redis.exists?("key:#{id}") - @redis.multi do - expire_at = Time.now.to_i + KEEPALIVE_TTL - @redis.zadd(AVAILABLE_SET, expire_at, id) - @redis.expire("key:#{id}", @redis.ttl("key:#{id}")+KEEPALIVE_TTL) - end + lua_script = <<-LUA + local id = KEYS[1] + local ttl = ARGV[1] + local available_set = ARGV[2] + + local current_ttl = redis.call('TTL', 'key:' .. id) + if current_ttl < 0 then + return 0 + end + + local expire_at = tonumber(redis.call('TIME')[1]) + tonumber(ttl) + + redis.call('ZADD', available_set, expire_at, id) + redis.call('EXPIRE', 'key:' .. id, current_ttl + tonumber(ttl)) + + return 1 + LUA + + result = @redis.eval(lua_script, keys: [id], argv: [KEEPALIVE_TTL, AVAILABLE_SET]) + + return false if result == 0 { message: "Keepalive refreshed for #{id}" } end def ttl(id) - return nil unless @redis.exists?("key:#{id}") @redis.ttl("key:#{id}") end -end +end \ No newline at end of file