server: fix tracing bug where some errors prevent spans from being emitted

PR-URL: https://github.com/hasura/graphql-engine-mono/pull/7450
GitOrigin-RevId: 23f6c9cfea8e7ca64b39866d15d2e6187aaaa0d9

server/src-lib/Hasura/Backends/DataConnector/Adapter/Metadata.hs

@@ -4,6 +4,7 @@
 module Hasura.Backends.DataConnector.Adapter.Metadata () where
 
 import Control.Arrow.Extended
+import Control.Monad.Trans.Control (MonadBaseControl)
 import Data.Aeson qualified as J
 import Data.Aeson.Key qualified as K
 import Data.Aeson.KeyMap qualified as KM
@@ -80,6 +81,7 @@ resolveBackendInfo' ::
     Inc.ArrowDistribute arr,
     ArrowWriter (Seq (Either InconsistentMetadata MetadataDependency)) arr,
     MonadIO m,
+    MonadBaseControl IO m,
     HasHttpManagerM m
   ) =>
   Logger Hasura ->
@@ -100,6 +102,7 @@ resolveBackendInfo' logger = proc (invalidationKeys, optionsMap) -> do
         Inc.ArrowCache m arr,
         ArrowWriter (Seq (Either InconsistentMetadata MetadataDependency)) arr,
         MonadIO m,
+        MonadBaseControl IO m,
         HasHttpManagerM m
       ) =>
       (Inc.Dependency (Maybe (HashMap DC.DataConnectorName Inc.InvalidationKey)), DC.DataConnectorName, DC.DataConnectorOptions) `arr` Maybe DC.DataConnectorInfo
@@ -114,7 +117,7 @@ resolveBackendInfo' logger = proc (invalidationKeys, optionsMap) -> do
         |) metadataObj
 
     getDataConnectorCapabilities ::
-      MonadIO m =>
+      (MonadIO m, MonadBaseControl IO m) =>
       DC.DataConnectorOptions ->
       HTTP.Manager ->
       m (Either QErr DC.DataConnectorInfo)
@@ -130,7 +133,7 @@ resolveBackendInfo' logger = proc (invalidationKeys, optionsMap) -> do
       capabilitiesCase defaultAction capabilitiesAction errorAction capabilitiesU
 
 resolveSourceConfig' ::
-  MonadIO m =>
+  (MonadIO m, MonadBaseControl IO m) =>
   Logger Hasura ->
   SourceName ->
   DC.ConnSourceConfig ->

server/src-lib/Hasura/Eventing/EventTrigger.hs

@@ -353,6 +353,7 @@ processEventQueue logger httpMgr getSchemaCache EventEngineCtx {..} LockedEvents
         Has (L.Logger L.Hasura) r,
         Tracing.HasReporter io,
         MonadMask io,
+        MonadBaseControl IO io,
         BackendEventTrigger b
       ) =>
       EventWithSource b ->

server/src-lib/Hasura/Eventing/ScheduledTrigger.hs

@@ -116,6 +116,7 @@ where
 import Control.Concurrent.Extended (Forever (..), sleep)
 import Control.Concurrent.STM
 import Control.Lens (view)
+import Control.Monad.Trans.Control (MonadBaseControl)
 import Data.Aeson qualified as J
 import Data.Environment qualified as Env
 import Data.Has
@@ -217,6 +218,7 @@ generateCronEventsFrom startTime CronTriggerInfo {..} =
 
 processCronEvents ::
   ( MonadIO m,
+    MonadBaseControl IO m,
     Tracing.HasReporter m,
     MonadMetadataStorage (MetadataStorageT m)
   ) =>
@@ -269,6 +271,7 @@ processCronEvents logger httpMgr prometheusMetrics cronEvents getSC lockedCronEv
 
 processOneOffScheduledEvents ::
   ( MonadIO m,
+    MonadBaseControl IO m,
     Tracing.HasReporter m,
     MonadMetadataStorage (MetadataStorageT m)
   ) =>
@@ -316,6 +319,7 @@ processOneOffScheduledEvents
 
 processScheduledTriggers ::
   ( MonadIO m,
+    MonadBaseControl IO m,
     Tracing.HasReporter m,
     MonadMetadataStorage (MetadataStorageT m)
   ) =>
@@ -349,6 +353,7 @@ processScheduledEvent ::
     Has HTTP.Manager r,
     Has (L.Logger L.Hasura) r,
     MonadIO m,
+    MonadBaseControl IO m,
     Tracing.HasReporter m,
     MonadMetadataStorage m
   ) =>

server/src-lib/Hasura/RQL/DDL/Schema/Cache.hs

@@ -438,6 +438,7 @@ buildSchemaCacheRule logger env = proc (metadataNoDefaults, invalidationKeys, st
         Inc.ArrowDistribute arr,
         ArrowWriter (Seq (Either InconsistentMetadata MetadataDependency)) arr,
         MonadIO m,
+        MonadBaseControl IO m,
         HasHttpManagerM m
       ) =>
       (BackendConfigWrapper b, Inc.Dependency (BackendMap BackendInvalidationKeysWrapper)) `arr` BackendCache
@@ -455,6 +456,7 @@ buildSchemaCacheRule logger env = proc (metadataNoDefaults, invalidationKeys, st
         Inc.ArrowDistribute arr,
         ArrowWriter (Seq (Either InconsistentMetadata MetadataDependency)) arr,
         MonadIO m,
+        MonadBaseControl IO m,
         HasHttpManagerM m
       ) =>
       (Inc.Dependency (BackendMap BackendInvalidationKeysWrapper), [AB.AnyBackend BackendConfigWrapper]) `arr` BackendCache
@@ -473,6 +475,7 @@ buildSchemaCacheRule logger env = proc (metadataNoDefaults, invalidationKeys, st
         Inc.ArrowCache m arr,
         ArrowWriter (Seq (Either InconsistentMetadata MetadataDependency)) arr,
         MonadIO m,
+        MonadBaseControl IO m,
         MonadResolveSource m,
         HasHttpManagerM m,
         BackendMetadata b

server/src-lib/Hasura/RQL/DDL/Schema/Source.hs

@@ -359,7 +359,8 @@ runGetSourceTables ::
     MonadReader r m,
     MonadError Error.QErr m,
     Metadata.MetadataM m,
-    MonadIO m
+    MonadIO m,
+    MonadBaseControl IO m
   ) =>
   Env.Environment ->
   GetSourceTables ->
@@ -410,7 +411,8 @@ runGetTableInfo ::
     MonadReader r m,
     MonadError Error.QErr m,
     Metadata.MetadataM m,
-    MonadIO m
+    MonadIO m,
+    MonadBaseControl IO m
   ) =>
   Env.Environment ->
   GetTableInfo ->
@@ -453,7 +455,15 @@ lookupDataConnectorOptions dcName bmap =
    in (InsOrdHashMap.lookup dcName =<< backendConfig)
         `onNothing` (Error.throw400 Error.DataConnectorError ("Data connector named " <> Text.E.toTxt dcName <> " was not found in the data connector backend config"))
 
-querySourceSchema :: (MonadIO m, MonadError QErr m) => L.Logger L.Hasura -> HTTP.Manager.Manager -> Maybe DC.Types.SourceTimeout -> BaseUrl -> SourceName -> API.Config -> m API.SchemaResponse
+querySourceSchema ::
+  (MonadIO m, MonadBaseControl IO m, MonadError QErr m) =>
+  L.Logger L.Hasura ->
+  HTTP.Manager.Manager ->
+  Maybe DC.Types.SourceTimeout ->
+  BaseUrl ->
+  SourceName ->
+  API.Config ->
+  m API.SchemaResponse
 querySourceSchema logger manager timeout uri sourceName transformedConfig =
   Tracing.ignoreTraceT
     . flip Agent.Client.runAgentClientT (Agent.Client.AgentClientContext logger uri manager (DC.Types.sourceTimeoutMicroseconds <$> timeout))

server/src-lib/Hasura/RQL/Types/Metadata/Backend.hs

@@ -74,6 +74,7 @@ class
       Inc.ArrowDistribute arr,
       ArrowWriter (Seq (Either InconsistentMetadata MetadataDependency)) arr,
       MonadIO m,
+      MonadBaseControl IO m,
       HasHttpManagerM m
     ) =>
     Logger Hasura ->
@@ -89,7 +90,7 @@ class
   -- | Function that resolves the connection related source configuration, and
   -- creates a connection pool (and other related parameters) in the process
   resolveSourceConfig ::
-    (MonadIO m, MonadResolveSource m) =>
+    (MonadIO m, MonadBaseControl IO m, MonadResolveSource m) =>
     Logger Hasura ->
     SourceName ->
     SourceConnConfiguration b ->

server/src-lib/Hasura/RemoteSchema/SchemaCache/Build.hs

@@ -8,6 +8,7 @@ where
 
 import Control.Arrow.Extended
 import Control.Arrow.Interpret
+import Control.Monad.Trans.Control (MonadBaseControl)
 import Data.Aeson
 import Data.ByteString.Lazy qualified as BL
 import Data.Environment qualified as Env
@@ -40,6 +41,7 @@ buildRemoteSchemas ::
     ArrowWriter (Seq (Either InconsistentMetadata MetadataDependency)) arr,
     Inc.ArrowCache m arr,
     MonadIO m,
+    MonadBaseControl IO m,
     HasHttpManagerM m,
     Eq remoteRelationshipDefinition,
     ToJSON remoteRelationshipDefinition,

server/src-lib/Hasura/Server/App.hs

@@ -313,7 +313,7 @@ mkSpockAction serverCtx@ServerCtx {..} qErrEncoder qErrModifier apiHandler = do
 
   let runTraceT ::
         forall m1 a1.
-        (MonadIO m1, Tracing.HasReporter m1) =>
+        (MonadIO m1, MonadBaseControl IO m1, Tracing.HasReporter m1) =>
         Tracing.TraceT m1 a1 ->
         m1 a1
       runTraceT = do

server/src-lib/Hasura/Tracing.hs

@@ -33,6 +33,7 @@ import Control.Monad.Trans.Control
 import Data.Aeson qualified as J
 import Data.Aeson.Lens qualified as JL
 import Data.ByteString.Char8 qualified as Char8
+import Data.IORef
 import Data.String (fromString)
 import Hasura.Prelude
 import Hasura.Tracing.TraceId
@@ -57,19 +58,21 @@ type TracingMetadata = [(Text, Text)]
 newtype Reporter = Reporter
   { runReporter ::
       forall io a.
-      MonadIO io =>
+      (MonadIO io, MonadBaseControl IO io) =>
       TraceContext ->
       -- the current trace context
       Text ->
       -- human-readable name for this block of code
-      io (a, TracingMetadata) ->
-      -- the action whose execution we want to report, returning
-      -- any metadata emitted
+      IO TracingMetadata ->
+      -- an IO action that gets all of the metadata logged so far by the action
+      -- being traced
+      io a ->
+      -- the action we want to trace
       io a
   }
 
 noReporter :: Reporter
-noReporter = Reporter \_ _ -> fmap fst
+noReporter = Reporter \_ _ _ -> id
 
 -- | A type class for monads which support some way to report execution traces.
 --
@@ -124,6 +127,7 @@ samplingStateFromHeader = \case
 data TraceTEnv = TraceTEnv
   { tteTraceContext :: TraceContext,
     tteReporter :: Reporter,
+    tteMetadataRef :: IORef TracingMetadata,
     tteSamplingDecision :: SamplingDecision
   }
 
@@ -185,14 +189,14 @@ sampleOneInN denominator
 
 -- | The 'TraceT' monad transformer adds the ability to keep track of
 -- the current trace context.
-newtype TraceT m a = TraceT {unTraceT :: ReaderT TraceTEnv (WriterT TracingMetadata m) a}
+newtype TraceT m a = TraceT {unTraceT :: ReaderT TraceTEnv m a}
   deriving (Functor, Applicative, Monad, MonadIO, MonadFix, MonadMask, MonadCatch, MonadThrow, MonadBase b, MonadBaseControl b)
 
 instance MonadTrans TraceT where
-  lift = TraceT . lift . lift
+  lift = TraceT . lift
 
 instance MFunctor TraceT where
-  hoist f (TraceT rwma) = TraceT (hoist (hoist f) rwma)
+  hoist f (TraceT rwma) = TraceT (hoist f rwma)
 
 instance MonadError e m => MonadError e (TraceT m) where
   throwError = lift . throwError
@@ -208,15 +212,27 @@ instance (HasHttpManagerM m) => HasHttpManagerM (TraceT m) where
 -- | Run an action in the 'TraceT' monad transformer.
 -- 'runTraceT' delimits a new trace with its root span, and the arguments
 -- specify a name and metadata for that span.
-runTraceT :: (HasReporter m, MonadIO m) => SamplingPolicy -> Text -> TraceT m a -> m a
+runTraceT ::
+  (HasReporter m, MonadIO m, MonadBaseControl IO m) =>
+  SamplingPolicy ->
+  Text ->
+  TraceT m a ->
+  m a
 runTraceT policy name tma = do
   rep <- askReporter
   runTraceTWithReporter rep policy name tma
 
 runTraceTWith ::
-  MonadIO m => TraceContext -> Reporter -> SamplingPolicy -> Text -> TraceT m a -> m a
+  (MonadIO m, MonadBaseControl IO m) =>
+  TraceContext ->
+  Reporter ->
+  SamplingPolicy ->
+  Text ->
+  TraceT m a ->
+  m a
 runTraceTWith ctx rep policy name tma = do
   samplingDecision <- liftIO $ decideSampling (tcSamplingState ctx) policy
+  metadataRef <- liftIO $ newIORef []
   let subCtx =
         ctx
           { tcSamplingState =
@@ -224,15 +240,21 @@ runTraceTWith ctx rep policy name tma = do
           }
       report =
         case samplingDecision of
-          SampleNever -> fmap fst
-          SampleAlways -> runReporter rep ctx name
-  report . runWriterT $
-    runReaderT (unTraceT tma) (TraceTEnv subCtx rep samplingDecision)
+          SampleNever -> id
+          SampleAlways -> do
+            runReporter rep ctx name (readIORef metadataRef)
+  report $
+    runReaderT (unTraceT tma) (TraceTEnv subCtx rep metadataRef samplingDecision)
 
 -- | Run an action in the 'TraceT' monad transformer in an
 -- existing context.
 runTraceTInContext ::
-  (MonadIO m, HasReporter m) => TraceContext -> SamplingPolicy -> Text -> TraceT m a -> m a
+  (MonadIO m, MonadBaseControl IO m, HasReporter m) =>
+  TraceContext ->
+  SamplingPolicy ->
+  Text ->
+  TraceT m a ->
+  m a
 runTraceTInContext ctx policy name tma = do
   rep <- askReporter
   runTraceTWith ctx rep policy name tma
@@ -240,7 +262,12 @@ runTraceTInContext ctx policy name tma = do
 -- | Run an action in the 'TraceT' monad transformer in an
 -- existing context.
 runTraceTWithReporter ::
-  MonadIO m => Reporter -> SamplingPolicy -> Text -> TraceT m a -> m a
+  (MonadIO m, MonadBaseControl IO m) =>
+  Reporter ->
+  SamplingPolicy ->
+  Text ->
+  TraceT m a ->
+  m a
 runTraceTWithReporter rep policy name tma = do
   ctx <-
     TraceContext
@@ -252,7 +279,7 @@ runTraceTWithReporter rep policy name tma = do
 
 -- | Run an action in the 'TraceT' monad transformer while suppressing all
 -- tracing-related side-effects.
-ignoreTraceT :: MonadIO m => TraceT m a -> m a
+ignoreTraceT :: (MonadIO m, MonadBaseControl IO m) => TraceT m a -> m a
 ignoreTraceT = runTraceTWithReporter noReporter sampleNever ""
 
 -- | Monads which support tracing. 'TraceT' is the standard example.
@@ -267,6 +294,9 @@ class Monad m => MonadTrace m where
   -- | Ask for the current tracing reporter
   currentReporter :: m Reporter
 
+  -- | Ask for the current handle on the tracing metadata
+  currentMetadataRef :: m (IORef TracingMetadata)
+
   -- | Ask for the current sampling decision
   currentSamplingDecision :: m SamplingDecision
 
@@ -296,49 +326,55 @@ class Monad m => MonadTrace m where
 -- Laws:
 --
 -- > interpTraceT id (hoist f (TraceT x)) = interpTraceT f (TraceT x)
-interpTraceT ::
-  MonadTrace n =>
-  (m (a, TracingMetadata) -> n (b, TracingMetadata)) ->
-  TraceT m a ->
-  n b
-interpTraceT f (TraceT rwma) = do
+interpTraceT :: MonadTrace n => (m a -> n b) -> TraceT m a -> n b
+interpTraceT f (TraceT rma) = do
   ctx <- currentContext
   rep <- currentReporter
+  metadataRef <- currentMetadataRef
   samplingDecision <- currentSamplingDecision
-  (b, meta) <- f (runWriterT (runReaderT rwma (TraceTEnv ctx rep samplingDecision)))
-  attachMetadata meta
-  pure b
+  f (runReaderT rma (TraceTEnv ctx rep metadataRef samplingDecision))
 
 -- | If the underlying monad can report trace data, then 'TraceT' will
 -- collect it and hand it off to that reporter.
-instance MonadIO m => MonadTrace (TraceT m) where
+instance (MonadIO m, MonadBaseControl IO m) => MonadTrace (TraceT m) where
   -- Note: this implementation is so awkward because we don't want to give the
   -- derived MonadReader/Writer instances to TraceT
-  trace name ma = TraceT . ReaderT $ \env@(TraceTEnv ctx rep samplingDecision) -> do
-    case samplingDecision of
-      SampleNever -> runReaderT (unTraceT ma) env
-      SampleAlways -> do
-        spanId <- liftIO randomSpanId
-        let subCtx =
-              ctx
-                { tcCurrentSpan = spanId,
-                  tcCurrentParent = Just (tcCurrentSpan ctx)
-                }
-        lift . runReporter rep subCtx name . runWriterT $
-          runReaderT (unTraceT ma) (TraceTEnv subCtx rep samplingDecision)
+  trace name ma =
+    TraceT $
+      ReaderT $ \env@(TraceTEnv ctx rep _ samplingDecision) -> do
+        case samplingDecision of
+          SampleNever -> runReaderT (unTraceT ma) env
+          SampleAlways -> do
+            spanId <- liftIO randomSpanId
+            let subCtx =
+                  ctx
+                    { tcCurrentSpan = spanId,
+                      tcCurrentParent = Just (tcCurrentSpan ctx)
+                    }
+            metadataRef <- liftIO $ newIORef []
+            runReporter rep subCtx name (readIORef metadataRef) $
+              runReaderT
+                (unTraceT ma)
+                (TraceTEnv subCtx rep metadataRef samplingDecision)
 
   currentContext = TraceT (asks tteTraceContext)
 
   currentReporter = TraceT (asks tteReporter)
 
+  currentMetadataRef = TraceT (asks tteMetadataRef)
+
   currentSamplingDecision = TraceT (asks tteSamplingDecision)
 
-  attachMetadata = TraceT . tell
+  attachMetadata metadata =
+    TraceT $
+      ReaderT $ \env ->
+        liftIO $ modifyIORef' (tteMetadataRef env) (metadata ++)
 
 instance MonadTrace m => MonadTrace (ReaderT r m) where
   trace = mapReaderT . trace
   currentContext = lift currentContext
   currentReporter = lift currentReporter
+  currentMetadataRef = lift currentMetadataRef
   currentSamplingDecision = lift currentSamplingDecision
   attachMetadata = lift . attachMetadata
 
@@ -346,6 +382,7 @@ instance MonadTrace m => MonadTrace (StateT e m) where
   trace = mapStateT . trace
   currentContext = lift currentContext
   currentReporter = lift currentReporter
+  currentMetadataRef = lift currentMetadataRef
   currentSamplingDecision = lift currentSamplingDecision
   attachMetadata = lift . attachMetadata
 
@@ -353,6 +390,7 @@ instance MonadTrace m => MonadTrace (ExceptT e m) where
   trace = mapExceptT . trace
   currentContext = lift currentContext
   currentReporter = lift currentReporter
+  currentMetadataRef = lift currentMetadataRef
   currentSamplingDecision = lift currentSamplingDecision
   attachMetadata = lift . attachMetadata