Ticket #1657: threadbug.hs

{-# OPTIONS_GHC -fglasgow-exts #-}
module Main where

import Control.Concurrent
import Control.Exception
import System.IO.Unsafe (unsafeInterleaveIO)
import Data.Dynamic (Typeable)
import Data.IORef

delay :: Int
delay = 10000

runFor :: Int
runFor = 1000

compute :: Int -> Int -> IO [Int]
compute max n
        | n > max = do
                putStrLn ("# complete!")
                return []
        | otherwise = do
                putStrLn ("# computing " ++ show n)
                --threadDelay delay
                ns <- unsafeInterleaveIO $ compute max (n+1)
                return (n:ns)

data BGState a = BGS
        { bgsIsRunning          :: MVar ()
        , bgsIsComplete         :: MVar () -- write-once
        , bgsResult             :: MVar [a]
        , bgsSrc                :: IORef [a] -- worker thread only
        , bgsIter               :: IORef Int -- main thread only
        }

data Suspend = Suspend Int deriving Typeable

newBGS :: [a] -> IO (BGState a)
newBGS computation = do
        isRunning <- newEmptyMVar
        isComplete <- newEmptyMVar
        result <- newEmptyMVar
        src <- newIORef computation
        iter <- newIORef 0
        
        return $ BGS isRunning isComplete result src iter

runBGSFor :: Int -> BGState a -> IO (Maybe [a])
runBGSFor us state = do
        done <- tryTakeMVar (bgsIsComplete state)
        iter <- readIORef (bgsIter state)
        
        case (iter, done) of
                (10,_) -> return Nothing
                (_,Just ()) -> return Nothing
                _ -> do
                        putStrLn "*** starting!"
                        thread <- forkIO $ runBGSThread state
                        
                        -- verify that we are inside "catch" block before
                        -- advancing further
                        takeMVar (bgsIsRunning state)
                        
                        -- wait
                        threadDelay us
                
                        -- suspend background computation
                        n <- readIORef (bgsIter state)
                        modifyIORef (bgsIter state) (+1)
                        putStrLn ("*** throwing " ++ show n)
                        throwDynTo thread (Suspend n)
                        result <- takeMVar (bgsResult state)
                        return $ Just result

runBGSThread :: BGState a -> IO ()
runBGSThread state = do
        computation <- readIORef (bgsSrc state)
        
        -- locally we store our results here
        ans <- newIORef ([], computation)
        
        let loop = do
                (out, inp) <- readIORef ans
                case inp of
                        [] -> putMVar (bgsIsComplete state) ()
                        (x:xs) -> x `seq` (writeIORef ans (x:out, xs) >> loop)
        let run = do
                putStrLn "*** running!"
                putMVar (bgsIsRunning state) ()
                loop
        let handler (Suspend n) = do
                putStrLn ("*** caught " ++ show n)
                return ()
        
        catchDyn run handler
        
        -- get the work we've done so far
        (result, rest) <- readIORef ans
        writeIORef (bgsSrc state) rest
        
        -- write output
        putMVar (bgsResult state) (reverse result)

main :: IO ()
main = do
        computation <- unsafeInterleaveIO $ compute (10^10) 0
        bgState <- newBGS computation
        
        let iter = do
                x <- runBGSFor runFor bgState
                case x of Nothing -> return ()
                          Just [] -> putStrLn "^ Got  0" >> iter
                          Just xs -> do
                                putStrLn ("^ Got  " ++ show (length xs)) 
                                putStrLn ("^ Last " ++ show (last xs))
                                iter
        iter