token-limiter-concurrent 0.0.0.0 → 0.1.0.0

raw patch · 5 files changed

+124/−70 lines, 5 filesdep +stmPVP ok

version bump matches the API change (PVP)

Dependencies added: stm

API changes (from Hackage documentation)

Files

• CHANGELOG.md +13/−0
• LICENSE +1/−1
• src/Control/Concurrent/TokenLimiter/Concurrent.hs +63/−35
• test/Control/Concurrent/TokenLimiter/ConcurrentSpec.hs +40/−31
• token-limiter-concurrent.cabal +7/−3

@@ -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`+

@@ -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

@@ -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

@@ -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

@@ -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