@@ -95,8 +95,8 @@ import Data.Word (Word)
9595
9696import Data.Typeable (Typeable)
9797
98−import Data.Conduit (ResourceT, runResourceT)
99−import qualified Data.Conduit as C
98+import Data.Conduit
99+import Data.Conduit.Internal (ResumableSource (..))
100100import qualified Data.Conduit.List as CL
101101import qualified Data.Conduit.Binary as CB
102102import Data.Conduit.Blaze (builderToByteString)
@@ -237,62 +237,60 @@ runSettingsConnection set getConn app = do
237237 ) `onException` (T.cancel th >> connClose conn >> onClose)
238238
239239-- | Contains a @Source@ and a byte count that is still to be read in.
240−newtype IsolatedBSSource = IsolatedBSSource (I.IORef (Int, C.Source (ResourceT IO) ByteString))
240+newtype IsolatedBSSource = IsolatedBSSource (I.IORef (Int, ResumableSource (ResourceT IO) ByteString))
241241
242242-- | Given an @IsolatedBSSource@ provide a @Source@ that only allows up to the
243243-- specified number of bytes to be passed downstream. All leftovers should be
244244-- retained within the @Source@. If there are not enough bytes available,
245245-- throws a @ConnectionClosedByPeer@ exception.
246−ibsIsolate :: IsolatedBSSource -> C.Source (ResourceT IO) ByteString
247−ibsIsolate ibs@(IsolatedBSSource ref) =
248− C.PipeM pull (return ())
249− where
250− pull = do
251− (count, src) <- liftIO $ I.readIORef ref
252− if count == 0
253− -- No more bytes wanted downstream, so we're done.
254− then return $ C.Done Nothing ()
255− else do
256− -- Get the next chunk (if available) and the updated source
257− (src', mbs) <- src C.$$+ CL.head
258−
259− -- If no chunk available, then there aren't enough bytes in the
260− -- stream. Throw a ConnectionClosedByPeer
261− bs <- maybe (liftIO $ throwIO ConnectionClosedByPeer) return mbs
262−
263− let -- How many of the bytes in this chunk to send downstream
264− toSend = min count (S.length bs)
265− -- How many bytes will still remain to be sent downstream
266− count' = count - toSend
267− case () of
268− ()
269− -- The expected count is greater than the size of the
270− -- chunk we just read. Send the entire chunk
271− -- downstream, and then loop on this function for the
272− -- next chunk.
273− | count' > 0 -> do
274− liftIO $ I.writeIORef ref (count', src')
275− return $ C.HaveOutput (ibsIsolate ibs) (return ()) bs
276−
277− -- The expected count is the total size of the chunk we
278− -- just read. Send this chunk downstream, and then
279− -- terminate the stream.
280− | count == S.length bs -> do
281− liftIO $ I.writeIORef ref (count', src')
282− return $ C.HaveOutput (C.Done Nothing ()) (return ()) bs
283−
284− -- Some of the bytes in this chunk should not be sent
285− -- downstream. Split up the chunk into the sent and
286− -- not-sent parts, add the not-sent parts onto the new
287− -- source, and send the rest of the chunk downstream.
288− | otherwise -> do
289− let (x, y) = S.splitAt toSend bs
290− liftIO $ I.writeIORef ref (count', C.HaveOutput src' (return ()) y)
291− return $ C.HaveOutput (C.Done Nothing ()) (return ()) x
246+ibsIsolate :: IsolatedBSSource -> Source (ResourceT IO) ByteString
247+ibsIsolate ibs@(IsolatedBSSource ref) = do
248+ (count, src) <- liftIO $ I.readIORef ref
249+ if count == 0
250+ -- No more bytes wanted downstream, so we're done.
251+ then return ()
252+ else do
253+ -- Get the next chunk (if available) and the updated source
254+ (src', mbs) <- lift $ src $$++ CL.head
255+
256+ -- If no chunk available, then there aren't enough bytes in the
257+ -- stream. Throw a ConnectionClosedByPeer
258+ bs <- maybe (liftIO $ throwIO ConnectionClosedByPeer) return mbs
259+
260+ let -- How many of the bytes in this chunk to send downstream
261+ toSend = min count (S.length bs)
262+ -- How many bytes will still remain to be sent downstream
263+ count' = count - toSend
264+ case () of
265+ ()
266+ -- The expected count is greater than the size of the
267+ -- chunk we just read. Send the entire chunk
268+ -- downstream, and then loop on this function for the
269+ -- next chunk.
270+ | count' > 0 -> do
271+ liftIO $ I.writeIORef ref (count', src')
272+ yield bs
273+ ibsIsolate ibs
274+
275+ -- The expected count is the total size of the chunk we
276+ -- just read. Send this chunk downstream, and then
277+ -- terminate the stream.
278+ | count == S.length bs -> do
279+ liftIO $ I.writeIORef ref (count', src')
280+ yield bs
281+
282+ -- Some of the bytes in this chunk should not be sent
283+ -- downstream. Split up the chunk into the sent and
284+ -- not-sent parts, add the not-sent parts onto the new
285+ -- source, and send the rest of the chunk downstream.
286+ | otherwise -> do
287+ let (x, y) = S.splitAt toSend bs
288+ liftIO $ I.writeIORef ref (count', fmapResume (yield y >>) src')
289+ yield x
292290
293291-- | Extract the underlying @Source@ from an @IsolatedBSSource@, which will not
294292-- perform any more isolation.
295−ibsDone :: IsolatedBSSource -> IO (C.Source (ResourceT IO) ByteString)
293+ibsDone :: IsolatedBSSource -> IO (ResumableSource (ResourceT IO) ByteString)
296294ibsDone (IsolatedBSSource ref) = fmap snd $ I.readIORef ref
297295
298296serveConnection :: Settings
@@ -315,22 +313,22 @@ serveConnection settings th port app conn remoteHost' =
315313 res <- app env
316314
317315 -- flush the rest of the request body
318− requestBody env C.$$ CL.sinkNull
319− fromClient' <- liftIO getSource
316+ requestBody env $$ CL.sinkNull
317+ ResumableSource fromClient' _ <- liftIO getSource
320318
321319 liftIO $ T.resume th
322320 keepAlive <- sendResponse th env conn res
323321 when keepAlive $ serveConnection'' fromClient'
324322 Just intercept -> do
325323 liftIO $ T.pause th
326− fromClient' <- liftIO getSource
324+ ResumableSource fromClient' _ <- liftIO getSource
327325 intercept fromClient' conn
328326
329327parseRequest :: Connection -> Port -> SockAddr
330− -> C.Source (ResourceT IO) S.ByteString
331− -> ResourceT IO (Request, IO (C.Source (ResourceT IO) ByteString))
328+ -> Source (ResourceT IO) S.ByteString
329+ -> ResourceT IO (Request, IO (ResumableSource (ResourceT IO) ByteString))
332330parseRequest conn port remoteHost' src1 = do
333− (src2, headers') <- src1 C.$$+ takeHeaders
331+ (src2, headers') <- src1 $$+ takeHeaders
334332 parseRequest' conn port headers' remoteHost' src2
335333
336334-- FIXME come up with good values here
@@ -367,8 +365,8 @@ parseRequest' :: Connection
367365 -> Port
368366 -> [ByteString]
369367 -> SockAddr
370− -> C.Source (ResourceT IO) S.ByteString -- FIXME was buffered
371− -> ResourceT IO (Request, IO (C.Source (ResourceT IO) ByteString))
368+ -> ResumableSource (ResourceT IO) S.ByteString -- FIXME was buffered
369+ -> ResourceT IO (Request, IO (ResumableSource (ResourceT IO) ByteString))
372370parseRequest' _ _ [] _ _ = throwIO $ NotEnoughLines []
373371parseRequest' conn port (firstLine:otherLines) remoteHost' src = do
374372 (method, rpath', gets, httpversion) <- parseFirst firstLine
@@ -417,24 +415,24 @@ data ChunkState = NeedLen
417415 | HaveLen Word
418416
419417chunkedSource :: MonadIO m
420− => I.IORef (C.Source m ByteString, ChunkState)
421− -> C.Source m ByteString
418+ => I.IORef (ResumableSource m ByteString, ChunkState)
419+ -> Source m ByteString
422420chunkedSource ipair = do
423421 (src, mlen) <- liftIO $ I.readIORef ipair
424422 go src mlen
425423 where
426424 go' src front = do
427− (src', (len, bs)) <- lift $ src C.$$+ front getLen
425+ (src', (len, bs)) <- lift $ src $$++ front getLen
428426 let src''
429427 | S.null bs = src'
430− | otherwise = C.yield bs >> src'
428+ | otherwise = fmapResume (yield bs >>) src'
431429 go src'' $ HaveLen len
432430
433431 go src NeedLen = go' src id
434432 go src NeedLenNewline = go' src (CB.take 2 >>)
435433 go src (HaveLen 0) = liftIO $ I.writeIORef ipair (src, HaveLen 0)
436434 go src (HaveLen len) = do
437− (src', mbs) <- lift $ src C.$$+ CL.head
435+ (src', mbs) <- lift $ src $$++ CL.head
438436 case mbs of
439437 Nothing -> liftIO $ I.writeIORef ipair (src', HaveLen 0)
440438 Just bs ->
@@ -445,15 +443,15 @@ chunkedSource ipair = do
445443 yield' src' mlen bs
446444 GT -> do
447445 let (x, y) = S.splitAt (fromIntegral len) bs
448− let src'' = C.yield y >> src'
446+ let src'' = fmapResume (yield y >>) src'
449447 yield' src'' NeedLenNewline x
450448
451449 yield' src mlen bs = do
452450 liftIO $ I.writeIORef ipair (src, mlen)
453− C.yield bs
451+ yield bs
454452 go src mlen
455453
456− getLen :: Monad m => C.Sink ByteString m (Word, ByteString)
454+ getLen :: Monad m => Sink ByteString m (Word, ByteString)
457455 getLen = do
458456 mbs <- CL.head
459457 case mbs of
@@ -621,31 +619,23 @@ sendResponse th req conn r = sendResponse' r
621619 sendResponse' (ResponseSource s hs bodyFlush)
622620 | hasBody s req = do
623621 let src = CL.sourceList [headers' needsChunked'] `mappend`
624− (if needsChunked' then body C.$= chunk else body)
625− src C.$$ builderToByteString C.=$ connSink conn th
622+ (if needsChunked' then body $= chunk else body)
623+ src $$ builderToByteString =$ connSink conn th
626624 return $ isKeepAlive hs
627625 | otherwise = liftIO $ do
628626 sendHeader $ headers' False
629627 T.tickle th
630628 return isPersist
631629 where
632− body = fmap2 (\x -> case x of
633− C.Flush -> flush
634− C.Chunk builder -> builder) bodyFlush
630+ body = mapOutput (\x -> case x of
631+ Flush -> flush
632+ Chunk builder -> builder) bodyFlush
635633 headers' = headers version s hs
636634 -- FIXME perhaps alloca a buffer per thread and reuse that in all
637635 -- functions below. Should lessen greatly the GC burden (I hope)
638636 needsChunked' = needsChunked hs
639− chunk :: C.Conduit Builder (ResourceT IO) Builder
640− chunk = C.NeedInput push close
641− push x = C.HaveOutput chunk (return ()) (chunkedTransferEncoding x)
642− close = C.HaveOutput (C.Done Nothing ()) (return ()) chunkedTransferTerminator
643−
644−fmap2 :: Functor m => (o1 -> o2) -> C.Pipe i o1 m r -> C.Pipe i o2 m r
645−fmap2 f (C.HaveOutput p c o) = C.HaveOutput (fmap2 f p) c (f o)
646−fmap2 f (C.NeedInput p c) = C.NeedInput (fmap2 f . p) (fmap2 f c)
647−fmap2 f (C.PipeM mp c) = C.PipeM (fmap (fmap2 f) mp) c
648−fmap2 _ (C.Done i x) = C.Done i x
637+ chunk :: Conduit Builder (ResourceT IO) Builder
638+ chunk = await >>= maybe (yield chunkedTransferTerminator) (\x -> yield (chunkedTransferEncoding x) >> chunk)
649639
650640parseHeaderNoAttr :: ByteString -> H.Header
651641parseHeaderNoAttr s =
@@ -653,31 +643,31 @@ parseHeaderNoAttr s =
653643 rest' = S.dropWhile (\c -> c == 32 || c == 9) $ S.drop 1 rest
654644 in (CI.mk k, rest')
655645
656−connSource :: Connection -> T.Handle -> C.Source (ResourceT IO) ByteString
646+connSource :: Connection -> T.Handle -> Source (ResourceT IO) ByteString
657647connSource Connection { connRecv = recv } th =
658648 src
659649 where
660− src = C.PipeM (do
650+ src = do
661651 bs <- liftIO recv
662652 if S.null bs
663− then return $ C.Done Nothing ()
653+ then return ()
664654 else do
665655 when (S.length bs >= 2048) $ liftIO $ T.tickle th
666− return (C.HaveOutput src (return ()) bs))
667− (return ())
656+ yield bs
657+ src
668658
669659-- | Use 'connSendAll' to send this data while respecting timeout rules.
670−connSink :: Connection -> T.Handle -> C.Sink B.ByteString (ResourceT IO) ()
660+connSink :: Connection -> T.Handle -> Sink B.ByteString (ResourceT IO) ()
671661connSink Connection { connSendAll = send } th =
672662 sink
673663 where
674− sink = C.NeedInput push close
664+ sink = await >>= maybe close push
675665 close = liftIO (T.resume th)
676− push x = C.PipeM (liftIO $ do
677− T.resume th
678− send x
679− T.pause th
680− return sink) (liftIO $ T.resume th)
666+ push x = do
667+ liftIO $ T.resume th
668+ liftIO $ send x
669+ liftIO $ T.pause th
670+ sink
681671 -- We pause timeouts before passing control back to user code. This ensures
682672 -- that a timeout will only ever be executed when Warp is in control. We
683673 -- also make sure to resume the timeout after the completion of user code
@@ -699,7 +689,7 @@ data Settings = Settings
699689 , settingsOnOpen :: IO () -- ^ What to do when a connection is open. Default: do nothing.
700690 , settingsOnClose :: IO () -- ^ What to do when a connection is close. Default: do nothing.
701691 , settingsTimeout :: Int -- ^ Timeout value in seconds. Default value: 30
702− , settingsIntercept :: Request -> Maybe (C.Source (ResourceT IO) S.ByteString -> Connection -> ResourceT IO ())
692+ , settingsIntercept :: Request -> Maybe (Source (ResourceT IO) S.ByteString -> Connection -> ResourceT IO ())
703693 , settingsManager :: Maybe Manager -- ^ Use an existing timeout manager instead of spawning a new one. If used, 'settingsTimeout' is ignored. Default is 'Nothing'
704694 }
705695
@@ -735,10 +725,11 @@ data THStatus = THStatus
735725 BSEndoList -- previously parsed lines
736726 BSEndo -- bytestrings to be prepended
737727
738−takeHeaders :: C.Sink ByteString (ResourceT IO) [ByteString]
728+takeHeaders :: Sink ByteString (ResourceT IO) [ByteString]
739729takeHeaders =
740− C.NeedInput (push (THStatus 0 id id)) close
730+ await >>= maybe close (push (THStatus 0 id id))
741731 where
732+ close :: Sink ByteString (ResourceT IO) a
742733 close = throwIO IncompleteHeaders
743734
744735 push (THStatus len lines prepend) bs
@@ -752,7 +743,7 @@ takeHeaders =
752743 let len' = len + bsLen
753744 prepend' = prepend . S.append bs
754745 status = THStatus len' lines prepend'
755− in C.NeedInput (push status) close
746+ in await >>= maybe close (push status)
756747 -- Found a newline, but next line continues as a multiline header
757748 Just (end, True) ->
758749 let rest = S.drop (end + 1) bs
@@ -776,7 +767,7 @@ takeHeaders =
776767 rest = if start < bsLen
777768 then Just (SU.unsafeDrop start bs)
778769 else Nothing
779− in C.Done rest lines'
770+ in maybe (return ()) leftover rest >> return lines'
780771 -- more headers
781772 else
782773 let len' = len + start
@@ -787,7 +778,7 @@ takeHeaders =
787778 then let bs' = SU.unsafeDrop start bs
788779 in push status bs'
789780 -- no more bytes in this chunk, ask for more
790− else C.NeedInput (push status) close
781+ else await >>= maybe close (push status)
791782 where
792783 bsLen = S.length bs
793784 mnl = do
@@ -831,3 +822,6 @@ serverHeader hdrs = case lookup key hdrs of
831822 key = "Server"
832823 ver = B.pack $ "Warp/" ++ warpVersion
833824 server = (key, ver)
825+
826+fmapResume :: (Source m o1 -> Source m o2) -> ResumableSource m o1 -> ResumableSource m o2
827+fmapResume f (ResumableSource src m) = ResumableSource (f src) m