Thread-local storage in Haskell that doesn't leak memory.

TLS.hs

{-# LANGUAGE MagicHash #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE UnliftedFFITypes #-}
{-# OPTIONS_GHC -Wall #-}

module TLS
  ( TLS,
    newTLS,
    freeTLS,
    getTLS,
    setTLS,
  )
where

import Control.Exception
import Data.Foldable
import Data.IORef
import qualified Data.IntMap.Lazy as IM
import Data.Tuple
import GHC.Conc
import GHC.Exts
import GHC.Types
import GHC.Weak

newtype TLS a
  = TLS (IORef (IM.IntMap (Weak a)))

newTLS :: IO (TLS a)
newTLS = do
  r <- newIORef IM.empty
  _ <- mkWeakIORef r $ freeTLS $ TLS r
  pure $ TLS r

freeTLS :: TLS a -> IO ()
freeTLS (TLS r) = do
  m <- atomicModifyIORef' r (IM.empty,)
  traverse_ markDead $ IM.elems m

getTLS :: TLS a -> IO (Maybe a)
getTLS (TLS r) = do
  ThreadId tid# <- myThreadId
  k <- evaluate $ c_rts_getThreadId tid#
  mw <- IM.lookup k <$> readIORef r
  case mw of
    Just w -> deRefWeak w
    _ -> pure Nothing

setTLS :: TLS a -> a -> IO ()
setTLS (TLS r) a = do
  ThreadId tid# <- myThreadId
  k <- evaluate $ c_rts_getThreadId tid#
  mw <- atomicModifyIORef' r (swap . IM.updateLookupWithKey (\_ _ -> Nothing) k)
  case mw of
    Just w -> markDead w
    _ -> pure ()
  w' <- IO $ \s0 ->
    case mkWeak#
      tid#
      a
      ( case atomicModifyIORef' r $ \m -> (IM.delete k m, ()) of
          IO fin -> fin
      )
      s0 of
      (# s1, w# #) -> (# s1, Weak w# #)
  atomicModifyIORef' r $ \m -> (IM.insert k w' m, ())

markDead :: Weak a -> IO ()
markDead (Weak w#) = IO $ \s0 -> case finalizeWeak# w# s0 of
  (# s1, _, _ #) -> (# s1, () #)

foreign import ccall unsafe "rts_getThreadId"
  c_rts_getThreadId ::
    ThreadId# -> Int