berater 0.1.2 → 0.4.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 04cc88a51db5eca83f7843c1d34cf86213a0ca412fd404e3fff88b3b6994ea8d
-  data.tar.gz: 4eb67771a03cd0f4179d8d80655efcad738cb877b2fb5813399f136dc743090c
+  metadata.gz: f57aeaa9319ae38bdd7d2468443fc5c9ea9ba44cecf47bc5a558f18ecd0ab274
+  data.tar.gz: ff14ef90856911e9e18cbfbad6f9799727c999697c5e4d96f47b5fbf665c725c
 SHA512:
-  metadata.gz: 9715ae524f7b0fdbbd60d352546ac86a0984f4a5e8966351ae00057e976c391076c270677b29f83c813043c24ad6fa7cae571f224df6400e98fa43bdb4f06cb6
-  data.tar.gz: 16999ff18ea843d12ce6b847b79ddf3bc2b4372cf01371a63d36213ccc27717e06e8794f6b4a519a59af7cba1afaed3cc3ca92a5baed60302e6f1fe114dd29ac
+  metadata.gz: c9d3abeeb9d35c608c194eaf505c5af1350a82583dd12257e70b9d4f96bd3a6cd30db802f57c3bf77d371b5deb5e856e9238eadd3990768205fd37ea3e2e3703
+  data.tar.gz: c98078a3c7d3a3eb4e2d523577f9d6dd7f1bbae76c9ecbd8d65961cc7aa1bab6e1a4b8fe84cfb30ee60f4d0ce11dcb58e3b94a3566cf50ee60d1f8e21bab82b1

data/lib/berater.rb

@@ -1,4 +1,5 @@
 require 'berater/version'
+require 'berater/lock'
 
 
 module Berater
@@ -8,41 +9,42 @@ module Berater
 
   MODES = {}
 
-  attr_accessor :redis, :mode
+  attr_accessor :redis
 
   def configure
-    self.mode = :unlimited # default
-
     yield self
   end
 
-  def new(mode, *args, **opts)
+  def new(key, mode = nil, *args, **opts, &block)
+    if mode.nil?
+      unless args.empty?
+        raise ArgumentError, '0 arguments expected with block'
+      end
+
+      unless block_given?
+        raise ArgumentError, 'expected either mode or block'
+      end
+
+      mode, *args = DSL.eval(&block)
+    else
+      if block_given?
+        raise ArgumentError, 'expected either mode or block, not both'
+      end
+    end
+
     klass = MODES[mode.to_sym]
 
     unless klass
       raise ArgumentError, "invalid mode: #{mode}"
     end
 
-    klass.new(*args, **opts)
+    klass.new(key, *args, **opts)
   end
 
   def register(mode, klass)
     MODES[mode.to_sym] = klass
   end
 
-  def mode=(mode)
-    unless MODES.include? mode.to_sym
-      raise ArgumentError, "invalid mode: #{mode}"
-    end
-
-    @mode = mode.to_sym
-  end
-
-  def limit(*args, **opts, &block)
-    mode = opts.delete(:mode) { self.mode }
-    new(mode, *args, **opts).limit(&block)
-  end
-
   def expunge
     redis.scan_each(match: "#{self.name}*") do |key|
       redis.del key
@@ -51,8 +53,13 @@ module Berater
 
 end
 
-# load and register limiters
-require 'berater/base_limiter'
+# convenience method
+def Berater(key, mode, *args, **opts, &block)
+  Berater.new(key, mode, *args, **opts).limit(&block)
+end
+
+# load limiters
+require 'berater/limiter'
 require 'berater/concurrency_limiter'
 require 'berater/inhibitor'
 require 'berater/rate_limiter'
@@ -62,3 +69,5 @@ Berater.register(:concurrency, Berater::ConcurrencyLimiter)
 Berater.register(:inhibited, Berater::Inhibitor)
 Berater.register(:rate, Berater::RateLimiter)
 Berater.register(:unlimited, Berater::Unlimiter)
+
+require 'berater/dsl'

data/lib/berater/concurrency_limiter.rb

@@ -1,18 +1,18 @@
 module Berater
-  class ConcurrencyLimiter < BaseLimiter
+  class ConcurrencyLimiter < Limiter
 
     class Incapacitated < Overloaded; end
 
     attr_reader :capacity, :timeout
 
-    def initialize(capacity, **opts)
-      super(**opts)
+    def initialize(key, capacity, **opts)
+      super(key, **opts)
 
       self.capacity = capacity
       self.timeout = opts[:timeout] || 0
     end
 
-    def capacity=(capacity)
+    private def capacity=(capacity)
       unless capacity.is_a? Integer
         raise ArgumentError, "expected Integer, found #{capacity.class}"
       end
@@ -22,7 +22,7 @@ module Berater
       @capacity = capacity
     end
 
-    def timeout=(timeout)
+    private def timeout=(timeout)
       unless timeout.is_a? Integer
         raise ArgumentError, "expected Integer, found #{timeout.class}"
       end
@@ -32,124 +32,62 @@ module Berater
       @timeout = timeout
     end
 
-    class Lock
-      attr_reader :limiter, :id, :contention
-
-      def initialize(limiter, id, contention)
-        @limiter = limiter
-        @id = id
-        @contention = contention
-        @released = false
-        @locked_at = Time.now
-      end
-
-      def release
-        raise 'lock already released' if released?
-        raise 'lock expired' if expired?
-
-        @released = limiter.release(self)
-      end
-
-      def released?
-        @released
-      end
-
-      def expired?
-        @locked_at + limiter.timeout < Time.now
-      end
-    end
-
-    LUA_SCRIPT = <<~LUA.gsub(/^\s*(--.*\n)?/, '')
+    LUA_SCRIPT = <<~LUA.gsub(/^\s*|\s*--.*/, '')
       local key = KEYS[1]
+      local lock_key = KEYS[2]
       local capacity = tonumber(ARGV[1])
       local ts = tonumber(ARGV[2])
       local ttl = tonumber(ARGV[3])
-
-      local exists
-      local count
       local lock
-      local ts = unpack(redis.call('TIME'))
 
-      -- check to see if key already exists
-      if ttl == 0 then
-        exists = redis.call('EXISTS', key)
-      else
-        -- and refresh TTL while we're at it
-        exists = redis.call('EXPIRE', key, ttl * 2)
+      -- purge stale hosts
+      if ttl > 0 then
+        redis.call('ZREMRANGEBYSCORE', key, '-inf', ts - ttl)
       end
 
-      if exists == 1 then
-        -- purge stale hosts
-        if ttl > 0 then
-          redis.call('ZREMRANGEBYSCORE', key, '-inf', ts - ttl)
-        end
-
-        -- check capacity (subtract one for next lock entry)
-        count = redis.call('ZCARD', key) - 1
-
-        if count < capacity then
-          -- yay, grab a lock
-
-          -- regenerate next lock entry, which has score inf
-          lock = unpack(redis.call('ZPOPMAX', key))
-          redis.call('ZADD', key, 'inf', (lock + 1) % 2^52)
-
-          count = count + 1
-        end
-      else
-        count = 1
-        lock = "1"
-
-        -- create structure to track locks and next id
-        redis.call('ZADD', key, 'inf', lock + 1)
+      -- check capacity
+      local count = redis.call('ZCARD', key)
 
-        if ttl > 0 then
-          redis.call('EXPIRE', key, ttl * 2)
-        end
-      end
-
-      if lock then
-        -- store lock and timestamp
+      if count < capacity then
+        -- grab a lock
+        lock = redis.call('INCR', lock_key)
         redis.call('ZADD', key, ts, lock)
+        count = count + 1
       end
 
       return { count, lock }
     LUA
 
-    def limit(**opts, &block)
-      unless opts.empty?
-        return self.class.new(
-          capacity,
-          **options.merge(opts)
-          # **options.merge(timeout: timeout).merge(opts)
-        ).limit(&block)
-      end
-
+    def limit
       count, lock_id = redis.eval(
         LUA_SCRIPT,
-        [ key ],
+        [ cache_key(key), cache_key('lock_id') ],
         [ capacity, Time.now.to_i, timeout ]
       )
 
       raise Incapacitated unless lock_id
 
-      lock = Lock.new(self, lock_id, count)
+      lock = Lock.new(self, lock_id, count, -> { release(lock_id) })
 
       if block_given?
         begin
           yield lock
         ensure
-          release(lock)
+          lock.release
         end
       else
         lock
       end
     end
 
-    def release(lock)
-      res = redis.zrem(key, lock.id)
+    private def release(lock_id)
+      res = redis.zrem(cache_key(key), lock_id)
       res == true || res == 1 # depending on which version of Redis
     end
 
+    def to_s
+      "#<#{self.class}(#{key}: #{capacity} at a time)>"
+    end
+
   end
 end

data/lib/berater/dsl.rb

@@ -0,0 +1,57 @@
+module Berater
+  module DSL
+    extend self
+
+    def eval &block
+      @keywords ||= Class.new do
+        # create a class where DSL keywords are methods
+        KEYWORDS.each do |keyword|
+          define_singleton_method(keyword) { keyword }
+        end
+      end
+
+      install
+      @keywords.class_eval &block
+    ensure
+      uninstall
+    end
+
+    private
+
+    def each &block
+      Berater::MODES.map do |mode, limiter|
+        next unless limiter.const_defined?(:DSL, false)
+        limiter.const_get(:DSL)
+      end.compact.each(&block)
+    end
+
+    KEYWORDS = [
+      :second, :minute, :hour,
+      :unlimited, :inhibited,
+    ].freeze
+
+    def install
+      Integer.class_eval do
+        def per(unit)
+          [ :rate, self, unit ]
+        end
+        alias every per
+
+        def at_once
+          [ :concurrency, self ]
+        end
+        alias concurrently at_once
+        alias at_a_time at_once
+      end
+    end
+
+    def uninstall
+      Integer.remove_method :per
+      Integer.remove_method :every
+
+      Integer.remove_method :at_once
+      Integer.remove_method :concurrently
+      Integer.remove_method :at_a_time
+    end
+  end
+end

data/lib/berater/inhibitor.rb

@@ -1,19 +1,13 @@
 module Berater
-  class Inhibitor < BaseLimiter
+  class Inhibitor < Limiter
 
     class Inhibited < Overloaded; end
 
-    def initialize(*args, **opts)
-      super(**opts)
+    def initialize(key = :inhibitor, *args, **opts)
+      super(key, **opts)
     end
 
-    def limit(**opts, &block)
-      unless opts.empty?
-        return self.class.new(
-          **options.merge(opts)
-        ).limit(&block)
-      end
-
+    def limit
       raise Inhibited
     end
 

data/lib/berater/limiter.rb

@@ -0,0 +1,30 @@
+module Berater
+  class Limiter
+
+    attr_reader :key, :options
+
+    def redis
+      options[:redis] || Berater.redis
+    end
+
+    def limit
+      raise NotImplementedError
+    end
+
+    def to_s
+      "#<#{self.class}>"
+    end
+
+    protected
+
+    def initialize(key, **opts)
+      @key = key
+      @options = opts
+    end
+
+    def cache_key(key)
+      "#{self.class}:#{key}"
+    end
+
+  end
+end

data/lib/berater/lock.rb

@@ -0,0 +1,36 @@
+module Berater
+  class Lock
+
+    attr_reader :limiter, :id, :contention
+
+    def initialize(limiter, id, contention, release_fn = nil)
+      @limiter = limiter
+      @id = id
+      @contention = contention
+      @locked_at = Time.now
+      @release_fn = release_fn
+      @released_at = nil
+    end
+
+    def locked?
+      @released_at.nil? && !expired?
+    end
+
+    def expired?
+      timeout ? @locked_at + timeout < Time.now : false
+    end
+
+    def release
+      raise 'lock expired' if expired?
+      raise 'lock already released' unless locked?
+
+      @released_at = Time.now
+      @release_fn ? @release_fn.call : true
+    end
+
+    def timeout
+      limiter.respond_to?(:timeout) ? limiter.timeout : nil
+    end
+
+  end
+end

data/lib/berater/rate_limiter.rb

@@ -1,18 +1,18 @@
 module Berater
-  class RateLimiter < BaseLimiter
+  class RateLimiter < Limiter
 
     class Overrated < Overloaded; end
 
     attr_accessor :count, :interval
 
-    def initialize(count, interval, **opts)
-      super(**opts)
+    def initialize(key, count, interval, **opts)
+      super(key, **opts)
 
       self.count = count
       self.interval = interval
     end
 
-    def count=(count)
+    private def count=(count)
       unless count.is_a? Integer
         raise ArgumentError, "expected Integer, found #{count.class}"
       end
@@ -22,12 +22,13 @@ module Berater
       @count = count
     end
 
-    def interval=(interval)
+    private def interval=(interval)
       @interval = interval.dup
 
       case @interval
       when Integer
         raise ArgumentError, "interval must be >= 0" unless @interval >= 0
+        @interval_sec = @interval
       when String
         @interval = @interval.to_sym
       when Symbol
@@ -38,47 +39,96 @@ module Berater
       if @interval.is_a? Symbol
         case @interval
         when :sec, :second, :seconds
-          @interval = 1
+          @interval = :second
+          @interval_sec = 1
         when :min, :minute, :minutes
-          @interval = 60
+          @interval = :minute
+          @interval_sec = 60
         when :hour, :hours
-          @interval = 60 * 60
+          @interval = :hour
+          @interval_sec = 60 * 60
         else
           raise ArgumentError, "unexpected interval value: #{interval}"
         end
       end
-
-      @interval
     end
 
-    def limit(**opts, &block)
-      unless opts.empty?
-        return self.class.new(
-          count,
-          interval,
-          options.merge(opts)
-        ).limit(&block)
+    LUA_SCRIPT = <<~LUA.gsub(/^\s*|\s*--.*/, '')
+      local key = KEYS[1]
+      local ts_key = KEYS[2]
+      local ts = tonumber(ARGV[1])
+      local capacity = tonumber(ARGV[2])
+      local usec_per_drip = tonumber(ARGV[3])
+      local count = 0
+
+      -- timestamp of last update
+      local last_ts = tonumber(redis.call('GET', ts_key))
+
+      if last_ts then
+        count = tonumber(redis.call('GET', key)) or 0
+
+        -- adjust for time passing
+        local drips = math.floor((ts - last_ts) / usec_per_drip)
+        count = math.max(0, count - drips)
       end
 
-      ts = Time.now.to_i
+      local allowed = count + 1 <= capacity
+
+      if allowed then
+        count = count + 1
 
-      # bucket into time slot
-      rkey = "%s:%d" % [ key, ts - ts % @interval ]
+        -- time for bucket to empty, in milliseconds
+        local ttl = math.ceil((count * usec_per_drip) / 1000)
 
-      count, _ = redis.multi do
-        redis.incr rkey
-        redis.expire rkey, @interval * 2
+        -- update count and last_ts, with expirations
+        redis.call('SET', key, count, 'PX', ttl)
+        redis.call('SET', ts_key, ts, 'PX', ttl)
       end
 
-      raise Overrated if count > @count
+      return { count, allowed }
+    LUA
+
+    def limit
+      usec_per_drip = (@interval_sec * 10**6) / @count
+
+      # timestamp in microseconds
+      ts = (Time.now.to_f * 10**6).to_i
+
+      count, allowed = redis.eval(
+        LUA_SCRIPT,
+        [ cache_key(key), cache_key("#{key}-ts") ],
+        [ ts, @count, usec_per_drip ]
+      )
+
+      raise Overrated unless allowed
+
+      lock = Lock.new(self, "#{ts}-#{count}", count)
 
       if block_given?
-        yield
+        begin
+          yield lock
+        ensure
+          lock.release
+        end
       else
-        count
+        lock
       end
     end
 
+    def to_s
+      msg = if @interval.is_a? Integer
+        if @interval == 1
+          "every second"
+        else
+          "every #{@interval} seconds"
+        end
+      else
+        "per #{@interval}"
+      end
+
+      "#<#{self.class}(#{key}: #{count} #{msg})>"
+    end
+
   end
 end