ratelimit.clj

(ns ratelimit
  "Implement rate limiting using a token bucket"
  (:require [clojure.core.async :as async :refer [<!! >! alts! go chan dropping-buffer close! timeout]]))

(defprotocol TokenBucket
  (take-token [this])
  (release [this]))

(defn- ->period-ms [p]
  (if (keyword? p)
    (condp = p
      :per-minute (* 60 1000)
      :per-hour   (* 60 60 1000)
      :per-day    (* 24 60 60 1000))
    p))

(defn rate-limiter
  "Creates a rate limiter. ``(take-token limiter)`` will block if n
   calls have been made in the previous period. Tokens are replenished
   regularly such that the call will unblock when doing so will not
   exceed n requests / period. period can either be a duration in
   milliseconds or one of :per-minute, :per-hour, :per-day" [n period]
   (let [buf       (dropping-buffer n)
         tchan     (chan buf)
         stop      (chan)
         period-ms (/ (->period-ms period) n)]
    (go (dotimes [_ n] (>! tchan ::token)) ;initial fill
        (loop []
          (let [[_ chan] (alts! [stop (timeout period-ms)] :priority true)]
            (when-not (= chan stop)
              (>! tchan ::token)
              (recur)))))
    (reify
      TokenBucket
      (take-token [this]
        (<!! tchan))
      (release [this]
        (go (>! stop ::stop)
            (close! tchan)))
      clojure.lang.Counted
      (count [this]
        (count buf)))))


(comment

  (def limit (rate-limiter 10 :per-minute))
  (dotimes [n 15] (take-token limit) (printf "%02d\t%s\n" n (java.util.Date.)))
   ;00  Tue Jul 16 12:08:46 BST 2013 (first 10 proceed immediately)
   ;01  Tue Jul 16 12:08:46 BST 2013
   ;02  Tue Jul 16 12:08:46 BST 2013
   ;03  Tue Jul 16 12:08:46 BST 2013
   ;04  Tue Jul 16 12:08:46 BST 2013
   ;05  Tue Jul 16 12:08:46 BST 2013
   ;06  Tue Jul 16 12:08:46 BST 2013
   ;07  Tue Jul 16 12:08:46 BST 2013
   ;08  Tue Jul 16 12:08:46 BST 2013
   ;09  Tue Jul 16 12:08:46 BST 2013
   ;10  Tue Jul 16 12:08:52 BST 2013 (rate limit kicks in, now proceeds every ~6s)
   ;11  Tue Jul 16 12:08:58 BST 2013
   ;12  Tue Jul 16 12:09:04 BST 2013
   ;13  Tue Jul 16 12:09:10 BST 2013
   ;14  Tue Jul 16 12:09:16 BST 2013
   ;nil
   (release limit)
   (dotimes [n 15] (take-token limit) (printf "%02d\t%s\n" n (java.util.Date.)))
   ;00  Tue Jul 16 12:09:31 BST 2013 (limit is released)
   ;01  Tue Jul 16 12:09:31 BST 2013
   ;02  Tue Jul 16 12:09:31 BST 2013
   ;03  Tue Jul 16 12:09:31 BST 2013
   ;04  Tue Jul 16 12:09:31 BST 2013
   ;05  Tue Jul 16 12:09:31 BST 2013
   ;06  Tue Jul 16 12:09:31 BST 2013
   ;07  Tue Jul 16 12:09:31 BST 2013
   ;08  Tue Jul 16 12:09:31 BST 2013
   ;09  Tue Jul 16 12:09:31 BST 2013
   ;10  Tue Jul 16 12:09:31 BST 2013
   ;11  Tue Jul 16 12:09:31 BST 2013
   ;12  Tue Jul 16 12:09:31 BST 2013
   ;13  Tue Jul 16 12:09:31 BST 2013
   ;14  Tue Jul 16 12:09:31 BST 2013
   ;nil
  )