Timeout manager for not threaded runtime. Work in progress

Author: knsd

Diff for: src/Control/Timeout.hs

@@ -23,7 +23,6 @@
 
 {-# LANGUAGE Trustworthy #-}
 {-# LANGUAGE RecursiveDo #-}
-{-# LANGUAGE DeriveDataTypeable #-}
 
 module Control.Timeout
     ( NominalDiffTime
@@ -33,30 +32,18 @@ module Control.Timeout
     , sleep
     ) where
 
-import Control.Exception (Exception)
-import Control.Concurrent (myThreadId, forkIO, killThread, threadDelay, throwTo,
+import Control.Concurrent (myThreadId, forkIO, killThread, throwTo,
                            rtsSupportsBoundThreads)
-import Data.Typeable (Typeable)
 import Data.Time.Clock (NominalDiffTime)
 import Data.Unique (newUnique)
-import GHC.Event (TimeoutKey, getSystemTimerManager, registerTimeout, unregisterTimeout)
+import GHC.Event (getSystemTimerManager, registerTimeout, unregisterTimeout)
 import Unsafe.Coerce (unsafeCoerce)
 
 import Control.Monad.Catch (MonadCatch(..), bracket, handleJust)
 import Control.Monad.Trans (MonadIO, liftIO)
 
--- | Exception used for timeout handling
-newtype Timeout = Timeout TimeoutKey
-    deriving (Eq, Typeable)
-
-instance Show Timeout where
-    show _ = "<<timeout>>"
-
-instance Exception Timeout
-
-timeToUsecs :: NominalDiffTime -> Int
-timeToUsecs t = floor $ (* 1000000) $ toRational t
-{-# INLINEABLE timeToUsecs #-}
+import Control.Timeout.EventManager.Types
+import Control.Timeout.Utils
 
 -- | Wrap an 'MonadIO' computation to time out and return @Nothing@ in case no result
 -- is available within @n@ seconds. In case a result
@@ -94,11 +81,11 @@ withTimeout t f | t <= 0 = return Nothing
     timer <- liftIO getSystemTimerManager
     ex@(Timeout key) <- liftIO $ mdo
         pid <- liftIO myThreadId
-        ex <- return . Timeout =<< (liftIO $ registerTimeout timer (timeToUsecs t) (throwTo pid ex))
+        ex <- fmap (Timeout . unsafeCoerce) (liftIO $ registerTimeout timer (timeToUsecs t) (throwTo pid ex))
         return ex
     handleJust (\e -> if e == ex then Just () else Nothing)
                (\_ -> return Nothing)
-               (f ex >>= \r -> (liftIO $ unregisterTimeout timer key) >> (return $ Just r))
+               (f ex >>= \r -> (liftIO $ unregisterTimeout timer (unsafeCoerce key)) >> (return $ Just r))
             | otherwise = do
     pid <- liftIO myThreadId
     ex  <- liftIO newUnique >>= return . Timeout . unsafeCoerce
@@ -107,9 +94,3 @@ withTimeout t f | t <= 0 = return Nothing
                (bracket (liftIO $ forkIO (sleep t >> throwTo pid ex))
                         (liftIO . killThread)
                         (\_ -> f ex >>= return . Just))
-
--- | Sleep for 'NominalDiffTime', example:
---
--- > sleep 5  -- Will sleep for 5 seconds
-sleep :: (MonadIO m) => NominalDiffTime -> m ()
-sleep = liftIO . threadDelay . timeToUsecs

Diff for: src/Control/Timeout/EventManager.hs

@@ -0,0 +1,26 @@
+module Control.Timeout.EventManager
+    ( Timeout(..)
+    , TimeoutKey
+    , registerTimeout
+    , unregisterTimeout
+    ) where
+
+import qualified GHC.Event as Event
+import Control.Concurrent (rtsSupportsBoundThreads)
+
+import Control.Timeout.EventManager.Types (Timeout)
+
+getTimerManager :: IO EventManager
+getTimerManager
+  | rtsSupportsBoundThreads = Event.getSystemTimerManager
+  | otherwise = undefined
+
+registerTimeout :: NominalDiffTime -> IO () -> IO Timeout
+registerTimeout t f = do
+    timer <- Event.getSystemTimerManager
+    fmap Timeout $ registerTimeout timer (timeToUsecs t) f
+
+unregisterTimeout :: Timeout -> IO ()
+unregisterTimeout (Timeout key) = do
+    timer <- Event.getSystemTimerManager
+    Event.unregisterTimeout timer key

Diff for: src/Control/Timeout/EventManager/Local.hs

@@ -0,0 +1,68 @@
+{-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE DeriveDataTypeable #-}
+
+module Control.Timeout.EventManager.Local () where
+
+import Control.Exception (Exception, handle)
+import Control.Concurrent (ThreadId, forkIO, throwTo, rtsSupportsBoundThreads)
+import Control.Monad (void)
+import Data.Time.Clock (NominalDiffTime, addUTCTime, diffUTCTime, getCurrentTime)
+import Data.Typeable (Typeable)
+import System.IO.Unsafe (unsafePerformIO)
+
+import Control.Monad.Trans (liftIO)
+import Control.Monad.Trans.State.Strict (evalStateT, get, put)
+
+import Control.Timeout.EventManager.Types (Timeout)
+import Control.Timeout.EventManager.PSQ (PSQ, Elem(..))
+import Control.Timeout.Utils (sleep)
+import qualified Control.Timeout.EventManager.PSQ as PSQ
+
+data Event = Register Timeout NominalDiffTime (IO ())
+           | Unregister Timeout
+    deriving (Typeable)
+
+instance Show Event where
+    show _ = "<<event>>"
+
+instance Exception Event
+
+type TimeoutQueue = PSQ (IO ())
+type EventManager = TimeoutQueue
+
+tick :: EventManager -> IO EventManager
+tick queue = handle eventHandler $ do
+    let mbNext = PSQ.findMin queue
+    case mbNext of
+        Just (E { prio, value }) -> do
+            now <- getCurrentTime
+            let diff = diffUTCTime prio now
+            sleep diff >> value
+            return $ PSQ.deleteMin queue
+        Nothing -> sleep 1 >> return queue
+  where
+    eventHandler (Register timeout time f) = do
+        now <- getCurrentTime
+        let time' = addUTCTime time now
+        return $ PSQ.insert timeout time' f queue
+    eventHandler (Unregister timeout) = return $ PSQ.delete timeout queue
+
+loop :: IO ()
+loop = void $ evalStateT go PSQ.empty
+  where
+    go = do
+        state <- get
+        new <- liftIO $ tick state
+        put new >> go
+
+managerThread :: ThreadId
+managerThread
+  | rtsSupportsBoundThreads = error "manager thread not running"
+  | otherwise = unsafePerformIO $ forkIO loop
+{-# NOINLINE managerThread #-}
+
+registerTimeout :: NominalDiffTime -> IO () -> IO Timeout
+registerTimeout = undefined
+
+unregisterTimeout :: Timeout -> IO ()
+unregisterTimeout t = throwTo managerThread $ Unregister t