From 0d138d097056a99d0b14640b7d652415b5c55430 Mon Sep 17 00:00:00 2001
From: Harendra Kumar <harendra.kumar@gmail.com>
Date: Mon, 16 Oct 2017 07:36:14 +0530
Subject: [PATCH] Gracefully deal with context EOF

When using a context asynchronously
---
 src/Asyncly/AsyncT.hs | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/src/Asyncly/AsyncT.hs b/src/Asyncly/AsyncT.hs
index 2eb9b909..614ea680 100644
--- a/src/Asyncly/AsyncT.hs
+++ b/src/Asyncly/AsyncT.hs
@@ -58,7 +58,7 @@ import           Control.Concurrent          (ThreadId, forkIO,
                                               myThreadId, threadDelay)
 import           Control.Concurrent.MVar     (MVar, newEmptyMVar, tryTakeMVar,
                                               tryPutMVar, takeMVar)
-import           Control.Exception           (SomeException (..))
+import           Control.Exception           (SomeException (..), Exception)
 import qualified Control.Exception.Lifted    as EL
 import           Control.Monad               (ap, liftM, MonadPlus(..), mzero,
                                               when)
@@ -390,10 +390,21 @@ sendWorkerWait ctx = dispatch >> void (liftIO $ takeMVar (doorBell ctx))
             when (not done) $ (pushWorker ctx) >> dispatch
 
 
+data ContextUsedAfterEOF = ContextUsedAfterEOF deriving Show
+instance Exception ContextUsedAfterEOF
+
 -- | Pull an AsyncT stream from a context
 {-# NOINLINE pullFromCtx #-}
 pullFromCtx :: MonadAsync m => Context m a -> AsyncT m a
 pullFromCtx ctx = AsyncT $ \_ stp yld -> do
+    -- When using an async handle to the context, one may keep using a stale
+    -- context even after it has been fullt drained. To detect it gracefully we
+    -- raise an explicit exception.
+    -- XXX if reading the IORef is costly we can use a flag in the context to
+    -- indicate we are done.
+    done <- allThreadsDone ctx
+    when done $ throwM ContextUsedAfterEOF
+
     res <- liftIO $ tryTakeMVar (doorBell ctx)
     when (isNothing res) $ sendWorkerWait ctx
     list <- liftIO $ atomicModifyIORefCAS (outputQueue ctx) $ \x -> ([], x)