STMRace.hs

module STMRace where

import           Control.Concurrent                  ( threadDelay )
import           Control.Concurrent.Async            ( async
                                                     , cancel
                                                     , wait
                                                     )
import           Control.Concurrent.MVar             ( MVar
                                                     , newMVar
                                                     , withMVar
                                                     )
import           Control.Concurrent.STM.TMVar        ( TMVar
                                                     , newEmptyTMVarIO
                                                     , newTMVarIO
                                                     , putTMVar
                                                     , takeTMVar
                                                     )
import           Control.Concurrent.STM.TVar         ( TVar
                                                     , modifyTVar'
                                                     , newTVarIO
                                                     , readTVar
                                                     )
import           Control.Exception                   ( bracket )
import           Control.Monad                       ( void )
import           Control.Monad.STM
import           Data.Sequence                       ( (|>)
                                                     , Seq( (:<|)
                                                          , Empty
                                                          )
                                                     )
import qualified Data.Sequence                as Seq
import           System.IO.Unsafe                    ( unsafePerformIO )

main :: IO ()
main = do
  signal   <- newEmptyTMVarIO
  queue    <- newTVarIO Empty
  timer    <- async $ wakeup signal
  worker   <- async $ work signal
  producer <- async $ produce queue 0
  consumer <- async $ consume signal timer queue
  wait producer
  where
    productonDelay   = 1
    consumptionDelay = 2
    timeout          = 5

    sleep n = threadDelay $ n * 1000 * 1000

    produce qVar n = do
      sleep productonDelay
      atomically $ modifyTVar' qVar (|> n)
      display $ "Enqueued: " ++ show n
      produce qVar (n+1)

    consume signal timer qVar = do
      sleep consumptionDelay
      result <- atomically $ do
        q <- readTVar qVar
        if Seq.null q
          then pure Nothing
          else do
            let x :<| xs = q
            modifyTVar' qVar (const xs)
            pure $ Just x
      case result of
        Nothing -> do
          cancel timer
          atomically $ putTMVar signal ()
          timer' <- async $ wakeup signal
          display "Nothing to consume so I kicked the worker and reset the timer."
          consume signal timer' qVar
        Just n  -> do
          display $ "Consumed: " ++ show n
          consume signal timer qVar

    wakeup signal = do
      sleep timeout
      atomically $ putTMVar signal ()
      wakeup signal

    work signal = do
      void . atomically $ takeTMVar signal
      display "I'm doing the thing."
      work signal

display :: String -> IO ()
display s = withMVar displayLock $ const (putStrLn s)

{-# NOINLINE displayLock #-}
displayLock :: MVar ()
displayLock = unsafePerformIO $ newMVar ()