Playing with Haskell Threads

Threads.hs

import Control.Concurrent
import Control.Exception
import Control.Monad
import System.IO
import Text.Printf
import Data.Maybe
import Data.Map as Map
import Data.Set as Set
import Data.String

main :: IO ()
main = do
  m <- initIdRef
  forever $ do
    s <- getLine
    let sl = words s in
      if length sl > 0 then
        if (sl !! 0) == "del" then do
          let time = (read $ sl !! 1 :: Int)
          count <- dropDelay m time
          printf "Dropped %d threads waiting for %d seconds\n" count time
        else do
          (num, id) <- setReminder m s
          trackThreadId m num id
      else
        return ()

-- ===========================================================
-- Track active thread ids

type IdMap = Map.Map Int (Set ThreadId)
data IdRef = IdRef (MVar IdMap)

initIdRef :: IO IdRef
initIdRef = do
  m <- newMVar $ Map.empty
  return $ IdRef m

trackThreadId :: IdRef -> Int -> ThreadId -> IO ()
trackThreadId (IdRef mv) num id = do
  m <- takeMVar mv
  let newm = Map.alter (\x -> if isNothing x then
                                Just $ Set.singleton id
                              else
                                Just $ Set.insert id $ fromJust x) num m
  putMVar mv newm

dropThreadId :: IdRef -> Int -> ThreadId -> IO ()
dropThreadId (IdRef mv) num id = do
  m <- takeMVar mv
  let newm = Map.alter (\x -> if isNothing x then
                                Nothing
                              else
                                Just $ Set.delete id $ fromJust x) num m
  printf "ThreadMap[%d]= %s\n" num $ show $ fromMaybe Set.empty $ Map.lookup num newm
  putMVar mv newm

dropDelay :: IdRef -> Int -> IO (Int)
dropDelay (IdRef mv) num = do
  m <- takeMVar mv
  case Map.lookup num m of
    Just existing -> do
      putMVar mv $ Map.delete num m
      mapM_ (\x -> throwTo x ThreadKilled) $ Set.toList existing
      return $ Set.size existing
    Nothing -> do
      putMVar mv m 
      return 0

-- ===========================================================
-- Fork delay

setReminder :: IdRef -> String -> IO (Int, ThreadId)
setReminder ref s = do
  let t = read s :: Int
  printf "Ok, I'll remind you in %d seconds\n" t
  id <- forkIO $ do
    threadDelay (10^6 * t)
    printf "%d seconds passed!\n" t
    mid <- myThreadId
    dropThreadId ref t mid
  return (t, id)