fix(apollo-forest-run): fix perf regression from new auto-eviction

change/@graphitation-apollo-forest-run-906abf62-5b94-4c5d-999d-2f429024c6f8.json

@@ -0,0 +1,7 @@
+{
+  "type": "patch",
+  "comment": "fix perf regression on auto eviction",
+  "packageName": "@graphitation/apollo-forest-run",
+  "email": "vrazuvaev@microsoft.com_msteamsmdb",
+  "dependentChangeType": "patch"
+}

packages/apollo-forest-run/src/ForestRun.ts

@@ -578,7 +578,7 @@ export class ForestRun<
       );
       logUpdateStats(this.env, activeTransaction.changelog, watchesToNotify);
     }
-    maybeEvictOldData(this.env, this.store, activeTransaction);
+    maybeEvictOldData(this.env, this.store);
 
     return result as T;
   }

packages/apollo-forest-run/src/ForestRunCompat.ts

@@ -47,7 +47,7 @@ export class ForestRunCompat extends ForestRun {
       );
       const operationResult = { data: write.result ?? {} };
       const tree = indexTree(this.env, operation, operationResult);
-      replaceTree(this.store.dataForest, tree);
+      replaceTree(this.env, this.store.dataForest, tree);
     }
     return this;
   }

packages/apollo-forest-run/src/__tests__/eviction.test.ts

@@ -152,7 +152,10 @@ it("allows partitioned eviction", () => {
     maxOperationCount: 1,
     autoEvict: false,
     partitionConfig: {
-      partitionKey: (o) => o.operation.debugName,
+      partitionKey: (o) =>
+        ["query a", "query b"].includes(o.operation.debugName)
+          ? o.operation.debugName
+          : null,
       partitions: {
         "query a": { maxOperationCount: 1 },
         "query b": { maxOperationCount: 2 },
@@ -260,9 +263,12 @@ it("allows partitioned eviction", () => {
 });
 
 it("should warn exactly once for the same warning", () => {
+  const calls: unknown[] = [];
   const mockConsole = {
     ...console,
-    warn: jest.fn(),
+    warn: (...args: unknown[]) => {
+      calls.push(args);
+    },
   };
 
   const cache = new ForestRun({
@@ -295,11 +301,11 @@ it("should warn exactly once for the same warning", () => {
 
   cache.gc();
 
-  expect(mockConsole.warn).toHaveBeenCalledTimes(1);
-  expect(mockConsole.warn).toHaveBeenCalledWith(
+  expect(calls.length).toBe(1);
+  expect(calls[0]).toEqual([
     "partition_not_configured",
     'Partition "query TestA" is not configured in partitionConfig. Using default partition instead.',
-  );
+  ]);
 
   // notably: "query TestB" is NOT logged
 });
@@ -505,6 +511,58 @@ it("gc() evicts all partitions regardless of per-partition autoEvict", () => {
   ).toEqual({ bar: 1 });
 });
 
+it("auto-evicting partition B does not evict partition A with autoEvict: false", () => {
+  const cache = new ForestRun({
+    maxOperationCount: 2,
+    autoEvict: true,
+    partitionConfig: {
+      partitionKey: (o) => o.operation.debugName,
+      partitions: {
+        "query a": { maxOperationCount: 2, autoEvict: false },
+        "query b": { maxOperationCount: 1, autoEvict: true },
+      },
+    },
+  });
+  const a = gql`
+    query a($i: Int) {
+      foo(i: $i)
+    }
+  `;
+  const b = gql`
+    query b($i: Int) {
+      bar(i: $i)
+    }
+  `;
+
+  // Fill partition A beyond maxOperationCount (but autoEvict: false)
+  cache.write({ query: a, variables: { i: 0 }, result: { foo: 0 } });
+  cache.write({ query: a, variables: { i: 1 }, result: { foo: 1 } });
+  cache.write({ query: a, variables: { i: 2 }, result: { foo: 2 } });
+
+  // Fill partition B to trigger auto-eviction (2 ops >= maxOperationCount * 2 = 2)
+  cache.write({ query: b, variables: { i: 0 }, result: { bar: 0 } });
+  cache.write({ query: b, variables: { i: 1 }, result: { bar: 1 } });
+
+  // Partition A should NOT be evicted (autoEvict: false)
+  expect(
+    cache.read({ query: a, variables: { i: 0 }, optimistic: true }),
+  ).toEqual({ foo: 0 });
+  expect(
+    cache.read({ query: a, variables: { i: 1 }, optimistic: true }),
+  ).toEqual({ foo: 1 });
+  expect(
+    cache.read({ query: a, variables: { i: 2 }, optimistic: true }),
+  ).toEqual({ foo: 2 });
+
+  // Partition B should be auto-evicted (autoEvict: true, over threshold)
+  expect(
+    cache.read({ query: b, variables: { i: 0 }, optimistic: true }),
+  ).toEqual(null);
+  expect(
+    cache.read({ query: b, variables: { i: 1 }, optimistic: true }),
+  ).toEqual({ bar: 1 });
+});
+
 it("unconfigured operations fall to default partition and inherit global autoEvict", () => {
   // autoEvict: false only applies to __default__ partition and partitions
   // that don't specify their own autoEvict

packages/apollo-forest-run/src/__tests__/helpers/forest.ts

@@ -36,6 +36,7 @@ export function createTestForest(): IndexedForest {
     operationsByNodes: new Map(),
     operationsWithErrors: new Set(),
     operationsByName: new Map(),
+    operationsByPartitions: new Map(),
     deletedNodes: new Set(),
   };
 }

packages/apollo-forest-run/src/cache/read.ts

@@ -264,7 +264,7 @@ function growOutputTree(
   }
   if (!dataTree) {
     dataTree = growDataTree(env, forest, operation);
-    addTree(forest, dataTree);
+    addTree(env, forest, dataTree);
   }
   const tree = applyTransformations(
     env,

packages/apollo-forest-run/src/cache/store.ts

@@ -3,8 +3,6 @@ import {
   DataForest,
   DataTree,
   OptimisticLayer,
-  ResolvedPartition,
-  ResolvedPartitionOptions,
   ResultTree,
   Store,
   Transaction,
@@ -17,51 +15,20 @@ import {
   TypeName,
 } from "../descriptor/types";
 import { assert } from "../jsutils/assert";
-import { IndexedTree } from "../forest/types";
+import { IndexedTree, DefaultPartition } from "../forest/types";
 import { NodeChunk } from "../values/types";
 import { operationCacheKey } from "./descriptor";
 
 const EMPTY_ARRAY = Object.freeze([]);
-
-const DEFAULT_PARTITION = "__default__";
-
-function resolvePartition(
-  env: CacheEnv,
-  store: Store,
-  tree: IndexedTree,
-): ResolvedPartition {
-  const cached = store.partitions.get(tree);
-  if (cached !== undefined) {
-    return cached;
-  }
-  const { partitions, partitionKey } = env.partitionConfig;
-  const key = partitionKey(tree);
-  let partition: string;
-  let options: ResolvedPartitionOptions;
-  if (key != null && partitions[key]) {
-    partition = key;
-    options = partitions[key];
-  } else {
-    partition = DEFAULT_PARTITION;
-    options = partitions[DEFAULT_PARTITION];
-    if (key != null) {
-      env.logger?.warnOnce(
-        "partition_not_configured",
-        `Partition "${key}" is not configured in partitionConfig. Using default partition instead.`,
-      );
-    }
-  }
-  const result: ResolvedPartition = { partition, options };
-  store.partitions.set(tree, result);
-  return result;
-}
+const DEFAULT_PARTITION: DefaultPartition = "__default__";
 
 export function createStore(_: CacheEnv): Store {
   const dataForest: DataForest = {
     trees: new Map(),
     operationsByNodes: new Map<NodeKey, Set<OperationId>>(),
-    operationsWithErrors: new Set<OperationDescriptor>(),
     operationsByName: new Map(),
+    operationsByPartitions: new Map(),
+    operationsWithErrors: new Set<OperationDescriptor>(),
     extraRootIds: new Map<NodeKey, TypeName>(),
     layerTag: null, // not an optimistic layer
     readResults: new Map(),
@@ -100,31 +67,33 @@ export function touchOperation(
   store.atime.set(operation.id, env.now());
 }
 
-function affectedPartitionsHaveAutoEvict(
+function shouldAutoEvictAtLeastOnePartition(
   env: CacheEnv,
   store: Store,
-  transaction: Transaction,
 ): boolean {
-  for (const entry of transaction.changelog) {
-    if ("incoming" in entry) {
-      const { options } = resolvePartition(env, store, entry.incoming);
-      if (options.autoEvict) {
-        return true;
-      }
+  const partitions = env.partitionConfig.partitions;
+
+  for (const [partition, ops] of store.dataForest.operationsByPartitions) {
+    const partitionConfig =
+      partitions[partition] ?? partitions[DEFAULT_PARTITION];
+
+    assert(partitionConfig);
+
+    if (
+      partitionConfig.autoEvict &&
+      ops.size >= partitionConfig.maxOperationCount * 2
+    ) {
+      return true;
     }
   }
   return false;
 }
 
-export function maybeEvictOldData(
-  env: CacheEnv,
-  store: Store,
-  transaction: Transaction,
-): void {
+export function maybeEvictOldData(env: CacheEnv, store: Store): void {
   if (store.pendingEviction != null) {
     return;
   }
-  if (!affectedPartitionsHaveAutoEvict(env, store, transaction)) {
+  if (!shouldAutoEvictAtLeastOnePartition(env, store)) {
     return;
   }
   store.pendingEviction = env.scheduleAutoEviction(() => {
@@ -145,48 +114,35 @@ export function evictOldData(
 ): OperationId[] {
   const { dataForest, atime } = store;
 
-  const partitionsOps = new Map<
-    string,
-    { options: ResolvedPartitionOptions; operations: OperationId[] }
-  >();
-
-  for (const tree of dataForest.trees.values()) {
-    if (!canEvict(env, store, tree)) {
-      continue; // Skip non-evictable operations entirely
-    }
+  const evictedIds: OperationId[] = [];
+  const partitions = env.partitionConfig?.partitions ?? {};
 
-    const { partition, options } = resolvePartition(env, store, tree);
-    let partitionOps = partitionsOps.get(partition);
-    if (!partitionOps) {
-      partitionOps = { options, operations: [] };
-      partitionsOps.set(partition, partitionOps);
+  // Process each partition
+  for (const [partition, ops] of dataForest.operationsByPartitions) {
+    if (!partitions[partition]) {
+      env.logger?.warnOnce(
+        "partition_not_configured",
+        `Partition "${partition}" is not configured in partitionConfig. Using default partition instead.`,
+      );
     }
 
-    partitionOps.operations.push(tree.operation.id);
-  }
+    const partitionConfig =
+      partitions[partition] ?? partitions[DEFAULT_PARTITION];
 
-  const toEvict: OperationId[] = [];
+    assert(partitionConfig);
 
-  // Process each partition
-  for (const {
-    options: partitionConf,
-    operations: evictableOperationIds,
-  } of partitionsOps.values()) {
-    // During auto-eviction, respect per-partition autoEvict settings
-    if (isAutoEvict) {
-      if (!partitionConf.autoEvict) {
-        continue;
-      }
+    if (isAutoEvict && !partitionConfig.autoEvict) {
+      continue;
     }
 
-    const maxCount = partitionConf.maxOperationCount;
+    const maxCount = partitionConfig.maxOperationCount;
 
-    if (!maxCount || evictableOperationIds.length <= maxCount) {
+    if (!maxCount || ops.size <= maxCount) {
       continue; // No eviction needed for this partition
     }
 
     // Sort by access time (LRU) and determine how many to evict
-    evictableOperationIds.sort(
+    const evictableOperationIds = [...ops].sort(
       (a, b) => (atime.get(a) ?? 0) - (atime.get(b) ?? 0),
     );
 
@@ -195,17 +151,20 @@ export function evictOldData(
     // Keep only the ones to evict without array copying w/ slice
     evictableOperationIds.length = evictCount;
 
-    toEvict.push(...evictableOperationIds);
-  }
+    // Remove evicted operations
+    for (const opId of evictableOperationIds) {
+      const tree = store.dataForest.trees.get(opId);
+      assert(tree);
+      if (!canEvict(env, store, tree)) {
+        continue; // Skip non-evictable operations entirely
+      }
+      removeDataTree(store, tree, partition);
+    }
 
-  // Remove evicted operations
-  for (const opId of toEvict) {
-    const tree = store.dataForest.trees.get(opId);
-    assert(tree);
-    removeDataTree(store, tree);
+    evictedIds.push(...evictableOperationIds);
   }
 
-  return toEvict;
+  return evictedIds;
 }
 
 function canEvict(env: CacheEnv, store: Store, resultTree: DataTree) {
@@ -234,6 +193,7 @@ export function createOptimisticLayer(
     operationsByNodes: new Map(),
     operationsWithErrors: new Set(),
     operationsByName: new Map(),
+    operationsByPartitions: new Map(),
     extraRootIds: new Map(),
     readResults: new Map(),
     mutations: new Set(),
@@ -394,11 +354,14 @@ function removeDataTree(
     atime,
   }: Store,
   { operation }: ResultTree,
+  partition: string,
 ) {
   assert(!watches.has(operation));
   dataForest.trees.delete(operation.id);
   dataForest.readResults.delete(operation);
   dataForest.operationsWithErrors.delete(operation);
+  dataForest.operationsByName.get(operation.name ?? "")?.delete(operation.id);
+  dataForest.operationsByPartitions.get(partition)?.delete(operation.id);
   optimisticReadResults.delete(operation);
   partialReadResults.delete(operation);
   operations.get(operation.document)?.delete(operationCacheKey(operation));

packages/apollo-forest-run/src/cache/write.ts

@@ -180,7 +180,7 @@ export function write(
     // Note: even with existingResult === undefined the tree for this operation may still exist in the cache
     //   (when existingResult is resolved with a different key descriptor due to key variables)
     // TODO: replace with addTree and add a proper check for keyVariables
-    replaceTree(targetForest, modifiedIncomingResult);
+    replaceTree(env, targetForest, modifiedIncomingResult);
   }
 
   appendAffectedOperationsFromOtherLayers(