Why wrapping the Data.Binary.Put monad creates a memory leak? (Part 2)

Question

As in my previous question, I\'m trying to wrap the Data.Binary.Put monad into another monad so that later I can ask it questions like \"how many bytes it\'s going to write\

Answer 1

It looks like the monad transformer is too lazy. You can create a heap profile (without having to build it specially) by running the program with:

$ ./myprog +RTS -hT
$ hp2ps myprog.hp
$ open hp2ps.ps    # Or whichever viewer you have

In this case it's not particularly helpful, because it only shows lots of PAPs, FUN_1_0s and FUN_2_0s. This means the heap is made up of lots of partially applied functions, and functions of one argument and two arguments. This usually means that something is not evaluated enough. Monad transformers are somewhat notorious for this.

The workaround is to use a more strict monad transformers using continuation passing style. (his requires {-# LANGUAGE Rank2Types #-}.

newtype MyStateT s m a =
  MyStateT { unMyStateT :: forall r. (s -> a -> m r) -> s -> m r }

Continuation passing style means that instead of returning a result directly, we call another function, the continuation, with our result, in this case s and a. The instance definitions look a bit funny. To understand it read the link above (Wikipedia).

instance Monad m => Monad (MyStateT s m) where
  return x = MyStateT (\k s -> k s x)
  MyStateT f >>= kk = MyStateT (\k s ->
    f (\s' a -> unMyStateT (kk a) k s') s)

runMyStateT :: Monad m => MyStateT s m a -> s -> m (a, s)
runMyStateT (MyStateT f) s0 = f (\s a -> return (a, s)) s0

instance MonadTrans (MyStateT s) where
  lift act = MyStateT (\k s -> do a <- act; k s a)

type Out = MyStateT Integer P.PutM ()

Running it now gives constant space (the "maximum residency" bit):

$ ./so1 +RTS -s 
begin
end
   8,001,343,308 bytes allocated in the heap
     877,696,096 bytes copied during GC
          46,628 bytes maximum residency (861 sample(s))
          33,196 bytes maximum slop
            2 MB total memory in use (0 MB lost due to fragmentation)

Generation 0: 14345 collections,     0 parallel,  3.32s,  3.38s elapsed
Generation 1:   861 collections,     0 parallel,  0.08s,  0.08s elapsed

The downside of using such strict transformers is that you can no longer define MonadFix instances and certain laziness tricks no longer work.

Answer 2

I started playing with this and realized what the bigger problem is -- your algorithm has terrible complexity. Rather than computing the size of each child tree once, you compute it once for each time you call getSize. And you call getSize recursively. For each leaf node, getSize is called once for each time getSize is called on its parent. And getSize is called on each parent once for itself + once for each time getSize is called on any of its parents. So getsize is called at least geometrically in the depth of the tree. You need to cache the sizes to get something resembling a reasonable runtime.

That said, here's a version of the core functions that appears to run properly without a leak, although it's really crawling along for the reasons stated above:

type MyPut = S (Offset,Size) P.PutM

peal_1 :: (Monad m, Num t, Num t1) => S (t, t1) m a -> m a
peal_1 put = unS put (\o -> return) (0,0)

writeToFile :: String -> MyPut () -> IO ()
writeToFile path put =
  BL.writeFile path $ P.runPut $ (peal_1 put) >> return ()

getSize :: MyPut a -> MyPut Int
getSize x = S $ \f os -> unS (x >> getCurrentSize) f os

getCurrentOffset :: MyPut Int
getCurrentOffset = S $ \f os -> f os (fst os)

getCurrentSize :: MyPut Int
getCurrentSize = S $ \f os -> f os (snd os)

I also have to say I'm not sure if your logic is correct in general. My code preserves the current behavior while fixing the leak. I tested this by running it and your code on a cut-down data set and producing files that are bit-for-bit identical.

But for your large test data, this code wrote 6.5G before I killed it (the provided code exhausted heap well before then). I suspect but have not tested that the underlying calls in the put monad are getting run once for each call to getSize, even though the result of getSize is getting thrown away.

My proposed proper solution is posted as an answer to your other question: How do you save a tree data structure to binary file in Haskell