rolling_counter.rb

require "securerandom"

class RollingCounter
  
  # :call-seq: RollingCounter.new(redis_client, max_window_seconds) -> counter
  # 
  # Create a new counter instance, that will store/retrieve counts with
  # +redis_client+ and return counts within the last +max_window_seconds+
  # seconds.
  # 
  # Do not attempt to use counters with different +max_window_seconds+ to access
  # the same keys, results will not be consistent. Instead set
  # +max_window_seconds+ to the max window needed, then pass the window required
  # to #get.
  # 
  def initialize(redis, max_window)
    @redis = redis
    @max_window = max_window
  end
  
  # :call-seq: counter.incr(key) -> count
  #            counter.incr(key, window_seconds) -> count
  # 
  # Increment the counter for +key+ by 1. Returns the updated count within
  # +window_seconds+, or the +max_window_seconds+ if +window_seconds+ is not
  # supplied.
  # 
  def incr(key, window=@max_window)
    key = key.to_s
    now = Time.now.to_f
    @redis.multi do
      @redis.zadd(key, now, SecureRandom.uuid)
      @redis.expire(key, @max_window.ceil)
      do_get(key, now, window)
    end.last
  end
  alias increment incr
  
  # :call-seq: counter.get(key) -> count
  #           counter.get(key, window_seconds) -> count
  # 
  # Returns the count for +key+ within +window_seconds+.
  # 
  def get(key, window=@max_window)
    @redis.multi do
      do_get(key.to_s, Time.now.to_f, window)
    end.last
  end
  
  private
  
  def do_get(key, now, window)
    @redis.zremrangebyscore(key, 0, now - @max_window)
    @redis.zcount(key, now - window, now)
  end
end