Subscriber.hs

{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}

{-# HLINT ignore "Avoid lambda" #-}
module Unclog.Subscriber
  ( -- * subscribing to a log channel
    Subscriber (..)
  , colourSubscriber
  , simpleSubscriber
  , fileSubscriber

    -- * run a logging situation
  , withLoggingWithSubscribers

    -- * creating, subscribing and reading from channels
  , newLogChan
  , newLogChanIO
  , subscribeLog
  , readLogEntry

    -- * helpers
  , withSubscriber
  , mkSubscriber
  , mkSubscriberSimple
  , bracketSubscriber
  , refSubscriber
  )
where

import Control.Concurrent.Async (cancelMany)
import Control.Monad (forever, when)
import Control.Monad.Codensity
import Data.ByteString.Builder qualified as BS
import Unclog.Common (LogEntry (level), LogLevel, PLogChan (..), SLogChan (..))
import Unclog.Render (colouredBuilder, simpleBuilder)
import UnliftIO

-- | a 'Subscriber' is something that can observe log entries and do something with them in the Monad @m@
newtype Subscriber m = MkSubscriber {runSubscriber :: Codensity m (LogEntry -> m ())}

-- | make a subscriber with a possible clean up (useful if you need to bracket your subscriber)
mkSubscriber :: (forall r. ((LogEntry -> m ()) -> m r) -> m r) -> Subscriber m
mkSubscriber k = MkSubscriber (Codensity k)
{-# INLINE mkSubscriber #-}

-- | make a subscriber from a simple log function
mkSubscriberSimple :: (LogEntry -> m ()) -> Subscriber m
mkSubscriberSimple k = MkSubscriber (pure k)
{-# INLINE mkSubscriberSimple #-}

-- | destruct a subscriber
withSubscriber :: Subscriber m -> ((LogEntry -> m ()) -> m b) -> m b
withSubscriber sub = runCodensity (runSubscriber sub)
{-# INLINE withSubscriber #-}

-- | write a log entry to a handle, printing all information, but use colour
colourSubscriber
  :: MonadIO m
  => LogLevel
  -- ^ the lowest loglevel to still log
  -> Handle
  -- ^ the handle to log to
  -> Subscriber m
colourSubscriber lvl hdl = mkSubscriberSimple \entry -> when (entry.level >= lvl) do
  liftIO $ BS.hPutBuilder hdl (colouredBuilder entry)
{-# INLINEABLE colourSubscriber #-}

-- | write a log entry to a handle, printing all information
simpleSubscriber
  :: MonadIO m
  => LogLevel
  -- ^ the lowest loglevel to still log
  -> Handle
  -- ^ the handle to log to
  -> Subscriber m
simpleSubscriber lvl hdl = mkSubscriberSimple \entry -> when (entry.level >= lvl) do
  liftIO $ BS.hPutBuilder hdl (simpleBuilder entry)
{-# INLINEABLE simpleSubscriber #-}

-- | create a simple subscriber that writes to a file
fileSubscriber
  :: MonadUnliftIO m
  => LogLevel
  -- ^ the lowest loglevel to still log
  -> FilePath
  -- ^ the file to log (append) to
  -> Subscriber m
fileSubscriber lvl fp = mkSubscriber \k ->
  withFile fp AppendMode \hdl ->
    withSubscriber (simpleSubscriber lvl hdl) k
{-# INLINEABLE fileSubscriber #-}

-- | build a subscriber that requires a resource
bracketSubscriber
  :: MonadUnliftIO m
  => m a
  -- ^ acquire the resource
  -> (a -> m b)
  -- ^ release the resource
  -> (a -> LogEntry -> m ())
  -- ^ how to write a log entry given a resource
  -> Subscriber m
bracketSubscriber acquire release entry = mkSubscriber \k ->
  bracket
    acquire
    release
    \a -> k (entry a)
{-# INLINEABLE bracketSubscriber #-}

-- | write log entries to an 'IORef'
refSubscriber :: MonadIO m => IORef [LogEntry] -> Subscriber m
refSubscriber ref = mkSubscriberSimple \entry -> atomicModifyIORef' ref \old -> (entry : old, ())
{-# INLINEABLE refSubscriber #-}

-- | it is not a priori unsafe to escape the chan out of the context, it will just mean that
--   all the subscribers are gone and nothing is logged anymore
withLoggingWithSubscribers
  :: forall m r
   . MonadUnliftIO m
  => [Subscriber m]
  -- ^ a list of subscribers
  -> (PLogChan -> m r)
  -- ^ a channel, a client can write to
  -> m r
withLoggingWithSubscribers subscribers k = do
  (chan, subchans) <- atomically do
    c <- newLogChan
    (c,) <$> traverse (\sub -> (,sub) <$> subscribeLog c) subscribers
  subAsyncs <- traverse (uncurry mkSubscriberAsync) subchans
  k chan `finally` liftIO (cancelMany subAsyncs)
 where
  mkSubscriberAsync :: SLogChan -> Subscriber m -> m (Async b)
  mkSubscriberAsync chan subscription = async $ runCodensity (runSubscriber subscription) \sub -> do
    let -- this mask is needed because if we receive a cancel in the middle of executing the subcriber, the
        -- subscriber is aborted and we're missing a log entry. This can happen when the subscribers are
        -- still reading but the main task as stopped writing and is trying to flush the subscribers
        runSubscription = forever $ mask_ do
          -- the normal case just reads the next log entry from the channel and runs the subscriber on it
          atomically (readLogEntry chan)
            >>= sub

        flushSubscription = do
          atomically (tryReadTChan (unSLogChan chan)) >>= \case
            Nothing -> pure ()
            Just a -> do
              sub a
              flushSubscription

    runSubscription `finally` flushSubscription
{-# INLINEABLE withLoggingWithSubscribers #-}

-- | create a new 'LogChan'
newLogChan :: STM PLogChan
newLogChan = MkPLogChan <$> newBroadcastTChan
{-# INLINE newLogChan #-}

-- | create a new 'LogChan' in 'IO'
newLogChanIO :: MonadIO m => m PLogChan
newLogChanIO = MkPLogChan <$> newBroadcastTChanIO
{-# INLINE newLogChanIO #-}

-- | subscribe to a publishing log chan to obtain a subscriber log chan
subscribeLog :: PLogChan -> STM SLogChan
subscribeLog = fmap MkSLogChan . dupTChan . unPLogChan
{-# INLINE subscribeLog #-}

-- | read one log entry from a subscriber log chan
readLogEntry :: SLogChan -> STM LogEntry
readLogEntry = readTChan . unSLogChan
{-# INLINE readLogEntry #-}