use new client capability of consumable notifications

integration/integration.cabal

@@ -179,6 +179,7 @@ library
     Testlib.App
     Testlib.Assertions
     Testlib.Cannon
+    Testlib.Cannon.ConsumableNotifications
     Testlib.Certs
     Testlib.Env
     Testlib.HTTP

integration/test/MLS/Util.hs

@@ -38,6 +38,7 @@ import System.IO.Temp
 import System.Posix.Files
 import System.Process
 import Testlib.Assertions
+import qualified Testlib.Cannon.ConsumableNotifications as CN
 import Testlib.HTTP
 import Testlib.JSON
 import Testlib.Prelude
@@ -570,6 +571,78 @@ createExternalCommit convId cid mgi = do
 data MLSNotificationTag = MLSNotificationMessageTag | MLSNotificationWelcomeTag
   deriving (Show, Eq, Ord)
 
+consumingMassagesViaNewCapability :: (HasCallStack) => MLSProtocol -> MessagePackage -> Codensity App ()
+consumingMassagesViaNewCapability mlsProtocol mp = Codensity $ \k -> do
+  conv <- getMLSConv mp.convId
+  -- clients that should receive the message itself
+  let oldClients = Set.delete mp.sender conv.members
+  -- clients that should receive a welcome message
+  let newClients = Set.delete mp.sender conv.newMembers
+  -- all clients that should receive some MLS notification, together with the
+  -- expected notification tag
+  let clients =
+        map (,MLSNotificationMessageTag) (toList oldClients)
+          <> map (,MLSNotificationWelcomeTag) (toList newClients)
+
+  let newUsers =
+        Set.delete mp.sender.user $
+          Set.difference
+            (Set.map (.user) newClients)
+            (Set.map (.user) oldClients)
+
+  let uidsWithClients =
+        fmap
+          ((\c -> (c.user, (object ["domain" .= c.domain, "id" .= c.user], c.client))) . fst)
+          clients
+
+  CN.withEventsWebSockets (fmap snd uidsWithClients) $ \chans -> do
+    r <- k ()
+
+    -- if the conversation is actually MLS (and not mixed), pick one client for
+    -- each new user and wait for its join event. In Mixed protocol, the user is
+    -- already in the conversation so they do not get a member-join
+    -- notification.
+    when (mlsProtocol == MLSProtocolMLS) $ do
+      let uidsWithChannels = zip uidsWithClients chans
+      let newUserChans = uidsWithChannels & filter (\((uid, _), _) -> Set.member uid newUsers) & fmap snd
+      let assertJoin e = do
+            eventType <- e %. "data.event.payload.0.type" & asString
+            pure $ eventType == "conversation.member-join"
+
+      traverse_
+        ( \(eventChan, ackChan) ->
+            CN.awaitMatch assertJoin eventChan
+              >>= CN.ackEvent ackChan
+        )
+        newUserChans
+
+    -- at this point we know that every new user has been added to the
+    -- conversation
+    for_ (zip clients chans) $ \((cid, t), (eventChan, ackChan)) -> case t of
+      MLSNotificationMessageTag -> do
+        event <-
+          CN.awaitMatch
+            ( \e -> do
+                eventType <- e %. "data.event.payload.0.type" & asString
+                pure $ eventType == "conversation.mls-message-add"
+            )
+            eventChan
+        CN.ackEvent ackChan event
+        eventData <- event %. "data.event.payload.0.data" & asByteString
+        void $ mlsCliConsume mp.convId conv.ciphersuite cid eventData
+      MLSNotificationWelcomeTag -> do
+        event <-
+          CN.awaitMatch
+            ( \e -> do
+                eventType <- e %. "data.event.payload.0.type" & asString
+                pure $ eventType == "conversation.mls-welcome"
+            )
+            eventChan
+        CN.ackEvent ackChan event
+        eventData <- event %. "data.event.payload.0.data" & asByteString
+        void $ fromWelcome mp.convId conv.ciphersuite cid eventData
+    pure r
+
 consumingMessages :: (HasCallStack) => MLSProtocol -> MessagePackage -> Codensity App ()
 consumingMessages mlsProtocol mp = Codensity $ \k -> do
   conv <- getMLSConv mp.convId
@@ -588,6 +661,7 @@ consumingMessages mlsProtocol mp = Codensity $ \k -> do
           Set.difference
             (Set.map (.user) newClients)
             (Set.map (.user) oldClients)
+
   withWebSockets (map fst clients) $ \wss -> do
     r <- k ()
 
@@ -863,3 +937,43 @@ getSubConvId user convId subConvName =
   getSubConversation user convId subConvName
     >>= getJSON 200
     >>= objConvId
+
+-- FUTUREWORK: we assume all clients in the conversation have the new consumable-notification capability
+-- to support both the legacy and the new capability,
+-- we need to add it to the client identity and store it in the local MLS state
+sendAndConsumeCommitBundleNew :: (HasCallStack) => MessagePackage -> App Value
+sendAndConsumeCommitBundleNew = sendAndConsumeCommitBundleWithProtocolNew MLSProtocolMLS
+
+-- | Send an MLS commit bundle, wait for clients to receive it, consume it, and
+-- update the test state accordingly.
+sendAndConsumeCommitBundleWithProtocolNew :: (HasCallStack) => MLSProtocol -> MessagePackage -> App Value
+sendAndConsumeCommitBundleWithProtocolNew protocol mp = do
+  lowerCodensity $ do
+    consumingMassagesViaNewCapability protocol mp
+    lift $ do
+      r <- postMLSCommitBundle mp.sender (mkBundle mp) >>= getJSON 201
+
+      -- if the sender is a new member (i.e. it's an external commit), then
+      -- process the welcome message directly
+      do
+        conv <- getMLSConv mp.convId
+        when (Set.member mp.sender conv.newMembers) $
+          traverse_ (fromWelcome mp.convId conv.ciphersuite mp.sender) mp.welcome
+
+      -- increment epoch and add new clients
+      modifyMLSState $ \mls ->
+        mls
+          { convs =
+              Map.adjust
+                ( \conv ->
+                    conv
+                      { epoch = conv.epoch + 1,
+                        members = conv.members <> conv.newMembers,
+                        newMembers = mempty
+                      }
+                )
+                mp.convId
+                mls.convs
+          }
+
+      pure r

integration/test/Performance/BigConversation.hs

@@ -3,6 +3,7 @@
 module Performance.BigConversation where
 
 import API.BrigCommon
+import API.Galley (getConversation)
 import qualified Data.ByteString as BS
 import qualified Data.ByteString.Base64 as B64
 import Data.List.Extra (chunksOf)
@@ -12,8 +13,9 @@ import MLS.Util
 import SetupHelpers
 import qualified System.CryptoBox as Cryptobox
 import Testlib.Prelude
-import UnliftIO (pooledMapConcurrentlyN)
+import UnliftIO (modifyIORef', newIORef, pooledMapConcurrentlyN, readIORef)
 import UnliftIO.Temporary
+import Prelude (writeFile)
 
 -- | A size saying how big an MLS conversation is. Each size is mapped to a
 -- number via the 'sizeToNumber' function.
@@ -38,27 +40,54 @@ batchForSize :: ConversationSize -> Word
 batchForSize Tiny = 10
 batchForSize Small = 20
 batchForSize Medium = 100
-batchForSize Big = 100
-batchForSize Large = 250
+batchForSize Big = 250
+batchForSize Large = 500
 batchForSize VeryLarge = 500
 
-testCreateBigMLSConversation :: ConversationSize -> App ()
-testCreateBigMLSConversation convSize = do
-  let teamSize = sizeToNumber convSize
-  let batchSize = fromIntegral . batchForSize $ convSize
-  totalTime <-
-    fmap snd $ timeIt do
-      (_, ownerClient, _, members, _) <- createTeamAndClients . fromIntegral $ teamSize
+testCreateBigMLSConversation :: App ()
+testCreateBigMLSConversation = do
+  domain <- OwnDomain & asString
+  let teamSize = 11
+  let batchSize = 20
+  let clientNotifCapability = Consumable
+  putStrLn $ "Creating a team with " <> show teamSize <> " members"
+  (owner, ownerClient, _, members, c1 : c2 : _) <- createTeamAndClients domain clientNotifCapability teamSize
+  putStrLn $ "Creating a conversation with " <> show teamSize <> " members in batches of " <> show batchSize
+  convSize <- liftIO $ newIORef (1 :: Int)
+  (convId, totalTime) <-
+    timeIt do
       convId <- createNewGroup def ownerClient
       let memberChunks = chunksOf batchSize members
       for_ memberChunks $ \chunk -> do
         (size, time) <- timeIt $ do
           msg <- createAddCommit ownerClient convId chunk
-          void $ sendAndConsumeCommitBundle msg
+          void $ case clientNotifCapability of
+            Legacy -> sendAndConsumeCommitBundle msg
+            Consumable -> sendAndConsumeCommitBundleNew msg
           pure (BS.length msg.message)
-        putStrLn $ "Sent " <> show size <> " bytes in " <> show time
+        cs <- liftIO $ readIORef convSize
+        putStrLn $ "Sent " <> show size <> " bytes in " <> show time <> ", adding " <> show (length chunk) <> " members to conv of size: " <> show cs
+        liftIO $ modifyIORef' convSize (+ (length chunk))
         pure (size, time)
+      pure convId
   putStrLn $ "Total time: " <> show totalTime
+  do
+    conv <- getConversation owner (convIdToQidObject convId) >>= getJSON 200
+    otherMembers <- conv %. "members.others" & asList
+    length otherMembers `shouldMatchInt` (teamSize - 1)
+  (bytes, timeRemoval) <- timeIt $ do
+    commit <- createRemoveCommit ownerClient convId [c1, c2]
+    -- m <- showMessage def ownerClient commit.message
+    -- prettyJSON m >>= liftIO . writeFile "removal.json"
+    case clientNotifCapability of
+      Legacy -> void $ sendAndConsumeCommitBundle commit
+      Consumable -> void $ sendAndConsumeCommitBundleNew commit
+    pure (BS.length commit.message)
+  putStrLn $ "Sent " <> show bytes <> " bytes in " <> show timeRemoval <> " for removing 2 members"
+  do
+    conv <- getConversation owner (convIdToQidObject convId) >>= getJSON 200
+    otherMembers <- conv %. "members.others" & asList
+    length otherMembers `shouldMatchInt` (teamSize - 3)
 
 timeIt :: App a -> App (a, NominalDiffTime)
 timeIt action = do
@@ -67,9 +96,9 @@ timeIt action = do
   end <- liftIO getCurrentTime
   pure (result, diffUTCTime end start)
 
-createTeamAndClients :: Int -> App (Value, ClientIdentity, String, [Value], [ClientIdentity])
-createTeamAndClients teamSize = do
-  (owner, tid, members) <- createTeam OwnDomain teamSize
+createTeamAndClients :: String -> ClientNotifCapability -> Int -> App (Value, ClientIdentity, String, [Value], [ClientIdentity])
+createTeamAndClients domain clientNotifCapability teamSize = do
+  (owner, tid, members) <- createTeam domain teamSize
   let genPrekeyInBox box i = do
         pk <- assertCrytoboxSuccess =<< liftIO (Cryptobox.newPrekey box i)
         pkBS <- liftIO $ Cryptobox.copyBytes pk.prekey
@@ -87,7 +116,10 @@ createTeamAndClients teamSize = do
                 { clientArgs =
                     def
                       { prekeys = Just [firstPrekey],
-                        lastPrekey = Just lastPrekey
+                        lastPrekey = Just lastPrekey,
+                        acapabilities = case clientNotifCapability of
+                          Legacy -> Nothing
+                          Consumable -> Just ["consumable-notifications"]
                       }
                 }
         createMLSClient def mlsClientOpts user
@@ -100,3 +132,5 @@ assertCrytoboxSuccess :: (Show a) => Cryptobox.Result a -> App a
 assertCrytoboxSuccess = \case
   Cryptobox.Success x -> pure x
   e -> assertFailure $ "Cryptobox exception: " <> show e
+
+data ClientNotifCapability = Legacy | Consumable