-- STM stress test

{-# OPTIONS -fffi #-}
module Main (main) where

import Control.Concurrent
import Control.Concurrent.STM
import System.Random
import Data.Array
import GHC.Conc 	( unsafeIOToSTM )
import Control.Monad	( when )
import System.IO
import System.IO.Unsafe
import System.Environment
import Foreign.C

-- | The number of bank accounts
n_accounts :: Int
n_accounts = 7

-- | The number of threads transferring money between accounts
n_actors :: Int
n_actors = 10

-- | The max initial number of monetary units in each account
init_credit :: Int
init_credit = 5

-- | The maximum size of a transfer
max_transfer :: Int
max_transfer = 3

-- | The maximum amount transferred by the source/sink thread
max_source :: Int
max_source = 3

max_transactions = 2000 :: Int

type Accounts = Array Int (TVar Int)

thread :: Int -> TVar Int -> Accounts -> IO ()
thread tid done accounts = loop max_transactions
 where loop 0 = atomically $ do x <- readTVar done; writeTVar done (x+1)
       loop n = do
	  src    <- randomRIO (1,n_accounts)
	  dst    <- randomRIO (1,n_accounts)
	  if (src == dst) then loop n else do
	  amount <- randomRIO (1,max_transfer)
	  start tid src dst amount
	  atomically_ tid  $ do
	    let src_acc = accounts ! src
		dst_acc = accounts ! dst
	    credit_src <- readTVar src_acc
	    when (credit_src < amount) retry
	    writeTVar src_acc (credit_src - amount)    
	    credit_dst <- readTVar dst_acc
	    writeTVar dst_acc (credit_dst + amount)    
	  loop (n-1)

start tid src dst amount = 
  puts ("start " ++ show tid ++ ' ':show src ++ ' ':show dst ++ ' ':show amount)

main = do
  hSetBuffering stdout LineBuffering

{-
  args <- getArgs
  case args of
   [n,m] -> let g = read (n ++ ' ':m) in setStdGen g
   []    -> do g <- getStdGen 
	       print g
-}

  -- for a deterministic run, we set the random seed explicitly:
  setStdGen (read "526454551 6356")

  -- HACK: the global commitVar requires atomically, so we want to seq it outside of
  -- an enclosing atomically (otherwise STM gets very confused).
  seq commitVar $ return ()

--  print n_actors
--  print n_accounts
  amounts <- sequence (take n_accounts (repeat (randomRIO (0,init_credit))))
--  mapM print amounts
  tvars <- atomically $ mapM newTVar amounts
  let accounts = listArray (1,n_accounts) tvars
  done <- atomically (newTVar 0)
  sequence [ forkIO (thread id done accounts) | id <- [1..n_actors] ]
  forkIO $ sourceSinkThread accounts
  atomically $ do
    x <- readTVar done
    when (x < n_actors) retry

sourceThreadId = 0 :: Int
sourceAccount  = 0 :: Int

-- A thread that alternates between dropping some cash into an account
-- (source), and removing some cash from an account (sink).
sourceSinkThread accounts = loop True
  where loop source = do
	   amount <- randomRIO (1,max_source)
	   acct   <- randomRIO (1,n_accounts)
	   if source
		then do start sourceThreadId sourceAccount acct amount
		        transfer acct amount
		else do start sourceThreadId acct sourceAccount amount
		        transfer acct (-amount)
	   loop (not source)

	transfer acct amount = do
	  let t = accounts ! acct
	  atomically_ sourceThreadId $ do
	    x <- readTVar t
	    writeTVar t (max 0 (x+amount)) -- never drop below zero, 
					   -- and don't block.

-- -----------------------------------------------------------------------------
-- Our tracing wrapper for atomically

{-# NOINLINE commitVar #-}
commitVar = unsafePerformIO $ atomically $ newTVar ([] :: [Int])

atomically_ :: Int -> STM a -> IO a
atomically_ tid stm = do
  r <- atomically $ do
    stmTrace ("execute " ++ show tid)
    r <- stm `orElse` do
		stmTrace ("retry " ++ show tid)
		retry
    c <- readTVar commitVar
    writeTVar commitVar (tid:c)
    return r

  atomically $ do
    c <- readTVar commitVar
    mapM stmTrace ["commit " ++ show tid | tid <- reverse c ]
    writeTVar commitVar []
  return r

stmTrace s = unsafeIOToSTM (puts s)

puts :: String -> IO ()
puts s = throwErrnoIfMinus1_ "puts" $ withCString s c_puts

foreign import ccall unsafe {-"puts"-} "strlen"
  c_puts :: CString -> IO CInt
