improvement: cancel handler receives next value

pmndrs · Sep 8, 2024 · ba0a731 · ba0a731
1 parent b56e632
commit ba0a731
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 22 deletions.
diff --git a/src/react/useAtomValue.ts b/src/react/useAtomValue.ts
@@ -56,10 +56,7 @@ const continuablePromiseMap = new WeakMap<
   Promise<unknown>
 >()
 
-const createContinuablePromise = <T>(
-  promise: PromiseLike<T>,
-  getLatest: () => PromiseLike<T> | T,
-) => {
+const createContinuablePromise = <T>(promise: PromiseLike<T>) => {
   let continuablePromise = continuablePromiseMap.get(promise)
   if (!continuablePromise) {
     continuablePromise = new Promise<T>((resolve, reject) => {
@@ -76,18 +73,18 @@ const createContinuablePromise = <T>(
       }
       const registerCancelHandler = (p: PromiseLike<T>) => {
         if ('onCancel' in p && typeof p.onCancel === 'function') {
-          p.onCancel(() => {
-            const nextValue = getLatest()
-            const nextPromise = isPromiseLike(nextValue)
-              ? nextValue
-              : Promise.resolve(nextValue)
-            if (import.meta.env?.MODE !== 'production' && nextPromise === p) {
+          p.onCancel((nextValue: PromiseLike<T> | T) => {
+            if (import.meta.env?.MODE !== 'production' && nextValue === p) {
               throw new Error('[Bug] p is not updated even after cancelation')
             }
-            continuablePromiseMap.set(nextPromise, continuablePromise!)
-            curr = nextPromise
-            nextPromise.then(onFulfilled(nextPromise), onRejected(nextPromise))
-            registerCancelHandler(nextPromise)
+            if (isPromiseLike(nextValue)) {
+              continuablePromiseMap.set(nextValue, continuablePromise!)
+              curr = nextValue
+              nextValue.then(onFulfilled(nextValue), onRejected(nextValue))
+              registerCancelHandler(nextValue)
+            } else {
+              resolve(nextValue)
+            }
           })
         }
       }
@@ -148,9 +145,7 @@ export function useAtomValue<Value>(atom: Atom<Value>, options?: Options) {
       if (typeof delay === 'number') {
         const value = store.get(atom)
         if (isPromiseLike(value)) {
-          attachPromiseMeta(
-            createContinuablePromise(value, () => store.get(atom)),
-          )
+          attachPromiseMeta(createContinuablePromise(value))
         }
         // delay rerendering to wait a promise possibly to resolve
         setTimeout(rerender, delay)
@@ -166,7 +161,7 @@ export function useAtomValue<Value>(atom: Atom<Value>, options?: Options) {
   // The use of isPromiseLike is to be consistent with `use` type.
   // `instanceof Promise` actually works fine in this case.
   if (isPromiseLike(value)) {
-    const promise = createContinuablePromise(value, () => store.get(atom))
+    const promise = createContinuablePromise(value)
     return use(promise)
   }
   return value as Awaited<Value>

diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts
@@ -23,19 +23,19 @@ const isActuallyWritableAtom = (atom: AnyAtom): atom is AnyWritableAtom =>
 // Cancelable Promise
 //
 
-type CancelHandler = () => void
+type CancelHandler = (nextValue: unknown) => void
 type PromiseState = [cancelHandlers: Set<CancelHandler>, settled: boolean]
 
 const cancelablePromiseMap = new WeakMap<PromiseLike<unknown>, PromiseState>()
 
 const isPendingPromise = (value: unknown): value is PromiseLike<unknown> =>
   isPromiseLike(value) && !cancelablePromiseMap.get(value)?.[1]
 
-const cancelPromise = <T>(promise: PromiseLike<T>) => {
+const cancelPromise = <T>(promise: PromiseLike<T>, nextValue: unknown) => {
   const promiseState = cancelablePromiseMap.get(promise)
   if (promiseState) {
     promiseState[1] = true
-    promiseState[0].forEach((fn) => fn())
+    promiseState[0].forEach((fn) => fn(nextValue))
   } else if (import.meta.env?.MODE !== 'production') {
     throw new Error('[Bug] cancelable promise not found')
   }
@@ -279,7 +279,7 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => {
     if (!hasPrevValue || !Object.is(prevValue, atomState.v)) {
       ++atomState.n
       if (pendingPromise) {
-        cancelPromise(pendingPromise)
+        cancelPromise(pendingPromise, valueOrPromise)
       }
     }
   }