{- |

If multiple threads are printing to the same handle at once, the messages can
interleave and become nonsensical. So instead of printing directly, we write
messages to a queue, and a single thread is responsible for reading from the
queue and writing to a handle.

-}
module SocketsAndPipes.Serve.Log ( withLogging, Write, writeException ) where

import Data.Functor                  ( ($>), (<&>), void )
import Control.Applicative           ( (<|>) )
import Control.Concurrent.Async      ( withAsync, wait )
import Control.Concurrent.STM.TQueue ( TQueue, newTQueue, writeTQueue, readTQueue, isEmptyTQueue )
import Control.Concurrent.STM.TVar   ( TVar, newTVar, writeTVar, readTVar )
import Control.Monad                 ( join, guard )
import Control.Monad.STM             ( STM, atomically )
import Control.Exception.Safe        ( SomeException, finally, displayException )
import System.IO                     ( Handle, hPutStrLn )

-- | The internal state of the logging system.
data Log =
   Log
     { Log -> TQueue Message
messageQueue :: TQueue Message
     , Log -> TVar RunState
runStateVar :: TVar RunState
     }

data RunState =
    Run      {- ^ Log state is initially 'Run'. -}
  | Shutdown {- ^ The state changes to 'Shutdown' when the rest of the program
                  has ended. In the 'Shutdown' state, the log-printing thread
                  continues printing until the queue is empty, then stops. -}
  deriving RunState -> RunState -> Bool
(RunState -> RunState -> Bool)
-> (RunState -> RunState -> Bool) -> Eq RunState
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: RunState -> RunState -> Bool
$c/= :: RunState -> RunState -> Bool
== :: RunState -> RunState -> Bool
$c== :: RunState -> RunState -> Bool
Eq

-- | A message that can be written to a log.
type Message = String

-- | Value returned after a message is written to the log.
data MessageWritten = MessageWritten

-- | Function that writes a message to the log.
type Write = Message -> IO MessageWritten

-- | Value returned at the very end when logging is shut down and the queue is empty.
data LoggingFinished = LoggingFinished

-- | A log is initially created with an empty queue, in the 'Run' state.
newLog :: STM Log
newLog :: STM Log
newLog = TQueue Message -> TVar RunState -> Log
Log (TQueue Message -> TVar RunState -> Log)
-> STM (TQueue Message) -> STM (TVar RunState -> Log)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> STM (TQueue Message)
forall a. STM (TQueue a)
newTQueue STM (TVar RunState -> Log) -> STM (TVar RunState) -> STM Log
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> RunState -> STM (TVar RunState)
forall a. a -> STM (TVar a)
newTVar RunState
Run

-- | This function typically encloses an entire `main` action.
withLogging ::
    Handle -- ^ What handle the log output shall be written to
    -> (Write -> IO ()) -- ^ Continuation provided with a thread-safe print function
    -> IO ()
withLogging :: Handle -> (Write -> IO ()) -> IO ()
withLogging Handle
h Write -> IO ()
go =
  do
    Log
l <- STM Log -> IO Log
forall a. STM a -> IO a
atomically STM Log
newLog
    IO () -> (Async () -> IO ()) -> IO ()
forall a b. IO a -> (Async a -> IO b) -> IO b
withAsync (Write -> IO ()
go (Log -> Write
writeToLog Log
l)) ((Async () -> IO ()) -> IO ()) -> (Async () -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Async ()
a1 ->
      IO LoggingFinished -> (Async LoggingFinished -> IO ()) -> IO ()
forall a b. IO a -> (Async a -> IO b) -> IO b
withAsync (Handle -> Log -> IO LoggingFinished
printFromLog Handle
h Log
l) ((Async LoggingFinished -> IO ()) -> IO ())
-> (Async LoggingFinished -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Async LoggingFinished
a2 ->
        do
          Async () -> IO ()
forall a. Async a -> IO a
wait Async ()
a1 IO () -> IO () -> IO ()
forall (m :: * -> *) a b. MonadMask m => m a -> m b -> m a
`finally` Log -> IO ()
requestLogStop Log
l
          LoggingFinished
LoggingFinished <- Async LoggingFinished -> IO LoggingFinished
forall a. Async a -> IO a
wait Async LoggingFinished
a2
          () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

requestLogStop :: Log -> IO ()
requestLogStop :: Log -> IO ()
requestLogStop Log
l = STM () -> IO ()
forall a. STM a -> IO a
atomically (Log -> STM ()
requestLogStopSTM Log
l)

requestLogStopSTM :: Log -> STM ()
requestLogStopSTM :: Log -> STM ()
requestLogStopSTM Log{ TVar RunState
runStateVar :: TVar RunState
runStateVar :: Log -> TVar RunState
runStateVar } = TVar RunState -> RunState -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar RunState
runStateVar RunState
Shutdown

writeToLog :: Log -> Write
writeToLog :: Log -> Write
writeToLog Log
l Message
x = STM MessageWritten -> IO MessageWritten
forall a. STM a -> IO a
atomically (Log -> Message -> STM MessageWritten
writeToLogSTM Log
l Message
x)

writeToLogSTM :: Log -> Message -> STM MessageWritten
writeToLogSTM :: Log -> Message -> STM MessageWritten
writeToLogSTM Log{ TQueue Message
messageQueue :: TQueue Message
messageQueue :: Log -> TQueue Message
messageQueue } Message
message =
    TQueue Message -> Message -> STM ()
forall a. TQueue a -> a -> STM ()
writeTQueue TQueue Message
messageQueue Message
message STM () -> MessageWritten -> STM MessageWritten
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> MessageWritten
MessageWritten

printFromLog :: Handle -> Log -> IO LoggingFinished
printFromLog :: Handle -> Log -> IO LoggingFinished
printFromLog Handle
h l :: Log
l@Log{ TQueue Message
messageQueue :: TQueue Message
messageQueue :: Log -> TQueue Message
messageQueue } = IO LoggingFinished
continue
  where
    continue :: IO LoggingFinished
    continue :: IO LoggingFinished
continue = IO (IO LoggingFinished) -> IO LoggingFinished
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (IO (IO LoggingFinished) -> IO LoggingFinished)
-> IO (IO LoggingFinished) -> IO LoggingFinished
forall a b. (a -> b) -> a -> b
$ STM (IO LoggingFinished) -> IO (IO LoggingFinished)
forall a. STM a -> IO a
atomically (STM (IO LoggingFinished)
a STM (IO LoggingFinished)
-> STM (IO LoggingFinished) -> STM (IO LoggingFinished)
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> STM (IO LoggingFinished)
b)

    a :: STM (IO LoggingFinished)
a = TQueue Message -> STM Message
forall a. TQueue a -> STM a
readTQueue TQueue Message
messageQueue STM Message
-> (Message -> IO LoggingFinished) -> STM (IO LoggingFinished)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \Message
x ->
          Handle -> Message -> IO ()
hPutStrLn Handle
h Message
x IO () -> IO LoggingFinished -> IO LoggingFinished
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> IO LoggingFinished
continue
    b :: STM (IO LoggingFinished)
b = (Bool -> STM ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> STM ()) -> STM Bool -> STM ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Log -> STM Bool
isReadyToStop Log
l) STM () -> IO LoggingFinished -> STM (IO LoggingFinished)
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$>
          LoggingFinished -> IO LoggingFinished
forall (m :: * -> *) a. Monad m => a -> m a
return LoggingFinished
LoggingFinished

shutdownRequested :: Log -> STM Bool
shutdownRequested :: Log -> STM Bool
shutdownRequested Log{ TVar RunState
runStateVar :: TVar RunState
runStateVar :: Log -> TVar RunState
runStateVar } =
    TVar RunState -> STM RunState
forall a. TVar a -> STM a
readTVar TVar RunState
runStateVar STM RunState -> (RunState -> Bool) -> STM Bool
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \RunState
s ->
        RunState
s RunState -> RunState -> Bool
forall a. Eq a => a -> a -> Bool
== RunState
Shutdown

isReadyToStop :: Log -> STM Bool
isReadyToStop :: Log -> STM Bool
isReadyToStop l :: Log
l@Log{ TQueue Message
messageQueue :: TQueue Message
messageQueue :: Log -> TQueue Message
messageQueue } =
    Bool -> Bool -> Bool
(&&) (Bool -> Bool -> Bool) -> STM Bool -> STM (Bool -> Bool)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Log -> STM Bool
shutdownRequested Log
l
         STM (Bool -> Bool) -> STM Bool -> STM Bool
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TQueue Message -> STM Bool
forall a. TQueue a -> STM Bool
isEmptyTQueue TQueue Message
messageQueue

writeException :: Write -> SomeException -> IO ()
writeException :: Write -> SomeException -> IO ()
writeException Write
w SomeException
e = IO MessageWritten -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (Write
w (SomeException -> Message
forall e. Exception e => e -> Message
displayException SomeException
e))