WPB-28483: remove serial per-user round trips from MLS commit-bundle path

changelog.d/5-internal/WPB-28483

@@ -0,0 +1 @@
+MLS commit-bundles are processed with less sequential I/O: proposal references are resolved from a single store read, client and client-store updates fan out concurrently, and welcome pushes no longer block the response.

libs/galley-types/src/Galley/Types/Error.hs

@@ -42,7 +42,7 @@ data InternalError
   | NoPrekeyForUser
   | CannotCreateManagedConv
   | InternalErrorWithDescription LText
-  deriving (Eq)
+  deriving (Eq, Show)
 
 internalErrorDescription :: InternalError -> LText
 internalErrorDescription = message . internalErrorToWai

libs/wire-subsystems/src/Wire/ConversationSubsystem/Federation.hs

@@ -41,6 +41,7 @@ import Galley.Types.Error
 import Imports
 import Network.Wai.Utilities.Exception
 import Polysemy
+import Polysemy.Async (Async)
 import Polysemy.Error
 import Polysemy.Input
 import Polysemy.Internal.Kind (Append)
@@ -603,7 +604,8 @@ sendMLSCommitBundle ::
     Member TeamCollaboratorsSubsystem r,
     Member E.MLSCommitLockStore r,
     Member FeaturesConfigSubsystem r,
-    Member (Input ConversationSubsystemConfig) r
+    Member (Input ConversationSubsystemConfig) r,
+    Member Async r
   ) =>
   Domain ->
   MLSMessageSendRequest ->
@@ -891,8 +893,6 @@ onMLSMessageSent domain rmm =
 mlsSendWelcome ::
   ( Member (Error InternalError) r,
     Member NotificationSubsystem r,
-    Member ExternalAccess r,
-    Member P.TinyLog r,
     Member (Input (Maybe (MLSKeysByPurpose MLSPrivateKeys))) r,
     Member (Input (Local ())) r,
     Member Now r

libs/wire-subsystems/src/Wire/ConversationSubsystem/Interpreter.hs

@@ -29,6 +29,7 @@ import Data.Qualified
 import Imports
 import Network.Wai.Utilities.JSONResponse (JSONResponse)
 import Polysemy
+import Polysemy.Async (Async)
 import Polysemy.Error
 import Polysemy.Input
 import Polysemy.Resource (Resource)
@@ -123,7 +124,8 @@ interpretConversationSubsystem ::
     Member (Input (Maybe (MLSKeysByPurpose MLSPrivateKeys))) r,
     Member UserClientIndexStore r,
     Member (Input FanoutLimit) r,
-    Member TinyLog r
+    Member TinyLog r,
+    Member Async r
   ) =>
   InterpreterFor ConversationSubsystem r
 interpretConversationSubsystem = interpret $ \case

libs/wire-subsystems/src/Wire/ConversationSubsystem/MLS/CheckClients.hs

@@ -29,8 +29,11 @@ import Data.Map qualified as Map
 import Data.Qualified
 import Data.Set qualified as Set
 import Data.Tuple.Extra
+import Galley.Types.Error (InternalError (..))
 import Imports
 import Polysemy
+import Polysemy.Async (Async)
+import Polysemy.Async qualified as P
 import Polysemy.Error
 import Wire.API.Error
 import Wire.API.Error.Galley
@@ -50,7 +53,9 @@ checkClients ::
     Member (FederationAPIAccess FederatorClient) r,
     Member (ErrorS MLSClientMismatch) r,
     Member (ErrorS MLSIdentityMismatch) r,
-    Member (Error MLSProtocolError) r
+    Member (Error MLSProtocolError) r,
+    Member (Error InternalError) r,
+    Member Async r
   ) =>
   Local ConvOrSubConv ->
   CipherSuiteTag ->
@@ -59,9 +64,33 @@ checkClients ::
 checkClients lConvOrSub ciphersuite newCM = do
   let convOrSub = tUnqualified lConvOrSub
       cm = convOrSub.members
-  fmap catMaybes . forM (Map.assocs (unClientMap newCM)) $
-    \(qtarget, newclients) -> do
-      mClientData <- getClientData lConvOrSub ciphersuite qtarget
+      assocs = Map.assocs (unClientMap newCM)
+  -- Fetch client data from brig concurrently. getClientData is total with
+  -- respect to 'FederationError' (hushed inside getClientData): an inner
+  -- 'Nothing' is a legitimate "user unreachable" result.
+  --
+  -- sequenceConcurrently attaches an outer 'Maybe' to every child result.
+  -- Under galley's production stack (asyncToIOFinal below pure
+  -- runError/mapError interpreters, cf. Galley.App), an 'Error'-effect
+  -- throw inside a spawned child (e.g. RpcException/ParseException from
+  -- interpretBrigAccess) has no interpreter inside the async boundary and
+  -- Polysemy collapses the child result to 'Nothing'. That is a crashed
+  -- child, not "no client data", and must abort the commit with an
+  -- internal error instead of being conflated with the unreachable case.
+  --
+  -- Validation below runs serially so that Error-effect throws abort the
+  -- whole commit exactly as in the fully serial implementation.
+  mClientDatas <-
+    P.sequenceConcurrently $
+      flip fmap assocs $ \(qtarget, _) ->
+        getClientData lConvOrSub ciphersuite qtarget
+  clientDatas <-
+    forM mClientDatas $
+      maybe
+        (throw (InternalErrorWithDescription "Concurrent brig client-data fetch failed while processing commit"))
+        pure
+  fmap catMaybes . forM (zip assocs clientDatas) $
+    \((qtarget, newclients), mClientData) -> do
       unreachable <- case (mClientData, cmLookup qtarget cm) of
         -- user is already present, skip check in this case
         (_, Just existingClients) -> do
@@ -103,9 +132,9 @@ checkClients lConvOrSub ciphersuite newCM = do
           pure False
 
       -- Check that new leaf nodes are using the registered signature keys.
-      for_ mClientData $ \clientData ->
+      for_ mClientData $ \cd ->
         for_ (Map.assocs newclients) $ \(cid, (_, mKp)) ->
-          checkSignatureKey (fmap (.leafNode) mKp) (Map.lookup cid clientData.infoMap)
+          checkSignatureKey (fmap (.leafNode) mKp) (Map.lookup cid cd.infoMap)
 
       pure $ guard unreachable $> qtarget
 

libs/wire-subsystems/src/Wire/ConversationSubsystem/MLS/Commit/Core.hs

@@ -18,6 +18,7 @@
 module Wire.ConversationSubsystem.MLS.Commit.Core
   ( getCommitData,
     incrementEpoch,
+    incrementEpochNoRead,
     getClientInfo,
     getSingleClientInfo,
     checkSignatureKey,
@@ -70,7 +71,7 @@ import Wire.ExternalAccess
 import Wire.FederationAPIAccess
 import Wire.LegalHoldStore (LegalHoldStore)
 import Wire.NotificationSubsystem
-import Wire.ProposalStore (ProposalStore)
+import Wire.ProposalStore (ProposalStore, StoredProposal, getAllPendingProposals)
 import Wire.Sem.Now (Now)
 import Wire.Sem.Random (Random)
 import Wire.TeamCollaboratorsSubsystem
@@ -114,22 +115,27 @@ getCommitData ::
   Epoch ->
   CipherSuiteTag ->
   IncomingBundle ->
-  Sem r (IndexMap, ProposalAction)
+  Sem r (IndexMap, ProposalAction, [StoredProposal])
 getCommitData senderIdentity lConvOrSub epoch ciphersuite bundle = do
   let convOrSub = tUnqualified lConvOrSub
       groupId = cnvmlsGroupId convOrSub.mlsMeta
 
-  runState convOrSub.indexMap $ do
-    creatorAction <-
-      if epoch == Epoch 0
-        then addProposedClient (Left . RegularClient $ senderIdentity.client)
-        else mempty
-    proposals <-
-      traverse
-        (derefOrCheckProposal epoch ciphersuite groupId)
-        bundle.commit.value.proposals
-    action <- applyProposals ciphersuite proposals
-    pure (creatorAction <> action)
+  -- Fetch all pending proposals once: used both for dereferencing commit
+  -- proposal refs and by checkReferences downstream.
+  storedProposals <- getAllPendingProposals groupId epoch
+  (creatorAction, action) <-
+    runState convOrSub.indexMap $ do
+      creatorAction <-
+        if epoch == Epoch 0
+          then addProposedClient (Left . RegularClient $ senderIdentity.client)
+          else mempty
+      proposals <-
+        traverse
+          (derefOrCheckProposalFrom storedProposals epoch ciphersuite groupId)
+          bundle.commit.value.proposals
+      action <- applyProposals ciphersuite proposals
+      pure (creatorAction <> action)
+  pure (creatorAction, action, storedProposals)
 
 incrementEpoch ::
   ( Member ConversationStore r,
@@ -149,6 +155,19 @@ incrementEpoch (SubConv c s) = do
     getSubConversation (mcId c) (scSubConvId s) >>= noteS @'ConvNotFound
   pure (SubConv c subconv)
 
+-- | Bump the MLS epoch without re-reading the conversation afterwards.
+-- Use when the caller discards the result; avoids 2-3 CQL round trips.
+incrementEpochNoRead ::
+  (Member ConversationStore r) =>
+  ConvOrSubConv ->
+  Sem r ()
+incrementEpochNoRead =
+  \case
+    Conv c ->
+      setConversationEpoch (mcId c) (succ (cnvmlsEpoch (mcMLSData c)))
+    SubConv _c s ->
+      setSubConversationEpoch (scParentConvId s) (scSubConvId s) (succ (cnvmlsEpoch (scMLSData s)))
+
 getClientInfo ::
   ( Member BrigAPIAccess r,
     Member (FederationAPIAccess FederatorClient) r,

libs/wire-subsystems/src/Wire/ConversationSubsystem/MLS/Commit/InternalCommit.hs

@@ -32,6 +32,8 @@ import Data.Tuple.Extra
 import Galley.Types.Error
 import Imports
 import Polysemy
+import Polysemy.Async (Async)
+import Polysemy.Async qualified as P
 import Polysemy.Error
 import Polysemy.Input (Input)
 import Polysemy.Resource (Resource)
@@ -78,6 +80,7 @@ processInternalCommit ::
     Member (ErrorS 'MissingLegalholdConsent) r,
     Member (ErrorS 'GroupIdVersionNotSupported) r,
     Member Resource r,
+    Member Async r,
     Member Random r,
     Member (ErrorS MLSInvalidLeafNodeSignature) r,
     Member MLSCommitLockStore r,
@@ -93,14 +96,14 @@ processInternalCommit ::
   Epoch ->
   ProposalAction ->
   Commit ->
+  [StoredProposal] ->
   Codensity (Sem r) [LocalConversationUpdate]
-processInternalCommit senderIdentity con lConvOrSub ciphersuite ciphersuiteUpdate epoch action commit = do
+processInternalCommit senderIdentity con lConvOrSub ciphersuite ciphersuiteUpdate epoch action commit storedProposals = do
   let convOrSub = tUnqualified lConvOrSub
       qusr = cidQualifiedUser senderIdentity.client
       cm = convOrSub.members
       newUserClients = Map.assocs (unClientMap (paAdd action))
-
-  lift $ checkReferences convOrSub epoch commit
+  lift $ checkReferences storedProposals commit
 
   -- check update path
   lift $ traverse_ (checkUpdatePath lConvOrSub senderIdentity ciphersuite) commit.path
@@ -244,9 +247,23 @@ processInternalCommit senderIdentity con lConvOrSub ciphersuite ciphersuiteUpdat
       removeMLSClients gid qtarget (Map.keysSet clients)
 
     -- add clients to the conversation state
-    for_ newUserClients $ \(qtarget, newClients) -> do
-      addMLSClients gid qtarget $
-        Set.fromList [(cid, idx) | (cid, (idx, _)) <- Map.assocs newClients]
+    -- Note: safe to run concurrently because the children only perform store
+    -- writes on disjoint rows. The store children fail via IO exceptions
+    -- (addMLSClients runs through embedClient, a pure IO embed), which the
+    -- Async interpretation rethrows. An 'Error'-effect throw in a child
+    -- would collapse to 'Nothing' (the error interpreters sit outside
+    -- asyncToIOFinal in Galley.App); the 'Nothing' guard below turns that
+    -- into a hard commit failure instead of a silently dropped write.
+    results <-
+      P.sequenceConcurrently $
+        flip fmap newUserClients $ \(qtarget, newClients) ->
+          addMLSClients gid qtarget $
+            Set.fromList [(cid, idx) | (cid, (idx, _)) <- Map.assocs newClients]
+    when (Nothing `elem` results) $
+      throw
+        ( InternalErrorWithDescription
+            "A concurrent client-store write failed while processing commit"
+        )
 
     for_ action.paHistoryClientAdd $ uncurry (addHistoryClient gid)
 
@@ -256,9 +273,8 @@ processInternalCommit senderIdentity con lConvOrSub ciphersuite ciphersuiteUpdat
     when ciphersuiteUpdate $ case convOrSub.id of
       Conv cid -> setConversationCipherSuite cid ciphersuite
       SubConv cid sub -> setSubConversationCipherSuite cid sub ciphersuite
-
     -- increment epoch number
-    for_ lConvOrSub incrementEpoch
+    for_ lConvOrSub incrementEpochNoRead
 
     pure events
 
@@ -330,12 +346,9 @@ existingMembers :: Local StoredConversation -> Set (Qualified UserId)
 existingMembers lconv = existingLocalMembers lconv <> existingRemoteMembers lconv
 
 checkReferences ::
-  ( Member ProposalStore r,
-    Member (ErrorS MLSCommitMissingReferences) r
-  ) =>
-  ConvOrSubConv -> Epoch -> Commit -> Sem r ()
-checkReferences convOrSub epoch commit = do
-  allPendingProposals <- getAllPendingProposals (cnvmlsGroupId convOrSub.mlsMeta) epoch
+  (Member (ErrorS MLSCommitMissingReferences) r) =>
+  [StoredProposal] -> Commit -> Sem r ()
+checkReferences allPendingProposals commit = do
   let referencedProposals = Set.fromList $ mapMaybe (\x -> preview _Ref x) commit.proposals
   let (includedProposals, missingProposals) =
         partition

libs/wire-subsystems/src/Wire/ConversationSubsystem/MLS/Message.hs

@@ -43,6 +43,7 @@ import Data.Tuple.Extra
 import Galley.Types.Error
 import Imports
 import Polysemy
+import Polysemy.Async (Async)
 import Polysemy.Error
 import Polysemy.Input
 import Polysemy.Output
@@ -183,7 +184,8 @@ postMLSCommitBundle ::
     Member FederationSubsystem r,
     Member TeamSubsystem r,
     Member (Input ConversationSubsystemConfig) r,
-    Member FeaturesConfigSubsystem r
+    Member FeaturesConfigSubsystem r,
+    Member Async r
   ) =>
   Local x ->
   Qualified UserId ->
@@ -220,7 +222,8 @@ postMLSCommitBundleFromLocalUser ::
     Member FederationSubsystem r,
     Member TeamSubsystem r,
     Member (Input ConversationSubsystemConfig) r,
-    Member FeaturesConfigSubsystem r
+    Member FeaturesConfigSubsystem r,
+    Member Async r
   ) =>
   Version ->
   Local UserId ->
@@ -257,7 +260,8 @@ postMLSCommitBundleToLocalConv ::
     Member FederationSubsystem r,
     Member TeamSubsystem r,
     Member (Input ConversationSubsystemConfig) r,
-    Member FeaturesConfigSubsystem r
+    Member FeaturesConfigSubsystem r,
+    Member Async r
   ) =>
   Qualified UserId ->
   ClientId ->
@@ -325,7 +329,7 @@ postMLSCommitBundleToLocalConv qusr c conn bundle ctype lConvOrSubId = do
     (events, newClients) <- case senderIdentity.index of
       Just _ -> do
         -- extract added/removed clients from bundle
-        (newIndexMap, action) <-
+        (newIndexMap, action, storedProposals) <-
           lift $
             getCommitData senderIdentity lConvOrSub bundle.epoch ciphersuite bundle
 
@@ -358,6 +362,7 @@ postMLSCommitBundleToLocalConv qusr c conn bundle ctype lConvOrSubId = do
             bundle.epoch
             action
             bundle.commit.value
+            storedProposals
         -- the sender client is included in the Add action on the first commit,
         -- but it doesn't need to get a welcome message, so we filter it out here
         let newClients = cmRemoveClient senderIdentity.client (paAdd action)

libs/wire-subsystems/src/Wire/ConversationSubsystem/MLS/Proposal.hs

@@ -17,7 +17,7 @@
 
 module Wire.ConversationSubsystem.MLS.Proposal
   ( -- * Proposal processing
-    derefOrCheckProposal,
+    derefOrCheckProposalFrom,
     checkProposal,
     processProposal,
     proposalProcessingStage,
@@ -145,24 +145,26 @@ type HasProposalEffects r =
     Member TeamCollaboratorsSubsystem r
   )
 
-derefOrCheckProposal ::
+-- | Dereference a commit proposal, looking up refs in a prefetched list of
+-- pending proposals instead of issuing one point-read per ref.
+derefOrCheckProposalFrom ::
   ( Member (Error MLSProtocolError) r,
     Member (ErrorS 'MLSInvalidLeafNodeIndex) r,
     Member (ErrorS 'MLSUnsupportedProposal) r,
-    Member ProposalStore r,
     Member (State IndexMap) r,
     Member (ErrorS 'MLSProposalNotFound) r,
     Member (ErrorS 'MLSInvalidLeafNodeSignature) r
   ) =>
+  [StoredProposal] ->
   Epoch ->
   CipherSuiteTag ->
   GroupId ->
   ProposalOrRef ->
   Sem r Proposal
-derefOrCheckProposal epoch _ciphersuite groupId (Ref ref) = do
-  p <- getProposal groupId epoch ref >>= noteS @'MLSProposalNotFound
-  pure p.value
-derefOrCheckProposal _epoch ciphersuite _ (Inline p) = do
+derefOrCheckProposalFrom stored _epoch _ciphersuite _groupId (Ref ref) =
+  noteS @'MLSProposalNotFound $
+    (.proposal.value) <$> find ((== ref) . (.ref)) stored
+derefOrCheckProposalFrom _stored _epoch ciphersuite _ (Inline p) = do
   im <- get
   checkProposal ciphersuite im p
   pure p