diff --git a/src/vanilla/utils/watch.ts b/src/vanilla/utils/watch.ts index 989e6366..c1022304 100644 --- a/src/vanilla/utils/watch.ts +++ b/src/vanilla/utils/watch.ts @@ -2,7 +2,9 @@ import { subscribe } from '../../vanilla.ts' type Cleanup = () => void type WatchGet = (proxyObject: T) => T -type WatchCallback = (get: WatchGet) => Cleanup | void | undefined +type WatchCallback = ( + get: WatchGet, +) => Cleanup | void | Promise | undefined type WatchOptions = { sync?: boolean } @@ -53,10 +55,12 @@ export function watch( } } - const revalidate = () => { + const revalidate = async () => { if (!alive) { return } + + // run own cleanups before re-subscribing cleanups.forEach((clean) => clean()) cleanups.clear() @@ -69,14 +73,28 @@ export function watch( // Ensures that the parent is reset if the callback throws an error. try { - const cleanupReturn = callback((proxyObject) => { + const promiseOrPossibleCleanup = callback((proxyObject) => { proxiesToSubscribe.add(proxyObject) + // in case the callback is a promise and the watch has ended + if (alive && !subscriptions.has(proxyObject)) { + // subscribe to new proxy immediately -> this fixes problems when Promises are used due to the callstack + const unsubscribe = subscribe(proxyObject, revalidate, options?.sync) + subscriptions.set(proxyObject, unsubscribe) + } return proxyObject }) + const couldBeCleanup = + promiseOrPossibleCleanup && promiseOrPossibleCleanup instanceof Promise + ? await promiseOrPossibleCleanup + : promiseOrPossibleCleanup // If there's a cleanup, we add this to the cleanups set - if (cleanupReturn) { - cleanups.add(cleanupReturn) + if (couldBeCleanup) { + if (alive) { + cleanups.add(couldBeCleanup) + } else { + cleanup() + } } } finally { currentCleanups = parent @@ -84,20 +102,11 @@ export function watch( // Unsubscribe old subscriptions subscriptions.forEach((unsubscribe, proxyObject) => { - if (proxiesToSubscribe.has(proxyObject)) { - // Already subscribed - proxiesToSubscribe.delete(proxyObject) - } else { + if (!proxiesToSubscribe.has(proxyObject)) { subscriptions.delete(proxyObject) unsubscribe() } }) - - // Subscribe to new proxies - proxiesToSubscribe.forEach((proxyObject) => { - const unsubscribe = subscribe(proxyObject, revalidate, options?.sync) - subscriptions.set(proxyObject, unsubscribe) - }) } // If there's a parent watch call, we attach this watch's diff --git a/tests/watch.test.tsx b/tests/watch.test.tsx index 7d2f9363..267368fa 100644 --- a/tests/watch.test.tsx +++ b/tests/watch.test.tsx @@ -1,8 +1,21 @@ -import { describe, expect, it, vi } from 'vitest' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { proxy } from 'valtio' import { watch } from 'valtio/utils' +const sleep = (ms: number) => + new Promise((resolve) => { + setTimeout(resolve, ms) + }) + describe('watch', () => { + beforeEach(() => { + vi.useFakeTimers() + }) + + afterEach(() => { + vi.useRealTimers() + }) + it('should re-run for individual proxy updates', async () => { const reference = proxy({ value: 'Example' }) @@ -96,4 +109,56 @@ describe('watch', () => { reference.value = 'Update' }) + it('should support promise watchers', async () => { + const reference = proxy({ value: 'Example' }) + + const callback = vi.fn() + + const waitPromise = sleep(10000) + watch(async (get) => { + await waitPromise + get(reference) + callback() + }) + + vi.runAllTimers() + await waitPromise + + expect(callback).toBeCalledTimes(1) + // listener will only be attached after one promise callback due to the await stack + await Promise.resolve() + reference.value = 'Update' + // wait for internal promise + await Promise.resolve() + // wait for next promise resolve call due to promise usage inside of callback + await Promise.resolve() + expect(callback).toBeCalledTimes(2) + }) + + it('should not subscribe if the watch is stopped before the promise completes', async () => { + const reference = proxy({ value: 'Example' }) + + const callback = vi.fn() + + const waitPromise = sleep(10000) + const stop = watch(async (get) => { + await waitPromise + get(reference) + callback() + }) + stop() + + vi.runAllTimers() + await waitPromise + + expect(callback).toBeCalledTimes(1) + // listener will only be attached after one promise callback due to the await stack + await Promise.resolve() + reference.value = 'Update' + // wait for internal promise + await Promise.resolve() + // wait for next promise resolve call due to promise usage inside of callback + await Promise.resolve() + expect(callback).toBeCalledTimes(1) + }) })