7afa715 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

track active workers in an SVar

7afa715Harendra Kumar8 years ago

src/Streamly/Core.hs+14−2

compare with master

unifiedsplit

@@ -154,6 +154,7 @@ data SVar m a =

154 , runqueue :: m ()

155 , runningThreads :: IORef (Set ThreadId)

156 , queueEmpty :: m Bool

157+ , activeWorkers :: IORef Int

157 , svarStyle :: SVarStyle

158 , svarStyle :: SVarStyle

158 }

159 }

159

160

@@ -401,7 +402,8 @@ handleChildException sv e = do

401

402

402{-# NOINLINE pushWorker #-}

403{-# NOINLINE pushWorker #-}

403pushWorker :: MonadParallel m => SVar m a -> m ()

404pushWorker :: MonadParallel m => SVar m a -> m ()

404−pushWorker sv =

405+pushWorker sv = do

406+ liftIO $ atomicModifyIORefCAS_ (activeWorkers sv) $ \n -> n + 1

405 doFork (runqueue sv) (handleChildException sv) >>= addThread sv

407 doFork (runqueue sv) (handleChildException sv) >>= addThread sv

406

408

407-- | In contrast to pushWorker which always happens only from the consumer

409-- | In contrast to pushWorker which always happens only from the consumer

@@ -455,6 +457,7 @@ fromStreamVar sv = Stream $ \_ stp sng yld -> do

455 where

457 where

456

458

457 handleException e tid = do

459 handleException e tid = do

460+ liftIO $ atomicModifyIORefCAS_ (activeWorkers sv) $ \n -> n - 1

458 modifyThread sv tid

461 modifyThread sv tid

459 -- XXX implement kill async exception handling

462 -- XXX implement kill async exception handling

460 -- liftIO $ readIORef (runningThreads sv) >>= mapM_ killThread

463 -- liftIO $ readIORef (runningThreads sv) >>= mapM_ killThread

@@ -475,13 +478,17 @@ fromStreamVar sv = Stream $ \_ stp sng yld -> do

475 ChildYield a -> yield a

478 ChildYield a -> yield a

476 ChildStop tid e ->

479 ChildStop tid e ->

477 case e of

480 case e of

481+ Nothing -> do

482+ let active = activeWorkers sv

483+ liftIO $ atomicModifyIORefCAS_ active $ \n -> n - 1

478− Nothing -> modifyThread sv tid >> continue

484+ modifyThread sv tid >> continue

479 Just ex -> handleException ex tid

485 Just ex -> handleException ex tid

480

486

481getFifoSVar :: MonadIO m => SVarStyle -> IO (SVar m a)

487getFifoSVar :: MonadIO m => SVarStyle -> IO (SVar m a)

482getFifoSVar ctype = do

488getFifoSVar ctype = do

483 outQ <- newIORef []

489 outQ <- newIORef []

484 outQMv <- newEmptyMVar

490 outQMv <- newEmptyMVar

491+ active <- newIORef 0

485 running <- newIORef S.empty

492 running <- newIORef S.empty

486 q <- newQ

493 q <- newQ

487 let sv =

494 let sv =

@@ -492,6 +499,7 @@ getFifoSVar ctype = do

492 , enqueue = pushL q

499 , enqueue = pushL q

493 , queueEmpty = liftIO $ nullQ q

500 , queueEmpty = liftIO $ nullQ q

494 , svarStyle = ctype

501 , svarStyle = ctype

502+ , activeWorkers = active

495 }

503 }

496 in return sv

504 in return sv

497

505

@@ -499,6 +507,7 @@ getLifoSVar :: MonadIO m => SVarStyle -> IO (SVar m a)

499getLifoSVar ctype = do

507getLifoSVar ctype = do

500 outQ <- newIORef []

508 outQ <- newIORef []

501 outQMv <- newEmptyMVar

509 outQMv <- newEmptyMVar

510+ active <- newIORef 0

502 running <- newIORef S.empty

511 running <- newIORef S.empty

503 q <- newIORef []

512 q <- newIORef []

504 let checkEmpty = null <$> liftIO (readIORef q)

513 let checkEmpty = null <$> liftIO (readIORef q)

@@ -510,6 +519,7 @@ getLifoSVar ctype = do

510 , enqueue = enqueueLIFO q

519 , enqueue = enqueueLIFO q

511 , queueEmpty = checkEmpty

520 , queueEmpty = checkEmpty

512 , svarStyle = ctype

521 , svarStyle = ctype

522+ , activeWorkers = active

513 }

523 }

514 in return sv

524 in return sv

515

525

@@ -517,6 +527,7 @@ getParSVar :: SVarStyle -> IO (SVar m a)

517getParSVar style = do

527getParSVar style = do

518 outQ <- newIORef []

528 outQ <- newIORef []

519 outQMv <- newEmptyMVar

529 outQMv <- newEmptyMVar

530+ active <- newIORef 0

520 running <- newIORef S.empty

531 running <- newIORef S.empty

521 let sv =

532 let sv =

522 SVar { outputQueue = outQ

533 SVar { outputQueue = outQ

@@ -526,6 +537,7 @@ getParSVar style = do

526 , enqueue = undefined

537 , enqueue = undefined

527 , queueEmpty = undefined

538 , queueEmpty = undefined

528 , svarStyle = style

539 , svarStyle = style

540+ , activeWorkers = active

529 }

541 }

530 in return sv

542 in return sv

531

543