Permalink
Browse files

Fix bug in lazyConsume

  • Loading branch information...
1 parent fd8f452 commit 87e890fe7ee58686d20cabba15dd37f18ba66620 @snoyberg committed Feb 21, 2012
Showing with 68 additions and 26 deletions.
  1. +50 −17 conduit/Control/Monad/Trans/Resource.hs
  2. +13 −8 conduit/Data/Conduit/Lazy.hs
  3. +5 −1 conduit/test/main.hs
@@ -43,12 +43,14 @@ module Control.Monad.Trans.Resource
, ResourceThrow (..)
-- ** Low-level
, HasRef (..)
+ , InvalidAccess (..)
+ , resourceActive
) where
import Data.Typeable
import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
-import Control.Exception (SomeException)
+import Control.Exception (SomeException, throw, Exception)
import Control.Monad.Trans.Control
( MonadTransControl (..), MonadBaseControl (..)
, ComposeSt, defaultLiftBaseWith, defaultRestoreM
@@ -260,6 +262,7 @@ type NextKey = Int
data ReleaseMap base =
ReleaseMap !NextKey !RefCount !(IntMap (base ()))
+ | ReleaseMapClosed
-- | The Resource transformer. This transformer keeps track of all registered
-- actions, and calls them upon exit (via 'runResourceT'). Actions may be
@@ -317,10 +320,25 @@ register' :: HasRef base
=> Ref base (ReleaseMap base)
-> base ()
-> base ReleaseKey
-register' istate rel = atomicModifyRef' istate $ \(ReleaseMap key rf m) ->
- ( ReleaseMap (key + 1) rf (IntMap.insert key rel m)
- , ReleaseKey key
- )
+register' istate rel = atomicModifyRef' istate $ \rm ->
+ case rm of
+ ReleaseMap key rf m ->
+ ( ReleaseMap (key + 1) rf (IntMap.insert key rel m)
+ , ReleaseKey key
+ )
+ ReleaseMapClosed -> throw $ InvalidAccess "register'"
+
+data InvalidAccess = InvalidAccess { functionName :: String }
+ deriving Typeable
+
+instance Show InvalidAccess where
+ show (InvalidAccess f) = concat
+ [ "Control.Monad.Trans.Resource."
+ , f
+ , ": The mutable state is being accessed after cleanup. Please contact the maintainers."
+ ]
+
+instance Exception InvalidAccess
-- | Call a release action early, and deregister it from the list of cleanup
-- actions to be performed.
@@ -344,25 +362,30 @@ release' istate (ReleaseKey key) = mask $ \restore -> do
( ReleaseMap next rf $ IntMap.delete key m
, Just action
)
+ lookupAction ReleaseMapClosed = throw $ InvalidAccess "release'"
stateAlloc :: HasRef m => Ref m (ReleaseMap m) -> m ()
stateAlloc istate = do
- atomicModifyRef' istate $ \(ReleaseMap nk rf m) ->
- (ReleaseMap nk (rf + 1) m, ())
+ atomicModifyRef' istate $ \rm ->
+ case rm of
+ ReleaseMap nk rf m ->
+ (ReleaseMap nk (rf + 1) m, ())
+ ReleaseMapClosed -> throw $ InvalidAccess "stateAlloc"
stateCleanup :: HasRef m => Ref m (ReleaseMap m) -> m ()
stateCleanup istate = mask_ $ do
- (rf, m) <- atomicModifyRef' istate $ \(ReleaseMap nk rf m) ->
- (ReleaseMap nk (rf - 1) m, (rf - 1, m))
- if rf == minBound
- then do
+ mm <- atomicModifyRef' istate $ \rm ->
+ case rm of
+ ReleaseMap nk rf m ->
+ let rf' = rf - 1
+ in if rf' == minBound
+ then (ReleaseMapClosed, Just m)
+ else (ReleaseMap nk rf' m, Nothing)
+ ReleaseMapClosed -> throw $ InvalidAccess "stateCleanup"
+ case mm of
+ Just m ->
mapM_ (\x -> try x >> return ()) $ IntMap.elems m
- -- Trigger an exception consistently for one race condition:
- -- let's put an undefined value in the state. If somehow
- -- another thread is still able to access it, at least we get
- -- clearer error messages.
- writeRef' istate $ error "Control.Monad.Trans.Resource.stateCleanup: There is a bug in the implementation. The mutable state is being accessed after cleanup. Please contact the maintainers."
- else return ()
+ Nothing -> return ()
-- | Unwrap a 'ResourceT' transformer, and call all registered release actions.
--
@@ -523,3 +546,13 @@ resourceForkIO (ResourceT f) = ResourceT $ \r -> L.mask $ \restore ->
(return ())
(stateCleanup r)
(restore $ f r))
+
+-- | Determine if the current @ResourceT@ is still active. This is necessary
+-- for such cases as lazy I\/O, where an unevaluated thunk may still refer to a
+-- closed @ResourceT@.
+resourceActive :: Resource m => ResourceT m Bool
+resourceActive = ResourceT $ \rmMap -> do
+ rm <- resourceLiftBase $ readRef' rmMap
+ case rm of
+ ReleaseMapClosed -> return False
+ _ -> return True
@@ -8,20 +8,25 @@ module Data.Conduit.Lazy
import Data.Conduit
import System.IO.Unsafe (unsafeInterleaveIO)
-import Control.Monad.Trans.Control
+import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_)
+import Control.Monad.Trans.Resource (resourceActive)
-- | Use lazy I\/O to consume all elements from a @Source@.
--
-- Since 0.2.0
-lazyConsume :: MonadBaseControl IO m => Source m a -> ResourceT m [a]
+lazyConsume :: (Resource m, MonadBaseControl IO m) => Source m a -> ResourceT m [a]
lazyConsume src0 = do
go src0
where
go src = liftBaseOp_ unsafeInterleaveIO $ do
- res <- sourcePull src
- case res of
- Closed -> return []
- Open src' x -> do
- y <- go src'
- return $ x : y
+ ra <- resourceActive
+ if ra
+ then do
+ res <- sourcePull src
+ case res of
+ Closed -> return []
+ Open src' x -> do
+ y <- go src'
+ return $ x : y
+ else return []
View
@@ -124,7 +124,7 @@ main = hspecX $ do
describe "zipping" $ do
it "zipping two small lists" $ do
res <- runResourceT $ CL.zip (CL.sourceList [1..10]) (CL.sourceList [11..12]) C.$$ CL.consume
- res @=? zip [1..10] [11..12]
+ res @=? zip [1..10 :: Int] [11..12 :: Int]
describe "Monad instance for Sink" $ do
it "binding" $ do
@@ -248,6 +248,10 @@ main = hspecX $ do
nums <- CLazy.lazyConsume $ mconcat $ map incr [1..10]
liftIO $ nums @?= [1..10]
+ it' "returns nothing outside ResourceT" $ do
+ bss <- runResourceT $ CLazy.lazyConsume $ CB.sourceFile "test/main.hs"
+ bss @?= []
+
describe "sequence" $ do
it "simple sink" $ do
let sumSink :: C.Resource m => C.Sink Int m Int

0 comments on commit 87e890f

Please sign in to comment.