Bug 1907212: Split large Close Tab commands into multiple sends r=lina

Differential Revision: https://phabricator.services.mozilla.com/D216240

services/fxaccounts/FxAccountsCommands.sys.mjs

@@ -26,6 +26,7 @@ ChromeUtils.defineESModuleGetters(lazy, {
   PushCrypto: "resource://gre/modules/PushCrypto.sys.mjs",
   getRemoteCommandStore: "resource://services-sync/TabsStore.sys.mjs",
   RemoteCommand: "resource://services-sync/TabsStore.sys.mjs",
+  Utils: "resource://services-sync/util.sys.mjs",
 });
 
 XPCOMUtils.defineLazyPreferenceGetter(
@@ -40,6 +41,7 @@ XPCOMUtils.defineLazyPreferenceGetter(
 );
 
 const TOPIC_TABS_CHANGED = "services.sync.tabs.changed";
+const COMMAND_MAX_PAYLOAD_SIZE = 16 * 1024;
 
 export class FxAccountsCommands {
   constructor(fxAccountsInternal) {
@@ -534,9 +536,8 @@ export class CloseRemoteTab extends Command {
    * @param {Device} target - Device object (typically returned by fxAccounts.getDevicesList()).
    * @param {String[]} urls - array of urls that should be closed on the remote device
    */
-  async sendCloseTabsCommand(target, urls) {
+  async sendCloseTabsCommand(target, urls, flowID) {
     log.info(`Sending tab closures to ${target.id} device.`);
-    const flowID = this._fxai.telemetry.generateFlowID();
     const encoder = new TextEncoder();
     try {
       const streamID = this._fxai.telemetry.generateFlowID();
@@ -703,39 +704,41 @@ export class CommandQueue {
           continue;
         }
       }
-      // this should be cleaned up a little more to better dispatch to the commands.
-      let toSendCloseTab = [];
-      for (let cmdToSend of toSend) {
-        if (cmdToSend.command instanceof lazy.RemoteCommand.CloseTab) {
-          toSendCloseTab.push(cmdToSend);
-        } else {
-          console.error("Unknown command", cmdToSend);
-          continue;
-        }
-      }
-      if (toSendCloseTab.length) {
-        let urlsToClose = toSendCloseTab.map(c => c.command.url);
-        // XXX - failure should cause a new error handling timer strategy (eg, ideally exponential backoff etc)
-        if (
-          await this._commands.closeTab.sendCloseTabsCommand(
-            device,
-            urlsToClose
-          )
-        ) {
-          // success! Mark them as sent.
-          for (let cmd of toSendCloseTab) {
-            log.trace(
-              `Setting pending command for device ${deviceId} as sent`,
-              cmd
-            );
-            await store.setPendingCommandSent(cmd);
-            didSend = true;
+
+      if (toSend.length) {
+        let urlsToClose = toSend.map(c => c.command.url);
+        // Generate a flowID to use for all chunked commands
+        const flowID = this._fxai.telemetry.generateFlowID();
+        // If we're dealing with large sets of urls, we should split them across
+        // multiple payloads to prevent breaking the issues for the user
+        let chunks = this.chunkUrls(urlsToClose, COMMAND_MAX_PAYLOAD_SIZE);
+        for (let chunk of chunks) {
+          if (
+            await this._commands.closeTab.sendCloseTabsCommand(
+              device,
+              chunk,
+              flowID
+            )
+          ) {
+            // We build a set from the sent urls for faster comparing
+            const urlChunkSet = new Set(chunk);
+            // success! Mark them as sent.
+            for (let cmd of toSend.filter(c =>
+              urlChunkSet.has(c.command.url)
+            )) {
+              log.trace(
+                `Setting pending command for device ${deviceId} as sent`,
+                cmd
+              );
+              await store.setPendingCommandSent(cmd);
+              didSend = true;
+            }
+          } else {
+            // We should investigate a better backoff strategy
+            // https://bugzilla.mozilla.org/show_bug.cgi?id=1899433
+            // For now just say 60s.
+            nextTime = Math.min(nextTime, now + 60000);
           }
-        } else {
-          // We should investigate a better backoff strategy
-          // https://bugzilla.mozilla.org/show_bug.cgi?id=1899433
-          // For now just say 60s.
-          nextTime = Math.min(nextTime, now + 60000);
         }
       }
     }
@@ -753,6 +756,31 @@ export class CommandQueue {
     }
   }
 
+  // Take a an array of urls and a max size and split them into chunks
+  // that are smaller than the passed in max size
+  // Note: This method modifies the passed in array
+  chunkUrls(urls, maxSize) {
+    let chunks = [];
+
+    // For optimal packing, we sort the array of urls from shortest-to-longest
+    urls.sort((a, b) => a.length - b.length);
+
+    while (urls.length) {
+      let chunk = lazy.Utils.tryFitItems(urls, maxSize);
+      if (!chunk.length) {
+        // None of the remaining URLs can fit into a single command
+        urls.forEach(url => {
+          log.warn(`Skipping oversized URL: ${url}`);
+        });
+        break;
+      }
+      chunks.push(chunk);
+      // Remove the processed URLs from the list
+      urls.splice(0, chunk.length);
+    }
+    return chunks;
+  }
+
   async _ensureTimer(timeout) {
     log.info(
       `Setting a new close-tab timer with delay=${timeout} with existing timer=${!!this

services/fxaccounts/tests/xpcshell/test_commands_closetab.js

@@ -362,3 +362,112 @@ add_task(async function test_telemetry_on_sendCloseTabsCommand() {
 
   commandQueue.shutdown();
 });
+
+// Should match the one in the FxAccountsCommands
+const COMMAND_MAX_PAYLOAD_SIZE = 16 * 1024;
+add_task(async function test_closetab_chunking() {
+  const targetDevice = { id: "dev1", name: "Device 1" };
+
+  const fxai = FxaInternalMock([targetDevice]);
+  let fxaCommands = {};
+  const closeTab = (fxaCommands.closeTab = new CloseRemoteTab(
+    fxaCommands,
+    fxai
+  ));
+  const commandQueue = (fxaCommands.commandQueue = new CommandQueue(
+    fxaCommands,
+    fxai
+  ));
+  let commandMock = sinon.mock(closeTab);
+  let queueMock = sinon.mock(commandQueue);
+
+  // freeze "now" to <= when the command was sent.
+  let now = Date.now();
+  commandQueue.now = () => now;
+
+  // Set the delay to 10ms
+  commandQueue.DELAY = 10;
+
+  // Generate a large number of commands to exceed the 16KB payload limit
+  const largeNumberOfCommands = [];
+  for (let i = 0; i < 300; i++) {
+    largeNumberOfCommands.push(
+      new RemoteCommand.CloseTab(
+        `https://example.com/addingsomeextralongstring/tab${i}`
+      )
+    );
+  }
+
+  // Add these commands to the store
+  const store = await getRemoteCommandStore();
+  for (let command of largeNumberOfCommands) {
+    await store.addRemoteCommandAt(targetDevice.id, command, now - 15);
+  }
+
+  const encoder = new TextEncoder();
+  // Calculate expected number of chunks
+  const totalPayloadSize = encoder.encode(
+    JSON.stringify(largeNumberOfCommands.map(cmd => cmd.url))
+  ).byteLength;
+  const expectedChunks = Math.ceil(totalPayloadSize / COMMAND_MAX_PAYLOAD_SIZE);
+
+  let flowIDUsed;
+  let chunksSent = 0;
+  commandMock
+    .expects("sendCloseTabsCommand")
+    .exactly(expectedChunks)
+    .callsFake((device, urls, flowID) => {
+      console.log(
+        "Chunk sent with size:",
+        encoder.encode(JSON.stringify(urls)).length
+      );
+      chunksSent++;
+      if (!flowIDUsed) {
+        flowIDUsed = flowID;
+      } else {
+        Assert.equal(
+          flowID,
+          flowIDUsed,
+          "FlowID should be consistent across chunks"
+        );
+      }
+
+      const chunkSize = encoder.encode(JSON.stringify(urls)).length;
+      Assert.ok(
+        chunkSize <= COMMAND_MAX_PAYLOAD_SIZE,
+        `Chunk size (${chunkSize}) should not exceed max payload size (${COMMAND_MAX_PAYLOAD_SIZE})`
+      );
+
+      return Promise.resolve(true);
+    });
+
+  await commandQueue.flushQueue();
+
+  // Check that all commands have been sent
+  Assert.equal((await store.getUnsentCommands()).length, 0);
+  Assert.equal(
+    chunksSent,
+    expectedChunks,
+    `Should have sent ${expectedChunks} chunks`
+  );
+
+  commandMock.verify();
+  queueMock.verify();
+
+  // Test edge case: URL exceeding max size
+  const oversizedCommand = new RemoteCommand.CloseTab(
+    "https://example.com/" + "a".repeat(COMMAND_MAX_PAYLOAD_SIZE)
+  );
+  await store.addRemoteCommandAt(targetDevice.id, oversizedCommand, now);
+
+  await commandQueue.flushQueue();
+
+  // The oversized command should still be unsent
+  Assert.equal((await store.getUnsentCommands()).length, 1);
+
+  commandMock.verify();
+  queueMock.verify();
+  commandQueue.shutdown();
+  commandMock.restore();
+  queueMock.restore();
+});