Tasks.hs (15381B)
1 {-# LANGUAGE ScopedTypeVariables #-} 2 -- Structured native task groups. Registry transactions are the linearization 3 -- points; ThreadKilled only wakes a body, never carries the cancellation reason. 4 module IOV2.Tasks 5 ( Outcome(..), Scheduler, Group, TaskRef(..), newScheduler, rootGroup 6 , scopeOf, runGroup, nested, spawn, join, discard, cancel, checkpoint, liveCounts 7 ) where 8 9 import Control.Concurrent hiding (yield) 10 import Control.Concurrent.STM 11 import Control.Exception 12 import Control.Monad (forM_, void) 13 import qualified Data.Map.Strict as M 14 import Runtime (T(..), ofString, ofList) 15 import qualified IOV2.Resources as R 16 17 data Outcome = Completed T | Failed T [R.Problem] | Cancelled T [R.Problem] deriving Show 18 data Phase = Open | Joining | Releasing | Closed deriving Eq 19 data Scheduler = Scheduler R.Runtime (TVar (M.Map Integer Task)) 20 (TVar (M.Map Integer Task)) (TVar (M.Map Integer [Integer])) 21 data Group = Group 22 { groupScheduler :: Scheduler 23 , scopeOf :: R.Scope 24 , taskIdentity :: Integer 25 , groupParent :: Maybe (Group, Bool) 26 , groupPhase :: TVar Phase 27 , groupReason :: TVar (Maybe Outcome) 28 , groupChildren :: TVar (M.Map Integer Task) 29 , groupCleanup :: TVar [R.Problem] 30 , groupLexical :: TVar (Maybe Group) 31 } 32 data Task = Task Group Group (TVar (Maybe Outcome)) 33 data TaskRef = TaskRef !Integer !Integer deriving (Eq, Show) 34 data Interrupted = Interrupted deriving Show 35 instance Exception Interrupted 36 37 schedulerOf :: Group -> Scheduler 38 schedulerOf = groupScheduler 39 runtimeOf :: Scheduler -> R.Runtime 40 runtimeOf (Scheduler runtime _ _ _) = runtime 41 42 newScheduler :: R.Runtime -> IO Scheduler 43 newScheduler runtime = Scheduler runtime <$> newTVarIO M.empty <*> newTVarIO M.empty <*> newTVarIO M.empty 44 newGroup :: Scheduler -> R.Scope -> Integer -> Maybe (Group, Bool) -> STM Group 45 newGroup scheduler scope task parent = Group scheduler scope task parent 46 <$> newTVar Open <*> newTVar Nothing <*> newTVar M.empty <*> newTVar [] <*> newTVar Nothing 47 rootGroup :: Scheduler -> IO Group 48 rootGroup scheduler = atomically (newGroup scheduler R.rootScope 0 Nothing) 49 50 problem :: String -> String -> R.Problem 51 problem category operation = R.Problem category operation "" "none" 52 failedException :: SomeException -> Outcome 53 failedException e = Failed (ofString (take 512 (displayException e))) [] 54 55 -- First recorded stop wins. Failure propagates across lexical groups and 56 -- uncontained child boundaries, but never across preinstalled supervision. 57 -- All siblings receive sticky requests in this same transaction; a slow OS 58 -- interruption of one worker cannot prevent cancellation of the others. 59 stop :: Group -> Outcome -> STM () 60 stop Group{ groupScheduler = scheduler, scopeOf = scope, groupParent = parent, 61 groupPhase = phase, groupReason = reason, groupChildren = children, groupLexical = lexical } outcome = do 62 state <- readTVar phase 63 previous <- readTVar reason 64 case (state, previous) of 65 (Closed, _) -> pure () 66 (_, Just _) -> pure () 67 _ -> do 68 writeTVar reason (Just (withoutCleanup outcome)) 69 R.beginCloseScopeSTM (runtimeOf scheduler) scope 70 inner <- readTVar lexical 71 forM_ inner (\child -> stop child (Cancelled (ofString "owning group stopped") [])) 72 active <- readTVar children 73 forM_ (M.elems active) $ \(Task _ child _) -> 74 stop child (Cancelled (ofString "owning group stopped") []) 75 case (outcome, parent) of 76 (Failed _ _, Just (outer, True)) -> stop outer (withoutCleanup outcome) 77 _ -> pure () 78 79 checkpoint :: Group -> IO () 80 checkpoint Group{ groupReason = reason } = atomically $ do 81 stopped <- readTVar reason 82 case stopped of Nothing -> pure (); Just _ -> throwSTM Interrupted 83 84 -- Propagation carries only the primary. Cleanup travels once, at settlement, 85 -- rather than being duplicated along early fail-fast notifications. 86 withoutCleanup :: Outcome -> Outcome 87 withoutCleanup (Failed diagnostic _) = Failed diagnostic [] 88 withoutCleanup (Cancelled reason _) = Cancelled reason [] 89 withoutCleanup completed = completed 90 cleanupOf :: Outcome -> [R.Problem] 91 cleanupOf (Failed _ errors) = errors 92 cleanupOf (Cancelled _ errors) = errors 93 cleanupOf _ = [] 94 95 attach :: Outcome -> [R.Problem] -> Outcome 96 attach (Failed diagnostic previous) errors = Failed diagnostic (R.combineErrors previous errors) 97 attach (Cancelled reason previous) errors = Cancelled reason (R.combineErrors previous errors) 98 attach completed [] = completed 99 attach (Completed _) (first:rest) = Failed (ofList (map ofString [R.category first, R.operation first, R.detail first, R.progress first])) rest 100 101 -- The controller never executes user code. Its body runs in a separate worker, 102 -- so cancellation cannot land inside the resource-release protocol. Lexical 103 -- nesting may interrupt a controller from outside; the exception boundary then 104 -- cancels/joins its entire subtree before rethrowing that host interruption. 105 -- Repeated external host interruption/process death is not a cleanup guarantee. 106 runGroup :: Group -> IO Outcome -> IO Outcome 107 runGroup group@Group{ groupScheduler = scheduler, scopeOf = scope, taskIdentity = taskId, 108 groupPhase = phase, groupReason = reason, groupChildren = children, 109 groupCleanup = childErrors, groupLexical = lexical } body = mask $ \restore -> do 110 bodyDone <- newEmptyTMVarIO 111 launched <- try $ forkIOWithUnmask $ \unmask -> do 112 result <- try (unmask (checkpoint group >> body >>= evaluate)) 113 atomically (putTMVar bodyDone result) 114 worker <- case launched of 115 Left (exception :: SomeException) -> atomically (putTMVar bodyDone (Left exception)) >> pure Nothing 116 Right tid -> pure (Just tid) 117 let interruptBody = do 118 -- A lexical controller is host teardown, not user work. Its sticky stop 119 -- was set by stop; let it settle naturally rather than injecting an 120 -- exception into its release protocol. Only interrupt actual body work. 121 inner <- readTVarIO lexical 122 case inner of Nothing -> forM_ worker (\tid -> throwTo tid ThreadKilled); Just _ -> pure () 123 closeAdmission = atomically $ do 124 writeTVar phase Joining 125 R.beginCloseScopeSTM (runtimeOf scheduler) scope 126 joinChildren = do 127 let Scheduler _ _ _ waits = scheduler 128 atomically $ do 129 active <- readTVar children 130 graph <- readTVar waits 131 let targets = M.keys active 132 if any (reaches graph taskId []) targets 133 then stop group (Failed (ofString "task join cycle") []) 134 else modifyTVar' waits (M.insert taskId targets) 135 atomically (readTVar children >>= check . M.null) 136 `finally` atomically (modifyTVar' waits (M.delete taskId)) 137 release = do 138 atomically (writeTVar phase Releasing) 139 R.closeScope (runtimeOf scheduler) scope 140 finish outcome errors = atomically $ do 141 prior <- readTVar reason 142 descendants <- readTVar childErrors 143 let diagnostics = R.combineErrors (cleanupOf outcome) (R.combineErrors descendants errors) 144 final = attach (maybe (withoutCleanup outcome) id prior) diagnostics 145 case final of Failed _ _ -> stop group final; _ -> pure () 146 writeTVar childErrors (cleanupOf final) 147 writeTVar phase Closed 148 let Scheduler _ _ tickets _ = scheduler 149 modifyTVar' tickets (M.filter (\(Task owner _ _) -> scopeOf owner /= scope)) 150 pure final 151 execute = do 152 signal <- restore $ atomically $ do 153 requested <- readTVar reason 154 case requested of 155 Just _ -> pure Nothing 156 Nothing -> Just <$> readTMVar bodyDone 157 result <- case signal of 158 Just value -> pure value 159 Nothing -> interruptBody >> atomically (readTMVar bodyDone) 160 let outcome = either failedException id result 161 atomically $ case outcome of Completed _ -> pure (); _ -> stop group outcome 162 closeAdmission 163 joinChildren 164 errors <- release 165 finish outcome errors 166 abort (exception :: SomeException) = do 167 atomically (stop group (Cancelled (ofString "host interruption") [])) 168 interruptBody 169 void (atomically (readTMVar bodyDone)) 170 closeAdmission 171 joinChildren 172 errors <- release 173 void (finish (Cancelled (ofString "host interruption") []) errors) 174 throwIO exception 175 execute `catch` abort 176 177 nested :: Group -> (Group -> IO Outcome) -> IO Outcome 178 nested parent@Group{ groupScheduler = scheduler, scopeOf = scope, taskIdentity = task, 179 groupCleanup = parentErrors, groupLexical = lexical } body = mask $ \restore -> do 180 allocated <- atomically $ do 181 checkpointSTM parent 182 reserved <- R.reserveChildSTM (runtimeOf scheduler) scope [] 183 case reserved of 184 Left e -> pure (Left e) 185 Right (childScope, _) -> do 186 child <- newGroup scheduler childScope task (Just (parent, True)) 187 writeTVar lexical (Just child) 188 pure (Right child) 189 case allocated of 190 Left e -> pure (Failed (ofString (R.category e)) []) 191 Right child@Group{ groupCleanup = diagnostics } -> do 192 result <- runGroup child (restore (body child)) `finally` atomically (do 193 errors <- readTVar diagnostics 194 modifyTVar' parentErrors (\previous -> R.combineErrors previous errors) 195 writeTVar lexical Nothing) 196 pure (withoutCleanup result) 197 198 checkpointSTM :: Group -> STM () 199 checkpointSTM Group{ groupPhase = phase, groupReason = reason } = do 200 state <- readTVar phase 201 stopped <- readTVar reason 202 case (state, stopped) of 203 (Open, Nothing) -> pure () 204 _ -> throwSTM Interrupted 205 206 -- A worker exists but cannot run before atomic resource transfer AND task 207 -- publication. Fork failure leaves the registry untouched. Failed validation 208 -- opens the gate with Nothing and joins that unused worker before returning. 209 spawn :: Group -> Bool -> Bool -> [R.Token] -> (Group -> [R.Token] -> IO Outcome) 210 -> IO (Either R.Problem TaskRef) 211 spawn owner@Group{ groupScheduler = scheduler, scopeOf = scope, 212 groupChildren = children, groupCleanup = ownerErrors } supervised keep tokens body = mask_ $ do 213 gate <- newEmptyTMVarIO 214 retired <- newEmptyMVar 215 let Scheduler runtime active tickets _ = scheduler 216 _ <- forkIOWithUnmask $ \unmask -> (do 217 entry <- atomically (takeTMVar gate) 218 case entry of 219 Nothing -> pure () 220 Just (task@(Task _ child completed), fresh) -> do 221 outcome <- unmask (runGroup child (body child fresh)) `catch` (pure . failedException) 222 atomically $ do 223 writeTVar completed (Just outcome) 224 if supervised then pure () else 225 modifyTVar' ownerErrors (\previous -> R.combineErrors previous (cleanupOf outcome)) 226 modifyTVar' children (M.delete (scopeOf child)) 227 modifyTVar' active (M.delete (scopeOf child)) 228 -- runGroup already propagated ordinary faults before release. This 229 -- also covers an unexpected controller fault at its outer boundary. 230 case outcome of Failed _ _ | not supervised -> stop owner outcome; _ -> pure () 231 task `seq` pure ()) `finally` putMVar retired () 232 let abandon = atomically (putTMVar gate Nothing) >> takeMVar retired 233 result <- (atomically $ do 234 checkpointSTM owner 235 reserved <- R.reserveChildSTM runtime scope tokens 236 case reserved of 237 Left e -> pure (Left e) 238 Right (childScope, fresh) -> do 239 child <- newGroup scheduler childScope childScope (Just (owner, not supervised)) 240 completed <- newTVar Nothing 241 let task = Task owner child completed 242 modifyTVar' children (M.insert childScope task) 243 modifyTVar' active (M.insert childScope task) 244 if keep then modifyTVar' tickets (M.insert childScope task) else pure () 245 putTMVar gate (Just (task, fresh)) 246 pure (Right childScope)) `onException` abandon 247 case result of 248 Left e -> abandon >> pure (Left e) 249 Right key -> do 250 identity <- R.runtimeIdentity runtime 251 pure (Right (TaskRef identity key)) 252 253 within :: Group -> Group -> Bool 254 within current@Group{ groupParent = parent } owner = 255 scopeOf current == scopeOf owner || maybe False (\(outer, _) -> within outer owner) parent 256 257 -- Registry lookup must remain inside the claiming transaction. Only the run 258 -- identity check is IO (it reads an immutable field). 259 checked :: Group -> TaskRef -> TVar (M.Map Integer Task) -> IO (STM (Either R.Problem Task)) 260 checked current (TaskRef run key) table = do 261 identity <- R.runtimeIdentity (runtimeOf (schedulerOf current)) 262 pure $ if identity /= run then pure (Left (problem "stale-task" "task")) else do 263 entries <- readTVar table 264 pure $ case M.lookup key entries of 265 Just task@(Task owner _ _) | within current owner -> Right task 266 Just _ -> Left (problem "out-of-scope" "task") 267 Nothing -> Left (problem "stale-task" "task") 268 269 -- Includes implicit scope joining as well as explicit ticket waits. Lexical 270 -- controllers share their body's task identity, so nested scopes cannot hide a 271 -- cycle. Edges are removed when a wait settles, never retained as history. 272 reaches :: M.Map Integer [Integer] -> Integer -> [Integer] -> Integer -> Bool 273 reaches graph wanted seen target 274 | target == wanted = True 275 | target `elem` seen = False 276 | otherwise = any (reaches graph wanted (target:seen)) (M.findWithDefault [] target graph) 277 278 join :: Group -> TaskRef -> IO (Either R.Problem Outcome) 279 join current@Group{ groupScheduler = Scheduler _ _ tickets waits, taskIdentity = taskId } ref@(TaskRef _ key) = mask $ \restore -> do 280 lookupEntry <- checked current ref tickets 281 claimed <- atomically $ do 282 checkpointSTM current 283 entry <- lookupEntry 284 graph <- readTVar waits 285 case entry of 286 Right task@(Task _ target _) | within current target || reaches graph taskId [] key -> pure (Left (problem "join-cycle" "join")) 287 | otherwise -> do 288 modifyTVar' tickets (M.delete key) 289 modifyTVar' waits (M.insert taskId [key]) 290 pure (Right task) 291 Left e -> pure (Left e) 292 case claimed of 293 Left e -> pure (Left e) 294 Right (Task _ _ completed) -> 295 (Right <$> restore (atomically (readTVar completed >>= maybe retry pure))) 296 `finally` atomically (modifyTVar' waits (M.delete taskId)) 297 298 discard :: Group -> TaskRef -> IO (Either R.Problem ()) 299 discard current ref@(TaskRef _ key) = do 300 let Scheduler _ _ tickets _ = schedulerOf current 301 lookupEntry <- checked current ref tickets 302 atomically $ do 303 checkpointSTM current 304 entry <- lookupEntry 305 case entry of Left e -> pure (Left e); Right _ -> modifyTVar' tickets (M.delete key) >> pure (Right ()) 306 307 cancel :: Group -> TaskRef -> T -> IO (Either R.Problem ()) 308 cancel current ref why = do 309 let Scheduler _ active _ _ = schedulerOf current 310 lookupEntry <- checked current ref active 311 atomically $ do 312 checkpointSTM current 313 entry <- lookupEntry 314 case entry of 315 Left e -> pure (Left e) 316 Right (Task _ child@Group{ groupPhase = phase } _) -> do 317 state <- readTVar phase 318 if state == Closed then pure (Left (problem "completed-task" "cancel")) 319 else stop child (Cancelled why []) >> pure (Right ()) 320 321 liveCounts :: Scheduler -> IO (Int, Int, Int) 322 liveCounts (Scheduler _ active tickets waits) = atomically $ (,,) 323 <$> (M.size <$> readTVar active) <*> (M.size <$> readTVar tickets) <*> (M.size <$> readTVar waits)