diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,11 @@ # Changelog +## [0.2.0.0] - 2025-05-12 + +### Changed + +* Fixed that `waitDebit` would not wait infinitely if tokens will never become available. + ## [0.1.0.0] - 2024-03-11 ### Changed diff --git a/src/Control/Concurrent/TokenLimiter/Concurrent.hs b/src/Control/Concurrent/TokenLimiter/Concurrent.hs --- a/src/Control/Concurrent/TokenLimiter/Concurrent.hs +++ b/src/Control/Concurrent/TokenLimiter/Concurrent.hs @@ -16,6 +16,7 @@ MonotonicDiffNanos (..), -- * Helper functions + computeMicrosecondsToWait, computeCurrentCount, ) where @@ -123,27 +124,48 @@ pure ((now, newCount), Nothing) else do let extraTokensNeeded = debit - currentCount - let microsecondsToWaitDouble :: Double - microsecondsToWaitDouble = - 1_000_000 - * (fromIntegral :: Word64 -> Double) extraTokensNeeded - / (fromIntegral :: Word64 -> Double) (tokenLimitConfigTokensPerSecond tokenLimiterConfig) + case computeMicrosecondsToWait (tokenLimitConfigTokensPerSecond tokenLimiterConfig) extraTokensNeeded of + Nothing -> do + -- We can't wait for this long, so we just wait forever. + -- This is a bit sad, but it is the best we can do. + -- We could also throw an exception, but that would be a bit rude. + -- So we just wait forever. + pure ((lastServiced, countThen), Nothing) + Just microsecondsToWait -> do + -- We can wait for this long, so we will. + -- + -- threadDelay guarantees that _at least_ the given number of microseconds will have passed. + threadDelay microsecondsToWait + -- However, it could be MUCH longer than that, so we will recalculate the time instead of + -- adding that number of microseconds to the old time. + nowAfterWaiting <- getMonotonicTimeNSec + let delta = MonotonicDiffNanos (nowAfterWaiting - now) + -- We do assume here that we will now have enough tokens and do not need to recalculate whether there will be enough. + -- (We would not know what to do if there weren't, anyway.) + -- BUT this assumption _should_ hold because _modifyMVar_ guarantees + -- atomicity if there are no other producers for this MVar, which there + -- aren't. + let currentCountAfterWaiting = computeCurrentCount tokenLimiterConfig lastServiced countThen nowAfterWaiting + let newCount = currentCountAfterWaiting - debit + pure ((nowAfterWaiting, newCount), Just delta) - let microsecondsToWait = ceiling microsecondsToWaitDouble - -- threadDelay guarantees that _at least_ the given number of microseconds will have passed. - threadDelay microsecondsToWait - -- However, it could be MUCH longer than that, so we will recalculate the time instead of - -- adding that number of microseconds to the old time. - nowAfterWaiting <- getMonotonicTimeNSec - let delta = MonotonicDiffNanos (nowAfterWaiting - now) - -- We do assume here that we will now have enough tokens and do not need to recalculate whether there will be enough. - -- (We would not know what to do if there weren't, anyway.) - -- BUT this assumption _should_ hold because _modifyMVar_ guarantees - -- atomicity if there are no other producers for this MVar, which there - -- aren't. - let currentCountAfterWaiting = computeCurrentCount tokenLimiterConfig lastServiced countThen nowAfterWaiting - let newCount = currentCountAfterWaiting - debit - pure ((nowAfterWaiting, newCount), Just delta) +computeMicrosecondsToWait :: Count -> Word64 -> Maybe Int +computeMicrosecondsToWait tokensPerSecond extraTokensNeeded = do + let microsecondsToWaitDouble :: Double + microsecondsToWaitDouble = + 1_000_000 + * (fromIntegral :: Word64 -> Double) extraTokensNeeded + / (fromIntegral :: Word64 -> Double) tokensPerSecond + + let maxBoundDouble :: Double + maxBoundDouble = fromIntegral (maxBound :: Int) + + guard $ microsecondsToWaitDouble < maxBoundDouble + + pure $ ceiling microsecondsToWaitDouble + +waitForever :: IO void +waitForever = forever $ threadDelay 1_000_000 -- | Compute the current number of tokens in a bucket purely. -- diff --git a/test/Control/Concurrent/TokenLimiter/ConcurrentSpec.hs b/test/Control/Concurrent/TokenLimiter/ConcurrentSpec.hs --- a/test/Control/Concurrent/TokenLimiter/ConcurrentSpec.hs +++ b/test/Control/Concurrent/TokenLimiter/ConcurrentSpec.hs @@ -64,6 +64,16 @@ -- 10 seconds later we should have no more than maxBound, eventhough the computed count would be maxBound + 9 computeCurrentCount config 0 maxBound 10_000_000_000 `shouldBe` maxBound + describe "computeMicrosecondsToWait" $ do + it "Instructs to wait for one second to get one token" $ do + computeMicrosecondsToWait 1 1 `shouldBe` Just 1_000_000 + + it "instructs to wait forever if no tokens are added" $ do + computeMicrosecondsToWait 0 1 `shouldBe` Nothing + + it "instructs to wait forever if no tokens will be available soon enough" $ do + computeMicrosecondsToWait 1 maxBound `shouldBe` Nothing + describe "makeTokenLimiter" $ do it "always succeeds" $ forAllValid $ \config -> do @@ -185,7 +195,7 @@ TokenLimitConfig { tokenLimitConfigInitialTokens = 0, tokenLimitConfigMaxTokens = 1, - tokenLimitConfigTokensPerSecond = 50 + tokenLimitConfigTokensPerSecond = 10 } limiter <- makeTokenLimiter config let l :: [Int] diff --git a/token-limiter-concurrent.cabal b/token-limiter-concurrent.cabal --- a/token-limiter-concurrent.cabal +++ b/token-limiter-concurrent.cabal @@ -1,57 +1,60 @@ cabal-version: 1.12 - --- This file has been generated from package.yaml by hpack version 0.35.2. --- --- see: https://github.com/sol/hpack - -name: token-limiter-concurrent -version: 0.1.0.0 -synopsis: A thread-safe concurrent token-bucket rate limiter that guarantees fairness -homepage: https://github.com/NorfairKing/token-limiter-concurrent#readme -bug-reports: https://github.com/NorfairKing/token-limiter-concurrent/issues -author: Tom Sydney Kerckhove -maintainer: syd@cs-syd.eu -copyright: Copyright (c) 2022-2024 Tom Sydney Kerckhove -license: MIT -license-file: LICENSE -build-type: Simple +name: token-limiter-concurrent +version: 0.2.0.0 +synopsis: A thread-safe concurrent token-bucket rate limiter that guarantees fairness +homepage: https://github.com/NorfairKing/token-limiter-concurrent#readme +bug-reports: https://github.com/NorfairKing/token-limiter-concurrent/issues +author: Tom Sydney Kerckhove +maintainer: syd@cs-syd.eu +copyright: Copyright (c) 2022-2024 Tom Sydney Kerckhove +license: MIT +license-file: LICENSE +build-type: Simple extra-source-files: - LICENSE - CHANGELOG.md + CHANGELOG.md + LICENSE source-repository head type: git location: https://github.com/NorfairKing/token-limiter-concurrent library - exposed-modules: - Control.Concurrent.TokenLimiter.Concurrent - other-modules: - Paths_token_limiter_concurrent + -- cabal-gild: discover src + exposed-modules: Control.Concurrent.TokenLimiter.Concurrent hs-source-dirs: - src + src + build-depends: - base >=4.7 && <5 + base >=4.7 && <5 + default-language: Haskell2010 test-suite token-limiter-concurrent-test type: exitcode-stdio-1.0 main-is: Spec.hs + -- cabal-gild: discover test --exclude=test/Spec.hs other-modules: - Control.Concurrent.TokenLimiter.ConcurrentSpec - Paths_token_limiter_concurrent + Control.Concurrent.TokenLimiter.ConcurrentSpec + hs-source-dirs: - test - ghc-options: -threaded -rtsopts -with-rtsopts=-N + test + + ghc-options: + -threaded + -rtsopts + -with-rtsopts=-N + build-tool-depends: - sydtest-discover:sydtest-discover + sydtest-discover:sydtest-discover + build-depends: - QuickCheck - , async - , base >=4.7 && <5 - , genvalidity - , genvalidity-sydtest - , stm - , sydtest - , token-limiter-concurrent + QuickCheck, + async, + base >=4.7 && <5, + genvalidity, + genvalidity-sydtest, + stm, + sydtest, + token-limiter-concurrent + default-language: Haskell2010