WPB-4888: Implement request tracing across federation

changelog.d/5-internal/WPB-4888

@@ -0,0 +1 @@
+Request tracing across federated requests

libs/types-common/default.nix

@@ -19,7 +19,6 @@
 , cryptohash-sha1
 , crypton
 , currency-codes
-, data-default
 , generic-random
 , gitignoreSource
 , hashable
@@ -75,7 +74,6 @@ mkDerivation {
     cryptohash-sha1
     crypton
     currency-codes
-    data-default
     generic-random
     hashable
     http-api-data

libs/types-common/src/Data/Id.hs

@@ -71,7 +71,6 @@ import Data.ByteString.Char8 qualified as B8
 import Data.ByteString.Conversion
 import Data.ByteString.Lazy qualified as L
 import Data.Char qualified as Char
-import Data.Default (Default (..))
 import Data.Hashable (Hashable)
 import Data.OpenApi qualified as S
 import Data.OpenApi.Internal.ParamSchema (ToParamSchema (..))
@@ -88,6 +87,7 @@ import Data.UUID qualified as UUID
 import Data.UUID.V4
 import Imports
 import Servant (FromHttpApiData (..), ToHttpApiData (..))
+import System.Logger (ToBytes)
 import Test.QuickCheck
 import Test.QuickCheck.Instances ()
 
@@ -413,18 +413,15 @@ newtype RequestId = RequestId
       ToByteString,
       Hashable,
       NFData,
-      Generic
+      Generic,
+      ToBytes
     )
 
 instance ToSchema RequestId where
   schema =
     RequestId . encodeUtf8
       <$> (decodeUtf8 . unRequestId) .= text "RequestId"
 
--- | Returns "N/A"
-instance Default RequestId where
-  def = RequestId "N/A"
-
 instance ToJSON RequestId where
   toJSON (RequestId r) = A.String (decodeUtf8 r)
 
@@ -437,6 +434,9 @@ instance EncodeWire RequestId where
 instance DecodeWire RequestId where
   decodeWire = fmap RequestId . decodeWire
 
+instance FromHttpApiData RequestId where
+  parseUrlPiece = Right . RequestId . encodeUtf8
+
 -- Rendering Id values in JSON objects -----------------------------------------
 
 newtype IdObject a = IdObject {fromIdObject :: a}

libs/types-common/types-common.cabal

@@ -105,7 +105,6 @@ library
     , cryptohash-sha1        >=0.11.7.2
     , crypton                >=0.26
     , currency-codes         >=3.0.0.1
-    , data-default           >=0.5
     , generic-random         >=1.4.0.0
     , hashable               >=1.2
     , http-api-data

libs/wire-api-federation/src/Wire/API/Federation/API.hs

@@ -93,7 +93,7 @@ fedQueueClient ::
   FedQueueClient (NotificationComponent k) ()
 fedQueueClient payload = do
   env <- ask
-  let notif = fedNotifToBackendNotif @tag env.originDomain payload
+  let notif = fedNotifToBackendNotif @tag env.requestId env.originDomain payload
       msg =
         newMsg
           { msgBody = encode notif,

libs/wire-api-federation/src/Wire/API/Federation/BackendNotifications.hs

@@ -7,6 +7,7 @@ import Control.Exception
 import Control.Monad.Except
 import Data.Aeson
 import Data.Domain
+import Data.Id (RequestId)
 import Data.Map qualified as Map
 import Data.Text qualified as Text
 import Data.Text.Lazy.Encoding qualified as TL
@@ -31,7 +32,8 @@ data BackendNotification = BackendNotification
     -- this body, which could be very large and completely useless to the
     -- pusher. This also makes development less clunky as we don't have to
     -- create a sum type here for all types of notifications that could exist.
-    body :: RawJson
+    body :: RawJson,
+    requestId :: Maybe RequestId
   }
   deriving (Show, Eq)
 
@@ -41,7 +43,8 @@ instance ToJSON BackendNotification where
       [ "ownDomain" .= notif.ownDomain,
         "targetComponent" .= notif.targetComponent,
         "path" .= notif.path,
-        "body" .= TL.decodeUtf8 notif.body.rawJsonBytes
+        "body" .= TL.decodeUtf8 notif.body.rawJsonBytes,
+        "requestId" .= notif.requestId
       ]
 
 instance FromJSON BackendNotification where
@@ -51,6 +54,7 @@ instance FromJSON BackendNotification where
       <*> o .: "targetComponent"
       <*> o .: "path"
       <*> (RawJson . TL.encodeUtf8 <$> o .: "body")
+      <*> o .:? "requestId"
 
 type BackendNotificationAPI = Capture "name" Text :> ReqBody '[JSON] RawJson :> Post '[JSON] EmptyResponse
 
@@ -70,8 +74,8 @@ sendNotification env component path body =
       runFederatorClient env . void $
         clientIn (Proxy @BackendNotificationAPI) (Proxy @(FederatorClient c)) (withoutFirstSlash path) body
 
-enqueue :: Q.Channel -> Domain -> Domain -> Q.DeliveryMode -> FedQueueClient c a -> IO a
-enqueue channel originDomain targetDomain deliveryMode (FedQueueClient action) =
+enqueue :: Q.Channel -> RequestId -> Domain -> Domain -> Q.DeliveryMode -> FedQueueClient c a -> IO a
+enqueue channel requestId originDomain targetDomain deliveryMode (FedQueueClient action) =
   runReaderT action FedQueueEnv {..}
 
 routingKey :: Text -> Text
@@ -127,7 +131,8 @@ data FedQueueEnv = FedQueueEnv
   { channel :: Q.Channel,
     originDomain :: Domain,
     targetDomain :: Domain,
-    deliveryMode :: Q.DeliveryMode
+    deliveryMode :: Q.DeliveryMode,
+    requestId :: RequestId
   }
 
 data EnqueueError = EnqueueError String

libs/wire-api-federation/src/Wire/API/Federation/Client.hs

@@ -45,6 +45,7 @@ import Data.ByteString.Builder
 import Data.ByteString.Conversion (toByteString')
 import Data.ByteString.Lazy qualified as LBS
 import Data.Domain
+import Data.Id
 import Data.Sequence qualified as Seq
 import Data.Set qualified as Set
 import Data.Text.Encoding qualified as Text
@@ -74,7 +75,8 @@ data FederatorClientEnv = FederatorClientEnv
   { ceOriginDomain :: Domain,
     ceTargetDomain :: Domain,
     ceFederator :: Endpoint,
-    ceHttp2Manager :: Http2Manager
+    ceHttp2Manager :: Http2Manager,
+    ceOriginRequestId :: RequestId
   }
 
 data FederatorClientVersionedEnv = FederatorClientVersionedEnv
@@ -215,6 +217,7 @@ withHTTP2StreamingRequest successfulStatus req handleResponse = do
         toList (requestHeaders req)
           <> [(originDomainHeaderName, toByteString' (ceOriginDomain env))]
           <> [(HTTP.hAccept, HTTP.renderHeader (toList $ req.requestAccept))]
+          <> [("Wire-Origin-Request-Id", toByteString' $ ceOriginRequestId env)]
       req' =
         HTTP2.requestBuilder
           (requestMethod req)

libs/wire-api-federation/src/Wire/API/Federation/HasNotificationEndpoint.hs

@@ -19,6 +19,7 @@ module Wire.API.Federation.HasNotificationEndpoint where
 
 import Data.Aeson
 import Data.Domain
+import Data.Id
 import Data.Kind
 import Data.Proxy
 import Data.Text qualified as T
@@ -46,10 +47,11 @@ fedNotifToBackendNotif ::
   KnownSymbol (NotificationPath tag) =>
   KnownComponent (NotificationComponent k) =>
   ToJSON (Payload tag) =>
+  RequestId ->
   Domain ->
   Payload tag ->
   BackendNotification
-fedNotifToBackendNotif ownDomain payload =
+fedNotifToBackendNotif rid ownDomain payload =
   let p = T.pack . symbolVal $ Proxy @(NotificationPath tag)
       b = RawJson . encode $ payload
    in toNotif p b
@@ -60,5 +62,6 @@ fedNotifToBackendNotif ownDomain payload =
         { ownDomain = ownDomain,
           targetComponent = componentVal @(NotificationComponent k),
           path = path,
-          body = body
+          body = body,
+          requestId = Just rid
         }

services/background-worker/src/Wire/BackendNotificationPusher.hs

@@ -7,6 +7,7 @@ import Control.Monad.Catch
 import Control.Retry
 import Data.Aeson qualified as A
 import Data.Domain
+import Data.Id
 import Data.Map.Strict qualified as Map
 import Data.Set qualified as Set
 import Data.Text qualified as Text
@@ -96,6 +97,7 @@ pushNotification runningFlag targetDomain (msg, envelope) = do
         ceHttp2Manager <- asks http2Manager
         let ceOriginDomain = notif.ownDomain
             ceTargetDomain = targetDomain
+            ceOriginRequestId = fromMaybe (RequestId "N/A") notif.requestId
             fcEnv = FederatorClientEnv {..}
         liftIO $ either throwM pure =<< sendNotification fcEnv notif.targetComponent notif.path notif.body
         lift $ ack envelope

services/background-worker/test/Test/Wire/BackendNotificationPusherSpec.hs

@@ -10,6 +10,7 @@ import Data.Aeson qualified as Aeson
 import Data.ByteString.Builder qualified as Builder
 import Data.ByteString.Lazy qualified as LBS
 import Data.Domain
+import Data.Id
 import Data.Range
 import Data.Sequence qualified as Seq
 import Data.Text qualified as Text
@@ -62,7 +63,8 @@ spec = do
               { targetComponent = Brig,
                 ownDomain = origDomain,
                 path = "/on-user-deleted-connections",
-                body = RawJson $ Aeson.encode notifContent
+                body = RawJson $ Aeson.encode notifContent,
+                requestId = Just $ RequestId "N/A"
               }
       envelope <- newMockEnvelope
       let msg =
@@ -128,7 +130,8 @@ spec = do
               { targetComponent = Brig,
                 ownDomain = origDomain,
                 path = "/on-user-deleted-connections",
-                body = RawJson $ Aeson.encode notifContent
+                body = RawJson $ Aeson.encode notifContent,
+                requestId = Just $ RequestId "N/A"
               }
       envelope <- newMockEnvelope
       let msg =

services/brig/brig.cabal

@@ -269,7 +269,6 @@ library
     , cql
     , cryptobox-haskell          >=0.1.1
     , currency-codes             >=2.0
-    , data-default               >=0.5
     , data-timeout               >=0.3
     , dns
     , dns-util

services/brig/default.nix

@@ -192,7 +192,6 @@ mkDerivation {
     cql
     cryptobox-haskell
     currency-codes
-    data-default
     data-timeout
     dns
     dns-util

services/brig/src/Brig/API/Handler.hs

@@ -42,13 +42,12 @@ import Brig.Options (setAllowlistEmailDomains, setAllowlistPhonePrefixes)
 import Brig.Phone (Phone, PhoneException (..))
 import Control.Error
 import Control.Exception (throwIO)
-import Control.Lens (set, view)
+import Control.Lens (view)
 import Control.Monad.Catch (catches, throwM)
 import Control.Monad.Catch qualified as Catch
 import Control.Monad.Except (MonadError, throwError)
 import Data.Aeson (FromJSON)
 import Data.Aeson qualified as Aeson
-import Data.Default (def)
 import Data.Text qualified as Text
 import Data.Text.Encoding qualified as Text
 import Data.ZAuth.Validation qualified as ZV
@@ -59,7 +58,7 @@ import Network.Wai.Predicate (Media)
 import Network.Wai.Routing (Continue)
 import Network.Wai.Utilities.Error ((!>>))
 import Network.Wai.Utilities.Error qualified as WaiError
-import Network.Wai.Utilities.Request (JsonRequest, lookupRequestId, parseBody)
+import Network.Wai.Utilities.Request (JsonRequest, parseBody)
 import Network.Wai.Utilities.Response (addHeader, json, setStatus)
 import Network.Wai.Utilities.Server qualified as Server
 import Servant qualified
@@ -80,11 +79,10 @@ runHandler ::
   Continue IO ->
   IO ResponseReceived
 runHandler e r h k = do
-  let e' = set requestId (maybe def RequestId (lookupRequestId r)) e
   a <-
-    runBrigToIO e' (runExceptT h)
+    runBrigToIO e (runExceptT h)
       `catches` brigErrorHandlers (view applog e) (unRequestId (view requestId e))
-  either (onError (view applog e') r k) pure a
+  either (onError (view applog e) r k) pure a
 
 toServantHandler :: Env -> (Handler BrigCanonicalEffects) a -> Servant.Handler a
 toServantHandler env action = do

services/brig/src/Brig/App.hs

@@ -117,7 +117,6 @@ import Control.Lens hiding (index, (.=))
 import Control.Monad.Catch
 import Control.Monad.Trans.Resource
 import Data.ByteString.Conversion
-import Data.Default (def)
 import Data.Domain
 import Data.GeoIP2 qualified as GeoIp
 import Data.IP
@@ -274,7 +273,7 @@ newEnv o = do
         _metrics = mtr,
         _applog = lgr,
         _internalEvents = eventsQueue,
-        _requestId = def,
+        _requestId = RequestId "N/A",
         _usrTemplates = utp,
         _provTemplates = ptp,
         _tmTemplates = ttp,

services/brig/src/Brig/Federation/Client.hs

@@ -18,7 +18,7 @@
 -- FUTUREWORK: Remove this module all together.
 module Brig.Federation.Client where
 
-import Brig.App
+import Brig.App as Brig
 import Control.Lens
 import Control.Monad
 import Control.Monad.Catch (MonadMask, throwM)
@@ -27,7 +27,7 @@ import Control.Retry
 import Control.Timeout
 import Data.Domain
 import Data.Handle
-import Data.Id (ClientId, UserId)
+import Data.Id
 import Data.Qualified
 import Data.Range (Range)
 import Data.Text qualified as T
@@ -161,22 +161,25 @@ notifyUserDeleted self remotes = do
           . Log.field "error" (show FederationNotConfigured)
 
 -- | Enqueues notifications in RabbitMQ. Retries 3 times with a delay of 1s.
-enqueueNotification :: (MonadIO m, MonadMask m, Log.MonadLogger m) => Domain -> Domain -> Q.DeliveryMode -> MVar Q.Channel -> FedQueueClient c () -> m ()
+enqueueNotification :: (MonadIO m, MonadMask m, Log.MonadLogger m, MonadReader Env m) => Domain -> Domain -> Q.DeliveryMode -> MVar Q.Channel -> FedQueueClient c () -> m ()
 enqueueNotification ownDomain remoteDomain deliveryMode chanVar action = do
   let policy = limitRetries 3 <> constantDelay 1_000_000
   recovering policy [logRetries (const $ pure True) logError] (const go)
   where
     logError willRetry (SomeException e) status = do
+      rid <- view Brig.requestId
       Log.err $
         Log.msg @Text "failed to enqueue notification in RabbitMQ"
           . Log.field "error" (displayException e)
           . Log.field "willRetry" willRetry
           . Log.field "retryCount" status.rsIterNumber
+          . Log.field "request" rid
     go = do
+      rid <- view Brig.requestId
       mChan <- timeout (1 :: Second) (readMVar chanVar)
       case mChan of
         Nothing -> throwM NoRabbitMqChannel
-        Just chan -> liftIO $ enqueue chan ownDomain remoteDomain deliveryMode action
+        Just chan -> liftIO $ enqueue chan rid ownDomain remoteDomain deliveryMode action
 
 data NoRabbitMqChannel = NoRabbitMqChannel
   deriving (Show)
@@ -192,12 +195,14 @@ runBrigFederatorClient targetDomain action = do
   ownDomain <- viewFederationDomain
   endpoint <- view federator >>= maybe (throwE FederationNotConfigured) pure
   mgr <- view http2Manager
+  rid <- view Brig.requestId
   let env =
         FederatorClientEnv
           { ceOriginDomain = ownDomain,
             ceTargetDomain = targetDomain,
             ceFederator = endpoint,
-            ceHttp2Manager = mgr
+            ceHttp2Manager = mgr,
+            ceOriginRequestId = rid
           }
   liftIO (runFederatorClient env action)
     >>= either (throwE . FederationCallFailure) pure