haskell logger handle pattern

Main.hs

{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE NoFieldSelectors #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE ImplicitParams #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE DuplicateRecordFields #-}

import Control.Concurrent.Async (async, wait, mapConcurrently_)
import Control.Concurrent.STM (newTQueueIO, newEmptyTMVarIO, isEmptyTMVar, atomically, retry, writeTQueue, flushTQueue, putTMVar)
import Control.Monad (when)
import Control.Exception (bracket)
import Data.Functor.Contravariant (Contravariant(contramap))
import Prelude hiding (log)

data Logger msg = Logger {log :: msg -> IO (), stop :: IO ()}

modify :: (msg' -> msg) -> Logger msg -> Logger msg'
modify = contramap

instance Contravariant Logger where
  contramap f logger = logger{log = logger.log . f}

makeSimpleLogger :: IO (Logger String)
makeSimpleLogger = pure Logger {log = putStrLn, stop = pure ()}

makeAsyncLogger :: IO (Logger String)
makeAsyncLogger = do
  messageBox <- newTQueueIO
  killMsg <- newEmptyTMVarIO
  let loop = do
        (messages, willContinue) <- atomically do
          messages <- flushTQueue messageBox
          willContinue <- isEmptyTMVar killMsg
          when (willContinue && null messages) retry
          pure (messages, willContinue)
        putStr (unlines messages)
        when willContinue loop
  pid <- async loop
  let log message = atomically do
        writeTQueue messageBox message
      stop = do
        atomically do putTMVar killMsg ()
        wait pid
  pure Logger {log, stop}


-- Example

type WithLogger msg = (?logger :: Logger msg)

withLogger :: IO (Logger msg) -> (WithLogger msg => IO a) -> IO a
withLogger mk k =
  bracket mk (.stop) \logger ->
    let ?logger = logger in k

log :: WithLogger msg => msg -> IO ()
log = ?logger.log

example :: WithLogger String => IO ()
example = mapConcurrently_ log ["hello", "world", "!!!"]

main = do
  withLogger (modify ("[async] " <>) <$> makeAsyncLogger) do
    example
  withLogger (modify ("[simple] " <>) <$> makeSimpleLogger) do
    example

-- ghci> main
-- [async] hello
-- [async] world
-- [async] !!!
-- [simple[][s sihimemplpllleoe]
-- ]  w!o!r!l
-- d

package.yaml

name: hspg

dependencies:
  - base >= 4.15
  - stm
  - async

executables:
  hspg-exe:
    main: Main.hs