Skip to content

Commit

Permalink
state needs to be implemented too
Browse files Browse the repository at this point in the history
  • Loading branch information
julik committed Feb 11, 2024
1 parent a26bb73 commit 769cf14
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 2 deletions.
24 changes: 22 additions & 2 deletions lib/pecorino/cached_throttle.rb
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ def request(n = 1)
def able_to_accept?(n = 1)
blocked_state = cached_blocked_state
return false if blocked_state&.blocked?

@throttle.able_to_accept?(n)
end

Expand All @@ -60,13 +61,32 @@ def throttled(&blk)
yield
end

# Returns the key of the throttle
#
# @see Pecorino::Throttle#key
def key
@throttle.key
end

# Returns `false` if there is a currently active block for that throttle in the cache. Otherwise forwards to underlying throttle.
#
# @see Pecorino::Throttle#able_to_accept?
def state
blocked_state = cached_blocked_state
return blocked_state if blocked_state&.blocked?

@throttle.state.tap do |state|
cache_blocked_state(state) if state.blocked?
end
end

private

def cache_blocked_state(state)
@cache_store.write("pcrn-cached-throttle-#{@throttle.key}", state, expires_after: state.blocked_until)
@cache_store.write("pecorino-cached-throttle-state-#{@throttle.key}", state, expires_after: state.blocked_until)
end

def cached_blocked_state
@cache_store.read("pcrn-cached-throttle-#{@throttle.key}")
@cache_store.read("pecorino-cached-throttle-state-#{@throttle.key}")
end
end
24 changes: 24 additions & 0 deletions test/cached_throttle_test.rb
Original file line number Diff line number Diff line change
Expand Up @@ -94,4 +94,28 @@ class << throttle

assert_nil cached_throttle.throttled { 345 }
end

test "caches results of state() and correctly returns cached state until the block is lifted" do
store = ActiveSupport::Cache::MemoryStore.new
throttle = Pecorino::Throttle.new(key: Random.uuid, capacity: 2, over_time: 1.second, block_for: 2.seconds)
cached_throttle = Pecorino::CachedThrottle.new(store, throttle)

cached_throttle.request(2)
blocked_state = cached_throttle.request(1)
blocked_state_from_cache = cached_throttle.state

assert_kind_of Pecorino::Throttle::State, blocked_state
assert_kind_of Pecorino::Throttle::State, blocked_state_from_cache
assert_predicate blocked_state, :blocked?
assert_predicate blocked_state_from_cache, :blocked?

# Delete the method on the actual throttle as it should not be called anymore until the block is lifted
class << throttle
undef :state
end

state_from_cache = cached_throttle.state
assert_kind_of Pecorino::Throttle::State, state_from_cache
assert_predicate state_from_cache, :blocked?
end
end

0 comments on commit 769cf14

Please sign in to comment.