Space leak in Pipes with RWST

11
votes

A memory analysis of the following program shows that the noleak functions runs in constant memory while the leak function leaks memory in a linear fashion. dflemstr indicated that this might be due to RWST causing an infinite chain of allocations. Is this the case and what other solutions exists? I actually dont need the Writer monad.

Environment:

GHC 7.8.3 on ARCH 64 bit

ghc Pipe.hs -o Pipe -prof

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

import Pipes
import Control.Monad.Trans.RWS.Strict

main = leak

effectLeak :: Effect (RWST () () () IO) ()
effectLeak =
  (forever $ do
      liftIO . threadDelay $ 10000 * 1
      yield "Space") >->
  (forever $ do
      text <- await
      yield $ text ++ (" leak" :: String)) >->
  (forever $ do
      text <- await
      liftIO . print $ text
  )

effectNoleak :: Effect IO ()
effectNoleak =
  (forever $ do
      lift . threadDelay $ 10000 * 1
      yield "Space") >->
  (forever $ do
      text <- await
      yield $ text ++ (" leak" :: String)) >->
  (forever $ do
      text <- await
      lift . print $ text
  )

leak = (\e -> runRWST e () ()) . runEffect $ effectLeak

noleak = runEffect $ effectNoleak
2
haskellhaskell-pipes
Cannot reproduce this behaviour on Ubuntu 14.04 (64bit), using GHC 7.6.3, pipes 4.1.2 and transformers 0.3.0.0. Both noleak and leak run in constant memory. Unless you compiled without optimization. So your solution is probably "use -O2". – Zeta
I notice that the space leak disappears with -O2. However, when I turn on optimizations on a larger program using some TH-generated RWST Pipes then the space leak is still existent. Is there any other solution to get state and reader-like behaviour in the base monad? – prinsen
Looking at the generated heap and stack profile reveals that the ARR_WORDS (ByteString) only grows to about 2 MB but the types (->>*) and (*) together with the stack keeps growing seemingly forever. The IO version uses less than 1 MB of memory while I've let the RWST run until around 32 MB before killing it – prinsen

2 Answers

14
votes

Zeta is right, and the space leak is because of WriterT. WriterT and RWST (both the "strict" and lazy versions) always leak space no matter what monoid you use.

I wrote up a longer explanation of this here, but here is the summary: the only way to not leak space is to simulate WriterT using a StateT monad where tell is simulated using a strict put, like this:

newtype WriterT w m a = WriterT { unWriterT :: w -> m (a, w) }

instance (Monad m, Monoid w) => Monad (WriterT w m) where
    return a = WriterT $ \w -> return (a, w)
    m >>= f  = WriterT $ \w -> do
        (a, w') <- unWriterT m w
        unWriterT (f a) w'

runWriterT :: (Monoid w) => WriterT w m a -> m (a, w)
runWriterT m = unWriterT m mempty

tell :: (Monad m, Monoid w) => w -> WriterT w m ()
tell w = WriterT $ \w' ->
    let wt = w `mappend` w'
     in wt `seq` return ((), wt)

That's basically equivalent to:

type WriterT = StateT

runWriterT m = runStateT m mempty

tell w = do
    w' <- get
    put $! mappend w w'
12
votes

It seems like the Writer part of RWST is actually the culprit:

instance (Monoid w, Monad m) => Monad (RWST r w s m) where
    return a = RWST $ \ _ s -> return (a, s, mempty)
    m >>= k  = RWST $ \ r s -> do
        (a, s', w)  <- runRWST m r s
        (b, s'',w') <- runRWST (k a) r s'
        return (b, s'', w `mappend` w') -- mappend
    fail msg = RWST $ \ _ _ -> fail msg

As you can see, the writer uses a plain mappend. Since (,,) isn't strict in its arguments, w `mappend` w' builds a series of thunks, even tough the Monoid instance of () is rather trivial:

instance Monoid () where
        -- Should it be strict?
        mempty        = ()
        _ `mappend` _ = ()
        mconcat _     = ()

In order to fix this, you need to add strictness to w `mappend` w' in the tuple: 

        let wt = w `mappend` w'
        wt `seq` return (b, s'', wt)

However, if you don't need the Writer anyway, you can simply use ReaderT r (StateT st m) instead:

import Control.Monad.Trans.Reader
import Control.Monad.Trans.State.Strict

type RST r st m = ReaderT r (StateT st m)

runRST :: Monad m => RST r st m a -> r -> st -> m (a,st)
runRST rst r st = flip runStateT st . flip runReaderT r $ rst

However, given that this will force you to lift the computations to the correct monad, you might want to use the mtl package instead. The code will stay the same, but the imports will be the following in this case

import Control.Monad.Reader
import Control.Monad.State.Strict