WIP: Capture request id in federator correctly

wireapp · May 21, 2024 · b8113da · b8113da
1 parent bd4911f
commit b8113da
Show file tree

Hide file tree

Showing 5 changed files with 51 additions and 41 deletions.
diff --git a/services/federator/src/Federator/ExternalServer.hs b/services/federator/src/Federator/ExternalServer.hs
@@ -62,13 +62,11 @@ import Servant.API.Extended.Endpath
 import Servant.Client.Core
 import Servant.Server (Tagged (..))
 import Servant.Server.Generic
-import System.Logger (msg, val, (.=), (~~))
 import System.Logger.Message qualified as Log
 import Wire.API.Federation.Component
 import Wire.API.Federation.Domain
 import Wire.API.Routes.FederationDomainConfig
 import Wire.API.VersionInfo
-import Wire.Sem.Logger (info)
 
 -- | Used to get PEM encoded certificate out of an HTTP header
 newtype CertHeader = CertHeader X509.Certificate
@@ -116,13 +114,16 @@ server ::
   ) =>
   Manager ->
   Word16 ->
-  (Sem r Wai.Response -> Codensity IO Wai.Response) ->
+  (RequestId -> Sem r Wai.Response -> Codensity IO Wai.Response) ->
   API AsServer
 server mgr intPort interpreter =
   API
     { status = Health.status mgr "internal server" intPort,
       externalRequest = \component rpc mReqId remoteDomain remoteCert ->
-        Tagged $ \req respond -> runCodensity (interpreter (callInward component rpc mReqId remoteDomain remoteCert req)) respond
+        Tagged $ \req respond -> do
+          -- TODO: Log generated request ID
+          rid <- maybe (RequestId . Text.encodeUtf8 . UUID.toText <$> UUID.nextRandom) pure mReqId
+          runCodensity (interpreter rid (callInward component rpc rid remoteDomain remoteCert req)) respond
     }
 
 -- FUTUREWORK(federation): Versioning of the federation API.
@@ -139,22 +140,22 @@ callInward ::
   ) =>
   Component ->
   RPC ->
-  Maybe RequestId ->
+  RequestId ->
   Domain ->
   CertHeader ->
   Wai.Request ->
   Sem r Wai.Response
-callInward component (RPC rpc) mReqId originDomain (CertHeader cert) wreq = do
-  rid <- case mReqId of
-    Just r -> pure r
-    Nothing -> do
-      localRid <- liftIO $ RequestId . Text.encodeUtf8 . UUID.toText <$> UUID.nextRandom
-      info $
-        "request-id" .= localRid
-          ~~ "method" .= Wai.requestMethod wreq
-          ~~ "path" .= Wai.rawPathInfo wreq
-          ~~ msg (val "generated a new request id for local request")
-      pure localRid
+callInward component (RPC rpc) rid originDomain (CertHeader cert) wreq = do
+  -- rid <- case mReqId of
+  --   Just r -> pure r
+  --   Nothing -> do
+  --     localRid <- liftIO $ RequestId . Text.encodeUtf8 . UUID.toText <$> UUID.nextRandom
+  --     info $
+  --       "request-id" .= localRid
+  --         ~~ "method" .= Wai.requestMethod wreq
+  --         ~~ "path" .= Wai.rawPathInfo wreq
+  --         ~~ msg (val "generated a new request id for local request")
+  --     pure localRid
   incomingCounterIncr originDomain
   -- only POST is supported
   when (Wai.requestMethod wreq /= HTTP.methodPost) $

diff --git a/services/federator/src/Federator/InternalServer.hs b/services/federator/src/Federator/InternalServer.hs
@@ -45,15 +45,15 @@ import Network.Wai qualified as Wai
 import Polysemy
 import Polysemy.Error
 import Polysemy.Input
+import Polysemy.TinyLog
 import Servant.API
 import Servant.API.Extended.Endpath
 import Servant.Server (Tagged (..))
 import Servant.Server.Generic
-import System.Logger (msg, val, (.=), (~~))
+import System.Logger (field, msg)
 import System.Logger.Class qualified as Log
 import Wire.API.Federation.Component
 import Wire.API.Routes.FederationDomainConfig
-import Wire.Sem.Logger (Logger, debug, info)
 
 data API mode = API
   { status ::
@@ -90,13 +90,27 @@ server ::
   ) =>
   Manager ->
   Word16 ->
-  (Sem r Wai.Response -> Codensity IO Wai.Response) ->
+  (RequestId -> Sem r Wai.Response -> Codensity IO Wai.Response) ->
   API AsServer
 server mgr extPort interpreter =
   API
     { status = Health.status mgr "external server" extPort,
       internalRequest = \mReqId remoteDomain component rpc ->
-        Tagged $ \req respond -> runCodensity (interpreter (callOutward mReqId remoteDomain component rpc req)) respond
+        Tagged $ \req respond -> do
+          -- TODO: Log generated request ID
+          rid <- maybe (RequestId . T.encodeUtf8 . UUID.toText <$> UUID.nextRandom) pure mReqId
+          -- rid <- case
+          -- rid <- case mReqId of
+          --   Just r -> pure r
+          --   Nothing -> do
+          --     localRid <- liftIO $ RequestId . T.encodeUtf8 . UUID.toText <$> UUID.nextRandom
+          --     info $
+          --       "request-id" .= localRid
+          --         ~~ "method" .= Wai.requestMethod req
+          --         ~~ "path" .= Wai.rawPathInfo req
+          --         ~~ msg (val "generated a new request id for local request")
+          --     pure localRid
+          runCodensity (interpreter rid (callOutward rid remoteDomain component rpc req)) respond
     }
 
 callOutward ::
@@ -108,23 +122,14 @@ callOutward ::
     Member Metrics r,
     Member (Logger (Log.Msg -> Log.Msg)) r
   ) =>
-  Maybe RequestId ->
+  RequestId ->
   Domain ->
   Component ->
   RPC ->
   Wai.Request ->
   Sem r Wai.Response
-callOutward mReqId targetDomain component (RPC path) req = do
-  rid <- case mReqId of
-    Just r -> pure r
-    Nothing -> do
-      localRid <- liftIO $ RequestId . T.encodeUtf8 . UUID.toText <$> UUID.nextRandom
-      info $
-        "request-id" .= localRid
-          ~~ "method" .= Wai.requestMethod req
-          ~~ "path" .= Wai.rawPathInfo req
-          ~~ msg (val "generated a new request id for local request")
-      pure localRid
+callOutward rid targetDomain component (RPC path) req = do
+  warn $ msg ("request id for new request" :: String) . field "req-id" rid
   -- only POST is supported
   when (Wai.requestMethod req /= HTTP.methodPost) $
     throw InvalidRoute

diff --git a/services/federator/src/Federator/Response.hs b/services/federator/src/Federator/Response.hs
@@ -28,6 +28,7 @@ where
 import Control.Lens
 import Control.Monad.Codensity
 import Data.ByteString.Builder
+import Data.Id
 import Data.Kind
 import Data.Text qualified as T
 import Federator.Discovery
@@ -151,11 +152,11 @@ type AllEffects =
 
 -- | Run Sem action containing HTTP handlers. All errors have to been handled
 -- already by this point.
-runFederator :: Env -> Sem AllEffects Wai.Response -> Codensity IO Wai.Response
-runFederator env =
+runFederator :: Env -> RequestId -> Sem AllEffects Wai.Response -> Codensity IO Wai.Response
+runFederator env rid =
   runM
     . runEmbedded @IO @(Codensity IO) liftIO
-    . loggerToTinyLogReqId (view requestId env) (view applog env)
+    . loggerToTinyLogReqId rid (view applog env)
     . runWaiErrors
       @'[ ValidationError,
           RemoteError,

diff --git a/services/federator/test/unit/Test/Federator/ExternalServer.hs b/services/federator/test/unit/Test/Federator/ExternalServer.hs
@@ -23,6 +23,7 @@ import Control.Monad.Codensity
 import Data.ByteString qualified as BS
 import Data.Default
 import Data.Domain
+import Data.Id
 import Data.Sequence as Seq
 import Data.String.Conversions
 import Data.Text.Encoding qualified as Text
@@ -147,7 +148,7 @@ requestBrigSuccess =
         . mockDiscoveryTrivial
         . runInputConst noClientCertSettings
         . runInputConst scaffoldingFederationDomainConfigs
-        $ callInward Brig (RPC "get-user-by-handle") Nothing aValidDomain (CertHeader cert) request
+        $ callInward Brig (RPC "get-user-by-handle") (RequestId "test") aValidDomain (CertHeader cert) request
     let expectedCall = Call Brig "/federation/get-user-by-handle" [("X-Wire-API-Version", "v0")] "\"foo\"" aValidDomain
     assertEqual "one call to brig should be made" [expectedCall] actualCalls
     Wai.responseStatus res @?= HTTP.status200
@@ -175,7 +176,7 @@ requestBrigFailure =
         . mockDiscoveryTrivial
         . runInputConst noClientCertSettings
         . runInputConst scaffoldingFederationDomainConfigs
-        $ callInward Brig (RPC "get-user-by-handle") Nothing aValidDomain (CertHeader cert) request
+        $ callInward Brig (RPC "get-user-by-handle") (RequestId "test") aValidDomain (CertHeader cert) request
 
     let expectedCall = Call Brig "/federation/get-user-by-handle" [] "\"foo\"" aValidDomain
     assertEqual "one call to brig should be made" [expectedCall] actualCalls
@@ -205,7 +206,7 @@ requestGalleySuccess =
           . mockDiscoveryTrivial
           . runInputConst noClientCertSettings
           . runInputConst scaffoldingFederationDomainConfigs
-          $ callInward Galley (RPC "get-conversations") Nothing aValidDomain (CertHeader cert) request
+          $ callInward Galley (RPC "get-conversations") (RequestId "test") aValidDomain (CertHeader cert) request
       let expectedCall = Call Galley "/federation/get-conversations" [] "\"foo\"" aValidDomain
       embed $ assertEqual "one call to galley should be made" [expectedCall] actualCalls
       embed $ Wai.responseStatus res @?= HTTP.status200
@@ -338,6 +339,7 @@ testMethod =
 
 testInterpretter ::
   IORef [Call] ->
+  RequestId ->
   Sem
     '[ Metrics,
        Input FederationDomainConfigs,
@@ -353,7 +355,7 @@ testInterpretter ::
      ]
     Wai.Response ->
   Codensity IO Wai.Response
-testInterpretter serviceCallsRef =
+testInterpretter serviceCallsRef _ =
   liftIO
     . runM @IO
     . runOutputMonoidIORef @Call serviceCallsRef (: [])

diff --git a/services/federator/test/unit/Test/Federator/InternalServer.hs b/services/federator/test/unit/Test/Federator/InternalServer.hs
@@ -23,6 +23,7 @@ import Data.ByteString.Builder
 import Data.ByteString.Conversion
 import Data.Default
 import Data.Domain
+import Data.Id
 import Federator.Error.ServerError
 import Federator.InternalServer (callOutward)
 import Federator.Metrics
@@ -102,7 +103,7 @@ federatedRequestSuccess =
         . runInputConst settings
         . runInputConst (FederationDomainConfigs AllowDynamic [FederationDomainConfig (Domain "target.example.com") FullSearch FederationRestrictionAllowAll] 10)
         . assertMetrics
-        $ callOutward Nothing targetDomain Brig (RPC "get-user-by-handle") request
+        $ callOutward (RequestId "test") targetDomain Brig (RPC "get-user-by-handle") request
     Wai.responseStatus res @?= HTTP.status200
     body <- Wai.lazyResponseBody res
     body @?= "\"bar\""
@@ -147,5 +148,5 @@ federatedRequestFailureAllowList =
         . runInputConst settings
         . runInputConst (FederationDomainConfigs AllowDynamic [FederationDomainConfig (Domain "hello.world") FullSearch FederationRestrictionAllowAll] 10)
         . interpretMetricsEmpty
-        $ callOutward Nothing targetDomain Brig (RPC "get-user-by-handle") request
+        $ callOutward (RequestId "test") targetDomain Brig (RPC "get-user-by-handle") request
     eith @?= Left (FederationDenied targetDomain)