5cde66c · streamly · git.iko.soy

23import qualified Control.Exception.Lifted as EL

24import Control.Monad.Catch (MonadThrow, throwM)

25import Control.Monad.IO.Class (MonadIO (..))

26−import Control.Monad.State (get, gets, modify, put,

26+import Control.Monad.State (gets, modify,

27− runStateT, when)

27+ runStateT, when, StateT)

28import Control.Monad.Trans.Class (MonadTrans (lift))

29import Control.Monad.Trans.Control (MonadBaseControl, liftBaseWith)

30import Data.Dynamic (Typeable)

69-- | Continue execution of the closure that we were executing when we migrated

70-- to a new thread.

71

72−resume :: MonadIO m => Context -> m ()

72+resume :: MonadIO m => Context -> m Context

73resume ctx = do

74 -- XXX rename to buildContext or buildState?

75 let s = runAsyncT (resumeContext ctx)

80 -- passing through all the parents? We can let the parent go away and

81 -- handle the ChildDone events as well in the root thread.

82 case parentChannel c of

83− Nothing -> return () -- TODO: yield the value for streaming here

83+ Nothing -> return ()

84 Just chan -> do

85 let r = accumResults c

86 when (length r /= 0) $

87 -- there is only one result in case of a non-root thread

88 -- XXX change the return type to 'a' instead of '[a]'

89 liftIO $ atomically $ writeTChan chan (PassOnResult (Right r))

90+ return c

90

91

91------------------------------------------------------------------------------

92------------------------------------------------------------------------------

92-- Thread Management (creation, reaping and killing)

93-- Thread Management (creation, reaping and killing)

227-- high.

228-- high.

228--

229--

229

230

230−forkFinally1 :: (MonadIO m, MonadBaseControl IO m) =>

231+forkFinally1 :: (MonadIO m, MonadBaseControl IO m)

231− m a -> (Either SomeException a -> IO ()) -> m ThreadId

232+ => Context -> (Either SomeException Context -> IO ()) -> m ThreadId

232−forkFinally1 action preExit =

233+forkFinally1 ctx preExit =

233 EL.mask $ \restore ->

234 EL.mask $ \restore ->

234 liftBaseWith $ \runInIO -> forkIO $ do

235 liftBaseWith $ \runInIO -> forkIO $ do

235− _ <- runInIO $ EL.try (restore action) >>= liftIO . preExit

236+ _ <- runInIO $ EL.try (restore (resume ctx)) >>= liftIO . preExit

236 return ()

237 return ()

237

238

238-- | Run a given context in a new thread.

239-- | Run a given context in a new thread.

239--

240--

240−forkContextWith :: (MonadBaseControl IO m, MonadIO m, MonadThrow m)

241+forkContext :: (MonadBaseControl IO m, MonadIO m, MonadThrow m)

241− => (Context -> m (Maybe a)) -> Context -> m ()

242+ => Context -> m ()

242−forkContextWith runCtx context = do

243+forkContext context = do

243 child <- childContext context

244 child <- childContext context

244− tid <- forkFinally1 (runCtx child) (beforeExit child)

245+ tid <- forkFinally1 child (beforeExit child)

245 updatePendingThreads context tid

246 updatePendingThreads context tid

246

247

247 where

248 where

272 dbg $ "beforeExit: " ++ show tid ++ " caught exception"

273 dbg $ "beforeExit: " ++ show tid ++ " caught exception"

273 liftIO $ killChildren ctx

274 liftIO $ killChildren ctx

274 return (Just e)

275 return (Just e)

275− Right Nothing -> return Nothing

276+ Right _ -> return Nothing

276− Right (Just _) -> error "Bug: should never happen"

277

278 e <- waitForChildren ctx exc

279

285 liftIO $ atomically $ writeTChan p

286 (ChildDone tid (maybe (Right []) Left e))

287

288−-- | A wrapper to first decide whether to run the context in the same thread or

288+-- | Decide whether to resume the context in the same thread or a new thread

289−-- a new thread and then run the context in that thread.

290--

289--

291−-- XXX create a canFork function

290+canFork :: Context -> IO Bool

292−resumeContextWith :: (MonadBaseControl IO m, MonadIO m, MonadThrow m)

291+canFork context = do

293− => (Context -> m (Maybe a)) -> Context -> m ()

294−resumeContextWith runCtx context = do

295 gotCredit <- liftIO $ waitQSemB (threadCredit context)

292 gotCredit <- liftIO $ waitQSemB (threadCredit context)

296− pending <- liftIO $ readIORef $ pendingThreads context

297 case gotCredit of

293 case gotCredit of

294+ False -> do

295+ pending <- liftIO $ readIORef $ pendingThreads context

298− False -> case pending of

296+ case pending of

299− [] -> do

300− _ <- runCtx context -- run synchronously

301− return ()

297+ [] -> return False

302 _ -> do

298 _ -> do

303 -- XXX If we have unreclaimable child threads e.g.

299 -- XXX If we have unreclaimable child threads e.g.

304 -- infinite loop, this is going to deadlock us. We need

300 -- infinite loop, this is going to deadlock us. We need

306 -- unreclaimable list? And always execute them in an

302 -- unreclaimable list? And always execute them in an

307 -- async thread, cannot use sync for those.

303 -- async thread, cannot use sync for those.

308 waitForOneEvent context

304 waitForOneEvent context

305+ canFork context

306+ True -> return True

307+

308+-- | resume a captured context with a given action. The context may be resumed

309+-- in the same thread or in a new thread depending on the synch parameter and

310+-- the current thread quota.

311+--

309− resumeContextWith runCtx context

312+resumeContextWith :: (MonadBaseControl IO m, MonadIO m, MonadThrow m)

310− True -> forkContextWith runCtx context

313+ => Context -- the context to resume

314+ -> Maybe Context -- context to copy the results from

315+ -> Bool -- force synchronous

316+ -> (Context -> AsyncT m (StreamData a)) -- the action to execute in the resumed context

317+ -> m (Maybe Context)

318+resumeContextWith context prevCtx synch action = do

319+ let ctx = setContextMailBox context (action context)

320+ can <- liftIO $ canFork context

321+ case can && (not synch) of

322+ False -> do

323+ -- transfer the results from the previous iteration to the next

324+ -- Only for top level thread, for others we would have already

325+ -- transferred the result to the parent via a channel

326+ let ctx' = case prevCtx of

327+ Nothing -> ctx

328+ Just x ->

329+ case parentChannel ctx of

330+ Nothing -> ctx { accumResults = accumResults x }

331+ Just _ -> ctx

332+ c <- resume ctx' -- run synchronously

333+ return (Just c)

334+ True -> do

335+ forkContext ctx

336+ return Nothing

311

337

312-- | 'StreamData' represents a task in a task stream being generated.

338-- | 'StreamData' represents a task in a task stream being generated.

313data StreamData a =

339data StreamData a =

332-- available we need a way to close the channel and wakeup all waiters so that

358-- available we need a way to close the channel and wakeup all waiters so that

333-- they can go away rather than waiting indefinitely.

359-- they can go away rather than waiting indefinitely.

334--

360--

361+

362+-- Housekeeping, invoked after spawning of all child tasks is done and the

363+-- parent task needs to terminate. Either the task is fully done or we handed

364+-- it over to another thread, in any case the current thread is done.

365+

366+spawningParentDone :: Monad m

367+ => (Maybe Context) -> StateT Context m (Maybe (StreamData a))

368+spawningParentDone ctx = do

369+ -- when running synchronously copy over the results from the returned

370+ -- context

371+ maybe (return ())

372+ (\x -> modify (\c -> c {accumResults = accumResults x}))

373+ ctx

374+

375+ loc <- getLocation

376+ when (loc /= RemoteNode) $ setLocation WaitingParent

377+ return Nothing

378+

379+-- | Captures the state of the current computation at this point, starts a new

380+-- thread, passing the captured state, runs the argument computation in the new

381+-- thread, returns its value and continues the computation from the capture

382+-- point. The end result is as if this function just returned the value

383+-- generated by the argument computation albeit in a new thread.

384+--

385+-- If a new thread cannot be created then the computation is run in the same

386+-- thread, but the functional behavior remains the same.

387+

388+spawnAsyncT :: (MonadIO m, MonadBaseControl IO m, MonadThrow m)

389+ => (Context -> AsyncT m (StreamData a))

390+ -> AsyncT m (StreamData a)

391+spawnAsyncT func = AsyncT $ do

392+ val <- takeContextMailBox

393+ case val of

394+ -- Child task

395+ Right x -> runAsyncT x

396+

397+ -- Spawning parent

398+ Left ctx -> do

399+ r <- lift $ resumeContextWith ctx Nothing False func

400+

401+ -- If we started the task asynchronously in a new thread then the

402+ -- parent thread reaches here, immediately after spawning the task.

403+ --

404+ -- Howver, if the task was executed synchronously then we will

405+ -- reach after it is completely done.

406+ spawningParentDone r

407+

335−-- | Execute the IO action, resume the saved context with the output of the io

408+-- | Execute the specified IO action, resume the saved context returning the

336−-- action.

409+-- output of the io action, continue this in a loop until the ioaction

410+-- indicates that its done.

411+

337loopContextWith :: (MonadIO m, MonadBaseControl IO m, MonadThrow m)

412loopContextWith :: (MonadIO m, MonadBaseControl IO m, MonadThrow m)

338− => IO (StreamData t) -> Context -> m (Maybe a)

413+ => IO (StreamData a) -> Context -> AsyncT m (StreamData a)

339−loopContextWith ioaction ctx = do

414+loopContextWith ioaction ctx = AsyncT $ loop Nothing

415+ where

416+

417+ loop prevCtx = do

340− streamData <- liftIO $ ioaction `catch`

418+ streamData <- liftIO $ ioaction `catch`

341− \(e :: SomeException) -> return $ SError e

419+ \(e :: SomeException) -> return $ SError e

342−

420+

343− let ctx' = setContextMailBox ctx streamData

421+ let resumeCtx synch =

422+ lift $ resumeContextWith ctx prevCtx synch $ \_ ->

423+ return streamData

424+

344− case streamData of

425+ case streamData of

345− SMore _ -> do

426+ SMore _ -> resumeCtx False >>= loop

346− resumeContextWith (\x -> resume x >> return Nothing) ctx'

427+ _ -> resumeCtx True >>= spawningParentDone

347− loopContextWith ioaction ctx

348− _ -> do

349− resume ctx' -- run synchronously

350− return Nothing -- we are done with the loop

351

428

352-- | Run an IO action one or more times to generate a stream of tasks. The IO

429-- | Run an IO action one or more times to generate a stream of tasks. The IO

353-- action returns a 'StreamData'. When it returns an 'SMore' or 'SLast' a new

430-- action returns a 'StreamData'. When it returns an 'SMore' or 'SLast' a new

358-- Unless the maximum number of threads (set with 'threads') has been reached,

435-- Unless the maximum number of threads (set with 'threads') has been reached,

359-- the task is generated in a new thread and the current thread returns a void

436-- the task is generated in a new thread and the current thread returns a void

360-- task.

437-- task.

361−parallel :: (Monad m, MonadIO m, MonadBaseControl IO m, MonadThrow m)

438+parallel :: (MonadIO m, MonadBaseControl IO m, MonadThrow m)

362 => IO (StreamData a) -> AsyncT m (StreamData a)

439 => IO (StreamData a) -> AsyncT m (StreamData a)

363−parallel ioaction = AsyncT $ do

364− -- We retrieve the context here and pass it on so that we can resume it

365− -- later. Control resumes at this place when we resume this context in a

366− -- new thread or in the same thread after generating the event data from

367− -- the ioaction.

368− mb <- takeContextMailBox

369− case mb of

370− -- We have already executed the ioaction in the parent thread and put

371− -- its result in the context and now we are continuing the context in a

372− -- new thread. Just return the result of the ioaction stored in the

373− -- 'event' field.

374− Right x -> return (Just x)

375−

376− -- We have to execute the io action, generate event data, put it in the

377− -- context mailbox and then continue from the point where this context

378− -- was retreieved, in this thread or a new thread.

379− Left ctx -> do

380− lift $ resumeContextWith (loopContextWith ioaction) ctx

440+parallel ioaction = spawnAsyncT (loopContextWith ioaction)

381−

382− -- We will never reach here if we continued the context in the same

383− -- thread. If we started a new thread then the parent thread

384− -- reaches here.

385− loc <- getLocation

386− when (loc /= RemoteNode) $ setLocation WaitingParent

387− return Nothing

388

441

389-- | An task stream generator that produces an infinite stream of tasks by

442-- | An task stream generator that produces an infinite stream of tasks by

390-- running an IO computation in a loop. A task is triggered carrying the output

443-- running an IO computation in a loop. A task is triggered carrying the output