WAI/Warp ResponseSource clean up

Question

Can't figure out how to do a clean up when a HTTP client drops connection (or other real world happens). I've tried to wrap my Source into a addCleanup, but it isn't getting called.

Here's my minimal example of an infinte Source streaming bytestrings:

{-# LANGUAGE OverloadedStrings #-}

module Main where

import Network.Wai
import Network.HTTP.Types
import Network.Wai.Handler.Warp (run)
import Data.ByteString.Lazy.Char8 ()

import Control.Monad
import Control.Monad.Trans
import Control.Concurrent (threadDelay)

import Data.Conduit
import Blaze.ByteString.Builder (Builder)
import qualified Blaze.ByteString.Builder.ByteString as BBBB
import qualified Data.ByteString.Char8 as BS

stream :: Source (ResourceT IO) (Flush Builder)
stream = addCleanup (\_ -> liftIO $ putStrLn "cleanup.") $ do
    liftIO $ putStrLn "source started."
    yield Flush

    forever $ do
        yield $ bchunk "whatever"
        yield Flush
        liftIO $ threadDelay 10000

app :: Application
app req = do
    liftIO $ putStrLn "in the handler."
    return $ ResponseSource status200 [("Content-Type", "text/plain")] stream

main :: IO ()
main = run 3000 app

bchunk = Chunk . BBBB.fromByteString . BS.pack

When i hit it with a http request, "startup" notices appear and the stream starts puring data. However after i close a connection, no "cleanup." message appears and no actions performed thus leaking resources in real code.

Answer 1

I think the preferred method for cleanup is using the allocate or register functions defined on Control.Monad.Trans.Resource.MonadResource. This will get your handler called when your ResponseSource terminates, exception or not.

From looking at the addCleanup code, it's only used for regular (non-exceptional) completion.

{-# LANGUAGE OverloadedStrings #-}

module Main where

import Network.Wai
import Network.HTTP.Types
import Network.Wai.Handler.Warp (run)
import Data.ByteString.Lazy.Char8 ()

import Control.Monad
import Control.Monad.Trans
import Control.Monad.Trans.Resource
import Control.Concurrent (threadDelay)

import Data.Conduit
import Blaze.ByteString.Builder (Builder)
import qualified Blaze.ByteString.Builder.ByteString as BBBB
import qualified Data.ByteString.Char8 as BS

stream :: MonadResource m => Source m (Flush Builder)
stream = do
  -- the release key can be used for early cleanup
  _releaseKey <- lift . register $ putStrLn "cleanup."

  liftIO $ putStrLn "source started."
  yield Flush

  forever $ do
    yield $ bchunk "whatever"
    yield Flush
    liftIO $ threadDelay 10000

app :: Application
app _ = do
  liftIO $ putStrLn "in the handler."
  return $ ResponseSource status200 [("Content-Type", "text/plain")] stream

main :: IO ()
main = run 3000 app

bchunk :: String -> Flush Builder
bchunk = Chunk . BBBB.fromByteString . BS.pack