Reference Counting Utility for ResourceT

Show and Tell
#1

This came up because I wanted to perform tasks asynchronously and then close the file descriptor when done. Specifically I was using fstatat and openat while walking a directory, and I wanted to make sure that the directory walk didn’t close the file descriptor while there was still something that was planning on opening the file relative to that directory descriptor.

I can imagine it would be useful for any asynchronous resource-using code.


{-# LANGUAGE OverloadedRecordDot #-}

module X where

import Control.Category ((>>>))
import Control.Concurrent.Counter qualified as Counter (Counter, add, new, sub)
import Control.Monad (when)
import Control.Monad.IO.Class (MonadIO (..))
import Control.Monad.Trans.Resource (MonadResource, ReleaseKey, allocate, release)

data RefCounted a = MkRefCounted
  { refCountedValue :: !a
  , refCount :: !Counter.Counter
  , refCountKey :: !ReleaseKey
  }

newRefCounted :: Int -> (ReleaseKey, a) -> IO (RefCounted a)
newRefCounted count (key, value) = do
  counter <- Counter.new count
  pure
    MkRefCounted
      { refCountedValue = value
      , refCount = counter
      , refCountKey = key
      }

refCountedAllocate
  :: (MonadResource m)
  => Int
  -> IO a
  -> (a -> IO ())
  -> m (RefCounted a)
refCountedAllocate initCount acquire unAcquire =
  allocate acquire unAcquire >>= (newRefCounted initCount >>> liftIO)

takeCount :: (MonadResource m) => Int -> forall a. RefCounted a -> m Int
takeCount count refCounted = do
  current <- liftIO $ Counter.add refCounted.refCount count
  pure $! current + count

freeCount :: (MonadResource m) => Int -> forall a. RefCounted a -> m Int
freeCount count refCounted = do
  remaining <- liftIO $ Counter.sub refCounted.refCount count
  when (remaining == count) (release refCounted.refCountKey)
  pure $! remaining - count
6 Likes