Memory leak when using Control.Concurrent.STM.TBQueue

Question

I have two threads: producer and consumer. Producer produces some (key,value) pairs and consumer inserts them into Map, wrapped into Data.IORef. I tried to use Control.Concurrent.BoundedChan for communication between producer and consumer, and it works fine (memory usage is constant), provided that I use BangPatterns where needed. The code is as follows:

{-# LANGUAGE BangPatterns #-}

module Main where

import qualified Data.Map.Strict as M
import Control.Concurrent
import Control.Concurrent.STM
import Control.Concurrent.STM.TBQueue
import qualified Control.Concurrent.BoundedChan as BC
import qualified Control.Concurrent.Chan.Unagi.Bounded as UBC
import qualified Data.Text as T
import System.Random
import Data.IORef
import Control.Monad

data Item = Item !Int !Int

data SinkState = SinkState {
  myMap :: !(M.Map Int Int)
}

testBCs = do
  chan <- BC.newBoundedChan 1000

  forkIO $ source chan
  sink chan
  where
    source chan = forever $ do
      threadDelay 500
      key <- getStdRandom (randomR (1,5000))
      value <- getStdRandom (randomR (1,1000000))
      BC.writeChan chan $ Item key value

    sink chan = do
      state <- newIORef SinkState {
          myMap = M.empty
        }
      forever $ do
        (Item key value) <- BC.readChan chan
        atomicModifyIORef' state (\s -> (s { myMap = myMap s `seq` M.insert key value (myMap s) }, ()))

Now, when I switch from BoundedChan to Control.Concurrent.STM.TBQueue, memory starts to leak:

testTBs = do
  chan <- atomically $ newTBQueue 1000

  forkIO $ source chan
  sink chan
  where
    source chan = forever $ do
      threadDelay 500
      key <- getStdRandom (randomR (1,5000))
      value <- getStdRandom (randomR (1,1000000))
      atomically $ writeTBQueue chan $ Item key value

    sink chan = do
      state <- newIORef SinkState {
          myMap = M.empty
        }
      forever $ do
        (Item key value) <- chan `seq` atomically $ readTBQueue chan
        atomicModifyIORef' state (\s -> (s { myMap = myMap s `seq` M.insert key value (myMap s) }, ()))

And the profiling results looks like this (ran with +RTS -hd):

So my questions are:

What's going on in the second case?
What are those stg_ap_2_upd_info and other symbols?
How to fix leak?

Switching out the IORef to a TVar alleviates the problem, though it's not obvious why. I prefer to keep all my atomic operations within one transaction whenever possible, so (imo) this is a preferrable change... however it's not an answer to why this particular case is problematic. Here's a version that works for me: https://gist.github.com/NathanHowell/b54e2805c10ff047dc46ec3e818afcd8 — Nathan Howell, Oct 27 '16 at 05:40
@NathanHowell, great, this works for me too! Please, make your comment an answer so I can properly upvote it, as it addresses an essential part of the problem. — user2155932, Oct 27 '16 at 11:08

score 1 · Accepted Answer · answered Oct 27 '16 at 13:58

Without digging in to the space leak itself, one solution is simply moving the consumer logic entirely into the STM. This is as simple as replacing the IORef with a TVar. But to fully take advantage of the STM the queue read and state updates should be placed into a single atomically block. Both operations will be executed in a transaction. A helpful side effect is that we've also gained exception safety.

import Control.Concurrent
import Control.Concurrent.STM
import Control.Monad
import System.Random

import qualified Data.Map.Strict as M

data Item = Item {-# UNPACK #-} !Int {-# UNPACK #-} !Int

data SinkState = SinkState {
  myMap :: !(M.Map Int Int)
}

main :: IO ()
main = do
  chan <- newTBQueueIO 1000

  forkIO . forever $ do
    threadDelay 500
    key <- getStdRandom $ randomR (1,5000)
    value <- getStdRandom $ randomR (1,1000000)
    atomically . writeTBQueue chan $ Item key value

  state <- newTVarIO SinkState {
      myMap = M.empty
    }

  forever . atomically $ do
    Item key value <- readTBQueue chan
    modifyTVar' state $ \s -> s { myMap = M.insert key value (myMap s) }

Memory leak when using Control.Concurrent.STM.TBQueue

1 Answers1