rate_limit.rb

require 'atomic'

class RateLimit
  def initialize(ops_per_second, max_bucket_size = ops_per_second) 
    @ops_per_second   = ops_per_second
    @max_bucket_size  = max_bucket_size
    @sec_per_op       = 1.0 / @ops_per_second

    @bucket           = Atomic.new(0)
    @incr_lock        = Atomic.new(0)

    @last_increment   = Time.now.to_f
  end 

  def use(ops = 1, sleep_wait = true)
    had_to_wait = false
    loop do
      # Increment ops in bucket

      if @incr_lock.compare_and_swap(0, 1)
        elapsed = Time.now.to_f - @last_increment

        if elapsed > @sec_per_op
          max_ops_to_add = (elapsed / @sec_per_op).to_i

          bucket_size = @bucket.value
          ops_to_add  = (bucket_size + max_ops_to_add) > @max_bucket_size ? (@max_bucket_size - bucket_size) : max_ops_to_add

          @bucket.update { |v| v + ops_to_add }

          @last_increment += (max_ops_to_add * @sec_per_op)
        end 
        @incr_lock.swap(0)
      end 

      if (v = @bucket.value) >= ops && @bucket.compare_and_swap(v, v-ops)
        return had_to_wait
      else
        had_to_wait = true
      end

      sleep @sec_per_op if sleep_wait
    end
  end
end

rate_limit_test.rb

require './rate_limit'
require 'benchmark'
require 'minitest/spec'
require 'minitest/autorun'

describe RateLimit do
  describe "100 operations per second" do
    it "takes a second for each 100 operations" do
      100.step(200, 20) do |ops|
        rl = RateLimit.new(100)
        elapsed = Benchmark.realtime { ops.times { rl.use } } 
        (ops / elapsed).must_be_within_epsilon 100, 0.01
      end
    end

    it "enforces max bucket size" do
      200.step(300, 25) do |ops|
        rl = RateLimit.new(100)
        sleep 2 
        elapsed = Benchmark.realtime { ops.times { rl.use } } 
        ((ops-100) / elapsed).must_be_within_epsilon 100, 0.01
      end
    end

    describe "concurrency" do
      it "takes a second for each 100 operations regardless of the number of threads" do
        total_ops = 1000

        100.step(200, 100) do |threads|
          rl = RateLimit.new(100)
          elapsed = Benchmark.realtime { 
            threads.times.collect { Thread.new { (total_ops/threads).times { rl.use } } } .each(&:join)
          }

          (total_ops / elapsed).must_be_within_epsilon 100, 0.01
        end
      end
    end
  end

end