File: test-tb.hs

package info (click to toggle)
haskell-token-bucket 0.1.0.1-11
  • links: PTS, VCS
  • area: main
  • in suites: bullseye, sid
  • size: 132 kB
  • sloc: haskell: 132; ansic: 11; makefile: 9
file content (55 lines) | stat: -rw-r--r-- 1,490 bytes parent folder | download | duplicates (2)
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)"