berater 0.0.1 → 0.1.4

checksums.yaml

@@ -1,7 +1,7 @@
 ---
-SHA1:
-  metadata.gz: 305ca5d94c6d92dd4b8bf73f2da7a03e15d419c7
-  data.tar.gz: 83a209c930780a32cf5cda81d1c932d1f8756f88
+SHA256:
+  metadata.gz: c8ca6fba7e38a014eed325df41dda7372b41ff62d9e78b61565f4d915d8b4677
+  data.tar.gz: 1fae864b93de87ea3dc70531f3f8697e80af526aea53d8ddfe2076eac9bbd61d
 SHA512:
-  metadata.gz: b0970141263af6e4b8461c0b974b01abd0873890d295e613e07db5fbb48ba0b736db44b67638e2eb8aa5f3d220d043ad3e7daaf17e58e00cabe999d67bd66f4d
-  data.tar.gz: d5d02d7331d13c355f3e3d2c56091552507a30e3ae7d978cbaa555984381c04033bd3d93441e8dd58bee9c019de7dd493a983ad28ad5b3bbd3c3fd23b6af8021
+  metadata.gz: 1b4b37134edece0657ea9be15119fe25b2208d554b9dbf56edd04f8e5acff77569022be025ec428b81ab1971171ae468e389a9568d7ad0163b0966e622a09f2e
+  data.tar.gz: bd927303eb1475e79a26161a184205008926f95c59a06eee43f453d5e464d15c205a91bffd36228ae28bc5d27f0fa32faebc7c5870ec5d1779bd6c57ecdc0fe2

data/lib/berater.rb

@@ -1,30 +1,64 @@
+require 'berater/version'
+
+
 module Berater
-  VERSION = '0.0.1'
+  extend self
 
-  class << self
+  class Overloaded < StandardError; end
 
-    def init redis
-      @@redis = redis
-    end
+  MODES = {}
+
+  attr_accessor :redis, :mode
 
+  def configure
+    self.mode = :unlimited # default
 
-    def incr key, limit, seconds
-      ts = Time.now.to_i
+    yield self
+  end
+
+  def new(mode, *args, **opts)
+    klass = MODES[mode.to_sym]
 
-      # bucket into time slot
-      rkey = "%s:%s:%d" % [ self.to_s, key, ts - ts % seconds ]
+    unless klass
+      raise ArgumentError, "invalid mode: #{mode}"
+    end
 
-      count = @@redis.multi do
-        @@redis.incr rkey
-        @@redis.expire rkey, seconds * 2
-      end.first
+    klass.new(*args, **opts)
+  end
 
-      raise LimitExceeded if count > limit
+  def register(mode, klass)
+    MODES[mode.to_sym] = klass
+  end
 
-      count
+  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
+    end
   end
 
-  class LimitExceeded < RuntimeError; end
 end
+
+# load and register limiters
+require 'berater/base_limiter'
+require 'berater/concurrency_limiter'
+require 'berater/inhibitor'
+require 'berater/rate_limiter'
+require 'berater/unlimiter'
+
+Berater.register(:concurrency, Berater::ConcurrencyLimiter)
+Berater.register(:inhibited, Berater::Inhibitor)
+Berater.register(:rate, Berater::RateLimiter)
+Berater.register(:unlimited, Berater::Unlimiter)

data/lib/berater/base_limiter.rb

@@ -0,0 +1,32 @@
+module Berater
+  class BaseLimiter
+
+    attr_reader :options
+
+    def initialize(**opts)
+      @options = opts
+    end
+
+    def key
+      if options[:key]
+        "#{self.class}:#{options[:key]}"
+      else
+        # default value
+        self.class.to_s
+      end
+    end
+
+    def redis
+      options[:redis] || Berater.redis
+    end
+
+    def limit(**opts)
+      raise NotImplementedError
+    end
+
+    def self.limit(*args, **opts, &block)
+      self.new(*args, **opts).limit(&block)
+    end
+
+  end
+end

data/lib/berater/concurrency_limiter.rb

@@ -0,0 +1,154 @@
+module Berater
+  class ConcurrencyLimiter < BaseLimiter
+
+    class Incapacitated < Overloaded; end
+
+    attr_reader :capacity, :timeout
+
+    def initialize(capacity, **opts)
+      super(**opts)
+
+      self.capacity = capacity
+      self.timeout = opts[:timeout] || 0
+    end
+
+    def capacity=(capacity)
+      unless capacity.is_a? Integer
+        raise ArgumentError, "expected Integer, found #{capacity.class}"
+      end
+
+      raise ArgumentError, "capacity must be >= 0" unless capacity >= 0
+
+      @capacity = capacity
+    end
+
+    def timeout=(timeout)
+      unless timeout.is_a? Integer
+        raise ArgumentError, "expected Integer, found #{timeout.class}"
+      end
+
+      raise ArgumentError, "timeout must be >= 0" unless timeout >= 0
+
+      @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?
+        limiter.timeout > 0 && @locked_at + limiter.timeout < Time.now
+      end
+    end
+
+    LUA_SCRIPT = <<~LUA.gsub(/^\s*(--.*\n)?/, '')
+      local key = KEYS[1]
+      local capacity = tonumber(ARGV[1])
+      local ts = tonumber(ARGV[2])
+      local ttl = tonumber(ARGV[3])
+
+      local exists
+      local count
+      local lock
+
+      -- 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)
+      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
+      elseif capacity > 0 then
+        count = 1
+        lock = "1"
+
+        -- create structure to track locks and next id
+        redis.call('ZADD', key, 'inf', lock + 1)
+
+        if ttl > 0 then
+          redis.call('EXPIRE', key, ttl * 2)
+        end
+      end
+
+      if lock then
+        -- store lock and timestamp
+        redis.call('ZADD', key, ts, lock)
+      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
+
+      count, lock_id = redis.eval(
+        LUA_SCRIPT,
+        [ key ],
+        [ capacity, Time.now.to_i, timeout ]
+      )
+
+      raise Incapacitated unless lock_id
+
+      lock = Lock.new(self, lock_id, count)
+
+      if block_given?
+        begin
+          yield lock
+        ensure
+          lock.release
+        end
+      else
+        lock
+      end
+    end
+
+    def release(lock)
+      res = redis.zrem(key, lock.id)
+      res == true || res == 1 # depending on which version of Redis
+    end
+
+  end
+end

data/lib/berater/inhibitor.rb

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

data/lib/berater/rate_limiter.rb

@@ -0,0 +1,84 @@
+module Berater
+  class RateLimiter < BaseLimiter
+
+    class Overrated < Overloaded; end
+
+    attr_accessor :count, :interval
+
+    def initialize(count, interval, **opts)
+      super(**opts)
+
+      self.count = count
+      self.interval = interval
+    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
+    end
+
+    def interval=(interval)
+      @interval = interval.dup
+
+      case @interval
+      when Integer
+        raise ArgumentError, "interval must be >= 0" unless @interval >= 0
+      when String
+        @interval = @interval.to_sym
+      when Symbol
+      else
+        raise ArgumentError, "unexpected interval type: #{interval.class}"
+      end
+
+      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}"
+        end
+      end
+
+      @interval
+    end
+
+    def limit(**opts, &block)
+      unless opts.empty?
+        return self.class.new(
+          count,
+          interval,
+          options.merge(opts)
+        ).limit(&block)
+      end
+
+      ts = Time.now.to_i
+
+      # bucket into time slot
+      rkey = "%s:%d" % [ key, ts - ts % @interval ]
+
+      count, _ = redis.multi do
+        redis.incr rkey
+        redis.expire rkey, @interval * 2
+      end
+
+      raise Overrated if count > @count
+
+      if block_given?
+        yield
+      else
+        count
+      end
+    end
+
+  end
+end
+

data/lib/berater/unlimiter.rb

@@ -0,0 +1,19 @@
+module Berater
+  class Unlimiter < BaseLimiter
+
+    def initialize(*args, **opts)
+      super(**opts)
+    end
+
+    def limit(**opts, &block)
+      unless opts.empty?
+        return self.class.new(
+          **options.merge(opts)
+        ).limit(&block)
+      end
+
+      yield if block_given?
+    end
+
+  end
+end

data/lib/berater/version.rb

@@ -0,0 +1,3 @@
+module Berater
+  VERSION = '0.1.4'
+end

data/spec/berater_spec.rb

@@ -0,0 +1,192 @@
+describe Berater do
+
+  it 'is connected to Redis' do
+    expect(Berater.redis.ping).to eq 'PONG'
+  end
+
+  it { is_expected.to respond_to :configure }
+
+  describe '.configure' do
+    it 'can be set via configure' do
+      Berater.configure do |c|
+        c.mode = :rate
+      end
+
+      expect(Berater.mode).to eq :rate
+    end
+  end
+
+  describe '.redis' do
+    it 'can be reset' do
+      expect(Berater.redis).not_to be_nil
+      Berater.redis = nil
+      expect(Berater.redis).to be_nil
+    end
+  end
+
+  describe '.mode' do
+    it 'validates inputs' do
+      expect { Berater.mode = :foo }.to raise_error(ArgumentError)
+    end
+  end
+
+  context 'unlimited mode' do
+    before { Berater.mode = :unlimited }
+
+    describe '.new' do
+      let(:limiter) { Berater.new(:unlimited) }
+
+      it 'instantiates an Unlimiter' do
+        expect(limiter).to be_a Berater::Unlimiter
+      end
+
+      it 'inherits redis' do
+        expect(limiter.redis).to be Berater.redis
+      end
+
+      it 'accepts options' do
+        redis = double('Redis')
+        limiter = Berater.new(:unlimited, key: 'key', redis: redis)
+        expect(limiter.key).to match(/key/)
+        expect(limiter.redis).to be redis
+      end
+    end
+
+    describe '.limit' do
+      it 'works' do
+        expect(Berater.limit).to be_nil
+      end
+
+      it 'yields' do
+        expect {|b| Berater.limit(&b) }.to yield_control
+      end
+
+      it 'never limits' do
+        10.times { expect(Berater.limit { 123 } ).to eq 123 }
+      end
+    end
+  end
+
+  context 'inhibited mode' do
+    before { Berater.mode = :inhibited }
+
+    describe '.new' do
+      let(:limiter) { Berater.new(:inhibited) }
+
+      it 'instantiates an Inhibitor' do
+        expect(limiter).to be_a Berater::Inhibitor
+      end
+
+      it 'inherits redis' do
+        expect(limiter.redis).to be Berater.redis
+      end
+
+      it 'accepts options' do
+        redis = double('Redis')
+        limiter = Berater.new(:inhibited, key: 'key', redis: redis)
+        expect(limiter.key).to match(/key/)
+        expect(limiter.redis).to be redis
+      end
+    end
+
+    describe '.limit' do
+      it 'always limits' do
+        expect { Berater.limit }.to be_inhibited
+      end
+    end
+  end
+
+  context 'rate mode' do
+    before { Berater.mode = :rate }
+
+    describe '.limiter' do
+      let(:limiter) { Berater.new(:rate, 1, :second) }
+
+      it 'instantiates a RateLimiter' do
+        expect(limiter).to be_a Berater::RateLimiter
+      end
+
+      it 'inherits redis' do
+        expect(limiter.redis).to be Berater.redis
+      end
+
+      it 'accepts options' do
+        redis = double('Redis')
+        limiter = Berater.new(:rate, 1, :second, key: 'key', redis: redis)
+        expect(limiter.key).to match(/key/)
+        expect(limiter.redis).to be redis
+      end
+    end
+
+    describe '.limit' do
+      it 'works' do
+        expect(Berater.limit(1, :second)).to eq 1
+      end
+
+      it 'yields' do
+        expect {|b| Berater.limit(2, :second, &b) }.to yield_control
+        expect(Berater.limit(2, :second) { 123 }).to eq 123
+      end
+
+      it 'limits excessive calls' do
+        expect(Berater.limit(1, :second)).to eq 1
+        expect { Berater.limit(1, :second) }.to be_overrated
+      end
+
+      it 'accepts options' do
+        redis = double('Redis')
+        expect(redis).to receive(:multi)
+
+        Berater.limit(1, :second, redis: redis) rescue nil
+      end
+    end
+  end
+
+  context 'concurrency mode' do
+    before { Berater.mode = :concurrency }
+
+    describe '.limiter' do
+      let(:limiter) { Berater.new(:concurrency, 1) }
+
+      it 'instantiates a ConcurrencyLimiter' do
+        expect(limiter).to be_a Berater::ConcurrencyLimiter
+      end
+
+      it 'inherits redis' do
+        expect(limiter.redis).to be Berater.redis
+      end
+
+      it 'accepts options' do
+        redis = double('Redis')
+        limiter = Berater.new(:concurrency, 1, key: 'key', redis: redis)
+        expect(limiter.key).to match(/key/)
+        expect(limiter.redis).to be redis
+      end
+    end
+
+    describe '.limit' do
+      it 'works (without blocks by returning a lock)' do
+        lock = Berater.limit(1)
+        expect(lock).to be_a Berater::ConcurrencyLimiter::Lock
+        expect(lock.release).to be true
+      end
+
+      it 'yields' do
+        expect {|b| Berater.limit(1, &b) }.to yield_control
+      end
+
+      it 'limits excessive calls' do
+        Berater.limit(1)
+        expect { Berater.limit(1) }.to be_incapacitated
+      end
+
+      it 'accepts options' do
+        redis = double('Redis')
+        expect(redis).to receive(:eval)
+
+        Berater.limit(1, redis: redis) rescue nil
+      end
+    end
+  end
+
+end