feat(rate-limiter): add Prism and RateLimiter
This commit is contained in:
1 parent
9b333d5046
commit
e75dc076e5
2 files changed
+119
No files matched your search
@@ -0,0 +1,57 @@
|
||||
require 'monitor'
|
||||
|
||||
module Prism
|
||||
module RateLimiter
|
||||
# In-process fixed-window counter + ban list, guarded by a Monitor so
|
||||
# concurrent requests (threaded server, parallel test threads) do not
|
||||
# race on the same key. Lost on process restart - use PStoreStore when
|
||||
# limits need to survive that.
|
||||
class MemoryStore
|
||||
include MonitorMixin
|
||||
|
||||
Window = Struct.new(:started_at, :count)
|
||||
|
||||
def initialize
|
||||
super()
|
||||
@windows = {}
|
||||
@bans = {}
|
||||
end
|
||||
|
||||
def count(key, period)
|
||||
synchronize do
|
||||
now = Time.now
|
||||
window = @windows[key]
|
||||
if window.nil? || now - window.started_at >= period
|
||||
window = Window.new(now, 0)
|
||||
@windows[key] = window
|
||||
end
|
||||
window.count += 1
|
||||
end
|
||||
end
|
||||
|
||||
def banned?(key)
|
||||
synchronize do
|
||||
expires_at = @bans[key]
|
||||
next false unless expires_at
|
||||
if Time.now >= expires_at
|
||||
@bans.delete(key)
|
||||
false
|
||||
else
|
||||
true
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
def ban!(key, duration)
|
||||
synchronize { @bans[key] = Time.now + duration }
|
||||
end
|
||||
|
||||
def reset!
|
||||
synchronize do
|
||||
@windows.clear
|
||||
@bans.clear
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
@@ -0,0 +1,62 @@
|
||||
require 'pstore'
|
||||
require 'fileutils'
|
||||
|
||||
module Prism
|
||||
module RateLimiter
|
||||
# PStore-backed fixed-window counter + ban list, same contract as
|
||||
# MemoryStore, but durable: limits and bans survive a process
|
||||
# restart, and multiple processes sharing one file share state
|
||||
# (each PStore#transaction takes an exclusive flock on the file).
|
||||
class PStoreStore
|
||||
def initialize(path)
|
||||
@path = path
|
||||
directory = File.dirname(path)
|
||||
FileUtils.mkdir_p(directory) unless directory == '.' || Dir.exist?(directory)
|
||||
@pstore = PStore.new(path)
|
||||
@pstore.transaction do
|
||||
@pstore[:windows] ||= {}
|
||||
@pstore[:bans] ||= {}
|
||||
end
|
||||
end
|
||||
|
||||
def count(key, period)
|
||||
@pstore.transaction do
|
||||
now = Time.now
|
||||
started_at, current = @pstore[:windows][key]
|
||||
if started_at.nil? || now - started_at >= period
|
||||
started_at = now
|
||||
current = 0
|
||||
end
|
||||
current += 1
|
||||
@pstore[:windows][key] = [started_at, current]
|
||||
current
|
||||
end
|
||||
end
|
||||
|
||||
def banned?(key)
|
||||
@pstore.transaction do
|
||||
expires_at = @pstore[:bans][key]
|
||||
if expires_at.nil?
|
||||
false
|
||||
elsif Time.now >= expires_at
|
||||
@pstore[:bans].delete(key)
|
||||
false
|
||||
else
|
||||
true
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
def ban!(key, duration)
|
||||
@pstore.transaction { @pstore[:bans][key] = Time.now + duration }
|
||||
end
|
||||
|
||||
def reset!
|
||||
@pstore.transaction do
|
||||
@pstore[:windows] = {}
|
||||
@pstore[:bans] = {}
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
Reference in new issue
Block a user