berater 0.1.3 → 0.5.0

data/lib/berater/lock.rb

@@ -0,0 +1,27 @@
+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?
+    end
+
+    def release
+      raise 'lock already released' unless locked?
+
+      @released_at = Time.now
+      @release_fn ? @release_fn.call : true
+    end
+
+  end
+end

data/lib/berater/lua_script.rb

@@ -0,0 +1,55 @@
+require 'digest'
+
+module Berater
+  class LuaScript
+
+    attr_reader :source
+
+    def initialize(source)
+      @source = source
+    end
+
+    def sha
+      @sha ||= Digest::SHA1.hexdigest(minify)
+    end
+
+    def eval(redis, *args)
+      redis.evalsha(sha, *args)
+    rescue Redis::CommandError => e
+      raise unless e.message.include?('NOSCRIPT')
+
+      # fall back to regular eval, which will trigger
+      # script to be cached for next time
+      redis.eval(minify, *args)
+    end
+
+    def load(redis)
+      redis.script(:load, minify).tap do |sha|
+        unless sha == self.sha
+          raise "unexpected script SHA: expected #{self.sha}, got #{sha}"
+        end
+      end
+    end
+
+    def loaded?(redis)
+      redis.script(:exists, sha)
+    end
+
+    def to_s
+      source
+    end
+
+    private
+
+    def minify
+      # trim comments (whole line and partial)
+      # and whitespace (prefix and empty lines)
+      @minify ||= source.gsub(/^\s*--.*\n|\s*--.*|^\s*|^$\n/, '').chomp
+    end
+
+  end
+
+  def LuaScript(source)
+    LuaScript.new(source)
+  end
+end

data/lib/berater/rate_limiter.rb

@@ -1,82 +1,101 @@
 module Berater
-  class RateLimiter < BaseLimiter
+  class RateLimiter < Limiter
 
     class Overrated < Overloaded; end
 
-    attr_accessor :count, :interval
+    attr_accessor :interval
 
-    def initialize(count, interval, **opts)
-      super(**opts)
-
-      self.count = count
+    def initialize(key, capacity, interval, **opts)
       self.interval = interval
+      super(key, capacity, @interval_usec, **opts)
     end
 
-    def count=(count)
-      unless count.is_a? Integer
-        raise ArgumentError, "expected Integer, found #{count.class}"
-      end
-
-      raise ArgumentError, "count must be >= 0" unless count >= 0
-
-      @count = count
+    private def interval=(interval)
+      @interval = interval
+      @interval_usec = Berater::Utils.to_usec(interval)
     end
 
-    def interval=(interval)
-      @interval = interval.dup
+    LUA_SCRIPT = Berater::LuaScript(<<~LUA
+      local key = KEYS[1]
+      local ts_key = KEYS[2]
+      local ts = tonumber(ARGV[1])
+      local capacity = tonumber(ARGV[2])
+      local interval_usec = tonumber(ARGV[3])
+      local cost = tonumber(ARGV[4])
+      local count = 0
+      local allowed
+      local usec_per_drip = interval_usec / capacity
+
+      -- 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
 
-      case @interval
-      when Integer
-        raise ArgumentError, "interval must be >= 0" unless @interval >= 0
-      when String
-        @interval = @interval.to_sym
-      when Symbol
+      if cost == 0 then
+        -- just check limit, ie. for .overlimit?
+        allowed = count < capacity
       else
-        raise ArgumentError, "unexpected interval type: #{interval.class}"
-      end
+        allowed = (count + cost) <= capacity
 
-      if @interval.is_a? Symbol
-        case @interval
-        when :sec, :second, :seconds
-          @interval = 1
-        when :min, :minute, :minutes
-          @interval = 60
-        when :hour, :hours
-          @interval = 60 * 60
-        else
-          raise ArgumentError, "unexpected interval value: #{interval}"
+        if allowed then
+          count = count + cost
+
+          -- time for bucket to empty, in milliseconds
+          local ttl = math.ceil((count * usec_per_drip) / 1000)
+
+          -- update count and last_ts, with expirations
+          redis.call('SET', key, count, 'PX', ttl)
+          redis.call('SET', ts_key, ts, 'PX', ttl)
         end
       end
 
-      @interval
-    end
+      return { count, allowed }
+    LUA
+    )
 
-    def limit(**opts, &block)
-      unless opts.empty?
-        return self.class.new(
-          count,
-          interval,
-          options.merge(opts)
-        ).limit(&block)
-      end
+    def limit(capacity: nil, cost: 1, &block)
+      capacity ||= @capacity
 
-      ts = Time.now.to_i
+      # timestamp in microseconds
+      ts = (Time.now.to_f * 10**6).to_i
 
-      # bucket into time slot
-      rkey = "%s:%d" % [ key, ts - ts % @interval ]
+      count, allowed = LUA_SCRIPT.eval(
+        redis,
+        [ cache_key(key), cache_key("#{key}-ts") ],
+        [ ts, capacity, @interval_usec, cost ]
+      )
 
-      count, _ = redis.multi do
-        redis.incr rkey
-        redis.expire rkey, @interval * 2
-      end
+      raise Overrated unless allowed
 
-      raise Overrated if count > @count
+      lock = Lock.new(self, ts, count)
+      yield_lock(lock, &block)
+    end
 
-      if block_given?
-        yield
+    def overloaded?
+      limit(cost: 0) { false }
+    rescue Overrated
+      true
+    end
+    alias overrated? overloaded?
+
+    def to_s
+      msg = if interval.is_a? Numeric
+        if interval == 1
+          "every second"
+        else
+          "every #{interval} seconds"
+        end
       else
-        count
+        "per #{interval}"
       end
+
+      "#<#{self.class}(#{key}: #{capacity} #{msg})>"
     end
 
   end

data/lib/berater/rspec.rb

@@ -0,0 +1,14 @@
+require 'berater'
+require 'berater/rspec/matchers'
+require 'berater/test_mode'
+require 'rspec'
+
+RSpec.configure do |config|
+  config.include(BeraterMatchers)
+
+  config.after do
+    Berater.expunge rescue nil
+    Berater.redis.script(:flush) rescue nil
+    Berater.reset
+  end
+end

data/lib/berater/rspec/matchers.rb

@@ -0,0 +1,62 @@
+module BeraterMatchers
+  class Overloaded
+    def initialize(type)
+      @type = type
+    end
+
+    def supports_block_expectations?
+      true
+    end
+
+    def matches?(obj)
+      begin
+        case obj
+        when Proc
+          # eg. expect { ... }.to be_overrated
+          res = obj.call
+
+          if res.is_a? Berater::Limiter
+            # eg. expect { Berater.new(...) }.to be_overloaded
+            res.overloaded?
+          else
+            # eg. expect { Berater(...)  }.to be_overloaded
+            # eg. expect { limiter.limit }.to be_overloaded
+            false
+          end
+        when Berater::Limiter
+          # eg. expect(Berater.new(...)).to be_overloaded
+          obj.overloaded?
+        end
+      rescue @type
+        true
+      end
+    end
+
+    # def description
+    # it { expect { Berater.new(:inhibitor) }.not_to be_overrated }
+
+    def failure_message
+      "expected #{@type} to be raised"
+    end
+
+    def failure_message_when_negated
+      "did not expect #{@type} to be raised"
+    end
+  end
+
+  def be_overloaded
+    Overloaded.new(Berater::Overloaded)
+  end
+
+  def be_overrated
+    Overloaded.new(Berater::RateLimiter::Overrated)
+  end
+
+  def be_incapacitated
+    Overloaded.new(Berater::ConcurrencyLimiter::Incapacitated)
+  end
+
+  def be_inhibited
+    Overloaded.new(Berater::Inhibitor::Inhibited)
+  end
+end

data/lib/berater/test_mode.rb

@@ -0,0 +1,52 @@
+require 'berater'
+
+module Berater
+  extend self
+
+  attr_reader :test_mode
+
+  def test_mode=(mode)
+    unless [ nil, :pass, :fail ].include?(mode)
+      raise ArgumentError, "invalid mode: #{Berater.test_mode}"
+    end
+
+    @test_mode = mode
+
+    # overload class methods
+    unless Berater::Limiter.singleton_class.ancestors.include?(TestMode)
+      Berater::Limiter.singleton_class.prepend(TestMode)
+    end
+  end
+
+  module TestMode
+    def new(*args, **opts)
+      return super unless Berater.test_mode
+
+      # chose a stub class with desired behavior
+      stub_klass = case Berater.test_mode
+      when :pass
+        Berater::Unlimiter
+      when :fail
+        Berater::Inhibitor
+      end
+
+      # don't stub self
+      return super if self < stub_klass
+
+      # swap out limit and overloaded? methods with stub
+      super.tap do |instance|
+        stub = stub_klass.allocate
+        stub.send(:initialize, *args, **opts)
+
+        instance.define_singleton_method(:limit) do |**opts, &block|
+          stub.limit(**opts, &block)
+        end
+
+        instance.define_singleton_method(:overloaded?) do
+          stub.overloaded?
+        end
+      end
+    end
+  end
+
+end

data/lib/berater/unlimiter.rb

@@ -1,18 +1,16 @@
 module Berater
-  class Unlimiter < BaseLimiter
+  class Unlimiter < Limiter
 
-    def initialize(*args, **opts)
-      super(**opts)
+    def initialize(key = :unlimiter, *args, **opts)
+      super(key, Float::INFINITY, **opts)
     end
 
     def limit(**opts, &block)
-      unless opts.empty?
-        return self.class.new(
-          **options.merge(opts)
-        ).limit(&block)
-      end
+      yield_lock(Lock.new(self, Float::INFINITY, 0), &block)
+    end
 
-      yield if block_given?
+    def overloaded?
+      false
     end
 
   end

data/lib/berater/utils.rb

@@ -0,0 +1,46 @@
+module Berater
+  module Utils
+    extend self
+
+    refine Object do
+      def to_usec
+        Berater::Utils.to_usec(self)
+      end
+    end
+
+    def to_usec(val)
+      res = val
+
+      if val.is_a? String
+        # naively attempt casting, otherwise maybe it's a keyword
+        res = Float(val) rescue val.to_sym
+      end
+
+      if res.is_a? Symbol
+        case res
+        when :sec, :second, :seconds
+          res = 1
+        when :min, :minute, :minutes
+          res = 60
+        when :hour, :hours
+          res = 60 * 60
+        end
+      end
+
+      unless res.is_a? Numeric
+        raise ArgumentError, "unexpected value: #{val}"
+      end
+
+      if res < 0
+        raise ArgumentError, "expected value >= 0, found: #{val}"
+      end
+
+      if res == Float::INFINITY
+        raise ArgumentError, "infinite values not allowed"
+      end
+
+      (res * 10**6).to_i
+    end
+
+  end
+end