Convert a strict ByteString Enumerator to a lazy ByteString

enumerator_to_lazy_bytestring.hs

{-# LANGUAGE PackageImports #-}
module Network.Wai.Util where

import Control.Concurrent
import Control.Concurrent.MVar
import Control.Exception (finally)
import "mtl" Control.Monad.Trans
import qualified Data.ByteString as B
import qualified Data.ByteString.Char8 as B8
import qualified Data.ByteString.Lazy as L
import Data.Enumerator (($$))
import qualified Data.Enumerator as E
import qualified Data.Enumerator.Binary as E
import Data.IORef
import Network.Wai
import System.IO.Unsafe

requestLBS :: (MVar Response -> L.ByteString -> IO ()) -> E.Iteratee B.ByteString IO Response
requestLBS f = do
  (ialive, mbs, mres) <- liftIO $ do
    mbs <- newEmptyMVar
    mres <- newEmptyMVar
    ialive <- newIORef True
    forkIO $ finally
      (runLBS mbs $ f mres)
      (writeIORef ialive False)
    return (ialive, mbs, mres)
  iterateLBS ialive mbs
  liftIO $ takeMVar mres

withLBS :: (L.ByteString -> IO ()) -> E.Iteratee B.ByteString IO ()
withLBS f = do
  (ialive, mbs) <- liftIO $ do
    mbs <- newEmptyMVar
    ialive <- newIORef True
    _ <- forkIO $ finally
      (runLBS mbs f)
      (writeIORef ialive False)
    return (ialive, mbs)
  iterateLBS ialive mbs

runLBS :: MVar [B.ByteString] -> (L.ByteString -> IO ()) -> IO ()
runLBS mbs f = f =<< evalLBS mbs

evalLBS :: MVar [B.ByteString] -> IO L.ByteString
evalLBS mbs = fmap L.fromChunks go
  where go = unsafeInterleaveIO $ do
               next <- takeMVar mbs
               case next of
                    [] -> return []
                    bs -> fmap (bs ++) go

iterateLBS :: IORef Bool -> MVar [B.ByteString] -> E.Iteratee B.ByteString IO ()
iterateLBS ialive mbs = E.continue go
  where go (E.Chunks []) = E.continue go

        go (E.Chunks cs) = do
          succ <- liftIO $ tryPut cs
          if succ
             then E.continue go
             else E.yield () $ E.Chunks cs

        go E.EOF = do
          liftIO $ tryPut []
          E.yield () E.EOF

        tryPut cs = do
          succ <- tryPutMVar mbs cs
          if succ
             then return True
             else waitPut cs

        waitPut cs = do
          alive <- readIORef ialive
          if alive
             then yield >> tryPut cs
             else return False

test = do
  let enum = E.enumFile "/tmp/in"
      out  = L.writeFile "/tmp/out"
  E.run_ (enum $$ withLBS out)

deadlock = do
  let enum = E.enumFile "/tmp/in"
      out lbs = print $ head $ L.toChunks lbs
  E.run_ (enum $$ withLBS out)