@@ -58,7 +58,7 @@ import Control.Concurrent (ThreadId, forkIO,
58 myThreadId, threadDelay)
58 myThreadId, threadDelay)
59import Control.Concurrent.MVar (MVar, newEmptyMVar, tryTakeMVar,
59import Control.Concurrent.MVar (MVar, newEmptyMVar, tryTakeMVar,
60 tryPutMVar, takeMVar)
60 tryPutMVar, takeMVar)
61−import Control.Exception (SomeException (..))
61+import Control.Exception (SomeException (..), Exception)
62import qualified Control.Exception.Lifted as EL
62import qualified Control.Exception.Lifted as EL
63import Control.Monad (ap, liftM, MonadPlus(..), mzero,
63import Control.Monad (ap, liftM, MonadPlus(..), mzero,
@@ -390,10 +390,21 @@ sendWorkerWait ctx = dispatch >> void (liftIO $ takeMVar (doorBell ctx))
390 when (not done) $ (pushWorker ctx) >> dispatch
390 when (not done) $ (pushWorker ctx) >> dispatch
393+data ContextUsedAfterEOF = ContextUsedAfterEOF deriving Show
394+instance Exception ContextUsedAfterEOF
393-- | Pull an AsyncT stream from a context
396-- | Pull an AsyncT stream from a context
394{-# NOINLINE pullFromCtx #-}
397{-# NOINLINE pullFromCtx #-}
395pullFromCtx :: MonadAsync m => Context m a -> AsyncT m a
398pullFromCtx :: MonadAsync m => Context m a -> AsyncT m a
396pullFromCtx ctx = AsyncT $ \_ stp yld -> do
399pullFromCtx ctx = AsyncT $ \_ stp yld -> do
400+ -- When using an async handle to the context, one may keep using a stale
401+ -- context even after it has been fullt drained. To detect it gracefully we
402+ -- raise an explicit exception.
403+ -- XXX if reading the IORef is costly we can use a flag in the context to
404+ -- indicate we are done.
405+ done <- allThreadsDone ctx
406+ when done $ throwM ContextUsedAfterEOF
397 res <- liftIO $ tryTakeMVar (doorBell ctx)
408 res <- liftIO $ tryTakeMVar (doorBell ctx)
398 when (isNothing res) $ sendWorkerWait ctx
409 when (isNothing res) $ sendWorkerWait ctx
399 list <- liftIO $ atomicModifyIORefCAS (outputQueue ctx) $ \x -> ([], x)
410 list <- liftIO $ atomicModifyIORefCAS (outputQueue ctx) $ \x -> ([], x)