From d765a5cab4177c467d913273e06429bc4934af61 Mon Sep 17 00:00:00 2001
From: Matteo Rossi <hello@matteorossi.dev>
Date: Wed, 29 Jan 2025 17:02:13 +0100
Subject: [PATCH] refactor: outer lock with locked!

---
 lib/rack/idempotency_key.rb              |  2 +-
 lib/rack/idempotency_key/memory_store.rb | 14 +++++++++-----
 lib/rack/idempotency_key/redis_store.rb  | 10 ++++++++--
 lib/rack/idempotency_key/request.rb      | 16 ++++++++++++++++
 4 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/lib/rack/idempotency_key.rb b/lib/rack/idempotency_key.rb
index 59a1a04..2049828 100644
--- a/lib/rack/idempotency_key.rb
+++ b/lib/rack/idempotency_key.rb
@@ -36,7 +36,7 @@ def handle_request!(request, env)
         cached_response = request.cached_response!
         return cached_response unless cached_response.nil?
 
-        app.call(env).tap { |response| request.cache!(response) }
+        request.locked! { app.call(env).tap { |response| request.cache!(response) } }
       end
   end
 end
diff --git a/lib/rack/idempotency_key/memory_store.rb b/lib/rack/idempotency_key/memory_store.rb
index 3857f5b..0198fd2 100644
--- a/lib/rack/idempotency_key/memory_store.rb
+++ b/lib/rack/idempotency_key/memory_store.rb
@@ -16,7 +16,7 @@ def get(key)
           value = store[key]
           return if value.nil?
 
-          if expired?(value[:added_at])
+          if expired?(value[:expires_at])
             store.delete(key)
             return
           end
@@ -25,21 +25,25 @@ def get(key)
         end
       end
 
-      def set(key, value)
+      def set(key, value, ttl: expires_in)
         mutex.synchronize do
           raise Rack::IdempotencyKey::ConflictError if store.key?(key)
 
-          store[key] ||= { value: value, added_at: Time.now.utc }
+          store[key] ||= { value: value, expires_at: Time.now.utc + ttl }
           store[key][:value]
         end
       end
 
+      def unset(key)
+        mutex.synchronize { store.delete(key) }
+      end
+
       private
 
         attr_reader :store, :expires_in, :mutex
 
-        def expired?(added_at)
-          Time.now.utc - added_at > expires_in
+        def expired?(expires_at)
+          Time.now.utc > expires_at
         end
     end
   end
diff --git a/lib/rack/idempotency_key/redis_store.rb b/lib/rack/idempotency_key/redis_store.rb
index 9f8a7bf..3eb1de2 100644
--- a/lib/rack/idempotency_key/redis_store.rb
+++ b/lib/rack/idempotency_key/redis_store.rb
@@ -20,9 +20,9 @@ def get(key)
         raise Rack::IdempotencyKey::StoreError, "#{self.class}: #{e.message}"
       end
 
-      def set(key, value)
+      def set(key, value, ttl: expires_in)
         with_redis do |redis|
-          result = redis.set(namespaced_key(key), value, nx: true, ex: expires_in)
+          result = redis.set(namespaced_key(key), value, nx: true, ex: ttl)
           raise Rack::IdempotencyKey::ConflictError unless result
         end
 
@@ -31,6 +31,12 @@ def set(key, value)
         raise Rack::IdempotencyKey::StoreError, "#{self.class}: #{e.message}"
       end
 
+      def unset(key)
+        with_redis { |redis| redis.del(key) }
+      rescue Redis::BaseError => e
+        raise Rack::IdempotencyKey::StoreError, "#{self.class}: #{e.message}"
+      end
+
       private
 
         attr_reader :store, :expires_in
diff --git a/lib/rack/idempotency_key/request.rb b/lib/rack/idempotency_key/request.rb
index c0e2899..2187bcc 100644
--- a/lib/rack/idempotency_key/request.rb
+++ b/lib/rack/idempotency_key/request.rb
@@ -3,6 +3,8 @@
 module Rack
   class IdempotencyKey
     class Request
+      DEFAULT_LOCK_TTL = 60 # seconds
+
       # @param request [Rack::Request]
       # @param store   [Store]
       def initialize(request, store)
@@ -23,6 +25,16 @@ def cached_response!
         end
       end
 
+      def locked!
+        store.set(lock_key, 1, ttl: DEFAULT_LOCK_TTL)
+
+        begin
+          yield
+        ensure
+          store.unset(lock_key)
+        end
+      end
+
       def cache!(response)
         status, = response
         store.set(cache_key, response) if status != 400
@@ -56,6 +68,10 @@ def cache_key
       private
 
         attr_reader :request, :store
+
+        def lock_key
+          "#{idempotency_key}_lock"
+        end
     end
   end
 end