⚡ [AI] Speed up resetting the budget cache (#9143)

* [AI] Speed up resetting the budget cache

Resetting the budget cache recomputed nearly the whole spreadsheet twice:
`loadUserBudgets` closed its own spreadsheet transaction, which sorted from
the budget input cells and queued a pass over everything downstream, and the
`await` let that pass drain before `recomputeAll` queued it all again. Wrap
both in a single transaction so there is one sort, one computation pass and
one cache write, and hold the cache barrier for the duration so an interrupted
reset leaves the cache marked dirty.

`recomputeAll` also marked every `sum-amount` cell dirty, each of which runs
its own `SELECT SUM(amount)` — the categories x months cost that #8140 removed
from cold builds by seeding those cells from one grouped query. Reuse that
query via `reseedSumAmounts` and let `recomputeAll` skip the seeded cells.

Co-Authored-By: Claude Opus 5 <[email protected]>

* [AI] Release the reset transactions and cache barrier on failure

`resetBudgetCache` now holds a spreadsheet transaction across two awaited
database calls, and `loadUserBudgets` already held its own across one. A
rejection in any of them left the transaction depth above zero, which stops
the spreadsheet queueing further computations and makes every later
`waitOnSpreadsheet` throw — including the one `closeBudget` awaits before
unloading the spreadsheet. Close both transaction scopes and the cache
barrier in `finally`.

Co-Authored-By: Claude Opus 5 <[email protected]>

---------

Co-authored-by: Claude Opus 5 <[email protected]>

packages/loot-core/src/server/budget/base.ts

@@ -38,6 +38,27 @@ export function getBudgetRange(start: string, end: string) {
   return { start, end, range: monthUtils.rangeInclusive(start, end) };
 }
 
+// The day range spanning the given months. They aren't guaranteed to be
+// sorted, so find the first and last explicitly ('YYYY-MM' strings compare
+// chronologically).
+function getMonthSpanBounds(months: string[]) {
+  let firstMonth = months[0];
+  let lastMonth = months[0];
+  for (const month of months) {
+    if (month < firstMonth) {
+      firstMonth = month;
+    }
+    if (month > lastMonth) {
+      lastMonth = month;
+    }
+  }
+
+  return {
+    rangeStart: monthUtils.bounds(firstMonth).start,
+    rangeEnd: monthUtils.bounds(lastMonth).end,
+  };
+}
+
 // Computes the spend total for every category in every month within the
 // given day range using a single grouped query. This is used to seed the
 // `sum-amount` cells on a cold build so we avoid running one
@@ -70,6 +91,45 @@ function getSumAmountsByMonth(
   return sums;
 }
 
+// Reseeds the `sum-amount` cells for the given months from the same single
+// grouped query a cold build uses, and returns the names of the cells that
+// were seeded. Resetting the budget cache recomputes every cell, which would
+// otherwise run one `SELECT SUM(amount)` per category per month; the caller
+// skips recomputing these and caches them directly instead.
+export async function reseedSumAmounts(months: string[]): Promise<string[]> {
+  if (months.length === 0) {
+    return [];
+  }
+
+  const { rangeStart, rangeEnd } = getMonthSpanBounds(months);
+  const sums = getSumAmountsByMonth(rangeStart, rangeEnd);
+
+  const { data: groups }: { data: CategoryGroupEntity[] } = await aqlQuery(
+    q('category_groups').select('*'),
+  );
+  const categories = groups.flatMap(group => group.categories);
+
+  const seededCells: string[] = [];
+  for (const month of months) {
+    const sheetName = monthUtils.sheetForMonth(month);
+    const dbMonth = parseInt(month.replace('-', ''));
+
+    for (const cat of categories) {
+      const name = resolveName(sheetName, `sum-amount-${cat.id}`);
+      // A cell the spreadsheet doesn't have stays out of the seed set and is
+      // left to be recomputed normally.
+      if (!sheet.get().hasCell(name)) {
+        continue;
+      }
+
+      sheet.get().load(name, sums.get(`${dbMonth}-${cat.id}`) || 0);
+      seededCells.push(name);
+    }
+  }
+
+  return seededCells;
+}
+
 export function createCategory(cat, sheetName, prevSheetName, start, end) {
   sheet.get().createDynamic(sheetName, 'sum-amount-' + cat.id, {
     initialValue: 0,
@@ -292,22 +352,8 @@ export async function createBudget(months) {
   let sumAmounts: Map<string, number> | null = null;
   const getSumAmounts = () => {
     if (!sumAmounts) {
-      // `monthsToCreate` isn't guaranteed to be sorted, so find the span
-      // explicitly ('YYYY-MM' strings compare chronologically).
-      let firstMonth = monthsToCreate[0];
-      let lastMonth = monthsToCreate[0];
-      for (const month of monthsToCreate) {
-        if (month < firstMonth) {
-          firstMonth = month;
-        }
-        if (month > lastMonth) {
-          lastMonth = month;
-        }
-      }
-      sumAmounts = getSumAmountsByMonth(
-        monthUtils.bounds(firstMonth).start,
-        monthUtils.bounds(lastMonth).end,
-      );
+      const { rangeStart, rangeEnd } = getMonthSpanBounds(monthsToCreate);
+      sumAmounts = getSumAmountsByMonth(rangeStart, rangeEnd);
     }
     return sumAmounts;
   };

packages/loot-core/src/server/budgetfiles/app.ts

@@ -146,10 +146,44 @@ async function getRemoteFiles() {
 }
 
 async function resetBudgetCache() {
-  // Recomputing everything will update the cache
-  await sheet.loadUserBudgets(db);
-  sheet.get().recomputeAll();
-  await sheet.waitOnSpreadsheet();
+  const spreadsheet = sheet.get();
+
+  // Hold the cache dirty for the whole reset so an interrupted one doesn't
+  // leave a half-written cache looking clean.
+  spreadsheet.startCacheBarrier();
+
+  try {
+    let seededCells: string[] = [];
+
+    // One transaction around the whole reset. `loadUserBudgets` would
+    // otherwise close its own, queueing a computation pass over everything
+    // downstream of the budget inputs that `recomputeAll` then repeats.
+    spreadsheet.startTransaction();
+    try {
+      await sheet.loadUserBudgets(db);
+
+      // Seeding the spend totals from a single grouped query keeps the
+      // recompute from running a `SELECT SUM(amount)` per category per month.
+      seededCells = await budget.reseedSumAmounts([
+        ...spreadsheet.meta().createdMonths,
+      ]);
+
+      // Recomputing everything will update the cache
+      spreadsheet.recomputeAll(new Set(seededCells));
+    } finally {
+      // A transaction left open would stop the spreadsheet from queueing any
+      // further computations, and make every later `waitOnSpreadsheet` throw.
+      spreadsheet.endTransaction();
+    }
+
+    // The seeded cells were loaded rather than computed, so they aren't part
+    // of the computation queue that normally gets cached.
+    spreadsheet.saveCachedCells(seededCells);
+
+    await sheet.waitOnSpreadsheet();
+  } finally {
+    spreadsheet.endCacheBarrier();
+  }
 }
 
 async function uploadBudget({ id }: { id?: Budget['id'] } = {}): Promise<{

packages/loot-core/src/server/main.test.ts

@@ -281,6 +281,77 @@ describe('Budget', () => {
       changed.sort((a, b) => (a > b ? 1 : a < b ? -1 : 0)),
     ).toMatchSnapshot();
   });
+
+  test('resetting the budget cache computes each cell at most once', async () => {
+    const spreadsheet = await sheet.loadSpreadsheet(db);
+
+    let catId;
+    await runMutator(async () => {
+      await db.insertCategoryGroup({
+        id: 'incomeGroup',
+        name: 'incomeGroup',
+        is_income: 1,
+      });
+      await db.insertCategoryGroup({ id: 'group1', name: 'group1' });
+      catId = await db.insertCategory({ name: 'foo', cat_group: 'group1' });
+      await db.insertCategory({ name: 'bar', cat_group: 'group1' });
+    });
+
+    // Creates the budget months, which is what populates `createdMonths`
+    await runHandler(handlers['get-budget-bounds']);
+
+    db.runQuery("INSERT INTO accounts (id, name) VALUES ('one', 'boa')");
+    await runHandler(handlers['transaction-add'], {
+      id: uuidv4(),
+      date: '2017-01-05',
+      amount: -3400,
+      account: 'one',
+      category: catId,
+    });
+    await runHandler(handlers['budget/budget-amount'], {
+      month: '2017-01',
+      category: catId,
+      amount: 5000,
+    });
+
+    await sheet.waitOnSpreadsheet();
+
+    const sheetName = monthUtils.sheetForMonth('2017-01');
+    const valuesOf = () =>
+      new Map(
+        [...spreadsheet.getNodes()].map(([name, node]) => [name, node.value]),
+      );
+    const before = valuesOf();
+    expect(before.get(`${sheetName}!sum-amount-${catId}`)).toBe(-3400);
+
+    let computed = [];
+    const remove = spreadsheet.addEventListener('change', ({ names }) => {
+      computed = computed.concat(names);
+    });
+    await runHandler(handlers['reset-budget-cache']);
+    remove();
+
+    // Resetting the cache recomputes values, it doesn't change them
+    expect(valuesOf()).toEqual(before);
+
+    // Spend totals are seeded from a single grouped query rather than each
+    // running its own `SELECT SUM(amount)`
+    expect(computed.filter(name => name.includes('!sum-amount-'))).toEqual([]);
+
+    // Everything downstream of the budget inputs used to be computed twice:
+    // once when `loadUserBudgets` closed its own transaction, and again for
+    // `recomputeAll`
+    const seen = new Set();
+    const computedTwice = computed.filter(name => {
+      const duplicate = seen.has(name);
+      seen.add(name);
+      return duplicate;
+    });
+    expect(computedTwice).toEqual([]);
+
+    // Everything else really was recomputed
+    expect(computed).toContain(`${sheetName}!to-budget`);
+  });
 });
 
 describe('Categories', () => {

packages/loot-core/src/server/sheet.ts

@@ -218,32 +218,39 @@ export async function loadUserBudgets(db: typeof DbModule): Promise<void> {
 
   sheet.startTransaction();
 
-  // Load all the budget amounts and carryover values
-  for (const budget of budgets) {
-    if (budget.month && budget.category) {
-      const sheetName = `budget${budget.month}`;
-      sheet.set(`${sheetName}!budget-${budget.category}`, budget.amount);
-      sheet.set(
-        `${sheetName}!carryover-${budget.category}`,
-        budget.carryover === 1 ? true : false,
+  try {
+    // Load all the budget amounts and carryover values
+    for (const budget of budgets) {
+      if (budget.month && budget.category) {
+        const sheetName = `budget${budget.month}`;
+        sheet.set(`${sheetName}!budget-${budget.category}`, budget.amount);
+        sheet.set(
+          `${sheetName}!carryover-${budget.category}`,
+          budget.carryover === 1 ? true : false,
+        );
+        sheet.set(`${sheetName}!goal-${budget.category}`, budget.goal);
+        sheet.set(
+          `${sheetName}!long-goal-${budget.category}`,
+          budget.long_goal,
+        );
+      }
+    }
+
+    // For zero-based budgets, load the buffered amounts
+    if (budgetType !== 'tracking') {
+      const budgetMonths = await db.all<DbZeroBudgetMonth>(
+        'SELECT * FROM zero_budget_months',
       );
-      sheet.set(`${sheetName}!goal-${budget.category}`, budget.goal);
-      sheet.set(`${sheetName}!long-goal-${budget.category}`, budget.long_goal);
+      for (const budgetMonth of budgetMonths) {
+        const sheetName = sheetForMonth(budgetMonth.id);
+        sheet.set(`${sheetName}!buffered`, budgetMonth.buffered);
+      }
     }
+  } finally {
+    // The query above runs inside the transaction, so a failure would
+    // otherwise leave it open and wedge the spreadsheet for the session.
+    sheet.endTransaction();
   }
-
-  // For zero-based budgets, load the buffered amounts
-  if (budgetType !== 'tracking') {
-    const budgetMonths = await db.all<DbZeroBudgetMonth>(
-      'SELECT * FROM zero_budget_months',
-    );
-    for (const budgetMonth of budgetMonths) {
-      const sheetName = sheetForMonth(budgetMonth.id);
-      sheet.set(`${sheetName}!buffered`, budgetMonth.buffered);
-    }
-  }
-
-  sheet.endTransaction();
 }
 
 export function getCell(sheet: string, name: string) {

packages/loot-core/src/server/spreadsheet/spreadsheet.ts

@@ -348,10 +348,12 @@ export class Spreadsheet {
     });
   }
 
-  recomputeAll(): void {
-    // Recompute everything!
+  // Recompute everything, except any cells in `skip` — those are expected to
+  // have already been given their final value with `load`.
+  recomputeAll(skip?: ReadonlySet<string>): void {
     this.transaction(() => {
-      this.dirtyCells = [...this.nodes.keys()];
+      const names = [...this.nodes.keys()];
+      this.dirtyCells = skip ? names.filter(name => !skip.has(name)) : names;
     });
   }
 

upcoming-release-notes/faster-reset-budget-cache.md

@@ -0,0 +1,6 @@
+---
+category: Enhancements
+authors: [youngcw]
+---
+
+Make resetting the budget cache much faster on large budgets