content: {-# LANGUAGE RecordWildCards #-}
  module Semantic.Telemetry.AsyncQueue
  (
    AsyncQueue(..)
  , newAsyncQueue
  , newAsyncQueue'
  , writeAsyncQueue
  , closeAsyncQueue
  )
  where
  
  import Control.Concurrent.Async as Async
  import Control.Concurrent.STM.TBMQueue
  import Control.Monad
  import GHC.Conc
  
  -- | 'AsyncQueue' represents a 'TBMQueue' that's drained from a separate thread.
  -- It is intended to be used to queue data from a pure function and then process
  -- that data in IO on a separate thread. 'AsyncQueue' is parameterized by:
  --   * 'a'     - the type of message stored on the queue.
  --   * 'extra' - any other type needed to process messages on the queue.
  data AsyncQueue a extra
    = AsyncQueue
    { asyncQueue      :: TBMQueue a -- ^ The underlying 'TBMQueue'.
    , asyncQueueSink  :: Async ()   -- ^ A sink that will drain the queue.
    , asyncQueueExtra :: extra      -- ^ Any exta data the queue needs to use.
    }
  
  -- | Create a new AsyncQueue with the given capacity using the defaultSink.
  newAsyncQueue :: Int -> (extra -> a -> IO ()) -> extra -> IO (AsyncQueue a extra)
  newAsyncQueue i = newAsyncQueue' i . defaultSink
  
  -- | Create a new AsyncQueue with the given capacity, specifying a custom sink.
  newAsyncQueue' :: Int -> (extra -> TBMQueue a -> IO ()) -> extra -> IO (AsyncQueue a extra)
  newAsyncQueue' i f extra = do
    q <- newTBMQueueIO i
    s <- Async.async (f extra q)
    pure (AsyncQueue q s extra)
  
  -- | Write a message to the queue.
  writeAsyncQueue :: AsyncQueue a extra -> a -> IO ()
  writeAsyncQueue AsyncQueue{..} = void . atomically . tryWriteTBMQueue asyncQueue
  
  -- | Drain messages from the queue, calling the specified function for each message.
  defaultSink :: (extra -> a -> IO ()) -> extra -> TBMQueue a -> IO ()
  defaultSink f extra q = do
    msg <- atomically (readTBMQueue q)
    maybe (pure ()) go msg
    where go msg = f extra msg >> defaultSink f extra q
  
  -- | Close the queue.
  closeAsyncQueue :: AsyncQueue a extra -> IO ()
  closeAsyncQueue AsyncQueue{..} = do
    atomically (closeTBMQueue asyncQueue)
    Async.wait asyncQueueSink
  