diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,13 @@ +# Changelog + +## [0.1.0.0] - 2024-03-11 + +### Changed + +* `waitDebit` now returns how long it waited, if it did. +* `tryDebit` no longer blocks if another `waitDebit` is already happening. + +### Removed + +* `canDebit` + diff --git a/LICENSE b/LICENSE --- a/LICENSE +++ b/LICENSE @@ -1,4 +1,4 @@ -Copyright 2021 Tom Sydney Kerckhove +Copyright 2021-2024 Tom Sydney Kerckhove Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in 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 @@ -3,14 +3,17 @@ {-# LANGUAGE RecordWildCards #-} module Control.Concurrent.TokenLimiter.Concurrent - ( Count, + ( -- * Create + Count, TokenLimitConfig (..), MonotonicTime, TokenLimiter (..), makeTokenLimiter, - canDebit, + + -- * Use tryDebit, waitDebit, + MonotonicDiffNanos (..), -- * Helper functions computeCurrentCount, @@ -18,6 +21,9 @@ where import Control.Concurrent +import Control.Exception +import Control.Monad +import Data.Maybe (fromMaybe) import Data.Word import GHC.Clock import GHC.Generics (Generic) @@ -42,6 +48,9 @@ -- This only exists because it is also a 'Word64' and would be too easy to confuse with a 'Count'. type MonotonicTime = Word64 +newtype MonotonicDiffNanos = MonotonicDiffNanos {unMonotonicDiffNanos :: Word64} + deriving (Show, Eq, Ord, Generic) + -- | A token bucket-based rate limiter -- -- This token limiter is thread-safe and guarantees that: @@ -67,56 +76,66 @@ tokenLimiterLastServiced <- newMVar (now, min (tokenLimitConfigInitialTokens tokenLimiterConfig) (tokenLimitConfigMaxTokens tokenLimiterConfig)) pure TokenLimiter {..} --- | Ask if we could debit a number of tokens, without actually doing it. --- --- Note that this information can become stale _very_ quickly. --- If you want to also actually debit a number of tokens, use 'tryDebit' instead. -canDebit :: TokenLimiter -> Word64 -> IO Bool -canDebit TokenLimiter {..} debit = withMVar tokenLimiterLastServiced $ \(lastServiced, countThen) -> do - now <- getMonotonicTimeNSec - let currentCount = computeCurrentCount tokenLimiterConfig lastServiced countThen now - let enoughAvailable = currentCount >= debit - pure enoughAvailable - -- | Check if we can debit a number of tokens, and do it if possible. -- -- The returned boolean represents whether the tokens were debited. +-- +-- Note that there is a small race-condition in which `tryDebit` sometimes +-- returns `False` eventhough it could (maybe) have debited because another +-- thread was currently `waitDebit`-ing without actually waiting (because it +-- didn't need to wait). tryDebit :: TokenLimiter -> Word64 -> IO Bool -tryDebit TokenLimiter {..} debit = modifyMVar tokenLimiterLastServiced $ \(lastServiced, countThen) -> do - now <- getMonotonicTimeNSec - let currentCount = computeCurrentCount tokenLimiterConfig lastServiced countThen now - let enoughAvailable = currentCount >= debit - if enoughAvailable - then do - let newCount = currentCount - debit - pure ((now, newCount), True) - else pure ((lastServiced, countThen), False) +tryDebit TokenLimiter {..} debit = + fmap (fromMaybe False) $ + tryModifyMVar tokenLimiterLastServiced $ \(lastServiced, countThen) -> do + now <- getMonotonicTimeNSec + let currentCount = computeCurrentCount tokenLimiterConfig lastServiced countThen now + let enoughAvailable = currentCount >= debit + pure $ + if enoughAvailable + then + let newCount = currentCount - debit + in ((now, newCount), True) + else ((lastServiced, countThen), False) --- | Wait until the given number of tokens can be debited -waitDebit :: TokenLimiter -> Word64 -> IO () -waitDebit TokenLimiter {..} debit = modifyMVar_ tokenLimiterLastServiced $ \(lastServiced, countThen) -> do +-- | Wait until the given number of tokens can be debited. +-- +-- Returns the time waited, for stats recording purposes. +-- +-- Note: only reports the time waited due to rate-limiting this specific action, +-- not the wall-clock time waited (that might include waiting for previous +-- actions limited by this rate limiter to finish). +-- +-- Note: debitor threads are serviced in FIFO order, so a request for a small +-- (and currently satisfiable) number of tokens can still be delayed by a debit +-- request for a larger amount of tokens. +-- +-- Note: the wait time reported can be inflated due to scheduling inaccuracy. +-- See https://gitlab.haskell.org/ghc/ghc/-/issues/16601. +waitDebit :: TokenLimiter -> Word64 -> IO (Maybe MonotonicDiffNanos) +waitDebit TokenLimiter {..} debit = modifyMVar tokenLimiterLastServiced $ \(lastServiced, countThen) -> do now <- getMonotonicTimeNSec let currentCount = computeCurrentCount tokenLimiterConfig lastServiced countThen now let enoughAvailable = currentCount >= debit if enoughAvailable then do let newCount = currentCount - debit - pure (now, newCount) + pure ((now, newCount), Nothing) else do let extraTokensNeeded = debit - currentCount let microsecondsToWaitDouble :: Double microsecondsToWaitDouble = 1_000_000 - -- fromIntegral :: Word64 -> Double - * fromIntegral extraTokensNeeded - -- fromIntegral :: Word64 -> Double - / fromIntegral (tokenLimitConfigTokensPerSecond tokenLimiterConfig) + * (fromIntegral :: Word64 -> Double) extraTokensNeeded + / (fromIntegral :: Word64 -> Double) (tokenLimitConfigTokensPerSecond tokenLimiterConfig) + 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 @@ -124,7 +143,7 @@ -- aren't. let currentCountAfterWaiting = computeCurrentCount tokenLimiterConfig lastServiced countThen nowAfterWaiting let newCount = currentCountAfterWaiting - debit - pure (nowAfterWaiting, newCount) + pure ((nowAfterWaiting, newCount), Just delta) -- | Compute the current number of tokens in a bucket purely. -- @@ -135,10 +154,8 @@ nanoDiff = now - lastServiced countToAddDouble :: Double countToAddDouble = - -- fromIntegral :: Word64 -> Double - fromIntegral nanoDiff - -- fromIntegral :: Word64 -> Double - * fromIntegral tokenLimitConfigTokensPerSecond + (fromIntegral :: Word64 -> Double) nanoDiff + * (fromIntegral :: Word64 -> Double) tokenLimitConfigTokensPerSecond / 1_000_000_000 countToAdd :: Word64 countToAdd = floor countToAddDouble @@ -149,3 +166,14 @@ in if willOverflow then tokenLimitConfigMaxTokens else min tokenLimitConfigMaxTokens totalCount + +tryModifyMVar :: MVar a -> (a -> IO (a, b)) -> IO (Maybe b) +tryModifyMVar m io = + mask $ \restore -> do + mA <- tryTakeMVar m + forM mA $ \a -> do + (a', b) <- + restore (io a) + `onException` putMVar m a + putMVar m a' + pure b 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 @@ -3,11 +3,14 @@ module Control.Concurrent.TokenLimiter.ConcurrentSpec (spec) where +import Control.Concurrent (threadDelay) import Control.Concurrent.Async +import Control.Concurrent.STM (atomically) +import Control.Concurrent.STM.TQueue import Control.Concurrent.TokenLimiter.Concurrent -import Control.Monad import Data.Word import GHC.Clock +import System.Timeout import Test.QuickCheck import Test.Syd import Test.Syd.Validity @@ -54,7 +57,7 @@ it "Can deal with overflow" $ do let config = TokenLimitConfig - { tokenLimitConfigInitialTokens = maxBound -1, + { tokenLimitConfigInitialTokens = maxBound - 1, tokenLimitConfigMaxTokens = maxBound, tokenLimitConfigTokensPerSecond = 1 } @@ -67,27 +70,6 @@ tokenLimiter <- makeTokenLimiter config tokenLimiterConfig tokenLimiter `shouldBe` config - describe "canDebit" $ do - it "is correct for this simple example" $ do - let config = - TokenLimitConfig - { tokenLimitConfigInitialTokens = 10, - tokenLimitConfigMaxTokens = 20, - tokenLimitConfigTokensPerSecond = 1 - } - limiter <- makeTokenLimiter config - canDebit limiter 10 `shouldReturn` True - - it "says true when tokens are available from the start and false otherwise" $ - forAllValid $ \config -> - forAllValid $ \needed -> do - -- We set the tokens per second to 1 because otherwise the generated - -- value can be large enough for the time difference between making - -- the token limiter and running 'canDebit' to matter. - limiter <- makeTokenLimiter config {tokenLimitConfigTokensPerSecond = 1} - couldDebit <- canDebit limiter needed - couldDebit `shouldBe` needed <= min (tokenLimitConfigInitialTokens config) (tokenLimitConfigMaxTokens config) - describe "tryDebit" $ do it "is correct for this simple example" $ do let config = @@ -109,12 +91,20 @@ didDebit <- tryDebit limiter needed didDebit `shouldBe` needed <= min (tokenLimitConfigInitialTokens config) (tokenLimitConfigMaxTokens config) - it "always says true if canDebit said true first and there are no other threads" $ - forAllValid $ \config -> - forAllValid $ \needed -> do - limiter <- makeTokenLimiter config - couldDebit <- canDebit limiter needed - when couldDebit $ tryDebit limiter needed `shouldReturn` True + it "does not block, even if another waitDebit is already happening" $ do + let config = + TokenLimitConfig + { tokenLimitConfigInitialTokens = 0, + tokenLimitConfigMaxTokens = 2, + tokenLimitConfigTokensPerSecond = 1 + } + limiter <- makeTokenLimiter config + concurrently_ (waitDebit limiter 2) $ do + threadDelay 100_000 -- Wait a bit to make sure tryDebit starts second. + mResult <- timeout 1_000_000 $ tryDebit limiter 1 + case mResult of + Nothing -> expectationFailure "tryDebit was blocking" + Just result -> result `shouldBe` False describe "waitDebit" $ do it "does not need to wait for this simple example" $ do @@ -175,10 +165,9 @@ nanos <- time_ $ concurrently_ (waitDebit limiter 1) (waitDebit limiter 1) nanos `shouldSatisfy` (>= 2_000_000_000) - modifyMaxSuccess (`div` 50) $ + modifyMaxSuccess (`div` 50) $ do it "Waits appropriately when there are many threads that want tokens at the same time in this example." $ do forAll (sized pure) $ \numberOfThreads -> do - print numberOfThreads let config = TokenLimitConfig { tokenLimitConfigInitialTokens = 0, @@ -188,6 +177,26 @@ limiter <- makeTokenLimiter config nanos <- time_ $ replicateConcurrently_ numberOfThreads (waitDebit limiter 1) nanos `shouldSatisfy` (>= fromIntegral numberOfThreads * 10_000_000) + + it "waits fairly" $ do + forAll (sized pure) $ \numberOfThreads -> do + queue <- newTQueueIO + let config = + TokenLimitConfig + { tokenLimitConfigInitialTokens = 0, + tokenLimitConfigMaxTokens = 1, + tokenLimitConfigTokensPerSecond = 50 + } + limiter <- makeTokenLimiter config + let l :: [Int] + l = [1 .. numberOfThreads] + -- The threads start waiting in order, so they need to wake up in order + forConcurrently_ l $ \ix -> do + threadDelay (10_000 * ix) + _ <- waitDebit limiter 1 + atomically $ writeTQueue queue ix + ixs <- atomically $ flushTQueue queue + ixs `shouldBe` l time_ :: IO a -> IO Word64 time_ func = snd <$> time func diff --git a/token-limiter-concurrent.cabal b/token-limiter-concurrent.cabal --- a/token-limiter-concurrent.cabal +++ b/token-limiter-concurrent.cabal @@ -1,20 +1,23 @@ cabal-version: 1.12 --- This file has been generated from package.yaml by hpack version 0.34.4. +-- 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.0.0.0 +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) 2021 Tom Sydney Kerckhove +copyright: Copyright (c) 2022-2024 Tom Sydney Kerckhove license: MIT license-file: LICENSE build-type: Simple +extra-source-files: + LICENSE + CHANGELOG.md source-repository head type: git @@ -48,6 +51,7 @@ , base >=4.7 && <5 , genvalidity , genvalidity-sydtest + , stm , sydtest , token-limiter-concurrent default-language: Haskell2010