refactor

Summary:
`traceMsg_` could not easily be used to trace asynchronous computations like:

https://www.internalfb.com/code/fbsource/[4a571cb146aa]/fbcode/glean/client/hs/Glean/Remote.hs?lines=541-547

Refactored it to be more flexible while keeping the existing api

Reviewed By: josefs

Differential Revision: D55980077

fbshipit-source-id: d10c77641223049dbf35b941960eea4695d07542

Author: Pepe Iborra

common/util/Control/Trace/Core.hs

@@ -19,15 +19,11 @@ module Control.Trace.Core (
   (>$<),
 ) where
 
-import Control.Exception (
-  Exception,
- )
 import Control.Monad.Catch (
   ExitCase (..),
   MonadCatch,
   MonadMask (generalBracket),
   MonadThrow,
-  try,
  )
 import Control.Monad.IO.Class (
   MonadIO (..),
@@ -45,43 +41,41 @@ import Data.Coerce
 -- | A contravariant tracing abstraction
 data Tracer msg = Tracer
   { -- | Log a message
-    logMsg_ :: forall m. (HasCallStack, MonadTrace m) => msg -> m ()
-  , -- | Trace the begin and end of a computation
-    traceMsg_ :: forall a m. (HasCallStack, MonadTrace m) => msg -> m a -> m a
+    logMsg_ :: msg -> IO ()
+  , -- | Starts a trace and returns an action to end it
+    traceMsg_
+      :: forall a. HasCallStack => msg -> IO (ExitCase a -> IO ())
   }
 
--- Explicit record accessors to preserve call stacks
-
-logMsg :: (HasCallStack, MonadTrace m) => Tracer msg -> msg -> m ()
-logMsg logger msg = withFrozenCallStack $ logMsg_ logger msg
+logMsg :: (HasCallStack, MonadIO m) => Tracer msg -> msg -> m ()
+logMsg logger msg = withFrozenCallStack $ liftIO $ logMsg_ logger msg
 
-traceMsg :: (HasCallStack, MonadTrace m) => Tracer msg -> msg -> m a -> m a
-traceMsg logger msg act = withFrozenCallStack $ traceMsg_ logger msg act
+traceMsg ::
+  (HasCallStack, MonadTrace m) => Tracer msg -> msg -> m a -> m a
+traceMsg logger msg act = withFrozenCallStack $
+  bracketM (traceMsg_ logger msg) id (const act)
 
 instance Contravariant Tracer where
   contramap f (Tracer logf traceF) = Tracer (logf . f) (traceF . f)
 
 instance Monoid (Tracer msg) where
-  mempty = Tracer (\_ -> pure ()) (const id)
+  mempty = Tracer (\_ -> pure ()) (const $ pure $ const $ pure ())
 
 instance Semigroup (Tracer msg) where
   l1 <> l2 =
     Tracer
       { logMsg_ = \m -> logMsg_ l1 m *> logMsg_ l2 m
-      , traceMsg_ = \msg -> traceMsg_ l1 msg . traceMsg_ l2 msg
+      , traceMsg_ = \msg -> do
+        end1 <- traceMsg_ l1 msg
+        end2 <- traceMsg_ l2 msg
+        return (\res -> end2 res >> end1 res)
       }
 
--------------------------------------------------------------------------------
--- Exceptions
-
 class MonadIO m => MonadTrace m where
-  tryM :: Exception e => m a -> m (Either e a)
   bracketM :: IO a -> (a -> ExitCase b -> IO ()) -> (a -> m b) -> m b
 
 -- deriving via (MonadMaskInstance IO) instance MonadTrace IO
 instance MonadTrace IO where
-  tryM :: forall e a . Exception e => IO a -> IO (Either e a)
-  tryM = coerce (tryM @(MonadMaskInstance IO) @e @a)
   bracketM
     :: forall a b . IO a -> (a -> ExitCase b -> IO ()) -> (a -> IO b) -> IO b
   bracketM = coerce (bracketM @(MonadMaskInstance IO) @a @b)
@@ -92,6 +86,5 @@ newtype MonadMaskInstance m a = MonadMaskInstance (m a)
     (Applicative, Functor, Monad, MonadCatch, MonadIO, MonadMask, MonadThrow)
 
 instance (MonadIO m, MonadMask m) => MonadTrace (MonadMaskInstance m) where
-  tryM = try
   bracketM acquire release =
     fmap fst . generalBracket (liftIO acquire) ((liftIO .) . release)

common/util/Control/Trace/Eventlog.hs

@@ -18,7 +18,6 @@ import Control.Monad.Catch (
   ),
  )
 import Control.Trace.Core (
-  MonadTrace (..),
   Tracer (traceMsg_),
  )
 import Data.ByteString (ByteString)
@@ -46,8 +45,7 @@ eventlogTracer
   | userTracingEnabled = mempty {traceMsg_ = trace}
   | otherwise = mempty
   where
-    trace :: MonadTrace m => Trace -> m b -> m b
-    trace Trace {..} act = bracketM acquire release (const act)
+    trace Trace {..} = release <$> acquire
       where
         acquire = beginSpan name
         release sp res = do

common/util/Control/Trace/VLog.hs

@@ -68,30 +68,24 @@ vlogTracerWithPriority = Tracer {..}
       S p s -> String.vlog p s
       Skip -> error "unreachable"
 
-    traceMsg_ :: (HasCallStack, MonadTrace m) => TraceWithPriority -> m b -> m b
-    traceMsg_ Skip act = act
-    traceMsg_ msg act = withFrozenCallStack $ do
+    traceMsg_ msg =
       case msg of
-        T p t ->
-          bracketM
-            (vlog p ("BEGIN " <> t))
-            ( \() res -> case res of
+        T p t -> do
+            vlog p ("BEGIN " <> t)
+            return ( \res -> case res of
                 ExitCaseSuccess {} -> vlog p ("END " <> t)
                 ExitCaseAbort {} -> vlog p ("ABORTED " <> t)
                 ExitCaseException e -> vlog p ("FAILED " <> t <> ": " <> showt e)
-            )
-            (\() -> act)
-        S p t ->
-          bracketM
-            (String.vlog p ("BEGIN " <> t))
-            ( \() res -> case res of
+              )
+        S p t -> do
+            String.vlog p ("BEGIN " <> t)
+            return ( \res -> case res of
                 ExitCaseSuccess {} -> String.vlog p ("END " <> t)
                 ExitCaseAbort {} -> String.vlog p ("ABORTED " <> t)
                 ExitCaseException e ->
                   String.vlog p ("FAILED " <> t <> ": " <> show e)
-            )
-            (\() -> act)
-        Skip -> error "unreachable"
+              )
+        Skip -> return $ const $ return ()
 
 vlogTracer ::
   forall a.
@@ -104,16 +98,15 @@ vlogTracer ::
   Tracer a
 vlogTracer beginend log_ prio = Tracer {..}
   where
-    logMsg_ :: (HasCallStack, MonadIO m) => a -> m ()
+    logMsg_ :: a -> IO ()
     logMsg_ msg =
-      withFrozenCallStack $
         let p = prio msg
          in when (p >= 0) $ vlog p $ log_ msg
 
-    traceMsg_ :: (HasCallStack, MonadTrace m) => a -> m b -> m b
-    traceMsg_ msg act = withFrozenCallStack $ do
+    traceMsg_ :: a -> IO (ExitCase b -> IO ())
+    traceMsg_ msg = do
       let p = prio msg
           (b, e) = beginend msg
       if p >= 0
-        then bracketM (vlog p b) (\() res -> vlog p (e $ mkSome res)) (\() -> act)
-        else act
+        then vlog p b >> return (\res -> vlog p (e $ mkSome res))
+        else return $ const $ return ()