1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
|
{-# LANGUAGE BangPatterns #-}
module Main where
import Control.Concurrent.TokenBucket
import Control.Concurrent
import Control.Exception
import Control.Monad
import Data.Time.Clock.POSIX (getPOSIXTime)
import System.Exit
import Data.Word
getPosixTime :: IO Double
getPosixTime = fmap realToFrac getPOSIXTime
toInvRate :: Double -> Word64
toInvRate r = round (1e6 / r)
timeIO :: IO a -> IO (Double, a)
timeIO act = do
ts0 <- getPosixTime
res <- act
ts1 <- getPosixTime
dt <- evaluate (ts1-ts0)
return (dt,res)
timeIO_ :: IO a -> IO Double
timeIO_ = fmap fst . timeIO
main :: IO ()
main = runInUnboundThread $ do
putStrLn "testing tocket-bucket..."
replicateM_ 3 $ do
check 10 10.0
check 20 20.0
check 50 50.0
check 100 100.0
check 200 200.0
check 500 500.0
check 1000 1000.0
putStrLn "============================================="
where
check :: Int -> Double -> IO ()
check n rate = do
-- threadDelay 100000
putStrLn $ "running "++show n++"+1 iterations with "++show rate++" Hz rate-limit..."
!tb <- newTokenBucket
dt <- timeIO_ (replicateM_ (n+1) $ (tokenBucketWait tb 1 (toInvRate rate)))
let rate' = fromIntegral n/dt
unless (rate' <= rate) $ do
putStrLn $ "...FAILED! (effective rate was " ++ show rate' ++ " Hz)"
exitFailure
putStrLn $ "...PASSED (effective rate was " ++ show rate' ++ " Hz)"
|
