GlobalSharedMutableSingleton.hs

module LightStep.GlobalSharedMutableSingleton where

import Control.Concurrent.Async
import Control.Concurrent.STM
import Control.Monad.Catch
import Control.Monad.Except
import LightStep.Internal.Debug
import LightStep.LowLevel
import System.IO.Unsafe

{-# NOINLINE globalSharedMutableSingletonState #-}
globalSharedMutableSingletonState :: TBQueue Span
globalSharedMutableSingletonState = unsafePerformIO $ newTBQueueIO 100

withSingletonLightStep :: LightStepConfig -> IO () -> IO ()
withSingletonLightStep cfg action =
  race_ (action >> waitUntilQueueIsEmpty) $ do
    let work client = do
          d_ "Getting more spans"
          spans <- liftIO $ atomically $ do
            x <- readTBQueue globalSharedMutableSingletonState
            xs <- flushTBQueue globalSharedMutableSingletonState
            pure (x : xs)
          d_ $ "Got " <> show (length spans) <> " spans"
          reportSpans client spans
          d_ $ "Reported " <> show (length spans) <> " spans"
        shutdown client = do
          d_ "Getting the last spans before shutdown"
          spans <- liftIO $ atomically $ flushTBQueue globalSharedMutableSingletonState
          d_ $ "Got " <> show (length spans) <> " spans"
          when (not $ null spans) $
            reportSpans client spans
          d_ $ "Reported " <> show (length spans) <> " spans"
          closeClient client
          d_ "Client closed"
    runExceptT
      $ bracket
        (mkClient cfg)
        shutdown
      $ \client -> do
        let loop = do
              work client
              loop
        loop
    pure ()

submitSpan :: Span -> IO ()
submitSpan = atomically . writeTBQueue globalSharedMutableSingletonState

-- TODO: handle span batches larger that queue size
submitSpans :: Foldable f => f Span -> IO ()
submitSpans = atomically . mapM_ (writeTBQueue globalSharedMutableSingletonState)

waitUntilQueueIsEmpty :: IO ()
waitUntilQueueIsEmpty = atomically $ do
  e <- isEmptyTBQueue globalSharedMutableSingletonState
  when (not e) retry