open source co-trace (#122)

Summary:
Pull Request resolved: #122

This is just a small library for tracing. Moving it under `common/hs/util` so that we can use it in Glean

Reviewed By: simonmar

Differential Revision: D51713343

fbshipit-source-id: a611441aee605861bdd57c1f4a4424ebe0e17166

common/util/Control/Trace.hs

@@ -0,0 +1,35 @@
+{-
+  Copyright (c) Meta Platforms, Inc. and affiliates.
+  All rights reserved.
+
+  This source code is licensed under the BSD-style license found in the
+  LICENSE file in the root directory of this source tree.
+-}
+
+-- | A tracing library inspired by dcoutts Contravariant logging talk[1]
+--   and the co-log package.
+--
+--   [1] -  https://www.youtube.com/watch?v=qzOQOmmkKEM
+--
+--  Example of usage:
+--
+--  > tracer :: Tracer Text
+--  > tracer = vlogTextTracer 1
+--  >
+--  > main = traceMsg tracer "main" $ do
+--  >   putStrLn "Hello world"
+--
+module Control.Trace
+  ( Tracer
+  , logMsg
+  , traceMsg
+  , (>$<)
+  , vlogTracer
+  , vlogTracerWithPriority
+  , TraceWithPriority(..)
+  , vlogShowTracer
+  , vlogTextTracer
+  ) where
+
+import Control.Trace.Core
+import Control.Trace.VLog

common/util/Control/Trace/Core.hs

@@ -0,0 +1,97 @@
+{-
+  Copyright (c) Meta Platforms, Inc. and affiliates.
+  All rights reserved.
+
+  This source code is licensed under the BSD-style license found in the
+  LICENSE file in the root directory of this source tree.
+-}
+
+{-# OPTIONS -Wno-orphans #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE InstanceSigs #-}
+module Control.Trace.Core (
+  Tracer (..),
+  MonadTrace (..),
+  MonadMaskInstance (..),
+  logMsg,
+  traceMsg,
+  Contravariant,
+  (>$<),
+) where
+
+import Control.Exception (
+  Exception,
+ )
+import Control.Monad.Catch (
+  ExitCase (..),
+  MonadCatch,
+  MonadMask (generalBracket),
+  MonadThrow,
+  try,
+ )
+import Control.Monad.IO.Class (
+  MonadIO (..),
+ )
+import Data.Functor.Contravariant (
+  Contravariant (contramap),
+  (>$<),
+ )
+import GHC.Stack (
+  HasCallStack,
+  withFrozenCallStack,
+ )
+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
+  }
+
+-- Explicit record accessors to preserve call stacks
+
+logMsg :: (HasCallStack, MonadTrace m) => Tracer msg -> msg -> m ()
+logMsg logger msg = withFrozenCallStack $ logMsg_ logger msg
+
+traceMsg :: (HasCallStack, MonadTrace m) => Tracer msg -> msg -> m a -> m a
+traceMsg logger msg act = withFrozenCallStack $ traceMsg_ logger msg 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)
+
+instance Semigroup (Tracer msg) where
+  l1 <> l2 =
+    Tracer
+      { logMsg_ = \m -> logMsg_ l1 m *> logMsg_ l2 m
+      , traceMsg_ = \msg -> traceMsg_ l1 msg . traceMsg_ l2 msg
+      }
+
+-------------------------------------------------------------------------------
+-- 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)
+
+-- | Deriving 'MonadTrace' via 'MonadMask'
+newtype MonadMaskInstance m a = MonadMaskInstance (m a)
+  deriving
+    (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

@@ -0,0 +1,62 @@
+{-
+  Copyright (c) Meta Platforms, Inc. and affiliates.
+  All rights reserved.
+
+  This source code is licensed under the BSD-style license found in the
+  LICENSE file in the root directory of this source tree.
+-}
+
+module Control.Trace.Eventlog (
+  eventlogTracer,
+  Trace (..),
+) where
+
+import Control.Monad.Catch (
+  ExitCase (
+    ExitCaseAbort,
+    ExitCaseException
+  ),
+ )
+import Control.Trace.Core (
+  MonadTrace (..),
+  Tracer (traceMsg_),
+ )
+import Data.ByteString (ByteString)
+import Data.String (fromString)
+import Debug.Trace.Flags (userTracingEnabled)
+import OpenTelemetry.Eventlog (
+  beginSpan,
+  endSpan,
+  setTag,
+ )
+
+data Trace = Trace
+  { name :: !ByteString
+  , tags :: ![(ByteString, ByteString)]
+  }
+
+{- | A tracer that emits opentelemetry spans to the GHC eventlog.
+
+   Compile with `-c fbcode.hs_eventlog=True`, run with `+RTS -l`
+   and visualize the traces with `//common/hs/ghc-chrome-trace` or
+   the opentelemetry-extra package.
+-}
+eventlogTracer :: Tracer Trace
+eventlogTracer
+  | userTracingEnabled = mempty {traceMsg_ = trace}
+  | otherwise = mempty
+  where
+    trace :: MonadTrace m => Trace -> m b -> m b
+    trace Trace {..} act = bracketM acquire release (const act)
+      where
+        acquire = beginSpan name
+        release sp res = do
+          mapM_ (uncurry $ setTag sp) tags
+          case res of
+            ExitCaseException e -> do
+              setTag sp "error" "1"
+              setTag sp "errorMsg" (fromString $ show e)
+            ExitCaseAbort ->
+              setTag sp "abort " "1"
+            _ -> pure ()
+          endSpan sp

common/util/Control/Trace/VLog.hs

@@ -0,0 +1,119 @@
+{-
+  Copyright (c) Meta Platforms, Inc. and affiliates.
+  All rights reserved.
+
+  This source code is licensed under the BSD-style license found in the
+  LICENSE file in the root directory of this source tree.
+-}
+
+{-# LANGUAGE ViewPatterns #-}
+module Control.Trace.VLog (
+  vlogTracer,
+  TraceWithPriority (..),
+  vlogTracerWithPriority,
+  vlogShowTracer,
+  vlogTextTracer,
+) where
+
+import Control.Monad (when)
+import Control.Monad.Catch (
+  ExitCase (
+    ExitCaseAbort,
+    ExitCaseException,
+    ExitCaseSuccess
+  ),
+ )
+import Control.Monad.IO.Class
+import Control.Trace.Core
+import Data.Some
+import Data.Text (Text)
+import GHC.Stack
+import TextShow (
+  TextShow,
+  showt,
+ )
+import qualified Util.Log.String as String
+import Util.Log.Text
+
+vlogShowTracer :: TextShow a => (a -> Int) -> Tracer a
+vlogShowTracer =
+  vlogTracer
+    (\(showt -> x) -> ("BEGIN " <> x, \e -> "END" <> renderExitCase e <> x))
+    showt
+
+renderExitCase :: Some ExitCase -> Text
+renderExitCase (Some ExitCaseAbort {}) = "(aborted) "
+renderExitCase (Some (ExitCaseException e)) = "(" <> showt e <> ") "
+renderExitCase (Some ExitCaseSuccess {}) = " "
+
+vlogTextTracer :: Int -> Tracer Text
+vlogTextTracer p =
+  vlogTracer
+    (\x -> ("BEGIN " <> x, \e -> "END" <> renderExitCase e <> x))
+    id
+    (const p)
+
+data TraceWithPriority
+  = Skip
+  | T !Int !Text
+  | S !Int !String
+
+vlogTracerWithPriority :: Tracer TraceWithPriority
+vlogTracerWithPriority = Tracer {..}
+  where
+    logMsg_ :: (HasCallStack, MonadIO m) => TraceWithPriority -> m ()
+    logMsg_ Skip = pure ()
+    logMsg_ x = withFrozenCallStack $ case x of
+      T p t -> vlog p t
+      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
+      case msg of
+        T p t ->
+          bracketM
+            (vlog p ("BEGIN " <> t))
+            ( \() 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
+                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"
+
+vlogTracer ::
+  forall a.
+  -- | render BEGIN and END messages
+  (a -> (Text, Some ExitCase -> Text)) ->
+  -- | render LOG message
+  (a -> Text) ->
+  -- | Priority (use -1 to skip)
+  (a -> Int) ->
+  Tracer a
+vlogTracer beginend log_ prio = Tracer {..}
+  where
+    logMsg_ :: (HasCallStack, MonadIO m) => a -> m ()
+    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
+      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

common/util/fb-util.cabal

@@ -74,6 +74,9 @@ library
 
     exposed-modules:
         Control.Concurrent.Stream
+        Control.Trace
+        Control.Trace.Core
+        Control.Trace.VLog
         Data.MovingAverageRateLimiter
         Data.RateLimiterMap
         Foreign.CPP.Addressable
@@ -202,6 +205,7 @@ library
         QuickCheck,
         scientific,
         haskell-src-exts,
+        some,
         stm,
         base >=4.11.1.0 && <4.15,
         containers >=0.5.11 && <0.7,