Haskell timeout with partial results based on CPU time

{-# LANGUAGE BangPatterns #-}

import Control.Concurrent
import Control.Monad
import Control.Exception
import System.CPUTime
import System.IO
import System.Mem


main :: IO ()
main = do
    hSetBuffering stdout NoBuffering
    hSetBuffering stderr NoBuffering
    loop 0

loop :: Int -> IO ()
loop n = do
    calculateAsMuchOfPiAsPossible
    when (n < 100) $ loop (n + 1)

calculateAsMuchOfPiAsPossible :: IO ()
calculateAsMuchOfPiAsPossible = do
    startCpu <- getCPUTime

    let timeLimitInMs = 40

    r <- runWithCpuTimeLimit timeLimitInMs $ \result -> do
        putMVar result []
        calculatePi (1, 0, 1, 1, 3, 3) result

    -- {--
    let showResult r = head digits : '.' : tail digits
            where digits = map (head . show) (reverse r)

    hPutStrLn stdout (showResult r)
    -- --}

    endCpu <- getCPUTime
    let elapsedCpuTime = (endCpu - startCpu) `div` (1000 * 1000 * 1000)

    when (elapsedCpuTime > timeLimitInMs) $ hPutStrLn stderr (show elapsedCpuTime ++ " ms")

-- Adapted from http://rosettacode.org/wiki/Pi#Haskell
calculatePi :: (Integer, Integer, Integer, Integer, Integer, Integer) -> MVar [Integer] -> IO ()
calculatePi (!q, !r, !t, !k, !n, !l) result = do
    {--
    Force rescheduling and faster timeout detection (if not,
    the default rescheduling timer has a 20 ms granularity).
    --}
    yield
    if 4*q+r-t < n*t
        then do
            ns <- takeMVar result
            putMVar result (n : ns)
            calculatePi (10*q, 10*(r-n*t), t, k, div (10*(3*q+r)) t - 10*n, l) result
        else
            calculatePi (q*k, (2*q+r)*l, t*l, k+1, div (q*(7*k+2)+r*l) (t*l), l+2) result

runWithCpuTimeLimit :: Integer -> (MVar a -> IO ()) -> IO a
runWithCpuTimeLimit delayInMs action = do

    result <- newEmptyMVar
    threadId <- forkIO $ action result

    done <- newEmptyMVar
    forkIO $ (getCPUTime >>= sleepUntilDelayElapsed delayInMs) `finally` (putMVar done ())
    takeMVar done

    r <- takeMVar result
    killThread threadId

    return r

sleepUntilDelayElapsed :: Integer -> Integer -> IO ()
sleepUntilDelayElapsed delayInMs t0 = do
    let
        sleepDuration = 1000 * 1000 * 1000 -- 1 ms
        picoToMicro x = x `div` (1000 * 1000)
        milliToPico x = x * 1000 * 1000 * 1000
    t <- getCPUTime
    when (t - t0 + sleepDuration < milliToPico delayInMs) $ do
        {--
        There is no guarantee that the thread will be rescheduled promptly
        when the delay has expired, but the thread will never continue to
        run earlier than specified.

        in http://hackage.haskell.org/package/base-4.8.1.0/docs/Control-Concurrent.html

        Could be a problem if we are rescheduled too late.
        --}
        threadDelay (fromIntegral $ picoToMicro sleepDuration)
        -- performGC
        sleepUntilDelayElapsed delayInMs t0