ResourceSpec.hs

{-# LANGUAGE ScopedTypeVariables #-}

module Z.IO.ResourceSpec where

import           Control.Concurrent
import           Control.Exception
import           Control.Monad
import           Z.Data.PrimRef.PrimIORef
import           Z.IO.Resource          as R
import           Test.Hspec
import           Test.HUnit

data WorkerException = WorkerException deriving (Show)

instance Exception WorkerException

spec :: Spec
spec = describe "resource tests" $ do
    it "resource pool" $ do
        resCounter <- newCounter 0
        workerCounter <- newCounter 0
        let res = initResource (atomicAddCounter_ resCounter 1)
                               (\ _ -> atomicSubCounter_ resCounter 1)
            resPool = initPool res 100 1
        R.withResource resPool $ \ pool -> do
            forM_ [1..200] $ \ k -> forkIO. R.withResourceInPool pool $ \ i -> do
                atomicAddCounter_ workerCounter 1
                r <- readPrimIORef resCounter
                assertEqual "pool should limit max usage" True (r <= 100)
                threadDelay 100000

            threadDelay 1000000
            -- first 100 worker quickly get resources
            -- then hold for 1s, rest 100 worker have to wait, and so on
            -- so here we wait for 5s to make sure every worker got a resource
            -- we used to use replicateConcurrently_ from async, but it's
            -- not really neccessary

            r <- readPrimIORef resCounter
            assertEqual "pool should keep returned resources alive" 100 r

            s <- poolStat pool
            assertEqual "pool should be scanning returned resources" PoolScanning s

            threadDelay 5000000  -- after 5s, 200 thread should release all resources

            w <- readPrimIORef workerCounter
            assertEqual "worker should be able to get resource" 200 w

            r <- readPrimIORef resCounter
            assertEqual "pool should reap unused resources" 0 r

            s <- poolStat pool
            assertEqual "pool should stop scanning returned resources" PoolEmpty s

            -- Let's test again

            writePrimIORef workerCounter 0

            forM_ [1..200] $ \ k -> forkIO. R.withResourceInPool pool $ \ i -> do
                atomicAddCounter_ workerCounter 1
                r <- readPrimIORef resCounter
                assertEqual "pool should limit max usage" True (r <= 100)
                threadDelay 100000


            threadDelay 1000000

            r <- readPrimIORef resCounter
            assertEqual "pool should keep returned resources alive" 100 r

            s <- poolStat pool
            assertEqual "pool should be scanning returned resources" PoolScanning s

            threadDelay 5000000

            w <- readPrimIORef workerCounter
            assertEqual "worker should be able to get resource" 200 w

            r <- readPrimIORef resCounter
            assertEqual "pool should reap unused resources" 0 r

            s <- poolStat pool
            assertEqual "pool should stop scanning returned resources" PoolEmpty s

    it "resource pool under exceptions" $ do
        resCounter <- newCounter 0
        let res = initResource (atomicAddCounter' resCounter 1)
                               (\ _ -> atomicSubCounter_ resCounter 1)
            resPool = initPool res 100 1
        R.withResource resPool $ \ pool -> do

            forM_ [1..200] $ \ k -> forkIO. R.withResourceInPool pool $ \ i -> do
                r <- readPrimIORef resCounter
                threadDelay 100000
                when (even i) (throwIO WorkerException)
                assertEqual "pool should limit max usage" True (r <= 100)

            threadDelay 1000000

            s <- poolStat pool
            assertEqual "pool should be scanning returned resources" PoolScanning s

            threadDelay 5000000

            r <- readPrimIORef resCounter
            assertEqual "pool should reap unused resources" 0 r

            s <- poolStat pool
            assertEqual "pool should stop scanning returned resources" PoolEmpty s