{-# LANGUAGE OverloadedStrings #-}

-- | Internal logging implementation for Arbiter.
--
-- This module is not part of the public API.
module Arbiter.Worker.Logger.Internal
  ( logMessage
  , withJobContext
  , runHook
  ) where

import Arbiter.Core.Job.Types qualified as Job
import Control.Monad (when)
import Control.Monad.Catch (MonadMask)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Control.Monad.Logger qualified as ML
import Control.Monad.Logger.Aeson qualified as MLA
import Data.Aeson (KeyValue (..), object)
import Data.Aeson.KeyMap qualified as KM
import Data.Aeson.Types (Pair)
import Data.List.NonEmpty (NonEmpty, toList)
import Data.Text (Text)
import Data.Text qualified as T
import UnliftIO (MonadUnliftIO, tryAny)

import Arbiter.Worker.Logger (LogConfig (..), LogDestination (..), LogLevel (..))

-- | Log a message using the given config.
logMessage :: LogConfig -> LogLevel -> Text -> IO ()
logMessage :: LogConfig -> LogLevel -> Text -> IO ()
logMessage LogConfig
config LogLevel
level Text
msg = Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (LogLevel
level LogLevel -> LogLevel -> Bool
forall a. Ord a => a -> a -> Bool
>= LogConfig -> LogLevel
minLogLevel LogConfig
config) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
  extraCtx <- LogConfig -> IO [Pair]
additionalContext LogConfig
config
  runWithDestination (logDestination config) extraCtx level msg

-- | Run an action with job context attached to the thread.
--
-- This sets up thread-local context that will be included in all log messages
-- within the action.
withJobContext
  :: (MonadIO m, MonadMask m)
  => NonEmpty (Job.JobRead payload)
  -> m a
  -> m a
withJobContext :: forall (m :: * -> *) payload a.
(MonadIO m, MonadMask m) =>
NonEmpty (JobRead payload) -> m a -> m a
withJobContext NonEmpty (JobRead payload)
jobs = [Pair] -> m a -> m a
forall (m :: * -> *) a.
(MonadIO m, MonadMask m) =>
[Pair] -> m a -> m a
MLA.withThreadContext (NonEmpty (JobRead payload) -> [Pair]
forall payload. NonEmpty (JobRead payload) -> [Pair]
buildJobContext NonEmpty (JobRead payload)
jobs)

-- | Build structured context for a batch of jobs.
buildJobContext :: NonEmpty (Job.JobRead payload) -> [Pair]
buildJobContext :: forall payload. NonEmpty (JobRead payload) -> [Pair]
buildJobContext NonEmpty (JobRead payload)
jobs = [Key
"jobs" Key -> [Value] -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= (JobRead payload -> Value) -> [JobRead payload] -> [Value]
forall a b. (a -> b) -> [a] -> [b]
map ([Pair] -> Value
object ([Pair] -> Value)
-> (JobRead payload -> [Pair]) -> JobRead payload -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. JobRead payload -> [Pair]
forall payload. JobRead payload -> [Pair]
mkContext) (NonEmpty (JobRead payload) -> [JobRead payload]
forall a. NonEmpty a -> [a]
toList NonEmpty (JobRead payload)
jobs)]
  where
    mkContext :: Job.JobRead payload -> [Pair]
    mkContext :: forall payload. JobRead payload -> [Pair]
mkContext JobRead payload
job =
      [ Key
"job_id" Key -> Int64 -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= JobRead payload -> Int64
forall payload key q insertedAt.
Job payload key q insertedAt -> key
Job.primaryKey JobRead payload
job
      , Key
"job_attempts" Key -> Int32 -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= JobRead payload -> Int32
forall payload key q insertedAt.
Job payload key q insertedAt -> Int32
Job.attempts JobRead payload
job
      , Key
"job_group_key" Key -> Maybe Text -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= JobRead payload -> Maybe Text
forall payload key q insertedAt.
Job payload key q insertedAt -> Maybe Text
Job.groupKey JobRead payload
job
      , Key
"job_queue" Key -> Text -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= JobRead payload -> Text
forall payload key q insertedAt. Job payload key q insertedAt -> q
Job.queueName JobRead payload
job
      ]

-- | Run the logger with the appropriate destination and context.
runWithDestination :: LogDestination -> [Pair] -> LogLevel -> Text -> IO ()
runWithDestination :: LogDestination -> [Pair] -> LogLevel -> Text -> IO ()
runWithDestination LogDestination
dest [Pair]
ctx LogLevel
level Text
msg = case LogDestination
dest of
  LogDestination
LogStdout -> LoggingT IO () -> IO ()
forall (m :: * -> *) a. LoggingT m a -> m a
MLA.runStdoutLoggingT (LoggingT IO () -> IO ()) -> LoggingT IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ [Pair] -> LoggingT IO () -> LoggingT IO ()
forall (m :: * -> *) a.
(MonadIO m, MonadMask m) =>
[Pair] -> m a -> m a
MLA.withThreadContext [Pair]
ctx (LoggingT IO () -> LoggingT IO ())
-> LoggingT IO () -> LoggingT IO ()
forall a b. (a -> b) -> a -> b
$ LogLevel -> Text -> LoggingT IO ()
forall (m :: * -> *). MonadLogger m => LogLevel -> Text -> m ()
logAt LogLevel
level Text
msg
  LogDestination
LogStderr -> LoggingT IO () -> IO ()
forall (m :: * -> *) a. LoggingT m a -> m a
MLA.runStderrLoggingT (LoggingT IO () -> IO ()) -> LoggingT IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ [Pair] -> LoggingT IO () -> LoggingT IO ()
forall (m :: * -> *) a.
(MonadIO m, MonadMask m) =>
[Pair] -> m a -> m a
MLA.withThreadContext [Pair]
ctx (LoggingT IO () -> LoggingT IO ())
-> LoggingT IO () -> LoggingT IO ()
forall a b. (a -> b) -> a -> b
$ LogLevel -> Text -> LoggingT IO ()
forall (m :: * -> *). MonadLogger m => LogLevel -> Text -> m ()
logAt LogLevel
level Text
msg
  LogFastLogger LoggerSet
loggerSet -> LoggerSet -> LoggingT IO () -> IO ()
forall (m :: * -> *) a. LoggerSet -> LoggingT m a -> m a
MLA.runFastLoggingT LoggerSet
loggerSet (LoggingT IO () -> IO ()) -> LoggingT IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ [Pair] -> LoggingT IO () -> LoggingT IO ()
forall (m :: * -> *) a.
(MonadIO m, MonadMask m) =>
[Pair] -> m a -> m a
MLA.withThreadContext [Pair]
ctx (LoggingT IO () -> LoggingT IO ())
-> LoggingT IO () -> LoggingT IO ()
forall a b. (a -> b) -> a -> b
$ LogLevel -> Text -> LoggingT IO ()
forall (m :: * -> *). MonadLogger m => LogLevel -> Text -> m ()
logAt LogLevel
level Text
msg
  LogCallback LogLevel -> Text -> [Pair] -> IO ()
cb -> do
    threadCtx <- KeyMap Value -> [Pair]
forall v. KeyMap v -> [(Key, v)]
KM.toList (KeyMap Value -> [Pair]) -> IO (KeyMap Value) -> IO [Pair]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (KeyMap Value)
forall (m :: * -> *). (MonadIO m, MonadThrow m) => m (KeyMap Value)
MLA.myThreadContext
    cb level msg (threadCtx <> ctx)
  LogDestination
LogDiscard -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  where
    logAt :: (ML.MonadLogger m) => LogLevel -> Text -> m ()
    logAt :: forall (m :: * -> *). MonadLogger m => LogLevel -> Text -> m ()
logAt LogLevel
Debug = Text -> m ()
forall (m :: * -> *). MonadLogger m => Text -> m ()
ML.logDebugN
    logAt LogLevel
Info = Text -> m ()
forall (m :: * -> *). MonadLogger m => Text -> m ()
ML.logInfoN
    logAt LogLevel
Warning = Text -> m ()
forall (m :: * -> *). MonadLogger m => Text -> m ()
ML.logWarnN
    logAt LogLevel
Error = Text -> m ()
forall (m :: * -> *). MonadLogger m => Text -> m ()
ML.logErrorN

-- | Run an observability hook, catching and logging any exceptions.
runHook
  :: (MonadUnliftIO m)
  => LogConfig
  -> Text
  -- ^ Hook name (for logging)
  -> m ()
  -- ^ Hook action
  -> m ()
runHook :: forall (m :: * -> *).
MonadUnliftIO m =>
LogConfig -> Text -> m () -> m ()
runHook LogConfig
cfg Text
hookName m ()
action = do
  result <- m () -> m (Either SomeException ())
forall (m :: * -> *) a.
MonadUnliftIO m =>
m a -> m (Either SomeException a)
tryAny m ()
action
  case result of
    Left SomeException
e -> IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ LogConfig -> LogLevel -> Text -> IO ()
logMessage LogConfig
cfg LogLevel
Warning (Text -> IO ()) -> Text -> IO ()
forall a b. (a -> b) -> a -> b
$ Text
"Observability hook '" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
hookName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"' failed: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (SomeException -> String
forall a. Show a => a -> String
show SomeException
e)
    Right ()
_ -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()