diff --git a/Control/RateLimit.hs b/Control/RateLimit.hs --- a/Control/RateLimit.hs +++ b/Control/RateLimit.hs @@ -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 diff --git a/rate-limit.cabal b/rate-limit.cabal --- a/rate-limit.cabal +++ b/rate-limit.cabal @@ -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