Recycle.hs

-- | Deferred, Vulkan-style resource reclamation.
module Fragr.Recycle
  ( RecycleQueue (..)
  , RetireItem
  , newRecycleQueue
  , mkRetireItem
  , retireItem
  , acquireItem
  , releaseItem
  , collect
  ) where

import Control.Exception (onException)
import Control.Monad.IO.Class (MonadIO (..))
import Data.IORef (IORef, modifyIORef', newIORef, readIORef, writeIORef)
import Data.IntMap.Strict qualified as IntMap
import Data.Word (Word64)

import Fragr.Types (QueueId (..))

{- |
A Vulkan-style recycle queue for deferred resource reclamation. Retired
items carry the per-queue timeline values that must be reached before they
may be reclaimed and an in-use refcount; 'collect' destroys (or, for a
pooling backend, could hand back) exactly those whose timelines are all met
and whose refcount is zero.

It is a plain single-threaded 'IORef' cell; the library uses one from
'executeQueued', but it is a standalone utility a backend may drive itself.
-}
newtype RecycleQueue = RecycleQueue (IORef [RetireItem])

{- |
One resource awaiting reclamation: its per-queue timeline requirements, an
in-use refcount (see 'acquireItem' / 'releaseItem') and the action that
actually frees it.
-}
data RetireItem = RetireItem
  { retireEntry :: Int
  , retireReqs :: [(QueueId, Word64)]
  , retireRefs :: IORef Int
  , retireDestroy :: IO ()
  }

-- | A fresh, empty recycle queue.
{-# INLINEABLE newRecycleQueue #-}
newRecycleQueue :: (MonadIO m) => m RecycleQueue
newRecycleQueue = liftIO (RecycleQueue <$> newIORef [])

{- |
Build a retire item: an identifying tag (the resource's entry id), the
per-queue timeline values that must be reached before it may be reclaimed,
and the destroy action to run when it is. Starts with a zero refcount.
-}
{-# INLINE mkRetireItem #-}
mkRetireItem :: (MonadIO m) => Int -> [(QueueId, Word64)] -> IO () -> m RetireItem
mkRetireItem tag reqs destroy = liftIO do
  refs <- newIORef 0
  pure RetireItem{retireEntry = tag, retireReqs = reqs, retireRefs = refs, retireDestroy = destroy}

-- | Hand an item to the recycle queue for eventual reclamation.
{-# INLINE retireItem #-}
retireItem :: (MonadIO m) => RecycleQueue -> RetireItem -> m ()
retireItem (RecycleQueue ref) item = liftIO $ modifyIORef' ref (item :)

-- | Add an in-use reference; 'collect' will not reclaim while any are held.
{-# INLINE acquireItem #-}
acquireItem :: (MonadIO m) => RetireItem -> m ()
acquireItem item = liftIO $ modifyIORef' item.retireRefs (+ 1)

-- | Drop an in-use reference added by 'acquireItem'.
{-# INLINE releaseItem #-}
releaseItem :: (MonadIO m) => RetireItem -> m ()
releaseItem item = liftIO $ modifyIORef' item.retireRefs (subtract 1)

{- |
Given the currently-reached timeline value per queue, reclaim every retired
item whose requirements are all met and whose in-use refcount is zero,
running its destroy action. Items not yet reclaimable stay queued. Returns
the tags (entry ids) of the items reclaimed.
-}
{-# INLINEABLE collect #-}
collect :: (MonadIO m) => RecycleQueue -> [(QueueId, Word64)] -> m [Int]
collect (RecycleQueue ref) done = liftIO do
  items <- readIORef ref
  verdicts <- traverse decide items
  let
    reclaimed = do
      (item, True) <- verdicts
      pure item
    kept = do
      (item, False) <- verdicts
      pure item
  -- Store the survivors /before/ running destroy actions: a pooling backend's
  -- destroy may itself 'retireItem' a follow-up onto this same queue, and a
  -- post-destroy 'writeIORef' would clobber it. Should a destroy throw, the
  -- items not yet destroyed go back onto the queue instead of leaking.
  writeIORef ref kept
  destroyAll reclaimed
  pure (map (.retireEntry) reclaimed)
  where
    doneMap = IntMap.fromList do
      (QueueId q, v) <- done
      pure (q, v)
    met item = all (\(QueueId q, v) -> IntMap.findWithDefault 0 q doneMap >= v) item.retireReqs
    decide item = do
      refs <- readIORef item.retireRefs
      pure (item, refs <= 0 && met item)
    destroyAll = \case
      [] -> pure ()
      item : rest -> do
        item.retireDestroy `onException` modifyIORef' ref (rest <>)
        destroyAll rest