Test.hs

Test.hs

{-# LANGUAGE OverloadedStrings #-}

------------------------------------------------------------------------------
import           Control.Applicative
import           Control.Monad
import           Control.Monad.Trans
import           Control.Monad.Trans.Resource as R
import           Control.Concurrent.Chan.Lifted
import           Control.Concurrent.MVar.Lifted
------------------------------------------------------------------------------
import qualified Data.Conduit               as C
import qualified Data.Conduit.List          as C
import qualified Data.ByteString.Lazy       as BL
import           Data.Monoid                      ((<>))
import qualified Blaze.ByteString.Builder   as B
------------------------------------------------------------------------------
import           Network.Wai (Application, Response(..), Request(..), responseLBS)
import           Network.Wai.Handler.Warp (runSettings, defaultSettings, settingsOnException)
import           Network.HTTP.Types (status200, status404)
------------------------------------------------------------------------------

main :: IO ()
main = do
    session <- newSession
    runSettings defaultSettings { settingsOnException = print } $ app session

data Session = Session
    { sessionMessages :: Chan BL.ByteString
    , sessionStatus :: MVar ()
    }

newSession :: IO Session
newSession = Session
    <$> newChan
    <*> newMVar ()

app :: Session -> Application
app session req =
    case pathInfo req of
         ["send"]   -> send session req
         ["stream"] -> stream session req
         _          -> return $ responseLBS status404 [] ""

send :: Session -> Application
send session@(Session messages status) req = do
    mres <- tryTakeMVar status
    case mres of
         Just () -> echo >> return (respond "send full\n")  <* putMVar status ()
         Nothing -> echo >> return (respond "send empty\n")
    where
      echo = do
        body <- BL.fromChunks <$> (requestBody req C.$$ C.consume)
        writeChan messages body

stream :: Session -> Application
stream session@(Session messages status) req = do
    mres <- tryTakeMVar status
    case mres of
         Just () -> return $ ResponseSource status200 [] source
         Nothing -> return $ respond "busy\n"
    where
      source = addCleanup session $ loop 0
      loop n key = do
          liftIO $ print n
          msg <- readChan messages
          yieldAndFlush $ msg <> "\n"
          if n < 10
             then loop (n + 1) key
             else lift (R.release key)

yieldAndFlush :: Monad m => BL.ByteString -> C.Pipe l i (C.Flush B.Builder) u m ()
yieldAndFlush load = C.yield (C.Chunk $ B.fromLazyByteString load) >> C.yield C.Flush

respond :: BL.ByteString -> Response
respond = responseLBS status200 []

addCleanup :: Session
           -> (R.ReleaseKey -> C.Source (C.ResourceT IO) (C.Flush B.Builder))
           -> C.Source (C.ResourceT IO) (C.Flush B.Builder)
addCleanup session@(Session messages status) fsrc = do
    key <- lift $ R.register $ putMVar status ()
    C.addCleanup (flip unless $! R.release key)
                 (fsrc key)