Merge pull request #15 from LegionIO/fleet/ws-00g-cache-set-nx

Fleet WS-00G: add Legion::Cache.set_nx

Author: Esity

lib/legion/cache.rb

@@ -190,6 +190,19 @@ def set(key, value, ttl: nil, async: true, phi: false)
         end
       end
 
+      def set_nx(key, value, ttl: nil)
+        effective_ttl = resolve_ttl(ttl)
+        return Legion::Cache::Memory.set_nx(key, value, ttl: effective_ttl) if using_memory?
+        return Legion::Cache::Local.set_nx(key, value, ttl: effective_ttl) if using_local?
+        return Legion::Cache::Local.set_nx(key, value, ttl: effective_ttl) if failback_to_local?
+
+        configure_shared_adapter!
+        super
+      rescue StandardError => e
+        handle_exception(e, level: :warn, handled: true, operation: :cache_set_nx, key: key)
+        false
+      end
+
       def set_sync(key, value, ttl: nil, **)
         return Legion::Cache::Memory.set_sync(key, value, ttl: ttl) if using_memory?
         return Legion::Cache::Local.set_sync(key, value, ttl: ttl) if using_local?

lib/legion/cache/memcached.rb

@@ -76,6 +76,16 @@ def fetch(key, ttl: nil, &)
         nil
       end
 
+      def set_nx(key, value, ttl: nil)
+        effective_ttl = ttl || default_ttl
+        result = client.with { |conn| conn.add(key, value, effective_ttl) == true }
+        log.debug { "[cache] SET_NX #{key} ttl=#{effective_ttl.inspect} result=#{result}" }
+        result
+      rescue StandardError => e
+        handle_exception(e, level: :error, handled: false, operation: :memcached_set_nx, key: key, ttl: effective_ttl)
+        raise
+      end
+
       def set(key, value, ttl: nil, **)
         set_sync(key, value, ttl: ttl, **)
       end

lib/legion/cache/memory.rb

@@ -55,6 +55,25 @@ def set(key, value, ttl: nil, async: true, phi: false) # rubocop:disable Lint/Un
         set_sync(key, value, ttl: ttl, phi: phi)
       end
 
+      def set_nx(key, value, ttl: nil)
+        @mutex.synchronize do
+          expire_if_needed(key)
+          return false if @store.key?(key)
+
+          @store[key] = value
+          if ttl&.positive?
+            @expiry[key] = Time.now + ttl
+          else
+            @expiry.delete(key)
+          end
+          log.debug { "[cache:memory] SET_NX #{key} ttl=#{ttl.inspect} result=true" }
+          true
+        end
+      rescue StandardError => e
+        handle_exception(e, level: :warn, handled: true, operation: :memory_set_nx)
+        false
+      end
+
       def set_sync(key, value, ttl: nil, phi: false)
         ttl = enforce_phi_ttl(ttl, phi: phi) if phi
         @mutex.synchronize do

lib/legion/cache/redis.rb

@@ -100,6 +100,17 @@ def fetch(key, ttl: nil)
         result
       end
 
+      def set_nx(key, value, ttl: nil)
+        effective_ttl = ttl || default_ttl
+        serialized = serialize_value(value)
+        result = client.with { |conn| conn.set(key, serialized, nx: true, ex: effective_ttl) == 'OK' }
+        log.debug { "[cache] SET_NX #{key} ttl=#{effective_ttl.inspect} result=#{result}" }
+        result
+      rescue StandardError => e
+        handle_exception(e, level: :error, handled: false, operation: :redis_set_nx, key: key, ttl: effective_ttl)
+        raise
+      end
+
       def set(key, value, ttl: nil, **)
         set_sync(key, value, ttl: ttl, **)
       end

lib/legion/cache/version.rb

@@ -2,6 +2,6 @@
 
 module Legion
   module Cache
-    VERSION = '1.4.1'
+    VERSION = '1.4.2'
   end
 end

spec/legion/cache/helper_spec.rb

@@ -459,7 +459,11 @@ def cache_default_ttl
 
       it 'serializes hash as JSON via cache set (merge)' do
         allow(Legion::Cache).to receive(:get).with('microsoft_teams:h').and_return(nil)
-        expect(Legion::Cache).to receive(:set).with('microsoft_teams:h', '{"f":"v"}', ttl: 3600, async: false)
+        expect(Legion::Cache).to receive(:set) do |key, json, **opts|
+          expect(key).to eq('microsoft_teams:h')
+          expect(opts).to eq(ttl: 3600, async: false)
+          expect(Legion::JSON.load(json)).to eq(f: 'v')
+        end
         subject.cache_hset(':h', { 'f' => 'v' })
       end
 

spec/legion/cache/set_nx_spec.rb

@@ -0,0 +1,116 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+require 'legion/cache/redis'
+require 'legion/cache/memcached'
+require 'legion/cache/memory'
+
+RSpec.describe 'Legion::Cache set_nx' do
+  describe Legion::Cache::Memory do
+    before { described_class.reset! }
+
+    describe '.set_nx' do
+      it 'returns true and stores value when key does not exist' do
+        result = described_class.set_nx('nx-key', 'value', ttl: 60)
+        expect(result).to be true
+        expect(described_class.get('nx-key')).to eq('value')
+      end
+
+      it 'returns false and does not overwrite when key already exists' do
+        described_class.set('nx-key', 'original', ttl: 60)
+        result = described_class.set_nx('nx-key', 'overwrite', ttl: 60)
+        expect(result).to be false
+        expect(described_class.get('nx-key')).to eq('original')
+      end
+
+      it 'returns true after an expired key has been purged' do
+        described_class.set('nx-expire', 'old', ttl: 0.05)
+        sleep 0.07
+        result = described_class.set_nx('nx-expire', 'new', ttl: 60)
+        expect(result).to be true
+        expect(described_class.get('nx-expire')).to eq('new')
+      end
+
+      it 'is atomic under concurrent access' do
+        winners = []
+        mutex = Mutex.new
+        threads = 10.times.map do |i|
+          Thread.new do
+            won = described_class.set_nx('race-key', "value-#{i}", ttl: 60)
+            mutex.synchronize { winners << i } if won
+          end
+        end
+        threads.each(&:join)
+        expect(winners.size).to eq(1)
+      end
+    end
+  end
+
+  describe Legion::Cache::Redis do
+    let(:cache) { described_class.dup }
+    let(:pool)  { instance_double(ConnectionPool) }
+    let(:redis) { instance_double(Redis) }
+
+    before do
+      cache.instance_variable_set(:@client, pool)
+      cache.instance_variable_set(:@connected, true)
+      allow(pool).to receive(:with).and_yield(redis)
+    end
+
+    describe '#set_nx' do
+      it 'returns true when Redis SET NX succeeds (returns "OK")' do
+        allow(redis).to receive(:set).with('nx-key', anything, nx: true, ex: 60).and_return('OK')
+        expect(cache.set_nx('nx-key', 'value', ttl: 60)).to be true
+      end
+
+      it 'returns false when Redis SET NX fails (key exists, returns nil)' do
+        allow(redis).to receive(:set).with('nx-key', anything, nx: true, ex: 60).and_return(nil)
+        expect(cache.set_nx('nx-key', 'value', ttl: 60)).to be false
+      end
+
+      it 'passes nx: true and ex: ttl to Redis SET' do
+        expect(redis).to receive(:set).with('nx-key', anything, nx: true, ex: 120).and_return('OK')
+        cache.set_nx('nx-key', 'value', ttl: 120)
+      end
+
+      it 'serializes the value before storing' do
+        captured = nil
+        allow(redis).to receive(:set) do |_key, val, **_opts|
+          captured = val
+          'OK'
+        end
+        cache.set_nx('nx-key', { data: 42 }, ttl: 60)
+        expect(captured).to be_a(String)
+      end
+    end
+  end
+
+  describe Legion::Cache::Memcached do
+    let(:cache) { described_class.dup }
+    let(:pool)  { instance_double(ConnectionPool) }
+    let(:dalli) { instance_double(Dalli::Client) }
+
+    before do
+      cache.instance_variable_set(:@client, pool)
+      cache.instance_variable_set(:@connected, true)
+      allow(pool).to receive(:with).and_yield(dalli)
+    end
+
+    describe '#set_nx' do
+      it 'returns true when Dalli#add succeeds (key did not exist)' do
+        allow(dalli).to receive(:add).with('nx-key', 'value', 60).and_return(true)
+        expect(cache.set_nx('nx-key', 'value', ttl: 60)).to be true
+      end
+
+      it 'returns false when Dalli#add fails (key already exists, returns nil/false)' do
+        allow(dalli).to receive(:add).with('nx-key', 'value', 60).and_return(nil)
+        expect(cache.set_nx('nx-key', 'value', ttl: 60)).to be false
+      end
+
+      it 'passes the ttl positionally to Dalli#add' do
+        expect(dalli).to receive(:add).with('nx-key', 'value', 90).and_return(true)
+        cache.set_nx('nx-key', 'value', ttl: 90)
+      end
+    end
+  end
+end

spec/legion/cache_interface_spec.rb

@@ -84,4 +84,14 @@
   it 'responds to enabled?' do
     expect(Legion::Cache).to respond_to(:enabled?)
   end
+
+  it 'has set_nx method' do
+    expect(Legion::Cache.method(:set_nx)).to be_a(Method)
+  end
+
+  it 'set_nx accepts keyword ttl' do
+    params = Legion::Cache.method(:set_nx).parameters
+    names = params.map(&:last)
+    expect(names).to include(:ttl)
+  end
 end