acme-missiles-0.3: Acme/Missiles/STM.hs

module Acme.Missiles.STM (
    -- * Launching missiles in the 'STM' monad
    withMissilesDo,
    launchMissilesSTM,
) where

import Acme.Missiles (launchMissiles)

import Control.Concurrent       (forkIO, killThread)
import Control.Concurrent.STM   (STM, atomically, retry,
                                 TVar, newTVarIO, readTVar, writeTVar)
import Control.Exception        (bracket)
import Control.Monad            (forever)
import System.IO.Unsafe         (unsafePerformIO)

-- | Perform initialization needed to launch missiles in the 'STM' monad.
withMissilesDo :: IO a -> IO a
withMissilesDo action =
        bracket (forkIO doLaunching)
                killThread
                (\_ -> action)
    where
        doLaunching = forever $ do
            atomically $ do
                n <- readTVar missileCommand
                if n > 0
                    then writeTVar missileCommand (n - 1)
                    else retry
            launchMissiles

-- | Launch missiles within an 'STM' computation.  Even if the memory
-- transaction is retried, only one salvo of missiles will be launched.
--
-- Example:
--
-- >import Acme.Missiles
-- >import Control.Concurrent
-- >import Control.Concurrent.STM
-- >
-- >main :: IO ()
-- >main = withMissilesDo $ do
-- >    xv <- atomically $ newTVar (2 :: Int)
-- >    yv <- atomically $ newTVar (1 :: Int)
-- >    atomically $ do
-- >        x <- readTVar xv
-- >        y <- readTVar yv
-- >        if x > y
-- >            then launchMissilesSTM
-- >            else return ()
-- >    threadDelay 100000
launchMissilesSTM :: STM ()
launchMissilesSTM = do
    n <- readTVar missileCommand
    writeTVar missileCommand $! n + 1

missileCommand :: TVar Integer
{-# NOINLINE missileCommand #-}
missileCommand = unsafePerformIO (newTVarIO 0)