WorkerPoolSpec.hs

module WorkerPoolSpec (spec) where

import           Control.Concurrent.STM
    ( atomically
    , check
    , modifyTVar'
    , newEmptyTMVarIO
    , newTQueueIO
    , newTVarIO
    , putTMVar
    , readTMVar
    , readTQueue
    , readTVar
    , readTVarIO
    , writeTQueue
    , writeTVar
    )
import           Control.Monad          (replicateM_, void, when)
import           Lib.WorkerPool         (startWorkerPool)
import           System.Timeout         (timeout)
import           Test.Hspec

spec :: Spec
spec = describe "WorkerPool" $ do
  it "limits concurrent jobs to the worker count" $ do
    queue <- newTQueueIO
    gate <- newTQueueIO
    stopVar <- newEmptyTMVarIO
    runningVar <- newTVarIO 0
    maxRunningVar <- newTVarIO 0

    let job = do
          atomically $ do
            running <- readTVar runningVar
            let running' = running + 1
            writeTVar runningVar running'
            maxRunning <- readTVar maxRunningVar
            when (running' > maxRunning) $
              writeTVar maxRunningVar running'
          atomically $ readTQueue gate
          atomically $ modifyTVar' runningVar (subtract 1)

    startWorkerPool 2 queue (void (readTMVar stopVar)) (const (pure ()))
    replicateM_ 4 (atomically $ writeTQueue queue job)

    started <- timeout 2000000 (atomically $ do
      maxRunning <- readTVar maxRunningVar
      check (maxRunning >= 2))
    started `shouldBe` Just ()

    replicateM_ 4 (atomically $ writeTQueue gate ())
    done <- timeout 2000000 (atomically $ do
      running <- readTVar runningVar
      check (running == 0))
    done `shouldBe` Just ()

    atomically $ putTMVar stopVar ()
    maxRunning <- readTVarIO maxRunningVar
    maxRunning `shouldBe` 2