Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 36 additions & 0 deletions packages/store/src/store.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -548,6 +548,42 @@ describe('batch', () => {
expect(useStore.getState()).toEqual({ a: 1, b: 2 })
})

it('async batch does not leak batch context to concurrent non-batched updates (issue 8866)', async () => {
const useStore = createStore(() => ({ x: 0, y: 0 }))
const spy = vi.fn()
useStore.subscribe(spy)

let resolveBatch: () => void = () => {}
const p = batch(async () => {
useStore.setState({ x: 1 })
await new Promise<void>((resolve) => {
resolveBatch = resolve
})
useStore.setState({ x: 2 })
})

// While the batch is suspended, perform a non-batched update
useStore.setState({ y: 10 })

// The non-batched update should apply immediately and trigger listeners,
// because the async batch should not leak its batching context to concurrent tasks
expect(useStore.getState().y).toBe(10)
expect(spy).toHaveBeenCalledOnce()
expect(spy.mock.calls[0][0]).toEqual({ x: 0, y: 10 }) // x is still 0 (deferred), y is 10 (applied)

spy.mockClear()

// Resolve the batch
resolveBatch()
await p
await new Promise(resolve => queueMicrotask(resolve))

// The batch should now complete and notify
expect(useStore.getState()).toEqual({ x: 2, y: 10 })
expect(spy).toHaveBeenCalledOnce()
expect(spy.mock.calls[0][0]).toEqual({ x: 2, y: 10 })
})

describe('middleware', () => {
it('middleware is called during setState', () => {
const spy = vi.fn((prev, update, next) => next(update))
Expand Down
159 changes: 101 additions & 58 deletions packages/store/src/store.ts
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ import { useState, useEffect, useRef } from '@termuijs/jsx';
import * as path from 'node:path';
import * as fs from 'node:fs';
import * as os from 'node:os';
import { AsyncLocalStorage } from 'node:async_hooks';
import type { EqualityFn } from './shallow.js'


Expand All @@ -37,11 +38,14 @@ interface BatchEntry<T> {
rollback: () => void;
}

let _batchDepth = 0;
let _batchEpoch = 0;
// Map store instance to batch entry. Using any for listener set type because
// the batch mechanism operates on the raw Set<Listener<T>> without knowing T at this level.
const _batchStores = new Map<Set<any>, BatchEntry<any>>();
interface BatchContext {
depth: number;
}

const batchLocalStorage = new AsyncLocalStorage<BatchContext>();
const globalBatchStores = new Map<Set<any>, BatchEntry<any>>();
let globalBatchEpoch = 0;
Comment on lines +45 to +47

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift

Scope pending entries and flush state to each BatchContext.

AsyncLocalStorage isolates only depth. globalBatchStores and globalBatchEpoch still join independent root batches.

If batch A queues its synchronous flush, then batch B starts before the microtask runs, Line 157 suppresses A's notification. If batch B remains suspended, A's update stays pending. If both batches update the same store, batch B can also commit or roll back batch A's entry.

Store the pending-entry map and flush generation in BatchContext. Pass that context to flushBatch. Add coverage for overlapping root batches where one batch is synchronous and the other is suspended.

Also applies to: 107-112, 147-159

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@packages/store/src/store.ts` around lines 45 - 47, Move the pending-entry map
and flush-generation state from the module-level globalBatchStores and
globalBatchEpoch variables into BatchContext so independent root batches remain
isolated. Update batch creation and the flush scheduling/processing flow,
including flushBatch and the logic around lines 147–159, to pass and use the
active context rather than shared globals. Add coverage for overlapping root
batches where one completes synchronously and the other is suspended, including
same-store updates and independent notifications/commit or rollback behavior.


/**
* Batch multiple state updates into a single render pass.
*
Expand All @@ -63,60 +67,96 @@ const _batchStores = new Map<Set<any>, BatchEntry<any>>();
* ```
*/
export function batch<T>(fn: () => T): T {
const isOutermost = _batchDepth === 0;
_batchDepth++;
if (isOutermost) _batchEpoch++;
let threw = false;
let res: any;
try {
res = fn();
} catch (err) {
threw = true;
_batchDepth--;
if (_batchDepth === 0) {
flushBatch(threw);
const parentContext = batchLocalStorage.getStore();

if (parentContext) {
parentContext.depth++;
try {
const res = fn();
// Cast res through unknown to Promise<unknown> to safely check if it is a thenable (Promise) without type errors on generic T
if (res && typeof (res as unknown as Promise<unknown>).then === 'function') {
// Cast through unknown to Promise<unknown> to chain the then callback before returning
return (res as unknown as Promise<unknown>).then(
(val) => {
parentContext.depth--;
if (parentContext.depth === 0) flushBatch(false, true);
return val;
},
(err) => {
parentContext.depth--;
if (parentContext.depth === 0) flushBatch(true, true);
throw err;
}
) as T;
} else {
parentContext.depth--;
if (parentContext.depth === 0) {
flushBatch(false);
}
return res;
}
} catch (err) {
parentContext.depth--;
if (parentContext.depth === 0) {
flushBatch(true);
}
throw err;
}
throw err;
}

if (res && typeof res.then === 'function') {
return (res as Promise<any>).then(
(val) => {
_batchDepth--;
if (_batchDepth === 0) flushBatch(false, true);
return val;
},
(err) => {
_batchDepth--;
if (_batchDepth === 0) flushBatch(true, true);
throw err;
globalBatchEpoch++;
const context: BatchContext = {
depth: 1
};

return batchLocalStorage.run(context, () => {
try {
const res = fn();
// Cast res through unknown to Promise<unknown> to safely check if it is a thenable (Promise) without type errors on generic T
if (res && typeof (res as unknown as Promise<unknown>).then === 'function') {
// Cast through unknown to Promise<unknown> to chain the then callback before returning
return (res as unknown as Promise<unknown>).then(
(val) => {
context.depth--;
if (context.depth === 0) flushBatch(false, true);
return val;
},
(err) => {
context.depth--;
if (context.depth === 0) flushBatch(true, true);
throw err;
}
) as T;
} else {
context.depth--;
if (context.depth === 0) {
flushBatch(false);
}
return res;
}
) as T;
} else {
_batchDepth--;
if (_batchDepth === 0) {
flushBatch(false);
} catch (err) {
context.depth--;
if (context.depth === 0) {
flushBatch(true);
}
throw err;
}
return res;
}
});
}

function flushBatch(threw: boolean, immediate = false) {
if (threw) {
for (const [, { rollback }] of _batchStores) {
for (const [, { rollback }] of globalBatchStores) {
rollback();
}
_batchStores.clear();
globalBatchStores.clear();
Comment on lines 147 to +152

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift

Do not restore a stale pre-batch snapshot on rollback.

A suspended batch captures prevState before an unrelated update can commit. If the batch later rejects, Line 150 calls rollback, and the closure at Line 461 or Line 574 restores that stale snapshot.

For example, a deferred { x: 1 } update followed by an immediate { y: 10 } update will revert y to its old value when the deferred batch rejects. Deferred batch changes have not been committed, so rollback must discard only the deferred entry. Final listener notifications must also use the state immediately before the deferred commit.

Add a rejection regression test with an independent update while the batch is suspended.

Also applies to: 449-467, 558-580

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@packages/store/src/store.ts` around lines 147 - 152, Update flushBatch and
the rollback closures registered in the deferred batch paths to discard only the
rejected batch entry rather than restoring a stale prevState snapshot. Ensure
final listener notifications use the state immediately before the deferred
commit, preserving unrelated immediate updates made while the batch was
suspended. Add a regression test covering an independent update followed by
rejection of the suspended batch.

} else {
if (_batchStores.size === 0) return;
// Snapshot the current epoch so the microtask can bail out if a new
// batch has started before it runs.
const epochAtFlush = _batchEpoch;
if (globalBatchStores.size === 0) return;
const epochAtFlush = globalBatchEpoch;
const notify = () => {
// A new batch started between flush and notify — skip.
if (_batchEpoch !== epochAtFlush) return;
const stores = Array.from(_batchStores.entries());
_batchStores.clear();
if (globalBatchEpoch !== epochAtFlush) return;
const stores = Array.from(globalBatchStores.entries());
globalBatchStores.clear();
const newStates = new Map<Set<any>, any>();
for (const [listeners, { commit }] of stores) {
newStates.set(listeners, commit());
Expand Down Expand Up @@ -387,16 +427,17 @@ export function createStore<T extends object>(

const setState: SetState<T> = (partial) => {
const prevState = state;
const context = batchLocalStorage.getStore();

// When in a batch, function updaters should see the pending batch state
const batchState: T = _batchDepth > 0 ? _batchStores.get(listeners)?.nextState ?? state : state;
const batchState: T = context && context.depth > 0 ? globalBatchStores.get(listeners)?.nextState ?? state : state;
const nextPartial = typeof partial === 'function'
? (partial as (state: T) => Partial<T>)(batchState)
: partial;

const applyUpdate = (finalPartial: Partial<T>): T => {
// When in a batch, compute nextState from pending batch state if available
const baseState: T = _batchDepth > 0 ? _batchStores.get(listeners)?.nextState ?? state : state;
const baseState: T = context && context.depth > 0 ? globalBatchStores.get(listeners)?.nextState ?? state : state;
const nextState = { ...baseState, ...finalPartial };

// Only notify if at least one key's value actually changed
Expand All @@ -405,14 +446,14 @@ export function createStore<T extends object>(
key => !Object.is((baseState as any)[key], (nextState as any)[key])
);
if (hasChanged) {
if (_batchDepth > 0) {
if (context && context.depth > 0) {
// We're in a batch: defer listener notifications and track the final state
const existing = _batchStores.get(listeners);
const existing = globalBatchStores.get(listeners);
if (!existing) {
// Track only the keys changed inside the batch so commit can merge
// onto any intermediate non-batched updates without overwriting them.
const changes: Partial<T> = { ...finalPartial };
_batchStores.set(listeners, {
globalBatchStores.set(listeners, {
prevState,
nextState,
changes,
Expand Down Expand Up @@ -461,8 +502,9 @@ export function createStore<T extends object>(
};

const getState: GetState<T> = () => {
if (_batchDepth > 0) {
const entry = _batchStores.get(listeners);
const context = batchLocalStorage.getStore();
if (context && context.depth > 0) {
const entry = globalBatchStores.get(listeners);
if (entry) return entry.nextState;
}
return state;
Expand Down Expand Up @@ -495,7 +537,7 @@ export function createStore<T extends object>(
};

const destroy = (): void => {
_batchStores.delete(listeners);
globalBatchStores.delete(listeners);
listeners.clear();
if (writeTimeout) {
clearTimeout(writeTimeout);
Expand All @@ -504,16 +546,17 @@ export function createStore<T extends object>(
};
const mutate = (recipe: (draft: T) => void): void => {
const prevState = state;
const context = batchLocalStorage.getStore();
// When in a batch, produce from pending batch state
const baseState: T = _batchDepth > 0 ? _batchStores.get(listeners)?.nextState ?? state : state;
const baseState: T = context && context.depth > 0 ? globalBatchStores.get(listeners)?.nextState ?? state : state;
const nextState = produce(baseState, (draft) => {
recipe(draft as T);
});
if (Object.is(baseState, nextState)) {
return;
}
if (_batchDepth > 0) {
const existing = _batchStores.get(listeners);
if (context && context.depth > 0) {
const existing = globalBatchStores.get(listeners);
// Compute which keys actually changed so commit can merge instead of replace
const changedKeys = Object.keys(nextState).filter(
k => !Object.is((nextState as any)[k], (baseState as any)[k])
Expand All @@ -523,7 +566,7 @@ export function createStore<T extends object>(
(changes as any)[k] = (nextState as any)[k];
}
if (!existing) {
_batchStores.set(listeners, {
globalBatchStores.set(listeners, {
prevState,
nextState,
changes,
Expand Down
Loading
Loading