token-limiter-concurrent (empty) → 0.0.0.0
raw patch · 5 files changed
+424/−0 lines, 5 filesdep +QuickCheckdep +asyncdep +base
Dependencies added: QuickCheck, async, base, genvalidity, genvalidity-sydtest, sydtest, token-limiter-concurrent
Files
- LICENSE +19/−0
- src/Control/Concurrent/TokenLimiter/Concurrent.hs +151/−0
- test/Control/Concurrent/TokenLimiter/ConcurrentSpec.hs +200/−0
- test/Spec.hs +1/−0
- token-limiter-concurrent.cabal +53/−0
+ LICENSE view
@@ -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.
+ src/Control/Concurrent/TokenLimiter/Concurrent.hs view
@@ -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:+--+-- * <https://en.wikipedia.org/wiki/Thundering_herd_problem There will be no thundering herd problem>+-- * <https://hackage.haskell.org/package/base-4.14.1.0/docs/Control-Concurrent-MVar.html#v:modifyMVar Fairness: Waiting processes will be serviced in a first-come first-service order.>+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
+ test/Control/Concurrent/TokenLimiter/ConcurrentSpec.hs view
@@ -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)
+ test/Spec.hs view
@@ -0,0 +1,1 @@+{-# OPTIONS_GHC -F -pgmF sydtest-discover #-}
+ token-limiter-concurrent.cabal view
@@ -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