Timer coalescing: handle exceptions in timer coordination thread

This corrects my (wrong) assumption that the timer coordination thread
will only fail if there's an error in the code, and in that case any
attempt to recover is futile. It turns out that the thread does fail
recoverably in one notable case: when running in the non-threaded RTS,
registerDelay fails immediately. And we probably still wish for xmobar
to support the non-threaded RTS.

One way to solve this issue is to add a bunch of #ifdefs and compile the
code only in the threaded case. This would double the number of
configurations that need to be tested, though.

Instead, let's make the code robust against all kinds of exceptions in
the timer coordination thread, and get non-threaded RTS support for
free.

1 files changed, 68 insertions, 47 deletions

diff --git a/src/Xmobar/App/Timer.hs b/src/Xmobar/App/Timer.hs
index 8464fb7..cf59630 100644
--- a/src/Xmobar/App/Timer.hs
+++ b/src/Xmobar/App/Timer.hs
@@ -23,30 +23,27 @@ import Control.Concurrent.Async (withAsync)
 import Control.Concurrent.STM
 import Control.Exception
 import Control.Monad (forever, forM, guard)
-import Data.Foldable (foldrM)
-import Data.IORef
+import Data.Foldable (foldrM, for_)
 import Data.Int (Int64)
 import Data.Map (Map)
 import qualified Data.Map as M
-import Data.Maybe (isJust)
+import Data.Maybe (isJust, fromJust)
 import Data.Time.Clock.POSIX (getPOSIXTime)
 import Data.Unique
 import System.IO.Unsafe (unsafePerformIO)
 
-newtype Timer = Timer (TVar Periods)
-
 type Periods = Map Unique Period
 
-data Tick = Tick (TMVar ()) | TimeOut
+data Tick = Tick (TMVar ()) | UnCoalesce
 
 data Period = Period { rate :: Int64, next :: Int64, tick :: TMVar Tick }
 
-data TimeOutException = TimeOutException deriving Show
-instance Exception TimeOutException
+data UnCoalesceException = UnCoalesceException deriving Show
+instance Exception UnCoalesceException
 
-{-# NOINLINE timer #-}
-timer :: IORef (Maybe Timer)
-timer = unsafePerformIO (newIORef Nothing)
+{-# NOINLINE periodsVar #-}
+periodsVar :: TVar (Maybe Periods)
+periodsVar = unsafePerformIO $ newTVarIO Nothing
 
 now :: IO Int64
 now = do
@@ -68,24 +65,27 @@ newPeriod r = do
 -- second or when the next timer is due), coalescing is disabled for it and it
 -- falls back to periodic sleep.
 doEveryTenthSeconds :: Int -> IO () -> IO ()
-doEveryTenthSeconds r action = do
-    Just t <- readIORef timer
-    doEveryTenthSecondsCoalesced t r action `catch` \TimeOutException ->
+doEveryTenthSeconds r action =
+    doEveryTenthSecondsCoalesced r action `catch` \UnCoalesceException ->
         doEveryTenthSecondsSleeping r action
 
 -- | Perform a given action every N tenths of a second,
 -- coalesce with other timers using a given Timer instance.
-doEveryTenthSecondsCoalesced :: Timer -> Int -> IO () -> IO ()
-doEveryTenthSecondsCoalesced (Timer periodsVar) r action = do
+doEveryTenthSecondsCoalesced :: Int -> IO () -> IO ()
+doEveryTenthSecondsCoalesced r action = do
     (u, p) <- newPeriod (fromIntegral r)
     bracket_ (push u p) (pop u) $ forever $ bracket (wait p) done $ const action
     where
-        push u p = atomically $ modifyTVar periodsVar (M.insert u p)
-        pop u = atomically $ modifyTVar periodsVar (M.delete u)
+        push u p = atomically $ modifyTVar' periodsVar $ \case
+            Just periods -> Just $ M.insert u p periods
+            Nothing -> throw UnCoalesceException
+        pop u = atomically $ modifyTVar' periodsVar $ \case
+            Just periods -> Just $ M.delete u periods
+            Nothing -> Nothing
 
         wait p = atomically (takeTMVar $ tick p) >>= \case
             Tick doneVar -> return doneVar
-            TimeOut -> throwIO TimeOutException
+            UnCoalesce -> throwIO UnCoalesceException
         done doneVar = atomically $ putTMVar doneVar ()
 
 -- | Perform a given action every N tenths of a second,
@@ -105,25 +105,39 @@ tenthSeconds s | s >= x = do threadDelay (x * 100000)
                | otherwise = threadDelay (s * 100000)
                where x = (maxBound :: Int) `div` 100000
 
--- | Start the timer coordination thread.
+-- | Start the timer coordination thread and perform a given IO action (this
+-- is meant to surround the entire xmobar execution), terminating the timer
+-- thread afterwards.
+--
+-- Additionally, if the timer thread fails, individual
+-- 'doEveryTenthSecondsCoalesced' invocations that are waiting to be
+-- coordinated by it are notified to fall back to periodic sleeping.
+--
+-- The timer thread _will_ fail immediately when running in a non-threaded
+-- RTS.
 withTimer :: (IO () -> IO ()) -> IO a -> IO a
-withTimer pauseRefresh action = do
-    periodsVar <- atomically $ newTVar M.empty
-    withAsync (timerLoop pauseRefresh periodsVar) $ \_ ->
-        bracket_
-            (writeIORef timer (Just (Timer periodsVar)))
-            (writeIORef timer Nothing) -- TODO: kill all periods?
-            action
-
-timerLoop :: (IO () -> IO ()) -> TVar Periods -> IO ()
-timerLoop pauseRefresh periodsVar = forever $ do
+withTimer pauseRefresh action =
+    withAsync (timerThread `finally` cleanup) $ const action
+    where
+        timerThread = do
+            atomically $ writeTVar periodsVar $ Just M.empty
+            timerLoop pauseRefresh
+
+        cleanup = atomically $ readTVar periodsVar >>= \case
+            Just periods -> do
+                for_ periods unCoalesceTimer'
+                writeTVar periodsVar Nothing
+            Nothing -> return ()
+
+timerLoop :: (IO () -> IO ()) -> IO ()
+timerLoop pauseRefresh = forever $ do
     tNow <- now
     (toFire, tMaybeNext) <- atomically $ do
-        periods <- readTVar periodsVar
+        periods <- fromJust <$> readTVar periodsVar
         let toFire = timersToFire tNow periods
         let periods' = advanceTimers tNow periods
         let tMaybeNext = nextFireTime periods'
-        writeTVar periodsVar periods'
+        writeTVar periodsVar $ Just periods'
         return (toFire, tMaybeNext)
     pauseRefresh $ do
         -- To avoid multiple refreshes, pause refreshing for up to 1 second,
@@ -135,8 +149,8 @@ timerLoop pauseRefresh periodsVar = forever $ do
             Nothing -> 1000000
         fired <- fireTimers toFire
         timeouted <- waitForTimers timeoutVar fired
-        timeoutTimers timeouted periodsVar
-    delayUntilNextFire periodsVar
+        unCoalesceTimers timeouted
+    delayUntilNextFire
 
 advanceTimers :: Int64 -> Periods -> Periods
 advanceTimers t = M.map advance
@@ -168,20 +182,26 @@ waitForTimers timeoutVar fired = atomically $ do
     return [u | (u, False) <- dones]
 
 -- | Handle slow timers (drop and signal them to stop coalescing).
-timeoutTimers :: [Unique] -> TVar Periods -> IO ()
-timeoutTimers timers periodsVar = atomically $ do
-    periods <- readTVar periodsVar
-    periods' <- foldrM timeoutTimer periods timers
-    writeTVar periodsVar periods'
-
-timeoutTimer :: Unique -> Periods -> STM Periods
-timeoutTimer u periods = do
-    putTMVar (tick (periods M.! u)) TimeOut
+unCoalesceTimers :: [Unique] -> IO ()
+unCoalesceTimers timers = atomically $ do
+    periods <- fromJust <$> readTVar periodsVar
+    periods' <- foldrM unCoalesceTimer periods timers
+    writeTVar periodsVar $ Just periods'
+
+unCoalesceTimer :: Unique -> Periods -> STM Periods
+unCoalesceTimer u periods = do
+    unCoalesceTimer' (periods M.! u)
     return $ u `M.delete` periods
 
-delayUntilNextFire :: TVar Periods -> IO ()
-delayUntilNextFire periodsVar = do
-    tMaybeNext <- fmap nextFireTime $ readTVarIO periodsVar
+unCoalesceTimer' :: Period -> STM ()
+unCoalesceTimer' p = do
+    _ <- tryTakeTMVar (tick p)
+    putTMVar (tick p) UnCoalesce
+
+delayUntilNextFire :: IO ()
+delayUntilNextFire = do
+    Just periods <- readTVarIO periodsVar
+    let tMaybeNext = nextFireTime periods
     tNow <- now
     delayVar <- case tMaybeNext of
         Just tNext -> do
@@ -195,6 +215,7 @@ delayUntilNextFire periodsVar = do
         Nothing -> atomically $ newTVar False
     atomically $ do
         delayOver <- readTVar delayVar
-        tMaybeNext' <- fmap nextFireTime $ readTVar periodsVar
+        periods' <- fromJust <$> readTVar periodsVar
+        let tMaybeNext' = nextFireTime periods'
         -- Return also if a new period is added (it may fire sooner).
         guard $ delayOver || tMaybeNext /= tMaybeNext'