Refactor IOSim primitives for register/threadDelay

bolt12 · coot · commit a1dd35e0c646 · 2022-11-09T21:08:41.000+01:00
diff --git a/io-classes/src/Control/Monad/Class/MonadTimer.hs b/io-classes/src/Control/Monad/Class/MonadTimer.hs
@@ -54,6 +54,9 @@ class Monad m => MonadDelay m where
   threadDelay d   = void . atomically . awaitTimeout =<< newTimeout d
 
 class (MonadSTM m, MonadDelay m) => MonadTimer m where
+  -- | The type of the timeout handle, used with 'newTimeout', 'readTimeout',
+  -- 'updateTimeout' and 'cancelTimeout'.
+  --
   data Timeout m :: Type
 
   -- | Create a new timeout which will fire at the given time duration in
diff --git a/io-sim/src/Control/Monad/IOSim/Internal.hs b/io-sim/src/Control/Monad/IOSim/Internal.hs
@@ -119,10 +119,13 @@ labelledThreads threadMap =
       [] threadMap
 
 
--- | Timers mutable variables.  First one supports 'newTimeout' api, the second
--- one 'registerDelay'.
+-- | Timers mutable variables. Supports 'newTimeout' api, the second
+-- one 'registerDelay', the third one 'threadDelay'.
 --
-data TimerVars s = TimerVars !(TVar s TimeoutState) !(TVar s Bool)
+data TimerCompletionInfo s =
+       Timer !(TVar s TimeoutState)
+     | TimerRegisterDelay !(TVar s Bool)
+     | TimerThreadDelay !ThreadId
 
 -- | Internal state.
 --
@@ -136,7 +139,7 @@ data SimState s a = SimState {
        -- | current time
        curTime  :: !Time,
        -- | ordered list of timers
-       timers   :: !(OrdPSQ TimeoutId Time (TimerVars s)),
+       timers   :: !(OrdPSQ TimeoutId Time (TimerCompletionInfo s)),
        -- | list of clocks
        clocks   :: !(Map ClockId UTCTime),
        nextVid  :: !TVarId,     -- ^ next unused 'TVarId'
@@ -346,30 +349,72 @@ schedule !thread@Thread{
     NewTimeout d k ->
       {-# SCC "schedule.NewTimeout.2" #-} do
       !tvar  <- execNewTVar nextVid
-                            (Just $ "<<timeout-state " ++ show (unTimeoutId nextTmid) ++ ">>")
-                            TimeoutPending
-      !tvar' <- execNewTVar (succ nextVid)
-                            (Just $ "<<timeout " ++ show (unTimeoutId nextTmid) ++ ">>")
-                            False
+                           (Just $ "<<timeout-state " ++ show (unTimeoutId nextTmid) ++ ">>")
+                           TimeoutPending
       let !expiry  = d `addTime` time
-          !t       = Timeout tvar tvar' nextTmid
-          !timers' = PSQ.insert nextTmid expiry (TimerVars tvar tvar') timers
+          !t       = Timeout tvar nextTmid
+          !timers' = PSQ.insert nextTmid expiry (Timer tvar) timers
           !thread' = thread { threadControl = ThreadControl (k t) ctl }
-      !trace <- schedule thread' simstate { timers   = timers'
-                                          , nextVid  = succ (succ nextVid)
-                                          , nextTmid = succ nextTmid }
+      trace <- schedule thread' simstate { timers   = timers'
+                                         , nextVid  = succ nextVid
+                                         , nextTmid = succ nextTmid }
       return (SimTrace time tid tlbl (EventTimerCreated nextTmid nextVid expiry) trace)
 
+    RegisterDelay d k | d < 0 ->
+      {-# SCC "schedule.NewRegisterDelay" #-} do
+      !tvar <- execNewTVar nextVid
+                          (Just $ "<<timeout " ++ show (unTimeoutId nextTmid) ++ ">>")
+                          True
+      let !expiry  = d `addTime` time
+          !thread' = thread { threadControl = ThreadControl (k tvar) ctl }
+      trace <- schedule thread' simstate { nextVid = succ nextVid }
+      return (SimTrace time tid tlbl (EventRegisterDelayCreated nextTmid nextVid expiry) $
+              SimTrace time tid tlbl (EventRegisterDelayFired nextTmid) $
+              trace)
+
+    RegisterDelay d k ->
+      {-# SCC "schedule.NewRegisterDelay" #-} do
+      !tvar <- execNewTVar nextVid
+                          (Just $ "<<timeout " ++ show (unTimeoutId nextTmid) ++ ">>")
+                          False
+      let !expiry  = d `addTime` time
+          !timers' = PSQ.insert nextTmid expiry (TimerRegisterDelay tvar) timers
+          !thread' = thread { threadControl = ThreadControl (k tvar) ctl }
+      trace <- schedule thread' simstate { timers   = timers'
+                                         , nextVid  = succ nextVid
+                                         , nextTmid = succ nextTmid }
+      return (SimTrace time tid tlbl
+                (EventRegisterDelayCreated nextTmid nextVid expiry) trace)
+
+    ThreadDelay d k | d < 0 ->
+      {-# SCC "schedule.NewThreadDelay" #-} do
+      let !expiry  = d `addTime` time
+          !thread' = thread { threadControl = ThreadControl k ctl }
+      trace <- schedule thread' simstate
+      return (SimTrace time tid tlbl (EventThreadDelay expiry) $
+              SimTrace time tid tlbl EventThreadDelayFired $
+              trace)
+
+    ThreadDelay d k ->
+      {-# SCC "schedule.NewThreadDelay" #-} do
+      let !expiry  = d `addTime` time
+          !timers' = PSQ.insert nextTmid expiry (TimerThreadDelay tid) timers
+          !thread' = thread { threadControl = ThreadControl k ctl }
+      !trace <- deschedule Blocked thread' simstate { timers   = timers'
+                                                    , nextTmid = succ nextTmid }
+      return (SimTrace time tid tlbl (EventThreadDelay expiry) trace)
+
+
     -- we do not follow `GHC.Event` behaviour here; updating a timer to the past
     -- effectively cancels it.
-    UpdateTimeout (Timeout _tvar _tvar' tmid) d k | d < 0 ->
+    UpdateTimeout (Timeout _tvar tmid) d k | d < 0 ->
       {-# SCC "schedule.UpdateTimeout" #-} do
       let !timers' = PSQ.delete tmid timers
           !thread' = thread { threadControl = ThreadControl k ctl }
       trace <- schedule thread' simstate { timers = timers' }
       return (SimTrace time tid tlbl (EventTimerCancelled tmid) trace)
 
-    UpdateTimeout (Timeout _tvar _tvar' tmid) d k ->
+    UpdateTimeout (Timeout _tvar tmid) d k ->
       {-# SCC "schedule.UpdateTimeout" #-} do
           -- updating an expired timeout is a noop, so it is safe
           -- to race using a timeout with updating or cancelling it
@@ -387,12 +432,12 @@ schedule !thread@Thread{
       let thread' = thread { threadControl = ThreadControl k ctl }
       schedule thread' simstate
 
-    CancelTimeout (Timeout tvar _tvar' tmid) k ->
+    CancelTimeout (Timeout tvar tmid) k ->
       {-# SCC "schedule.CancelTimeout" #-} do
       let !timers' = PSQ.delete tmid timers
           !thread' = thread { threadControl = ThreadControl k ctl }
       !written <- execAtomically' (runSTM $ writeTVar tvar TimeoutCancelled)
-      (!wakeup, wokeby) <- threadsUnblockedByWrites written
+      (wakeup, wokeby) <- threadsUnblockedByWrites written
       mapM_ (\(SomeTVar var) -> unblockAllThreadsFromTVar var) written
       let (unblocked,
            simstate') = unblockThreads wakeup simstate
@@ -525,7 +570,7 @@ schedule !thread@Thread{
                                                (runIOSim action')
                                                (MaskFrame k maskst ctl)
                            , threadMasking = maskst' }
-      !trace <-
+      trace <-
         case maskst' of
           -- If we're now unmasked then check for any pending async exceptions
           Unmasked -> SimTrace time tid tlbl (EventDeschedule Interruptable)
@@ -723,30 +768,43 @@ reschedule !simstate@SimState{ threads, timers, curTime = time } =
 
         -- Reuse the STM functionality here to write all the timer TVars.
         -- Simplify to a special case that only reads and writes TVars.
-        !written <- execAtomically' (runSTM $ mapM_ timeoutAction fired)
-        (wakeup, wokeby) <- threadsUnblockedByWrites written
+        !written <- execAtomically' (runSTM $ mapM_ timeoutSTMAction fired)
+        (wakeupSTM, wokeby) <- threadsUnblockedByWrites written
         !_ <- mapM_ (\(SomeTVar tvar) -> unblockAllThreadsFromTVar tvar) written
 
-        let (unblocked,
-             simstate') = unblockThreads wakeup simstate
+            -- Check all fired threadDelays
+        let wakeupThreadDelay = [ tid | TimerThreadDelay tid <- fired ]
+            wakeup            = wakeupThreadDelay ++ wakeupSTM
+            (_, !simstate')   = unblockThreads wakeup simstate
+
         !trace <- reschedule simstate' { curTime = time'
                                        , timers  = timers' }
         return $
-          traceMany ([ (time', ThreadId [-1], Just "timer", EventTimerExpired tmid)
-                     | tmid <- tmids ]
+          traceMany ([ ( time', ThreadId [-1], Just "timer"
+                       , EventTimerFired tmid)
+                     | (tmid, Timer _) <- zip tmids fired ]
+                  ++ [ ( time', ThreadId [-1], Just "register delay timer"
+                       , EventRegisterDelayFired tmid)
+                     | (tmid, TimerRegisterDelay _) <- zip tmids fired ]
                   ++ [ (time', tid', tlbl', EventTxWakeup vids)
-                     | tid' <- unblocked
+                     | tid' <- wakeupSTM
                      , let tlbl' = lookupThreadLabel tid' threads
-                     , let Just vids = Set.toList <$> Map.lookup tid' wokeby ])
+                     , let Just vids = Set.toList <$> Map.lookup tid' wokeby ]
+                  ++ [ ( time', tid, Just "thread delay timer"
+                       , EventThreadDelayFired)
+                     | tid <- wakeupThreadDelay ])
                     trace
   where
-    timeoutAction (TimerVars var bvar) = do
+    timeoutSTMAction (Timer var) = do
       x <- readTVar var
       case x of
-        TimeoutPending   -> writeTVar var  TimeoutFired
-                         >> writeTVar bvar True
+        TimeoutPending   -> writeTVar var TimeoutFired
         TimeoutFired     -> error "MonadTimer(Sim): invariant violation"
         TimeoutCancelled -> return ()
+    timeoutSTMAction (TimerRegisterDelay var) = writeTVar var True
+    -- Note that 'threadDelay' is not handled via STM style wakeup, but rather
+    -- it's handled directly above with 'wakeupThreadDelay' and 'unblockThreads'
+    timeoutSTMAction (TimerThreadDelay _)     = return ()
 
 unblockThreads :: [ThreadId] -> SimState s a -> ([ThreadId], SimState s a)
 unblockThreads !wakeup !simstate@SimState {runqueue, threads} =
diff --git a/io-sim/src/Control/Monad/IOSim/Types.hs b/io-sim/src/Control/Monad/IOSim/Types.hs
@@ -150,9 +150,13 @@ data SimA s a where
   SetWallTime  ::  UTCTime -> SimA s b  -> SimA s b
   UnshareClock :: SimA s b -> SimA s b
 
-  NewTimeout   :: DiffTime -> (Timeout (IOSim s) -> SimA s b) -> SimA s b
-  UpdateTimeout:: Timeout (IOSim s) -> DiffTime -> SimA s b -> SimA s b
-  CancelTimeout:: Timeout (IOSim s) -> SimA s b -> SimA s b
+  RegisterDelay :: DiffTime -> (TVar s Bool -> SimA s b) -> SimA s b
+
+  ThreadDelay :: DiffTime -> SimA s b -> SimA s b
+
+  NewTimeout    :: DiffTime -> (Timeout (IOSim s) -> SimA s b) -> SimA s b
+  UpdateTimeout :: Timeout (IOSim s) -> DiffTime -> SimA s b -> SimA s b
+  CancelTimeout :: Timeout (IOSim s) -> SimA s b -> SimA s b
 
   Throw        :: Thrower -> SomeException -> SimA s a
   Catch        :: Exception e =>
@@ -542,17 +546,17 @@ unshareClock :: IOSim s ()
 unshareClock = IOSim $ oneShot $ \k -> UnshareClock (k ())
 
 instance MonadDelay (IOSim s) where
-  -- Use default in terms of MonadTimer
+  -- Use optimized IOSim primitive
+  threadDelay d = IOSim $ oneShot $ \k -> ThreadDelay d (k ())
 
 instance MonadTimer (IOSim s) where
-  data Timeout (IOSim s) = Timeout !(TVar s TimeoutState) !(TVar s Bool) !TimeoutId
-                         -- ^ a timeout; we keep both 'TVar's to support
-                         -- `newTimer` and 'registerTimeout'.
+  data Timeout (IOSim s) = Timeout !(TVar s TimeoutState) !TimeoutId
+                         -- ^ a timeout
                          | NegativeTimeout !TimeoutId
                          -- ^ a negative timeout
 
-  readTimeout (Timeout var _bvar _key) = MonadSTM.readTVar var
-  readTimeout (NegativeTimeout _key)   = pure TimeoutCancelled
+  readTimeout (Timeout var _key)     = MonadSTM.readTVar var
+  readTimeout (NegativeTimeout _key) = pure TimeoutCancelled
 
   newTimeout      d = IOSim $ oneShot $ \k -> NewTimeout      d k
   updateTimeout t d = IOSim $ oneShot $ \k -> UpdateTimeout t d (k ())
@@ -563,7 +567,7 @@ instance MonadTimer (IOSim s) where
     | d == 0    = return Nothing
     | otherwise = do
         pid <- myThreadId
-        t@(Timeout _ _ tid) <- newTimeout d
+        t@(Timeout _ tid) <- newTimeout d
         handleJust
           (\(TimeoutException tid') -> if tid' == tid
                                          then Just ()
@@ -579,7 +583,7 @@ instance MonadTimer (IOSim s) where
                   throwTo pid' AsyncCancelled)
             (\_ -> Just <$> action)
 
-  registerDelay d = IOSim $ oneShot $ \k -> NewTimeout d (\(Timeout _var bvar _) -> k bvar)
+  registerDelay d = IOSim $ oneShot $ \k -> RegisterDelay d k
 
 newtype TimeoutException = TimeoutException TimeoutId deriving Eq
 
@@ -806,10 +810,16 @@ data SimEventType
                        (Maybe Effect)    -- effect performed (only for `IOSimPOR`)
   | EventTxWakeup      [Labelled TVarId] -- changed vars causing retry
 
-  | EventTimerCreated   TimeoutId TVarId Time
-  | EventTimerUpdated   TimeoutId        Time
-  | EventTimerCancelled TimeoutId
-  | EventTimerExpired   TimeoutId
+  | EventThreadDelay        Time
+  | EventThreadDelayFired
+
+  | EventRegisterDelayCreated TimeoutId TVarId Time
+  | EventRegisterDelayFired TimeoutId
+
+  | EventTimerCreated         TimeoutId TVarId Time
+  | EventTimerUpdated         TimeoutId        Time
+  | EventTimerCancelled       TimeoutId
+  | EventTimerFired           TimeoutId
 
   -- the following events are inserted to mark the difference between
   -- a failed trace and a similar passing trace of the same action
diff --git a/io-sim/src/Control/Monad/IOSimPOR/Internal.hs b/io-sim/src/Control/Monad/IOSimPOR/Internal.hs
@@ -16,6 +16,7 @@
 {-# LANGUAGE StandaloneDeriving         #-}
 {-# LANGUAGE TypeApplications           #-}
 {-# LANGUAGE TypeFamilies               #-}
+{-# LANGUAGE StandaloneDeriving         #-}
 
 {-# OPTIONS_GHC -Wno-orphans            #-}
 -- incomplete uni patterns in 'schedule' (when interpreting 'StmTxCommitted')
@@ -456,7 +457,7 @@ schedule thread@Thread{
                            False
       modifySTRef (tvarVClock tvar') (leastUpperBoundVClock vClock)
       let expiry  = d `addTime` time
-          t       = Timeout tvar tvar' nextTmid
+          t       = Timeout tvar nextTmid
           timers' = PSQ.insert nextTmid expiry (TimerVars tvar tvar') timers
           thread' = thread { threadControl = ThreadControl (k t) ctl }
       !trace <- schedule thread' simstate { timers   = timers'
@@ -466,13 +467,13 @@ schedule thread@Thread{
 
     -- we do not follow `GHC.Event` behaviour here; updating a timer to the past
     -- effectively cancels it.
-    UpdateTimeout (Timeout _tvar _tvar' tmid) d k | d < 0 -> do
+    UpdateTimeout (Timeout _tvar tmid) d k | d < 0 -> do
       let timers' = PSQ.delete tmid timers
           thread' = thread { threadControl = ThreadControl k ctl }
       trace <- schedule thread' simstate { timers = timers' }
       return (SimPORTrace time tid tstep tlbl (EventTimerCancelled tmid) trace)
 
-    UpdateTimeout (Timeout _tvar _tvar' tmid) d k -> do
+    UpdateTimeout (Timeout _tvar tmid) d k -> do
           -- updating an expired timeout is a noop, so it is safe
           -- to race using a timeout with updating or cancelling it
       let updateTimeout_  Nothing       = ((), Nothing)
@@ -488,7 +489,7 @@ schedule thread@Thread{
       let thread' = thread { threadControl = ThreadControl k ctl }
       schedule thread' simstate
 
-    CancelTimeout (Timeout tvar tvar' tmid) k -> do
+    CancelTimeout (Timeout tvar tmid) k -> do
       let timers' = PSQ.delete tmid timers
       written <- execAtomically' (runSTM $ writeTVar tvar TimeoutCancelled)
       (wakeup, wokeby) <- threadsUnblockedByWrites written
@@ -503,7 +504,6 @@ schedule thread@Thread{
           (unblocked,
            simstate') = unblockThreads vClock wakeup simstate
       modifySTRef (tvarVClock tvar)  (leastUpperBoundVClock vClock)
-      modifySTRef (tvarVClock tvar') (leastUpperBoundVClock vClock)
       !trace <- deschedule Yield thread' simstate' { timers = timers' }
       return $ SimPORTrace time tid tstep tlbl (EventTimerCancelled tmid)
              $ traceMany
@@ -946,7 +946,7 @@ reschedule simstate@SimState{ threads, timers, curTime = time, races } =
         !trace <- reschedule simstate'' { curTime = time'
                                         , timers  = timers' }
         let traceEntries =
-                     [ (time', ThreadId [-1], (-1), Just "timer", EventTimerExpired tmid)
+                     [ (time', ThreadId [-1], (-1), Just "timer", EventTimerFired tmid)
                      | tmid <- tmids ]
                   ++ [ (time', tid', (-1), tlbl', EventTxWakeup vids)
                      | tid' <- unblocked
