diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,19 @@ +Copyright 2021 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 +the Software without restriction, including without limitation the rights to +use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies +of the Software, and to permit persons to whom the Software is furnished to do +so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/src/Control/Concurrent/TokenLimiter/Concurrent.hs b/src/Control/Concurrent/TokenLimiter/Concurrent.hs new file mode 100644 --- /dev/null +++ b/src/Control/Concurrent/TokenLimiter/Concurrent.hs @@ -0,0 +1,151 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE NumericUnderscores #-} +{-# LANGUAGE RecordWildCards #-} + +module Control.Concurrent.TokenLimiter.Concurrent + ( Count, + TokenLimitConfig (..), + MonotonicTime, + TokenLimiter (..), + makeTokenLimiter, + canDebit, + tryDebit, + waitDebit, + + -- * Helper functions + computeCurrentCount, + ) +where + +import Control.Concurrent +import Data.Word +import GHC.Clock +import GHC.Generics (Generic) +import Numeric.Natural + +-- | An amount of tokens +type Count = Word64 + +-- | A configuration for 'TokenLimiter' +data TokenLimitConfig = TokenLimitConfig + { -- | How many tokens should be in the bucket when it's created + tokenLimitConfigInitialTokens :: !Count, + -- | Maximum number of tokens the bucket can hold at any one time + tokenLimitConfigMaxTokens :: !Count, + -- | How many tokens are added to the bucket per second + tokenLimitConfigTokensPerSecond :: !Count + } + deriving (Show, Eq, Generic) + +-- | A type synonym for a number of "monotonic time" nanoseconds. +-- +-- This only exists because it is also a 'Word64' and would be too easy to confuse with a 'Count'. +type MonotonicTime = Word64 + +-- | A token bucket-based rate limiter +-- +-- This token limiter is thread-safe and guarantees that: +-- +-- * +-- * +data TokenLimiter = TokenLimiter + { tokenLimiterConfig :: !TokenLimitConfig, + -- | The last time the limiter was used, and what the token count was at that time. + -- + -- Not that this library assumes that you never put anything into this mvar + -- yourself and only use the functions in this library to interact with it. + tokenLimiterLastServiced :: !(MVar (MonotonicTime, Count)) + } + deriving (Eq, Generic) + +-- | Make a token limiter +-- +-- The initial number of tokens will be the minimum of the 'tokenLimitConfigInitialTokens' and the 'tokenLimitConfigMaxTokens', +makeTokenLimiter :: TokenLimitConfig -> IO TokenLimiter +makeTokenLimiter tokenLimiterConfig = do + now <- getMonotonicTimeNSec + 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. +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) + +-- | Wait until the given number of tokens can be debited +waitDebit :: TokenLimiter -> Word64 -> IO () +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) + else do + let extraTokensNeeded = debit - currentCount + let microsecondsToWaitDouble :: Double + microsecondsToWaitDouble = + 1_000_000 + -- fromIntegral :: Word64 -> Double + * fromIntegral extraTokensNeeded + -- fromIntegral :: Word64 -> Double + / fromIntegral (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 + -- 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) + +-- | Compute the current number of tokens in a bucket purely. +-- +-- You should not need this function. +computeCurrentCount :: TokenLimitConfig -> MonotonicTime -> Count -> MonotonicTime -> Count +computeCurrentCount TokenLimitConfig {..} lastServiced countThen now = + let nanoDiff :: Word64 + nanoDiff = now - lastServiced + countToAddDouble :: Double + countToAddDouble = + -- fromIntegral :: Word64 -> Double + fromIntegral nanoDiff + -- fromIntegral :: Word64 -> Double + * fromIntegral tokenLimitConfigTokensPerSecond + / 1_000_000_000 + countToAdd :: Word64 + countToAdd = floor countToAddDouble + totalPrecise :: Natural + totalPrecise = fromIntegral countThen + fromIntegral countToAdd + willOverflow = totalPrecise > fromIntegral (maxBound :: Word64) + totalCount = countThen + countToAdd + in if willOverflow + then tokenLimitConfigMaxTokens + else min tokenLimitConfigMaxTokens totalCount diff --git a/test/Control/Concurrent/TokenLimiter/ConcurrentSpec.hs b/test/Control/Concurrent/TokenLimiter/ConcurrentSpec.hs new file mode 100644 --- /dev/null +++ b/test/Control/Concurrent/TokenLimiter/ConcurrentSpec.hs @@ -0,0 +1,200 @@ +{-# LANGUAGE NumericUnderscores #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} + +module Control.Concurrent.TokenLimiter.ConcurrentSpec (spec) where + +import Control.Concurrent.Async +import Control.Concurrent.TokenLimiter.Concurrent +import Control.Monad +import Data.Word +import GHC.Clock +import Test.QuickCheck +import Test.Syd +import Test.Syd.Validity + +instance Validity TokenLimitConfig + +instance GenValid TokenLimitConfig where + genValid = genValidStructurallyWithoutExtraChecking + shrinkValid = shrinkValidStructurallyWithoutExtraFiltering + +spec :: Spec +spec = do + describe "computeCurrentCount" $ do + it "works in this super simple case" $ do + let config = + TokenLimitConfig + { tokenLimitConfigInitialTokens = 0, + tokenLimitConfigMaxTokens = 1, + tokenLimitConfigTokensPerSecond = 1 + } + -- One nanosecond later we shouldn't have any more tokens + computeCurrentCount config 0 0 1 `shouldBe` 0 + + it "Does not need to wait whole seconds to work" $ do + let config = + TokenLimitConfig + { tokenLimitConfigInitialTokens = 0, + tokenLimitConfigMaxTokens = 1000, + tokenLimitConfigTokensPerSecond = 1000 + } + -- After half a second we should have 500 tokens + computeCurrentCount config 0 0 500_000_000 `shouldBe` 500 + + it "Does not go over the maximum" $ do + let config = + TokenLimitConfig + { tokenLimitConfigInitialTokens = 0, + tokenLimitConfigMaxTokens = 10, + tokenLimitConfigTokensPerSecond = 10 + } + -- One nanosecond later we shouldn't have any more tokens + computeCurrentCount config 0 0 2_000_000_000 `shouldBe` 10 + + it "Can deal with overflow" $ do + let config = + TokenLimitConfig + { tokenLimitConfigInitialTokens = maxBound -1, + tokenLimitConfigMaxTokens = maxBound, + tokenLimitConfigTokensPerSecond = 1 + } + -- 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 "makeTokenLimiter" $ do + it "always succeeds" $ + forAllValid $ \config -> do + 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 = + TokenLimitConfig + { tokenLimitConfigInitialTokens = 10, + tokenLimitConfigMaxTokens = 20, + tokenLimitConfigTokensPerSecond = 1 + } + limiter <- makeTokenLimiter config + tryDebit 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 'tryDebit' to matter. + limiter <- makeTokenLimiter config {tokenLimitConfigTokensPerSecond = 1} + 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 + + describe "waitDebit" $ do + it "does not need to wait for this simple example" $ do + let config = + TokenLimitConfig + { tokenLimitConfigInitialTokens = 10, + tokenLimitConfigMaxTokens = 20, + tokenLimitConfigTokensPerSecond = 1 + } + limiter <- makeTokenLimiter config + nanos <- time_ $ waitDebit limiter 10 + nanos `shouldSatisfy` (< 1_000_000_000) + + it "Does not need to wait if tokens are available" $ + forAllValid $ \initial -> + forAllValid $ \maxTokens -> do + let config = + TokenLimitConfig + { tokenLimitConfigInitialTokens = initial, + tokenLimitConfigMaxTokens = maxTokens, + tokenLimitConfigTokensPerSecond = 1 + } + limiter <- makeTokenLimiter config + let needed = min initial maxTokens + nanos <- time_ $ waitDebit limiter needed + nanos `shouldSatisfy` (< 1_000_000_000) + + it "Waits appropriately when there is one threads that want tokens in this example." $ do + let config = + TokenLimitConfig + { tokenLimitConfigInitialTokens = 0, + tokenLimitConfigMaxTokens = 1, + tokenLimitConfigTokensPerSecond = 1 + } + limiter <- makeTokenLimiter config + nanos <- time_ $ waitDebit limiter 1 + nanos `shouldSatisfy` (>= 1_000_000_000) + + it "does not need to wait a whole number of seconds" $ do + let config = + TokenLimitConfig + { tokenLimitConfigInitialTokens = 0, + tokenLimitConfigMaxTokens = 10, + tokenLimitConfigTokensPerSecond = 10 + } + limiter <- makeTokenLimiter config + nanos <- time_ $ waitDebit limiter 1 + nanos `shouldSatisfy` (<= 500_000_000) + + it "Waits appropriately when there are multiple threads that want tokens at the same time in this example." $ do + let config = + TokenLimitConfig + { tokenLimitConfigInitialTokens = 0, + tokenLimitConfigMaxTokens = 1, + tokenLimitConfigTokensPerSecond = 1 + } + limiter <- makeTokenLimiter config + nanos <- time_ $ concurrently_ (waitDebit limiter 1) (waitDebit limiter 1) + nanos `shouldSatisfy` (>= 2_000_000_000) + + modifyMaxSuccess (`div` 50) $ + 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, + tokenLimitConfigMaxTokens = 1, + tokenLimitConfigTokensPerSecond = 100 + } + limiter <- makeTokenLimiter config + nanos <- time_ $ replicateConcurrently_ numberOfThreads (waitDebit limiter 1) + nanos `shouldSatisfy` (>= fromIntegral numberOfThreads * 10_000_000) + +time_ :: IO a -> IO Word64 +time_ func = snd <$> time func + +time :: IO a -> IO (a, Word64) +time func = do + begin <- getMonotonicTimeNSec + result <- func + end <- getMonotonicTimeNSec + pure (result, end - begin) diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,1 @@ +{-# OPTIONS_GHC -F -pgmF sydtest-discover #-} diff --git a/token-limiter-concurrent.cabal b/token-limiter-concurrent.cabal new file mode 100644 --- /dev/null +++ b/token-limiter-concurrent.cabal @@ -0,0 +1,53 @@ +cabal-version: 1.12 + +-- This file has been generated from package.yaml by hpack version 0.34.4. +-- +-- see: https://github.com/sol/hpack + +name: token-limiter-concurrent +version: 0.0.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 +license: MIT +license-file: LICENSE +build-type: Simple + +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 + hs-source-dirs: + src + build-depends: + base >=4.7 && <5 + default-language: Haskell2010 + +test-suite token-limiter-concurrent-test + type: exitcode-stdio-1.0 + main-is: Spec.hs + other-modules: + Control.Concurrent.TokenLimiter.ConcurrentSpec + Paths_token_limiter_concurrent + hs-source-dirs: + test + ghc-options: -threaded -rtsopts -with-rtsopts=-N + build-tool-depends: + sydtest-discover:sydtest-discover + build-depends: + QuickCheck + , async + , base >=4.7 && <5 + , genvalidity + , genvalidity-sydtest + , sydtest + , token-limiter-concurrent + default-language: Haskell2010