diff --git a/src/diehard/core.clj b/src/diehard/core.clj index 625b5a5..c129814 100644 --- a/src/diehard/core.clj +++ b/src/diehard/core.clj @@ -462,8 +462,13 @@ You can always check circuit breaker state with (defmacro ^{:doc "Create a rate limiter with options. -* `:rate` execution permits per second. -* `:max-cached-tokens` the max size of permits we can cache when idle"} +* `:rate` execution permits per second (may be a floating point number, e.g. + 0.5 <=> 1 req every 2 sec) +* `:max-cached-tokens` the max size of permit tokens that the bucket can cache + when it's idle +* `:sleep-fn` a unary fn of millis to sleep for, allowing for custom sleep + semantics; by default, sleeps interruptedly; pass `uninterruptible-sleep` + to sleep uninterruptedly"} defratelimiter [name opts] `(def ~name (rl/rate-limiter (u/verify-opt-map-keys-with-spec :rate-limiter/rate-limiter-new ~opts)))) @@ -482,10 +487,9 @@ to given rate. Use `defratelimiter` to define a ratelimiter and use it as option By default it will wait forever until there is permits available. You can also specify a `max-wait-ms` to wait for a given time. If there's no permits in this period, this block -will throw a Clojure `ex-info`, with `ex-data` as +will throw a Clojure exception with a `:throttled true` entry in `ex-data`, as follows: ```clojure - (try (with-rate-limiter {:ratelimiter myfl :max-wait-ms 1000} diff --git a/src/diehard/rate_limiter.clj b/src/diehard/rate_limiter.clj index 70f4f1f..73aab77 100644 --- a/src/diehard/rate_limiter.clj +++ b/src/diehard/rate_limiter.clj @@ -12,41 +12,26 @@ "Try to acquire given number of permits, allows blocking for at most `wait-ms` milliseconds. Return true if there are enough permits in permitted time.")) -(declare refill acquire-sleep-ms try-acquire-sleep-ms) - -(defn- do-acquire [rate-limiter permits] - (refill rate-limiter) - (acquire-sleep-ms rate-limiter permits)) - -(defn- do-try-acquire [rate-limiter permits max-wait-ms] - (refill rate-limiter) - (try-acquire-sleep-ms rate-limiter permits max-wait-ms)) +(declare refill do-acquire! do-try-acquire) (defrecord TokenBucketRateLimiter [rate max-tokens ;; internal state - state] + state sleep-fn] IRateLimiter (acquire! [this] (acquire! this 1)) (acquire! [this permits] - (let [sleep (do-acquire this permits)] - (when (> sleep 0) - (Thread/sleep ^long sleep)))) + (refill this) + (do-acquire! this permits)) (try-acquire [this] (try-acquire this 1)) (try-acquire [this permits] (try-acquire this permits 0)) (try-acquire [this permits wait-ms] - (let [sleep (do-try-acquire this permits wait-ms)] - (if (false? sleep) - false - (do - (when (> sleep 0) - (Thread/sleep ^long sleep)) - true))))) + (refill this) + (do-try-acquire this permits wait-ms))) (defn- refill [^TokenBucketRateLimiter rate-limiter] - ;; refill (let [now (System/currentTimeMillis)] (swap! (.-state rate-limiter) (fn [state] @@ -59,46 +44,77 @@ %)) (assoc :last-refill-ts now)))))) -(defn- acquire-sleep-ms [^TokenBucketRateLimiter rate-limiter permits] - (let [{pending-tokens :reserved-tokens} (swap! (.-state rate-limiter) - update :reserved-tokens + permits)] - (if (<= pending-tokens 0) - 0 - ;; time as milliseconds - (long (/ pending-tokens (.-rate rate-limiter)))))) +(defn- ->sleep-ms ^long [pending-tokens rate] + (if (<= pending-tokens 0) 0 (long (/ pending-tokens rate)))) + +(defn- do-acquire! + [^TokenBucketRateLimiter rate-limiter permits] + (let [state (swap! (.-state rate-limiter) update :reserved-tokens + permits) + sleep-ms (->sleep-ms (:reserved-tokens state) (.-rate rate-limiter))] + ((.-sleep-fn rate-limiter) sleep-ms))) -(defn- try-acquire-sleep-ms [^TokenBucketRateLimiter rate-limiter permits max-wait-ms] +(defn- do-try-acquire + [^TokenBucketRateLimiter rate-limiter permits max-wait-ms] (try - (let [{pending-tokens :reserved-tokens} - (swap! (.-state rate-limiter) - (fn [state] - (update state :reserved-tokens - (fn [pending-tokens] - ;; test if we can pass in wait period - (if (<= (- (+ pending-tokens permits) - (* max-wait-ms (.-rate rate-limiter))) - 0) - (+ pending-tokens permits) - (throw (ex-info "Not enough permits." {:rate-limiter true})))))))] - (if (<= pending-tokens 0) - 0 - (long (/ pending-tokens (.-rate rate-limiter))))) - (catch clojure.lang.ExceptionInfo _ - false))) + (let [state (swap! (.-state rate-limiter) + (fn [state] + (update state :reserved-tokens + (fn [pending-tokens] + ;; test if we can pass in wait period + (if (<= (- (+ pending-tokens permits) + (* max-wait-ms (.-rate rate-limiter))) + 0) + (+ pending-tokens permits) + (throw (ex-info "Not enough permits" + {:rate-limiter true}))))))) + sleep-ms (->sleep-ms (:reserved-tokens state) (.-rate rate-limiter))] + ((.-sleep-fn rate-limiter) sleep-ms) + true) + (catch Exception e + (if-not (:rate-limiter (ex-data e)) + (throw e) + false)))) + +(defn interruptible-sleep [^long ms] + (when (pos? ms) + (Thread/sleep ms))) + +(def ^:private nanos-in-ms 1000000) -(def ^{:const true :no-doc true} - allowed-rate-limiter-option-keys - #{:rate :max-cached-tokens}) +(defn uninterruptible-sleep [^long ms] + (when (pos? ms) + (let [end-time-ns (+ (System/nanoTime) (* nanos-in-ms ms))] + (with-local-vars [interrupted? false] + (try (loop [] + (let [remaining-ns (- end-time-ns (System/nanoTime))] + ;; required sleep duration fully consumed — exit + (when (< 0 remaining-ns) + (try + (Thread/sleep (quot remaining-ns nanos-in-ms) + (mod remaining-ns nanos-in-ms)) + ;; successful sleep — exit + (catch InterruptedException _ + (var-set interrupted? true))) + ;; can only recur from tail position + (when @interrupted? (recur))))) + (finally + (when @interrupted? + (Thread/.interrupt (Thread/currentThread))))))))) (defn rate-limiter "Create a default rate limiter with: * `rate`: permits per second (may be a floating point, e.g. 0.5 <=> 1 req every 2 sec) - * `max-cached-tokens`: the max size of tokens that the bucket can cache when it's idle" - [{:keys [rate max-cached-tokens] :as _opts}] + * `max-cached-tokens`: the max size of tokens that the bucket can cache when it's idle + * `sleep-fn`: a unary fn of millis to sleep for, allowing for custom sleep semantics; + by default, sleeps interruptedly; pass `uninterruptible-sleep` to sleep + uninterruptedly" + [{:keys [rate max-cached-tokens sleep-fn] :as _opts}] (if (some? rate) - (let [max-cached-tokens (or max-cached-tokens (int rate))] + (let [max-cached-tokens (or max-cached-tokens (int rate)) + sleep-fn (or sleep-fn interruptible-sleep)] (TokenBucketRateLimiter. (/ (double rate) 1000) max-cached-tokens (atom {:reserved-tokens (double 0) - :last-refill-ts (long -1)}))) + :last-refill-ts (long -1)}) + sleep-fn)) (throw (IllegalArgumentException. ":rate is required for rate-limiter")))) diff --git a/src/diehard/spec.clj b/src/diehard/spec.clj index b47f87d..88186a1 100644 --- a/src/diehard/spec.clj +++ b/src/diehard/spec.clj @@ -135,14 +135,15 @@ :timeout/on-failure])) ;; rate limiter -; allow floating point numbers, so that we can pass numbers such as 0.5, to signify 1 req / 2 seconds -;(s/def :rate-limiter/rate int?) + (s/def :rate-limiter/rate number?) (s/def :rate-limiter/max-cached-tokens int?) +(s/def :rate-limiter/sleep-fn fn?) (s/def :rate-limiter/rate-limiter-new (only-keys :req-un [:rate-limiter/rate] - :opt-un [:rate-limiter/max-cached-tokens])) + :opt-un [:rate-limiter/max-cached-tokens + :rate-limiter/sleep-fn])) (s/def :rate-limiter/ratelimiter #(satisfies? dr/IRateLimiter %)) (s/def :rate-limiter/max-wait-ms int?) diff --git a/test/diehard/core_test.clj b/test/diehard/core_test.clj index ca34a61..6810c70 100644 --- a/test/diehard/core_test.clj +++ b/test/diehard/core_test.clj @@ -2,7 +2,8 @@ (:require [clojure.test :refer :all] [diehard.circuit-breaker :as cb] [diehard.core :refer :all]) - (:import [dev.failsafe CircuitBreakerOpenException TimeoutExceededException] + (:import [clojure.lang IExceptionInfo] + [dev.failsafe CircuitBreakerOpenException TimeoutExceededException] [java.util.concurrent CountDownLatch])) (deftest test-retry @@ -390,11 +391,12 @@ (let [counter0 (atom 0)] (try (while (< @counter0 200) - (with-rate-limiter {:ratelimiter my-rl + (with-rate-limiter {:ratelimiter my-rl3 :max-wait-ms 1} (my-fn counter0))) (is false) (catch Exception e + (is (instance? IExceptionInfo e)) (is (:throttled (ex-data e)))))))) (testing "permits" diff --git a/test/diehard/rate_limiter_test.clj b/test/diehard/rate_limiter_test.clj index a82e820..658adc2 100644 --- a/test/diehard/rate_limiter_test.clj +++ b/test/diehard/rate_limiter_test.clj @@ -1,42 +1,216 @@ (ns diehard.rate-limiter-test - (:require [clojure.test :as t] - [diehard.rate-limiter :as r]) - (:import [java.util.concurrent Executors])) + (:require [clojure.test :refer [deftest testing is]] + [diehard.rate-limiter :as rl]) + (:import [java.util.concurrent ExecutorService Executors Future TimeUnit] + [java.util.concurrent.atomic AtomicBoolean])) -(t/deftest rate-limiter-test - (t/testing "base case" +(def default-error + "Relative error in percent. 3% by default." + 0.03) + +(defn- approx== + ([expected actual] + (approx== expected actual (* default-error expected))) + ([expected actual tolerance] + (< (abs (- expected actual)) tolerance))) + +(defn- total-block-time + [rate n-threads] + (long (* (/ 1.0 rate) n-threads 1000))) + +(defn- submit + ^Future [^ExecutorService executor ^Callable f] + (ExecutorService/.submit executor f)) + +(defn- run-rate-limited-counting + [& {:keys [proceed-run? do-shutdown!]}] + (fn [rl-opts n-threads run-time-sec] + (let [pool (Executors/newFixedThreadPool n-threads) + stop? (AtomicBoolean. false) + rate-limiter (rl/rate-limiter rl-opts) + counter (atom 0)] + (doseq [_ (range n-threads)] + (submit pool (fn [] + (try + (loop [] + (rl/acquire! rate-limiter) + (when (AtomicBoolean/.get stop?) + (throw (ex-info "Stopped" {:type ::stop}))) + (when (proceed-run?) + (swap! counter inc) + (recur))) + (catch Exception ex + (when-not (= ::stop (:type (ex-data ex))) + (throw ex))))))) + (Thread/sleep ^long (* 1000 run-time-sec)) + (do-shutdown! pool) + {:await-fn (fn [timeout-ms] + (ExecutorService/.awaitTermination + pool timeout-ms TimeUnit/MILLISECONDS)) + :stop-fn #(AtomicBoolean/.set stop? true) + :counter counter}))) + +(defn- run-rate-limited-counting:interruption [] + (run-rate-limited-counting + :proceed-run? (constantly true) + :do-shutdown! ExecutorService/.shutdownNow)) + +(defn- run-rate-limited-counting:custom-flag [] + (let [running? (AtomicBoolean. true)] + (run-rate-limited-counting + :proceed-run? #(AtomicBoolean/.get running?) + :do-shutdown! (fn [pool] + (AtomicBoolean/.set running? false) + (ExecutorService/.shutdown pool))))) + +(defn- run-rate-limited-counting:exception [] + (let [throw? (AtomicBoolean. false)] + (run-rate-limited-counting + :proceed-run? #(if (AtomicBoolean/.get throw?) + (throw (Exception. "Task failed!")) + true) + :do-shutdown! (fn [pool] + (AtomicBoolean/.set throw? true) + (ExecutorService/.shutdown pool))))) + +(deftest rate-limiter-at-high-rates-test + (testing "high rates" (let [rate 1000 - threads 32 - pool (Executors/newFixedThreadPool threads) - rl (r/rate-limiter {:rate rate}) - counter (atom 0) - time-secs 2] - (doseq [_ (range threads)] - (.submit pool (cast Runnable (fn [] - (while true - (r/acquire! rl) - (swap! counter inc)))))) - (Thread/sleep (* time-secs 1000)) - (.shutdown pool) - (t/is (< (Math/abs (- @counter (* rate time-secs))) - ;; error tolerance 3% - (* 0.03 (* rate time-secs)))))) - - (t/testing "rate limits less than 1.0" + run-time-sec 2 + + exp-count (* rate run-time-sec) + + n-threads 32 + ;; tenfold be enough to cover up thread switching costs + term-timeout-ms (* 10 (total-block-time rate n-threads))] + + (testing "with interruptible sleep" + (testing "and interruption signal to stop tasks" + (let [{:keys [await-fn counter]} ((run-rate-limited-counting:interruption) + {:rate rate} n-threads run-time-sec)] + (is (await-fn term-timeout-ms) + "Terminates successfully, without timeout") + (is (approx== exp-count @counter) + "Count must be close to an expected value"))) + (testing "and custom running flag to stop tasks" + (let [{:keys [await-fn counter]} ((run-rate-limited-counting:custom-flag) + {:rate rate} n-threads run-time-sec)] + (is (await-fn term-timeout-ms) + "Terminates successfully, without timeout") + (is (approx== exp-count @counter) + "Count must be close to an expected value"))) + (testing "when a task logic throws an exception" + (let [{:keys [await-fn counter]} ((run-rate-limited-counting:exception) + {:rate rate} n-threads run-time-sec)] + (is (await-fn term-timeout-ms) + "Terminates successfully, without timeout") + (is (approx== exp-count @counter) + "Count must be close to an expected value")))) + + (testing "with uninterruptible sleep" + (testing "and interruption signal to stop tasks" + (let [{:keys [await-fn stop-fn counter]} ((run-rate-limited-counting:interruption) + {:rate rate + :sleep-fn rl/uninterruptible-sleep} + n-threads run-time-sec)] + ;; here we are forced to check the count before it's too late, + ;; since the task pool termination won't happen straight away + (is (approx== exp-count @counter) + "Count must be close to an expected value") + + (is (false? (await-fn term-timeout-ms)) + "Cannot terminate due to (some) tasks still running") + ;; the only way to actually stop all running tasks (isolation) + (stop-fn) + (is (await-fn term-timeout-ms) + "Terminates successfully, without timeout THIS TIME"))) + (testing "and custom running flag to stop tasks" + (let [{:keys [await-fn counter]} ((run-rate-limited-counting:custom-flag) + {:rate rate + :sleep-fn rl/uninterruptible-sleep} + n-threads run-time-sec)] + (is (await-fn term-timeout-ms) + "Terminates successfully, without timeout") + (is (approx== exp-count @counter) + "Count must be close to an expected value"))) + (testing "when a task logic throws an exception" + (let [{:keys [await-fn counter]} ((run-rate-limited-counting:exception) + {:rate rate + :sleep-fn rl/uninterruptible-sleep} + n-threads run-time-sec)] + (is (await-fn term-timeout-ms) + "Terminates successfully, without timeout") + (is (approx== exp-count @counter) + "Count must be close to an expected value"))))))) + +(deftest rate-limiter-at-low-rates-test + (testing "low rates" ; less than 1.0 (let [rate 0.5 - threads 1 - pool (Executors/newFixedThreadPool threads) - rl (r/rate-limiter {:rate rate}) - counter (atom 0) - time-secs 5] - (doseq [_ (range threads)] - (.submit pool (cast Runnable (fn [] - (while true - (r/acquire! rl) - (swap! counter inc)))))) - (Thread/sleep (* time-secs 1000)) - (.shutdown pool) - (let [variance (Math/abs (- @counter (* rate time-secs)))] - (t/is (<= variance - ;; error tolerance: 0.5 -- larger, because we are dealing with slower ticks - 0.5)))))) + run-time-sec 5 + + exp-count (* rate run-time-sec) + ;; exact error tolerance, since we are dealing with much slower ticks + tolerance 1 + + n-threads 2 + ;; each thread will have to sleep for ≈2s (that's why we have just 2) + term-timeout-ms (total-block-time rate n-threads)] + + (testing "with interruptible sleep" + (testing "and interruption signal to stop tasks" + (let [{:keys [await-fn counter]} ((run-rate-limited-counting:interruption) + {:rate rate} n-threads run-time-sec)] + (is (await-fn term-timeout-ms) + "Terminates successfully, without timeout") + (is (approx== exp-count @counter tolerance) + "Count must be close to an expected value"))) + (testing "and custom running flag to stop tasks" + (let [{:keys [await-fn counter]} ((run-rate-limited-counting:custom-flag) + {:rate rate} n-threads run-time-sec)] + (is (await-fn term-timeout-ms) + "Terminates successfully, without timeout") + (is (approx== exp-count @counter tolerance) + "Count must be close to an expected value"))) + (testing "when a task logic throws an exception" + (let [{:keys [await-fn counter]} ((run-rate-limited-counting:exception) + {:rate rate} n-threads run-time-sec)] + (is (await-fn term-timeout-ms) + "Terminates successfully, without timeout") + (is (approx== exp-count @counter tolerance) + "Count must be close to an expected value")))) + + (testing "with uninterruptible sleep" + (testing "and interruption signal to stop tasks" + (let [{:keys [await-fn stop-fn counter]} ((run-rate-limited-counting:interruption) + {:rate rate + :sleep-fn rl/uninterruptible-sleep} + n-threads run-time-sec)] + ;; here we are forced to check the count before it's too late, + ;; since the task pool termination won't happen straight away + (is (approx== exp-count @counter tolerance) + "Count must be close to an expected value") + + (is (false? (await-fn term-timeout-ms)) + "Cannot terminate due to (some) tasks still running") + ;; the only way to actually stop all running tasks (isolation) + (stop-fn) + (is (await-fn term-timeout-ms) + "Terminates successfully, without timeout THIS TIME"))) + (testing "and custom running flag to stop tasks" + (let [{:keys [await-fn counter]} ((run-rate-limited-counting:custom-flag) + {:rate rate + :sleep-fn rl/uninterruptible-sleep} + n-threads run-time-sec)] + (is (await-fn term-timeout-ms) + "Terminates successfully, without timeout") + (is (approx== exp-count @counter tolerance) + "Count must be close to an expected value"))) + (testing "when a task logic throws an exception" + (let [{:keys [await-fn counter]} ((run-rate-limited-counting:exception) + {:rate rate + :sleep-fn rl/uninterruptible-sleep} + n-threads run-time-sec)] + (is (await-fn term-timeout-ms) + "Terminates successfully, without timeout") + (is (approx== exp-count @counter tolerance) + "Count must be close to an expected value")))))))