解决并发 xhr 的 session rule 数量问题

packages/chrome-extension-mock/declarativ_net_request.ts

@@ -1,4 +1,8 @@
 export default class DeclarativeNetRequest {
+  MAX_NUMBER_OF_SESSION_RULES = 5000;
+
+  private _sessionRules: chrome.declarativeNetRequest.Rule[] = [];
+
   HeaderOperation = {
     APPEND: "append",
     SET: "set",
@@ -30,9 +34,44 @@ export default class DeclarativeNetRequest {
     OTHER: "other",
   };
 
-  updateSessionRules() {
+  updateSessionRules(arg1: any, arg2: any): Promise<void> {
+    let options: {
+      addRules?: chrome.declarativeNetRequest.Rule[];
+      removeRuleIds?: number[];
+    } = {};
+    let callback: undefined | ((...args: any) => any) = undefined;
+
+    if (typeof arg1 === "function") {
+      callback = arg1;
+    } else if (typeof arg2 === "function") {
+      callback = arg2;
+    }
+    if (typeof arg1 === "object" && arg1) options = arg1;
+
     return new Promise<void>((resolve) => {
+      const { addRules = [], removeRuleIds = [] } = options;
+
+      // Remove rules by ID
+      if (removeRuleIds.length > 0) {
+        this._sessionRules = this._sessionRules.filter((rule) => !removeRuleIds.includes(rule.id));
+      }
+
+      // Add or update rules (upsert by ID)
+      for (const newRule of addRules) {
+        const existingIndex = this._sessionRules.findIndex((rule) => rule.id === newRule.id);
+        if (existingIndex !== -1) {
+          this._sessionRules[existingIndex] = newRule; // update
+        } else {
+          this._sessionRules.push(newRule); // add
+        }
+      }
+
       resolve();
+      callback?.();
     });
   }
+
+  getSessionRules(): Promise<chrome.declarativeNetRequest.Rule[]> {
+    return Promise.resolve([...this._sessionRules]);
+  }
 }

src/app/service/service_worker/gm_api/dnr_id_controller.test.ts

@@ -0,0 +1,200 @@
+import { sleep } from "@App/pkg/utils/utils";
+import { describe, expect, it } from "vitest";
+import {
+  getSessionRuleIds,
+  LIMIT_SESSION_RULES,
+  nextSessionRuleId,
+  removeSessionRuleIdEntry,
+} from "./dnr_id_controller";
+
+describe("getSessionRuleIds", () => {
+  it("initializes from existing chrome session rules", async () => {
+    const ids = await getSessionRuleIds();
+    expect(ids.size).lessThan(100);
+    await nextSessionRuleId();
+    expect(ids.size).greaterThanOrEqual(1);
+    await nextSessionRuleId();
+    expect(ids.size).greaterThanOrEqual(2);
+  });
+});
+
+describe("nextSessionRuleId", () => {
+  it("returns unique incrementing IDs on each call", async () => {
+    const id1 = await nextSessionRuleId();
+    const id2 = await nextSessionRuleId();
+    const id3 = await nextSessionRuleId();
+
+    expect(id2).toBeGreaterThan(id1);
+    expect(id3).toBeGreaterThan(id2);
+  });
+
+  it("skips IDs that already exist in session rules", async () => {
+    const ids = await getSessionRuleIds();
+    const next = await nextSessionRuleId();
+
+    ids.add(next + 1);
+    const skipped = await nextSessionRuleId();
+
+    expect(skipped).toBeGreaterThan(next + 1);
+  });
+});
+
+describe("removeSessionRuleIdEntry", () => {
+  it("removes the given ID from the tracked set", async () => {
+    const ids = await getSessionRuleIds();
+    const id = await nextSessionRuleId();
+
+    ids.add(id);
+    removeSessionRuleIdEntry(id);
+
+    expect(ids.has(id)).toBe(false);
+  });
+
+  it("is a no-op when called before sessionRuleIds is initialized", () => {
+    expect(() => removeSessionRuleIdEntry(10404)).not.toThrow();
+  });
+
+  it("throws when ruleId is <= 10000", () => {
+    expect(() => removeSessionRuleIdEntry(10000)).toThrow();
+    expect(() => removeSessionRuleIdEntry(1)).toThrow();
+  });
+
+  it("rewinds SESSION_RULE_ID_BEGIN so the removed id gets reused next", async () => {
+    const ids = await getSessionRuleIds();
+
+    const id1 = await nextSessionRuleId();
+    const id2 = await nextSessionRuleId();
+    const id3 = await nextSessionRuleId();
+
+    ids.add(id1);
+    ids.add(id2);
+    ids.add(id3);
+
+    // Remove the latest — counter should rewind and reuse it
+    removeSessionRuleIdEntry(id3);
+    const reused = await nextSessionRuleId();
+    expect(reused).toBe(id3);
+  });
+
+  it("rewinds SESSION_RULE_ID_BEGIN when removed id is behind the counter", async () => {
+    const ids = await getSessionRuleIds();
+
+    const id1 = await nextSessionRuleId();
+    const id2 = await nextSessionRuleId();
+    const id3 = await nextSessionRuleId();
+
+    ids.add(id1);
+    ids.add(id2);
+    ids.add(id3);
+
+    // Remove an older id — counter rewinds to id1 - 1, so id1 gets reused
+    removeSessionRuleIdEntry(id1);
+    const reused = await nextSessionRuleId();
+    expect(reused).toBe(id1);
+  });
+
+  it("does not rewind SESSION_RULE_ID_BEGIN when removed id was pre-existing (ahead of counter)", async () => {
+    const id1 = await nextSessionRuleId();
+    const id2 = await nextSessionRuleId();
+    const id3 = await nextSessionRuleId();
+    const _id4 = await nextSessionRuleId();
+    const id5 = await nextSessionRuleId();
+    const id6 = await nextSessionRuleId();
+
+    removeSessionRuleIdEntry(id1);
+    removeSessionRuleIdEntry(id5);
+    removeSessionRuleIdEntry(id3);
+    // removeSessionRuleIdEntry(id4);
+    removeSessionRuleIdEntry(id2);
+    const nextBefore = await nextSessionRuleId(); // e.g. 10001
+    removeSessionRuleIdEntry(id6);
+
+    expect(await nextSessionRuleId()).toBe(nextBefore + 1); // counter unchanged, just increments normally
+    expect(await nextSessionRuleId()).toBe(nextBefore + 2); // counter unchanged, just increments normally
+    // expect(await nextSessionRuleId()).toBe(nextBefore + 3); // counter unchanged, just increments normally
+    expect(await nextSessionRuleId()).toBe(nextBefore + 4); // counter unchanged, just increments normally
+    expect(await nextSessionRuleId()).toBe(nextBefore + 5); // counter unchanged, just increments normally
+  });
+});
+
+describe("nextSessionRuleId limit control", () => {
+  it("locks when session rules reach the limit and unlocks when an entry is removed", async () => {
+    const ids = await getSessionRuleIds();
+    expect(ids.size).toBeLessThan(100);
+
+    const added = [];
+    for (let w = ids.size; w < LIMIT_SESSION_RULES; w++) {
+      const j = await nextSessionRuleId();
+      added.push(j);
+      ids.add(j);
+    }
+    expect(ids.size).toBeGreaterThan(1000);
+
+    const lockedPromise = nextSessionRuleId();
+
+    const raceResult1 = await Promise.race([lockedPromise.then(() => "resolved"), sleep(5).then(() => "pending")]);
+    expect(raceResult1).toBe("pending");
+
+    const m1 = Math.floor(Math.random() * (added.length - 9));
+    const p1 = added[m1];
+    const p2 = added[m1 + 6];
+    removeSessionRuleIdEntry(p1);
+    removeSessionRuleIdEntry(p2);
+
+    const raceResult2 = await Promise.race([lockedPromise.then(() => "resolved"), sleep(5).then(() => "pending")]);
+    expect(raceResult2).toBe("resolved");
+
+    const id1 = await lockedPromise;
+    const id2 = await nextSessionRuleId();
+    expect(id1).toBe(p1);
+    expect(id2).toBe(p2);
+
+    for (const k of added) {
+      removeSessionRuleIdEntry(k);
+    }
+
+    const res = await getSessionRuleIds();
+    expect(res).toBe(ids);
+    expect(res.size).toBeLessThan(100);
+  });
+
+  it("only one lock is created even with concurrent nextSessionRuleId calls", async () => {
+    const ids = await getSessionRuleIds();
+    expect(ids.size).toBeLessThan(100);
+
+    const added = [];
+    for (let w = ids.size; w < LIMIT_SESSION_RULES; w++) {
+      const j = await nextSessionRuleId();
+      added.push(j);
+      ids.add(j);
+    }
+    expect(ids.size).toBeGreaterThan(1000);
+
+    // Fire multiple concurrent calls while locked
+    const p1 = nextSessionRuleId();
+    const p2 = nextSessionRuleId();
+    const p3 = nextSessionRuleId();
+
+    const raceResult = await Promise.race([
+      Promise.all([p1, p2, p3]).then(() => "resolved"),
+      sleep(5).then(() => "pending"),
+    ]);
+    expect(raceResult).toBe("pending");
+
+    // Single removal should unlock all waiters sequentially
+    const toRemove = [...ids].find((id) => id > 10000)!;
+    removeSessionRuleIdEntry(toRemove);
+
+    // All three should eventually resolve
+    const results = await Promise.race([Promise.all([p1, p2, p3]).then((ids) => ids), sleep(50).then(() => null)]);
+    expect(results).not.toBeNull();
+
+    for (const k of added) {
+      removeSessionRuleIdEntry(k);
+    }
+
+    const res = await getSessionRuleIds();
+    expect(res).toBe(ids);
+    expect(res.size).toBeLessThan(100);
+  });
+});

src/app/service/service_worker/gm_api/dnr_id_controller.ts

@@ -0,0 +1,57 @@
+import type { Deferred } from "@App/pkg/utils/utils";
+import { deferred } from "@App/pkg/utils/utils";
+
+let sessionRuleIdsPromise: Promise<Set<number>> | null = null;
+let sessionRuleIds: Set<number> | null = null;
+
+let SESSION_RULE_ID_BEGIN = 10000;
+export const LIMIT_SESSION_RULES =
+  process.env.VI_TESTING === "true" ? 1234 : chrome.declarativeNetRequest.MAX_NUMBER_OF_SESSION_RULES - 300;
+let lockSessionRuleCreation: Deferred<void> | null = null;
+
+export const getSessionRuleIds = async (): Promise<Set<number>> => {
+  if (!sessionRuleIdsPromise) {
+    sessionRuleIdsPromise = chrome.declarativeNetRequest
+      .getSessionRules()
+      .then((rules) => {
+        sessionRuleIds = new Set(rules.map((rule) => rule.id).filter(Boolean));
+        return sessionRuleIds;
+      })
+      .catch((e) => {
+        console.warn(e);
+        sessionRuleIds = new Set<number>([]);
+        return sessionRuleIds;
+      });
+  }
+  const ruleIds = sessionRuleIds || (await sessionRuleIdsPromise);
+  return ruleIds;
+};
+
+export const removeSessionRuleIdEntry = (ruleId: number) => {
+  if (ruleId <= 10000) {
+    throw new Error("removeSessionRuleIdEntry cannot remove ids not created by nextSessionRuleId");
+  }
+  if (sessionRuleIds) {
+    if (sessionRuleIds.delete(ruleId) === true) {
+      if (ruleId <= SESSION_RULE_ID_BEGIN + 1) {
+        SESSION_RULE_ID_BEGIN = ruleId - 1;
+      }
+      if (sessionRuleIds.size < LIMIT_SESSION_RULES) {
+        lockSessionRuleCreation?.resolve();
+        lockSessionRuleCreation = null;
+      }
+    }
+  }
+};
+
+export const nextSessionRuleId = async () => {
+  const ruleIds = await getSessionRuleIds();
+  if (!lockSessionRuleCreation && ruleIds.size + 1 > LIMIT_SESSION_RULES) lockSessionRuleCreation = deferred<void>();
+  if (lockSessionRuleCreation) await lockSessionRuleCreation.promise;
+  let id;
+  do {
+    id = ++SESSION_RULE_ID_BEGIN;
+  } while (ruleIds.has(id));
+  ruleIds.add(id);
+  return id;
+};

src/app/service/service_worker/gm_api/gm_api.ts

@@ -53,6 +53,7 @@ import { headerModifierMap, headersReceivedMap } from "./gm_xhr";
 import { BgGMXhr } from "@App/pkg/utils/xhr/bg_gm_xhr";
 import { mightPrepareSetClipboard, setClipboard } from "../clipboard";
 import { nativePageWindowOpen } from "../../offscreen/gm_api";
+import { nextSessionRuleId, removeSessionRuleIdEntry } from "./dnr_id_controller";
 
 let generatedUniqueMarkerIDs = "";
 let generatedUniqueMarkerIDWhen = "";
@@ -685,8 +686,7 @@ export default class GMApi {
       }
       const redirectNotManual = params.redirect !== "manual";
 
-      // 使用 cacheInstance 避免SW重启造成重复 DNR Rule ID
-      const ruleId = 10000 + (await cacheInstance.incr("gmXhrRequestId", 1));
+      const ruleId = await nextSessionRuleId();
       const rule = {
         id: ruleId,
         action: {
@@ -1609,6 +1609,7 @@ export default class GMApi {
                   if (lastError) {
                     console.error("chrome.declarativeNetRequest.updateSessionRules:", lastError);
                   }
+                  removeSessionRuleIdEntry(rule.id);
                 }
               );
             }

tests/vitest.setup.ts

@@ -9,6 +9,8 @@ import { mockFetch } from "./mocks/fetch";
 
 vi.stubGlobal("chrome", chromeMock);
 chromeMock.init();
+// vi.spyOn(chromeMock.declarativeNetRequest, "getSessionRules");
+// vi.spyOn(chromeMock.declarativeNetRequest, "updateSessionRules");
 initTestEnv();
 
 chromeMock.runtime.getURL = vi.fn().mockImplementation((path: string) => {