{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}

module Plow.Logging.Async (withAsyncHandleTracer) where

import qualified Control.Monad.IO.Class
import Data.Conduit ((.|))
import qualified Data.Conduit as Conduit
import qualified Data.Conduit.TMChan as Conduit.TMChan
import Data.String (fromString)
import Data.Text (Text)
import qualified Data.Text.IO as T
import Data.Time (getCurrentTime)
import Plow.Logging (IOTracer (..), Tracer (..), traceWith)
import System.IO (Handle, hFlush)
import UnliftIO (MonadUnliftIO)
import UnliftIO.Async (link, withAsync)
import qualified UnliftIO.STM as STM

-- | Returns (in CPS) a 'IOTracer' that pushes messages to a thread-safe queue.
-- This 'IOTracer' won't block unless the queue is full (size is configurable with
-- queueSize)
--
-- An async thread that continuously consumes traces in queue by printing them to a
-- 'Handle' will be launched. Any exceptions thrown inside (or to) the thread will be
-- rethrown in the caller of this function
--
-- Example use
--
-- main =
--   withAsyncHandleTracer stdout 100 $ \tracer' -> do
--     -- We use contramap to convert the tracer to a tracer that accepts
--     -- domain-specic trace types and displays them as Text
--     let tracer = contramap displaySomeTrace tracer'
--     traceWith tracer (SomeTrace a b c)
--     ...
withAsyncHandleTracer :: MonadUnliftIO m => Handle -> Int -> (IOTracer Text -> m a) -> m a
withAsyncHandleTracer :: forall (m :: * -> *) a.
MonadUnliftIO m =>
Handle -> Int -> (IOTracer Text -> m a) -> m a
withAsyncHandleTracer Handle
handle Int
queueSize IOTracer Text -> m a
f = do
  TBMChan (UTCTime, Text)
chan <- forall (m :: * -> *) a. MonadIO m => STM a -> m a
STM.atomically forall a b. (a -> b) -> a -> b
$ forall a. Int -> STM (TBMChan a)
Conduit.TMChan.newTBMChan Int
queueSize
  forall (m :: * -> *) a b.
MonadUnliftIO m =>
m a -> (Async a -> m b) -> m b
withAsync (forall {m :: * -> *} {a}.
(MonadIO m, Show a) =>
TBMChan (a, Text) -> m ()
logConsumer TBMChan (UTCTime, Text)
chan) forall a b. (a -> b) -> a -> b
$ \Async ()
logConsumerThread -> do
    let tracer :: IOTracer Text
tracer = forall {a}. TBMChan (UTCTime, a) -> IOTracer a
asyncTracer TBMChan (UTCTime, Text)
chan
    forall (m :: * -> *) a. MonadIO m => Async a -> m ()
link Async ()
logConsumerThread
    a
res <- IOTracer Text -> m a
f IOTracer Text
tracer
    forall (x :: * -> *) (m :: * -> *) a.
TraceWith x m =>
x a -> a -> m ()
traceWith IOTracer Text
tracer Text
"exit" forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall {m :: * -> *} {a}. MonadIO m => TBMChan a -> m ()
waitUntilEmpty TBMChan (UTCTime, Text)
chan forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return a
res
  where
    logConsumer :: TBMChan (a, Text) -> m ()
logConsumer TBMChan (a, Text)
chan =
      forall (m :: * -> *) r. Monad m => ConduitT () Void m r -> m r
Conduit.runConduit forall a b. (a -> b) -> a -> b
$
        ( forall (m :: * -> *) a.
MonadIO m =>
TBMChan a -> ConduitT () a m ()
Conduit.TMChan.sourceTBMChan TBMChan (a, Text)
chan
            forall (m :: * -> *) a b c r.
Monad m =>
ConduitT a b m () -> ConduitT b c m r -> ConduitT a c m r
.| forall (m :: * -> *) i o r.
Monad m =>
(i -> ConduitT i o m r) -> ConduitT i o m ()
Conduit.awaitForever
              ( \(a
time, Text
msg) ->
                  forall (m :: * -> *) a. MonadIO m => IO a -> m a
Control.Monad.IO.Class.liftIO forall a b. (a -> b) -> a -> b
$ do
                    Handle -> Text -> IO ()
T.hPutStrLn Handle
handle forall a b. (a -> b) -> a -> b
$ forall a. IsString a => String -> a
fromString (forall a. Show a => a -> String
show a
time forall a. Semigroup a => a -> a -> a
<> String
": ") forall a. Semigroup a => a -> a -> a
<> Text
msg
                    Handle -> IO ()
hFlush Handle
handle
              )
        )

    asyncTracer :: TBMChan (UTCTime, a) -> IOTracer a
asyncTracer TBMChan (UTCTime, a)
chan = forall a.
(forall (m :: * -> *). MonadIO m => Tracer m a) -> IOTracer a
IOTracer forall a b. (a -> b) -> a -> b
$
      forall (m :: * -> *) a. (a -> m ()) -> Tracer m a
Tracer forall a b. (a -> b) -> a -> b
$ \a
msg -> do
        UTCTime
time <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
Control.Monad.IO.Class.liftIO forall a b. (a -> b) -> a -> b
$ IO UTCTime
getCurrentTime
        forall (m :: * -> *) r. Monad m => ConduitT () Void m r -> m r
Conduit.runConduit forall a b. (a -> b) -> a -> b
$
          forall (m :: * -> *) o i. Monad m => o -> ConduitT i o m ()
Conduit.yield (UTCTime
time, a
msg)
            forall (m :: * -> *) a b c r.
Monad m =>
ConduitT a b m () -> ConduitT b c m r -> ConduitT a c m r
.| forall (m :: * -> *) a z.
MonadIO m =>
TBMChan a -> ConduitT a z m ()
Conduit.TMChan.sinkTBMChan TBMChan (UTCTime, a)
chan

    waitUntilEmpty :: TBMChan a -> m ()
waitUntilEmpty TBMChan a
chan =
      forall (m :: * -> *) a. MonadIO m => STM a -> m a
STM.atomically forall a b. (a -> b) -> a -> b
$
        forall a. TBMChan a -> STM Bool
Conduit.TMChan.isEmptyTBMChan TBMChan a
chan forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          Bool
True -> forall (m :: * -> *) a. Monad m => a -> m a
return ()
          Bool
False -> forall a. STM a
STM.retrySTM