Merge pull request #5030 from DataDog/no-ticket-extract-lru-cache-into-core-utils

[NO-TICKET] Extract AppSec LRUCache implementation into Core::Utils

Author: Strech

lib/datadog/appsec/api_security/sampler.rb

@@ -1,9 +1,9 @@
 # frozen_string_literal: true
 
 require 'zlib'
-require_relative 'lru_cache'
 require_relative 'route_extractor'
 require_relative '../../core/utils/time'
+require_relative '../../core/utils/lru_cache'
 
 module Datadog
   module AppSec
@@ -37,7 +37,7 @@ def sample_delay
         def initialize(sample_delay)
           raise ArgumentError, 'sample_delay must be an Integer' unless sample_delay.is_a?(Integer)
 
-          @cache = LRUCache.new(MAX_CACHE_SIZE)
+          @cache = Core::Utils::LRUCache.new(MAX_CACHE_SIZE)
           @sample_delay_seconds = sample_delay
         end
 
@@ -51,7 +51,7 @@ def sample?(request, response)
 
           return false if current_timestamp - cached_timestamp <= @sample_delay_seconds
 
-          @cache.store(key, current_timestamp)
+          @cache[key] = current_timestamp
           true
         end
       end

lib/datadog/appsec/api_security/lru_cache.rb renamed to lib/datadog/core/utils/lru_cache.rb

@@ -3,13 +3,13 @@
 require 'forwardable'
 
 module Datadog
-  module AppSec
-    module APISecurity
+  module Core
+    module Utils
       # An LRU (Least Recently Used) cache implementation that relies on the
       # Ruby 1.9+ `Hash` implementation that guarantees insertion order.
       #
       # WARNING: This implementation is NOT thread-safe and should be used
-      #          in a single-threaded context.
+      #          in a single-threaded context or guarded by Mutex.
       class LRUCache
         extend Forwardable
 
@@ -30,25 +30,14 @@ def [](key)
           end
         end
 
-        def store(key, value)
-          return @store[key] = value if @store.delete(key)
-
-          # NOTE: evict the oldest entry if store reached the maximum allowed size
-          @store.shift if @store.size >= @max_size
-          @store[key] = value
-        end
-
-        # NOTE: If the key exists, it's moved to the end of the list and
-        #       if does not, the given block will be executed and the result
-        #       will be stored (which will add it to the end of the list).
-        def fetch_or_store(key)
-          if (entry = @store.delete(key))
-            return @store[key] = entry
+        def []=(key, value)
+          if @store.delete(key)
+            @store[key] = value
+          else
+            # NOTE: evict the oldest entry if store reached the maximum allowed size
+            @store.shift if @store.size >= @max_size
+            @store[key] = value
           end
-
-          # NOTE: evict the oldest entry if store reached the maximum allowed size
-          @store.shift if @store.size >= @max_size
-          @store[key] = yield
         end
       end
     end

sig/datadog/appsec/api_security/lru_cache.rbs

@@ -6,7 +6,7 @@ module Datadog
 
         MAX_CACHE_SIZE: Integer
 
-        @cache: LRUCache
+        @cache: Core::Utils::LRUCache
 
         @sample_delay_seconds: Integer
 

sig/datadog/appsec/api_security/sampler.rbs

@@ -0,0 +1,23 @@
+module Datadog
+  module Core
+    module Utils
+      class LRUCache
+        extend Forwardable
+
+        @store: ::Hash[untyped, untyped]
+
+        @max_size: ::Integer
+
+        def initialize: (::Integer max_size) -> void
+
+        def []: (untyped key) -> untyped?
+
+        def []=: (untyped key, untyped value) -> untyped
+
+        def clear: () -> void
+
+        def empty?: () -> bool
+      end
+    end
+  end
+end

sig/datadog/core/utils/lru_cache.rbs

@@ -1,9 +1,9 @@
 # frozen_string_literal: true
 
 require 'spec_helper'
-require 'datadog/appsec/api_security/lru_cache'
+require 'datadog/core/utils/lru_cache'
 
-RSpec.describe Datadog::AppSec::APISecurity::LRUCache do
+RSpec.describe Datadog::Core::Utils::LRUCache do
   describe '#initialize' do
     it { expect(described_class.new(3)).to be_empty }
     it { expect { described_class.new('0') }.to raise_error(ArgumentError) }
@@ -19,19 +19,19 @@
     end
 
     context 'when key exists' do
-      before { lru_cache.fetch_or_store(:key) { 'value' } }
+      before { lru_cache[:key] = 'value' }
 
       it { expect(lru_cache[:key]).to eq('value') }
     end
 
     context 'when key exists and is accessed' do
       it 'updates the key order' do
-        lru_cache.fetch_or_store(:key1) { 'value1' }
-        lru_cache.fetch_or_store(:key2) { 'value2' }
-        lru_cache.fetch_or_store(:key3) { 'value3' }
+        lru_cache[:key1] = 'value1'
+        lru_cache[:key2] = 'value2'
+        lru_cache[:key3] = 'value3'
 
         lru_cache[:key1] # NOTE: as key accessed, it's moved to the end of the list
-        lru_cache.fetch_or_store(:key4) { 'value4' }
+        lru_cache[:key4] = 'value4'
 
         aggregate_failures 'verifying LRU cache state after key access and eviction' do
           expect(lru_cache[:key2]).to be_nil
@@ -43,30 +43,30 @@
     end
   end
 
-  describe '#store' do
+  describe '#[]=' do
     let(:lru_cache) { described_class.new(3) }
 
     it 'stores a key-value pair' do
-      expect { lru_cache.store(:key, 'value') }.to change { lru_cache[:key] }
+      expect { lru_cache[:key] = 'value' }.to change { lru_cache[:key] }
         .from(nil).to('value')
     end
 
     it 'overwrites existing values' do
-      lru_cache.store(:key, 'old-value')
+      lru_cache[:key] = 'old-value'
 
-      expect { lru_cache.store(:key, 'new-value') }.to change { lru_cache[:key] }
+      expect { lru_cache[:key] = 'new-value' }.to change { lru_cache[:key] }
         .from('old-value').to('new-value')
     end
 
     context 'when cache is full and an existing key is updated' do
       it 'does not remove other entries' do
-        lru_cache.store(:key2, 'value2')
-        lru_cache.store(:key3, 'value3')
-        lru_cache.store(:key1, 'value1')
+        lru_cache[:key2] = 'value2'
+        lru_cache[:key3] = 'value3'
+        lru_cache[:key1] = 'value1'
 
-        lru_cache.store(:key1, 'new-value1')
-        lru_cache.store(:key3, 'new-value3')
-        lru_cache.store(:key2, 'new-value2')
+        lru_cache[:key1] = 'new-value1'
+        lru_cache[:key3] = 'new-value3'
+        lru_cache[:key2] = 'new-value2'
 
         aggregate_failures 'verifying LRU cache contents' do
           expect(lru_cache[:key1]).to eq('new-value1')
@@ -78,47 +78,10 @@
 
     context 'when maximum size is reached' do
       it 'evicts the least recently used item' do
-        lru_cache.store(:key1, 'value1')
-        lru_cache.store(:key2, 'value2')
-        lru_cache.store(:key3, 'value3')
-        lru_cache.store(:key4, 'value4')
-
-        aggregate_failures 'verifying LRU cache state after eviction' do
-          expect(lru_cache[:key1]).to be_nil
-          expect(lru_cache[:key2]).to eq('value2')
-          expect(lru_cache[:key3]).to eq('value3')
-          expect(lru_cache[:key4]).to eq('value4')
-        end
-      end
-    end
-  end
-
-  describe '#fetch_or_store' do
-    context 'when key does not exist' do
-      let(:lru_cache) { described_class.new(3) }
-
-      it 'computes and stores the value' do
-        expect(lru_cache.fetch_or_store(:key) { 'value' }).to eq('value')
-        expect(lru_cache[:key]).to eq('value')
-      end
-
-      it 'computes the missing key only once' do
-        expect { lru_cache.fetch_or_store(:key) { 'value' } }
-          .to change { lru_cache[:key] }.from(nil).to('value')
-
-        expect { lru_cache.fetch_or_store(:key) { 'new-value' } }
-          .not_to(change { lru_cache[:key] })
-      end
-    end
-
-    context 'when maximum size is reached' do
-      let(:lru_cache) { described_class.new(3) }
-
-      it 'evicts the least recently used item' do
-        lru_cache.fetch_or_store(:key1) { 'value1' }
-        lru_cache.fetch_or_store(:key2) { 'value2' }
-        lru_cache.fetch_or_store(:key3) { 'value3' }
-        lru_cache.fetch_or_store(:key4) { 'value4' }
+        lru_cache[:key1] = 'value1'
+        lru_cache[:key2] = 'value2'
+        lru_cache[:key3] = 'value3'
+        lru_cache[:key4] = 'value4'
 
         aggregate_failures 'verifying LRU cache state after eviction' do
           expect(lru_cache[:key1]).to be_nil
@@ -134,8 +97,8 @@
     let(:lru_cache) { described_class.new(3) }
 
     it 'removes all items from the cache' do
-      lru_cache.fetch_or_store(:key1) { 'value1' }
-      lru_cache.fetch_or_store(:key2) { 'value2' }
+      lru_cache[:key1] = 'value1'
+      lru_cache[:key2] = 'value2'
 
       expect { lru_cache.clear }.to change { lru_cache[:key1] }.from('value1').to(nil)
         .and change { lru_cache[:key2] }.from('value2').to(nil)
@@ -147,7 +110,7 @@
     let(:lru_cache) { described_class.new(3) }
 
     it 'returns false when cache has items' do
-      expect { lru_cache.fetch_or_store(:key) { 'value' } }
+      expect { lru_cache[:key] = 'value' }
         .to change { lru_cache.empty? }.from(true).to(false)
     end
   end