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