Skip to content
Browse files

Add an 'escape valve' to safe buffer iteratee

  • Loading branch information...
1 parent a737d42 commit c79652f5e5978eee3c3aad06c2fad5a97c60045d @gregorycollins gregorycollins committed May 30, 2010
Showing with 38 additions and 15 deletions.
  1. +13 −2 src/Snap/Iteratee.hs
  2. +25 −13 test/suite/Snap/Iteratee/Tests.hs
View
15 src/Snap/Iteratee.hs
@@ -126,11 +126,22 @@ countBytes = go 0
-- Our enumerators produce a lot of little strings; rather than spending all
-- our time doing kernel context switches for 4-byte write() calls, we buffer
-- the iteratee to send 8KB at a time.
-bufferIteratee :: (Monad m) => Enumerator m a
-bufferIteratee = return . go (D.empty,0)
+bufferIteratee :: Iteratee IO a -> IO (Iteratee IO a, IORef Bool)
+bufferIteratee iteratee = do
+ esc <- newIORef False
+ return $ (start esc iteratee, esc)
+
where
blocksize = 8192
+ start esc iter = IterateeG $! checkRef esc iter
+
+ checkRef esc iter ch = do
+ quit <- readIORef esc
+ if quit
+ then runIter iter ch
+ else f (D.empty,0) iter ch
+
--go :: (DList ByteString, Int) -> Iteratee m a -> Iteratee m a
go (!dl,!n) iter = IterateeG $! f (dl,n) iter
View
38 test/suite/Snap/Iteratee/Tests.hs
@@ -78,46 +78,58 @@ testEnumLBS = testProperty "enumLBS" prop
testBuffer :: Test
-testBuffer = testProperty "testBuffer" prop
+testBuffer = testProperty "testBuffer" $
+ monadicIO $ forAllM arbitrary prop
where
- prop s = s /= L.empty ==> fromWrap (runIdentity (run iter)) == s'
+ prop s = do
+ pre (s /= L.empty)
+
+ (i,_) <- liftQ $ bufferIteratee stream2stream
+ iter <- liftQ $ enumLBS s' i
+ x <- liftQ $ run iter
+
+ QC.assert $ fromWrap x == s'
where
s' = L.take 20000 $ L.cycle s
- i = runIdentity $ bufferIteratee stream2stream
- iter = runIdentity $ enumLBS s' i
testBuffer2 :: Test
testBuffer2 = testCase "testBuffer2" prop
where
prop = do
- i <- bufferIteratee $ drop 4 >> stream2stream
+ (i,_) <- bufferIteratee $ drop 4 >> stream2stream
s <- enumLBS "abcdefgh" i >>= run >>= return . fromWrap
H.assertEqual "s == 'efgh'" "efgh" s
testBuffer3 :: Test
-testBuffer3 = testProperty "testBuffer3" prop
+testBuffer3 = testProperty "testBuffer3" $
+ monadicIO $ forAllM arbitrary prop
where
- prop s = s /= L.empty ==> fromWrap (runIdentity (run iter)) == (L.take 19999 s')
+ prop s = do
+ pre (s /= L.empty)
+ (i,_) <- liftQ $ bufferIteratee (ss >>= \x -> drop 1 >> return x)
+ iter <- liftQ $ enumLBS s' i
+ x <- liftQ $ run iter
+
+ QC.assert $ fromWrap x == (L.take 19999 s')
where
s' = L.take 20000 $ L.cycle s
ss = joinI $ take 19999 stream2stream
- i = runIdentity $ bufferIteratee (ss >>= \x -> drop 1 >> return x)
- iter = runIdentity $ enumLBS s' i
+
testBuffer4 :: Test
testBuffer4 = testProperty "testBuffer4" $
monadicIO $ forAllM arbitrary prop
where
prop s = do
- i <- liftQ $ bufferIteratee (stream2stream >> throwErr (Err "foo"))
- i' <- liftQ $ enumLBS s i
+ (i,_) <- liftQ $ bufferIteratee (stream2stream >> throwErr (Err "foo"))
+ i' <- liftQ $ enumLBS s i
expectException $ run i'
- j <- liftQ $ bufferIteratee (throwErr (Err "foo") >> stream2stream)
- j' <- liftQ $ enumLBS s j
+ (j,_) <- liftQ $ bufferIteratee (throwErr (Err "foo") >> stream2stream)
+ j' <- liftQ $ enumLBS s j
expectException $ run j'
k <- liftQ $ enumErr "foo" j

0 comments on commit c79652f

Please sign in to comment.
Something went wrong with that request. Please try again.