Add yield to MonadFork

io-classes/src/Control/Monad/Class/MonadFork.hs

@@ -38,6 +38,8 @@ class MonadThread m => MonadFork m where
   killThread       :: ThreadId m -> m ()
   killThread tid = throwTo tid ThreadKilled
 
+  yield            :: m ()
+
 fork :: MonadFork m => m () -> m (ThreadId m)
 fork = forkIO
 {-# DEPRECATED fork "use forkIO" #-}
@@ -57,6 +59,7 @@ instance MonadFork IO where
   forkIOWithUnmask = IO.forkIOWithUnmask
   throwTo          = IO.throwTo
   killThread       = IO.killThread
+  yield            = IO.yield
 
 instance MonadThread m => MonadThread (ReaderT r m) where
   type ThreadId (ReaderT r m) = ThreadId m
@@ -70,6 +73,7 @@ instance MonadFork m => MonadFork (ReaderT e m) where
                            restore' (ReaderT f) = ReaderT $ restore . f
                        in runReaderT (k restore') e
   throwTo e t = lift (throwTo e t)
+  yield       = lift yield
 
 -- | Apply the label to the current thread
 labelThisThread :: MonadThread m => String -> m ()

io-sim/src/Control/Monad/IOSim/Internal.hs

@@ -562,6 +562,10 @@ schedule !thread@Thread{
           return $ SimTrace time tid tlbl (EventThrowTo e tid')
                  $ trace
 
+    YieldSim k -> do
+      let thread' = thread { threadControl = ThreadControl k ctl }
+      deschedule Yield thread' simstate
+
     -- ExploreRaces is ignored by this simulator
     ExploreRaces k ->
       {-# SCC "schedule.ExploreRaces" #-}

io-sim/src/Control/Monad/IOSim/Types.hs

@@ -164,6 +164,8 @@ data SimA s a where
   SetMaskState :: MaskingState  -> IOSim s a -> (a -> SimA s b) -> SimA s b
   GetMaskState :: (MaskingState -> SimA s b) -> SimA s b
 
+  YieldSim     :: SimA s a -> SimA s a
+
   ExploreRaces :: SimA s b -> SimA s b
 
   Fix          :: (x -> IOSim s x) -> (x -> SimA s r) -> SimA s r
@@ -371,6 +373,7 @@ instance MonadFork (IOSim s) where
   forkIO task        = IOSim $ oneShot $ \k -> Fork task k
   forkIOWithUnmask f = forkIO (f unblock)
   throwTo tid e      = IOSim $ oneShot $ \k -> ThrowTo (toException e) tid (k ())
+  yield              = IOSim $ oneShot $ \k -> YieldSim (k ())
 
 instance MonadTest (IOSim s) where
   exploreRaces       = IOSim $ oneShot $ \k -> ExploreRaces (k ())

io-sim/src/Control/Monad/IOSimPOR/Internal.hs

@@ -696,6 +696,11 @@ schedule thread@Thread{
           return $ SimPORTrace time tid tstep tlbl (EventThrowTo e tid')
                  $ trace
 
+    -- intentionally a no-op (at least for now)
+    YieldSim k -> do
+      let thread' = thread { threadControl = ThreadControl k ctl }
+      schedule thread' simstate
+
 
 threadInterruptible :: Thread s a -> Bool
 threadInterruptible thread =

io-sim/test/Test/IOSim.hs

@@ -899,9 +899,6 @@ unit_async_10 =
           )
  ===
    ["child 1", "child 2", "child 1 running", "parent done"]
-  where
-    yield :: IOSim s ()
-    yield = atomically (return ())  -- yield, go to end of runqueue
 
 
 unit_async_11 =
@@ -934,9 +931,6 @@ unit_async_11 =
           )
  ===
    ["child 1", "child 2", "child 1 running", "parent done"]
-  where
-    yield :: IOSim s ()
-    yield = atomically (return ())  -- yield, go to end of runqueue
 
 
 unit_async_12 =

ouroboros-consensus/src/Ouroboros/Consensus/Util/EarlyExit.hs

@@ -198,6 +198,7 @@ instance MonadFork m => MonadFork (WithEarlyExit m) where
                              unmask' = earlyExit . unmask . withEarlyExit
                          in collapse <$> withEarlyExit (f unmask')
   throwTo            = lift .: throwTo
+  yield              = lift yield
 
 instance MonadST m => MonadST (WithEarlyExit m) where
   withLiftST f = lowerLiftST $ \(_proxy :: Proxy s) liftST ->