rate-limit 1.1.1 → 1.2.0

raw patch · 2 files changed

+99/−68 lines, 2 filesdep +stmdep ~basePVP ok

version bump matches the API change (PVP)

Dependencies added: stm

Dependency ranges changed: base

API changes (from Hackage documentation)

Files

• Control/RateLimit.hs +95/−66
• rate-limit.cabal +4/−2

@@ -1,4 +1,6 @@--- |This module implements rate-limiting functionality for Haskell programs.+{-# LANGUAGE ScopedTypeVariables #-}++-- | This module implements rate-limiting functionality for Haskell programs. -- Rate-limiting is useful when trying to control / limit access to a -- particular resource over time. For example, you might want to limit the -- rate at which you make requests to a server, as an administrator may block@@ -28,43 +30,43 @@ -- limiting functionality, but specialized versions of it are also exported -- for convenience. ---module Control.RateLimit(-         generateRateLimitedFunction-       , RateLimit(..)-       , ResultsCombiner-       , dontCombine-       , rateLimitInvocation-       , rateLimitExecution-       )- where+module Control.RateLimit (+    generateRateLimitedFunction+  , RateLimit(..)+  , ResultsCombiner+  , dontCombine+  , rateLimitInvocation+  , rateLimitExecution+  ) where  import Control.Concurrent-import Control.Concurrent.Chan-import Control.Monad(when)+import Control.Concurrent.STM+import Control.Monad (void, when)+import Data.Functor (($>)) import Data.Time.Units --- |The rate at which to limit an action.-data TimeUnit a => RateLimit a =-    PerInvocation a -- ^Rate limit the action to invocation once per time-                    --  unit. With this option, the time it takes for the-                    --  action to take place is not taken into consideration-                    --  when computing the rate, only the time between -                    --  invocations of the action. This may cause the action-                    --  to execute concurrently, as an invocation may occur-                    --  while an action is still running.-  | PerExecution a  -- ^Rate limit the action to execution once per time-                    --  unit. With this option, the time it takes for the-                    --  action to take plase is taken into account, and all-                    --  actions will necessarily occur sequentially. However,-                    --  if your action takes longer than the time unit given,-                    --  then the rate of execution will be slower than the-                    --  given unit of time.+-- | The rate at which to limit an action.+data RateLimit a+  = PerInvocation a -- ^ Rate limit the action to invocation once per time+                    -- unit. With this option, the time it takes for the+                    -- action to take place is not taken into consideration+                    -- when computing the rate, only the time between+                    -- invocations of the action. This may cause the action+                    -- to execute concurrently, as an invocation may occur+                    -- while an action is still running.+  | PerExecution a  -- ^ Rate limit the action to execution once per time+                    -- unit. With this option, the time it takes for the+                    -- action to take plase is taken into account, and all+                    -- actions will necessarily occur sequentially. However,+                    -- if your action takes longer than the time unit given,+                    -- then the rate of execution will be slower than the+                    -- given unit of time. --- |In some cases, if two requests are waiting to be run, it may be possible+-- | In some cases, if two requests are waiting to be run, it may be possible -- to combine them into a single request and thus increase the overall -- bandwidth. The rate limit system supports this, but requires a little -- additional information to make everything work out right. You may also--- need to do something a bit wonky with your types to make this work ... +-- need to do something a bit wonky with your types to make this work ... -- sorry. -- -- The basic idea is this: Given two requests, you can either return Nothing@@ -80,80 +82,107 @@ dontCombine :: ResultsCombiner a b dontCombine _ _ = Nothing --- |Rate limit the invocation of a given action. This is equivalent to calling+-- | Rate limit the invocation of a given action. This is equivalent to calling -- 'generateRateLimitedFunction' with a 'PerInvocation' rate limit and the -- 'dontCombine' combining function.-rateLimitInvocation :: TimeUnit t => -                       t -> (req -> IO resp) ->-                       IO (req -> IO resp)+rateLimitInvocation :: TimeUnit t+                    => t+                    -> (req -> IO resp)+                    -> IO (req -> IO resp) rateLimitInvocation pertime action =   generateRateLimitedFunction (PerInvocation pertime) action dontCombine --- |Rate limit the execution of a given action. This is equivalent to calling+-- | Rate limit the execution of a given action. This is equivalent to calling -- 'generateRateLimitedFunction' with a 'PerExecution' rate limit and the -- 'dontCombine' combining function.-rateLimitExecution :: TimeUnit t =>-                      t -> (req -> IO resp) ->-                      IO (req -> IO resp)+rateLimitExecution :: TimeUnit t+                   => t+                   -> (req -> IO resp)+                   -> IO (req -> IO resp) rateLimitExecution pertime action =   generateRateLimitedFunction (PerExecution pertime) action dontCombine --- |The most generic way to rate limit an invocation.-generateRateLimitedFunction :: TimeUnit t =>-  RateLimit t -- ^What is the rate limit for this action-  -> (req -> IO resp) -- ^What is the action you want to rate limit, given as an-                      --  a MonadIO function from requests to responses?-  -> ResultsCombiner req resp -- ^A function that can combine requests if rate-                              -- limiting happens. If you cannot combine two-                              -- requests into one request, we suggest using-                              -- 'dontCombine'.-  -> IO (req -> IO resp)-generateRateLimitedFunction ratelimit action combiner-  = do chan <- newChan-       forkIO $ runner (-42) chan-       return $ resultFunction chan- where+-- | The most generic way to rate limit an invocation.+generateRateLimitedFunction :: forall req resp t+                             . TimeUnit t+                            => RateLimit t+                               -- ^ What is the rate limit for this action+                            -> (req -> IO resp)+                               -- ^ What is the action you want to rate limit,+                               -- given as an a MonadIO function from requests+                               -- to responses?+                            -> ResultsCombiner req resp+                               -- ^ A function that can combine requests if+                               -- rate limiting happens. If you cannot combine+                               -- two requests into one request, we suggest+                               -- using 'dontCombine'.+                            -> IO (req -> IO resp)+generateRateLimitedFunction ratelimit action combiner = do+  chan <- atomically newTChan+  void $ forkIO $ runner (-42) chan+  return $ resultFunction chan++  where+  currentMicros :: IO Integer+  currentMicros = toMicroseconds `fmap` (getCPUTimeWithUnit :: IO Microsecond)++  runner :: Integer -> TChan (req, MVar resp) -> IO a   runner lastTime chan = do     -- should we wait for some amount of time?-    now <- toMicroseconds `fmap` (getCPUTimeWithUnit :: IO Microsecond)+    now <- currentMicros     when (now - lastTime < toMicroseconds (getRate ratelimit)) $ do       let delay = toMicroseconds (getRate ratelimit) - (now - lastTime)       threadDelay (fromIntegral delay)     -- OK, we're ready for the next item-    (req, respMV) <- readChan chan+    (req, respMV) <- atomically $ readTChan chan     let baseHandler resp = putMVar respMV resp     -- can we combine this with any other requests on the pipe?     (req', finalHandler) <- updateRequestWithFollowers chan req baseHandler     if shouldFork ratelimit       then forkIO (action req' >>= finalHandler) >> return ()       else action req' >>= finalHandler-    nextTime <- toMicroseconds `fmap` (getCPUTimeWithUnit :: IO Microsecond)+    nextTime <- currentMicros     runner nextTime chan+   -- updateRequestWithFollowers: We have one request. Can we combine it with   -- some other requests into a cohesive whole?+  updateRequestWithFollowers :: TChan (req, MVar resp)+                             -> req+                             -> (resp -> IO ())+                             -> IO (req, (resp -> IO ()))   updateRequestWithFollowers chan req handler = do-    isEmpty <- isEmptyChan chan+    isEmpty <- atomically $ isEmptyTChan chan     if isEmpty       then return (req, handler)-      else do (next, nextRespMV) <- readChan chan-              case combiner req next of-                Nothing -> do-                  unGetChan chan (next, nextRespMV)+      else do mCombinedAndMV <- atomically $ do+                tup@(next, nextRespMV) <- readTChan chan+                case combiner req next of+                  Nothing -> unGetTChan chan tup $> Nothing+                  Just combined -> return $ Just (combined, nextRespMV)++              case mCombinedAndMV of+                Nothing ->                   return (req, handler)-                Just (req', splitResponse) ->-                  updateRequestWithFollowers chan req' $ \ resp -> do+                Just ((req', splitResponse), nextRespMV) ->+                  updateRequestWithFollowers chan req' $ \resp -> do                     let (theirs, mine) = splitResponse resp                     putMVar nextRespMV mine                     handler theirs+   -- shouldFork: should we fork or execute the action in place?+  shouldFork :: RateLimit t -> Bool   shouldFork (PerInvocation _) = True   shouldFork (PerExecution _)  = False+   -- getRate: what is the rate of this action?-  getRate (PerInvocation x)    = x-  getRate (PerExecution  x)    = x+  getRate :: RateLimit t -> t+  getRate (PerInvocation x) = x+  getRate (PerExecution  x) = x+   -- resultFunction: the function (partially applied on the channel) that will   -- be returned from this monstrosity.+  resultFunction :: TChan (req, MVar resp) -> req -> IO resp   resultFunction chan req = do     respMV <- newEmptyMVar-    writeChan chan (req, respMV)+    atomically $ writeTChan chan (req, respMV)     takeMVar respMV

@@ -1,5 +1,5 @@ Name: rate-limit-Version: 1.1.1+Version: 1.2.0 Build-Type: Simple Cabal-Version: >= 1.6 License: BSD3@@ -17,7 +17,9 @@   is intended as a general-purpose mechanism for rate-limiting IO actions.   Library-  Build-Depends: base >= 4.0 && < 5.0, time-units >= 1.0 && < 2.0+  Build-Depends: base       >= 4.0 && < 5.0+               , stm        >= 2.4 && < 2.5+               , time-units >= 1.0 && < 2.0   Exposed-Modules: Control.RateLimit  source-repository head