ratelimiter (empty) → 0.1.0

raw patch · 6 files changed

+228/−0 lines, 6 filesdep +basedep +containersdep +extrasetup-changed

Dependencies added: base, containers, extra, mtl, ratelimiter, time, timespan, vector

Files

• LICENSE +30/−0
• README.md +24/−0
• Setup.hs +2/−0
• ratelimiter.cabal +63/−0
• src/Control/RateLimiter.hs +107/−0
• test/Spec.hs +2/−0

@@ -0,0 +1,30 @@+Copyright Alexander Thiemann <mail@thiemann.at> (c) 2020++All rights reserved.++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions are met:++    * Redistributions of source code must retain the above copyright+      notice, this list of conditions and the following disclaimer.++    * Redistributions in binary form must reproduce the above+      copyright notice, this list of conditions and the following+      disclaimer in the documentation and/or other materials provided+      with the distribution.++    * Neither the name of Author name here nor the names of other+      contributors may be used to endorse or promote products derived+      from this software without specific prior written permission.++THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT+OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

@@ -0,0 +1,24 @@+# Haskell: ratelimiter++A simple in-memory rate-limiter library.++## Usage++``` haskell+import Control.RateLimiter+import qualified Data.Vector as V++main :: IO+main =+  -- one rate limiter can have multiple rules+  do limiter <- +        newRateLimiter $ V.fromList+        [ RateLimitConfig (RollingWindow 60) 200 -- 200 per minute+        , RateLimitConfig (RollingWindow 3600) 400 -- 400 per hour+        ]+     let myRateLimitedFunction =+           do isLimited <- isRateLimited () limiter+              if isLimitd then pure Nothing else Just <$> someExpensiveWork+     -- ... use myRateLimitedFunction+```+

@@ -0,0 +1,2 @@+import Distribution.Simple+main = defaultMain

@@ -0,0 +1,63 @@+cabal-version: 1.12++-- This file has been generated from package.yaml by hpack version 0.33.0.+--+-- see: https://github.com/sol/hpack+--+-- hash: 8abbf5dee172728a76890aceb5dd74e24de46d5f252563c5a1a1a80e06444b12++name:           ratelimiter+version:        0.1.0+synopsis:       In-memory rate limiter+description:    An in-memory rate limiter implementation+category:       Web+homepage:       https://github.com/agrafix/ratelimiter#readme+bug-reports:    https://github.com/agrafix/ratelimiter/issues+author:         Alexander Thiemann <mail@thiemann.at>+maintainer:     Alexander Thiemann <mail@thiemann.at>+copyright:      2020 Alexander Thiemann <mail@thiemann.at>+license:        BSD3+license-file:   LICENSE+build-type:     Simple+extra-source-files:+    README.md++source-repository head+  type: git+  location: https://github.com/agrafix/ratelimiter++library+  exposed-modules:+      Control.RateLimiter+  other-modules:+      Paths_ratelimiter+  hs-source-dirs:+      src+  build-depends:+      base >=4.7 && <5+    , containers+    , extra+    , mtl+    , time+    , timespan+    , vector+  default-language: Haskell2010++test-suite ratelimiter-test+  type: exitcode-stdio-1.0+  main-is: Spec.hs+  other-modules:+      Paths_ratelimiter+  hs-source-dirs:+      test+  ghc-options: -threaded -rtsopts -with-rtsopts=-N+  build-depends:+      base >=4.7 && <5+    , containers+    , extra+    , mtl+    , ratelimiter+    , time+    , timespan+    , vector+  default-language: Haskell2010

@@ -0,0 +1,107 @@+{-# LANGUAGE OverloadedStrings #-}+module Control.RateLimiter+    ( WindowSize(..), RateLimitMode(..), RateLimitConfig(..)+    , RateLimiter, newRateLimiter+    , isRateLimited, withRetryRateLimiter+    )+where++import Control.Monad+import Control.Monad.Extra+import Control.Monad.Trans+import Data.IORef+import Data.Time+import Data.Time.TimeSpan+import qualified Data.Sequence as Seq+import qualified Data.Vector as V++data WindowSize+    = WsMinute+    | WsHour+    | WsSecond+    | WsDay+    deriving (Show, Eq, Enum, Bounded)++data RateLimitMode+    = RollingWindow !DiffTime+    | FixedWindow !WindowSize+    deriving (Show, Eq)++data RateLimitConfig+    = RateLimitConfig+    { rlc_mode :: !RateLimitMode+    , rlc_maximum :: !Int+    }++data Entry k+  = Entry+  { eTime :: !UTCTime+  , eKey :: !k+  }++data RateLimiterImpl k+    = RateLimiterImpl+    { rl_config :: !RateLimitConfig+    , rl_state :: !(IORef (Seq.Seq (Entry k)))+    }++newtype RateLimiter k+    = RateLimiter { _unRateLimiter :: V.Vector (RateLimiterImpl k) }++-- | Create a new rate limiter with a list of configurations+newRateLimiter :: V.Vector RateLimitConfig -> IO (RateLimiter k)+newRateLimiter cfgs =+    liftM RateLimiter $ forM cfgs $ \cfg ->+    do ref <- newIORef mempty+       pure $ RateLimiterImpl { rl_config = cfg, rl_state = ref }++-- | Check if a given key is rate limited. Use `()` if you don't need multiple keys+isRateLimited :: (Eq k, MonadIO m) => k -> RateLimiter k -> m Bool+isRateLimited key (RateLimiter rls) =+    anyM (isRateLimitedImpl key) $ V.toList rls++isRateLimitedImpl :: Eq k => MonadIO m => k -> RateLimiterImpl k -> m Bool+isRateLimitedImpl key rl =+    do now <- liftIO getCurrentTime+       (_, isLimited) <-+           liftIO $+           atomicModifyIORef' (rl_state rl) $ \xs ->+           do let (xs', bucketSize) =+                      currentBucketSize now key (rlc_mode $ rl_config rl) xs+              if bucketSize < rlc_maximum (rl_config rl)+                 then (xs' Seq.|> Entry now key, (bucketSize, False))+                 else (xs', (bucketSize, True))+       pure isLimited++-- | Retry action if rate limited after 1 second+withRetryRateLimiter :: Eq k => MonadIO m => k -> RateLimiter k -> m a -> m a+withRetryRateLimiter key rl action =+    do limited <- isRateLimited key rl+       if not limited+          then action+          else do liftIO $ sleepTS (seconds 1)+                  withRetryRateLimiter key rl action++currentBucketSize ::+    Eq k => UTCTime -> k -> RateLimitMode -> Seq.Seq (Entry k) -> (Seq.Seq (Entry k), Int)+currentBucketSize now key rlm times =+    let timeOfDay =+          timeToTimeOfDay (utctDayTime now)+        takeUntil =+            case rlm of+              RollingWindow dt ->+                  addUTCTime (fromRational . toRational $ (-1) * dt) now+              FixedWindow ws ->+                  case ws of+                    WsMinute ->+                        now {utctDayTime = timeOfDayToTime $ timeOfDay { todSec = 0}}+                    WsHour ->+                        now {utctDayTime = timeOfDayToTime $ timeOfDay { todMin = 0, todSec = 0}}+                    WsSecond ->+                        now { utctDayTime = fromInteger $ truncate (utctDayTime now) }+                    WsDay ->+                        now { utctDayTime = 0 }+        newSeq = Seq.takeWhileR (\el -> eTime el >= takeUntil) times+    in ( newSeq+       , Seq.length $ Seq.filter (\el -> eKey el == key) newSeq+       )

@@ -0,0 +1,2 @@+main :: IO ()+main = putStrLn "Test suite not yet implemented"