fix(scheduler): drain after promise + refactor

wmertens · wmertens · commit bddac0550fe1 · 2025-01-21T18:53:22.000+01:00
diff --git a/packages/qwik/src/core/shared/scheduler.ts b/packages/qwik/src/core/shared/scheduler.ts
@@ -80,15 +80,7 @@
  *   declaration order within component.
  */
 
-import { assertEqual } from './error/assert';
-import type { QRLInternal } from './qrl/qrl-class';
-import type { JSXOutput } from './jsx/types/jsx-node';
-import { Task, TaskFlags, cleanupTask, runTask, type TaskFn } from '../use/use-task';
-import { runResource, type ResourceDescriptor } from '../use/use-resource';
-import { logWarn } from './utils/log';
-import { isPromise, maybeThenPassError, retryOnPromise, safeCall } from './utils/promises';
-import type { ValueOrPromise } from './utils/types';
-import { isDomContainer } from '../client/dom-container';
+import { isDomContainer, type DomContainer } from '../client/dom-container';
 import {
   ElementVNodeProps,
   VNodeFlags,
@@ -97,21 +89,28 @@ import {
   type ElementVNode,
   type VirtualVNode,
 } from '../client/types';
-import { vnode_isVNode, vnode_setAttr, VNodeJournalOpCode } from '../client/vnode';
+import { VNodeJournalOpCode, vnode_isVNode, vnode_setAttr } from '../client/vnode';
 import { vnode_diff } from '../client/vnode-diff';
-import { executeComponent } from './component-execution';
-import type { Container, HostElement } from './types';
+import { triggerEffects, type ComputedSignal, type WrappedSignal } from '../signal/signal';
 import { isSignal, type Signal } from '../signal/signal.public';
-import { type DomContainer } from '../client/dom-container';
-import { serializeAttribute } from './utils/styles';
+import type { TargetType } from '../signal/store';
+import type { ISsrNode } from '../ssr/ssr-types';
+import { runResource, type ResourceDescriptor } from '../use/use-resource';
+import { Task, TaskFlags, cleanupTask, runTask, type TaskFn } from '../use/use-task';
+import { executeComponent } from './component-execution';
 import type { OnRenderFn } from './component.public';
+import { assertEqual } from './error/assert';
 import type { Props } from './jsx/jsx-runtime';
+import type { JSXOutput } from './jsx/types/jsx-node';
+import type { QRLInternal } from './qrl/qrl-class';
+import { ssrNodeDocumentPosition, vnode_documentPosition } from './scheduler-document-position';
+import type { Container, HostElement } from './types';
+import { logWarn } from './utils/log';
 import { QScopedStyle } from './utils/markers';
+import { isPromise, retryOnPromise, safeCall } from './utils/promises';
 import { addComponentStylePrefix } from './utils/scoped-styles';
-import { type WrappedSignal, type ComputedSignal, triggerEffects } from '../signal/signal';
-import type { TargetType } from '../signal/store';
-import { ssrNodeDocumentPosition, vnode_documentPosition } from './scheduler-document-position';
-import type { ISsrNode } from '../ssr/ssr-types';
+import { serializeAttribute } from './utils/styles';
+import type { ValueOrPromise } from './utils/types';
 
 // Turn this on to get debug output of what the scheduler is doing.
 const DEBUG: boolean = false;
@@ -150,7 +149,7 @@ export interface Chore {
   $target$: ChoreTarget | null;
   $payload$: unknown;
   $resolve$: (value: any) => void;
-  $promise$: Promise<any>;
+  $promise$?: Promise<any>;
   $returnValue$: any;
   $executed$: boolean;
 }
@@ -168,6 +167,10 @@ export type Scheduler = ReturnType<typeof createScheduler>;
 
 type ChoreTarget = HostElement | QRLInternal<(...args: unknown[]) => unknown> | Signal | TargetType;
 
+const getPromise = (chore: Chore) => {
+  return (chore.$promise$ ||= new Promise((resolve) => (chore.$resolve$ = resolve)));
+};
+
 export const createScheduler = (
   container: Container,
   scheduleDrain: () => void,
@@ -273,7 +276,6 @@ export const createScheduler = (
       $returnValue$: null,
       $executed$: false,
     };
-    chore.$promise$ = new Promise((resolve) => (chore.$resolve$ = resolve));
     DEBUG && debugTrace('schedule', chore, currentChore, choreQueue);
     chore = sortedInsert(choreQueue, chore, (container as DomContainer).rootVNode || null);
     if (!journalFlushScheduled && runLater) {
@@ -283,7 +285,7 @@ export const createScheduler = (
       scheduleDrain();
     }
     if (runLater) {
-      return chore.$promise$;
+      return getPromise(chore);
     } else {
       return drainUpTo(chore, (container as DomContainer).rootVNode || null);
     }
@@ -301,7 +303,7 @@ export const createScheduler = (
     }
     if (currentChore) {
       // Already running chore, which means we're waiting for async completion
-      return runUptoChore.$promise$;
+      return getPromise(currentChore).then(() => drainUpTo(runUptoChore, rootVNode));
     }
     while (choreQueue.length) {
       const nextChore = choreQueue[0];
@@ -327,9 +329,10 @@ export const createScheduler = (
         continue;
       }
       const returnValue = executeChore(nextChore);
-      if (isPromise(returnValue)) {
-        const promise = returnValue.then(() => drainUpTo(runUptoChore, rootVNode));
-        return promise;
+      if (currentChore) {
+        // We're waiting for the returnValue promise
+        // Continue draining after it is resolved
+        return returnValue.then(() => drainUpTo(runUptoChore, rootVNode));
       }
     }
     return runUptoChore.$returnValue$;
@@ -376,12 +379,18 @@ export const createScheduler = (
         );
         break;
       case ChoreType.RESOURCE:
-        // Don't await the return value of the resource, because async resources should not be awaited.
-        // The reason for this is that we should be able to update for example a node with loading
-        // text. If we await the resource, the loading text will not be displayed until the resource
-        // is loaded.
-        const result = runResource(chore.$payload$ as ResourceDescriptor<TaskFn>, container, host);
-        returnValue = isDomContainer(container) ? null : result;
+        {
+          // Don't await the return value of the resource, because async resources should not be awaited.
+          // The reason for this is that we should be able to update for example a node with loading
+          // text. If we await the resource, the loading text will not be displayed until the resource
+          // is loaded.
+          const result = runResource(
+            chore.$payload$ as ResourceDescriptor<TaskFn>,
+            container,
+            host
+          );
+          returnValue = isDomContainer(container) ? null : result;
+        }
         break;
       case ChoreType.RUN_QRL:
         {
@@ -396,66 +405,84 @@ export const createScheduler = (
         returnValue = runTask(chore.$payload$ as Task<TaskFn, TaskFn>, container, host);
         break;
       case ChoreType.CLEANUP_VISIBLE:
-        const task = chore.$payload$ as Task<TaskFn, TaskFn>;
-        cleanupTask(task);
+        {
+          const task = chore.$payload$ as Task<TaskFn, TaskFn>;
+          cleanupTask(task);
+        }
         break;
       case ChoreType.NODE_DIFF:
-        const parentVirtualNode = chore.$target$ as VirtualVNode;
-        let jsx = chore.$payload$ as JSXOutput;
-        if (isSignal(jsx)) {
-          jsx = jsx.value as any;
+        {
+          const parentVirtualNode = chore.$target$ as VirtualVNode;
+          let jsx = chore.$payload$ as JSXOutput;
+          if (isSignal(jsx)) {
+            jsx = jsx.value as any;
+          }
+          returnValue = retryOnPromise(() =>
+            vnode_diff(container as DomContainer, jsx, parentVirtualNode, null)
+          );
         }
-        returnValue = retryOnPromise(() =>
-          vnode_diff(container as DomContainer, jsx, parentVirtualNode, null)
-        );
         break;
       case ChoreType.NODE_PROP:
-        const virtualNode = chore.$host$ as unknown as ElementVNode;
-        const payload = chore.$payload$ as NodePropPayload;
-        let value: Signal<any> | string = payload.$value$;
-        if (isSignal(value)) {
-          value = value.value as any;
-        }
-        const isConst = payload.$isConst$;
-        const journal = (container as DomContainer).$journal$;
-        const property = chore.$idx$ as string;
-        const serializedValue = serializeAttribute(property, value, payload.$scopedStyleIdPrefix$);
-        if (isConst) {
-          const element = virtualNode[ElementVNodeProps.element] as Element;
-          journal.push(VNodeJournalOpCode.SetAttribute, element, property, serializedValue);
-        } else {
-          vnode_setAttr(journal, virtualNode, property, serializedValue);
+        {
+          const virtualNode = chore.$host$ as unknown as ElementVNode;
+          const payload = chore.$payload$ as NodePropPayload;
+          let value: Signal<any> | string = payload.$value$;
+          if (isSignal(value)) {
+            value = value.value as any;
+          }
+          const isConst = payload.$isConst$;
+          const journal = (container as DomContainer).$journal$;
+          const property = chore.$idx$ as string;
+          const serializedValue = serializeAttribute(
+            property,
+            value,
+            payload.$scopedStyleIdPrefix$
+          );
+          if (isConst) {
+            const element = virtualNode[ElementVNodeProps.element] as Element;
+            journal.push(VNodeJournalOpCode.SetAttribute, element, property, serializedValue);
+          } else {
+            vnode_setAttr(journal, virtualNode, property, serializedValue);
+          }
         }
         break;
       case ChoreType.QRL_RESOLVE: {
-        const target = chore.$target$ as QRLInternal<any>;
-        returnValue = !target.resolved ? target.resolve() : null;
+        {
+          const target = chore.$target$ as QRLInternal<any>;
+          returnValue = !target.resolved ? target.resolve() : null;
+        }
         break;
       }
       case ChoreType.RECOMPUTE_AND_SCHEDULE_EFFECTS: {
-        const target = chore.$target$ as ComputedSignal<unknown> | WrappedSignal<unknown>;
-        const forceRunEffects = target.$forceRunEffects$;
-        target.$forceRunEffects$ = false;
-        if (!target.$effects$?.length) {
-          break;
-        }
-        returnValue = retryOnPromise(() => {
-          if (target.$computeIfNeeded$() || forceRunEffects) {
-            triggerEffects(container, target, target.$effects$);
+        {
+          const target = chore.$target$ as ComputedSignal<unknown> | WrappedSignal<unknown>;
+          const forceRunEffects = target.$forceRunEffects$;
+          target.$forceRunEffects$ = false;
+          if (!target.$effects$?.length) {
+            break;
           }
-        });
+          returnValue = retryOnPromise(() => {
+            if (target.$computeIfNeeded$() || forceRunEffects) {
+              triggerEffects(container, target, target.$effects$);
+            }
+          });
+        }
         break;
       }
     }
-    return maybeThenPassError(returnValue, (value) => {
-      if (currentChore) {
-        currentChore.$executed$ = true;
-        currentChore.$resolve$?.(value);
-      }
-      DEBUG && debugTrace('execute.DONE', null, currentChore, choreQueue);
+    const after = (value: any) => {
       currentChore = null;
-      return (chore.$returnValue$ = value);
-    });
+      chore.$executed$ = true;
+      chore.$returnValue$ = value;
+      DEBUG && debugTrace('execute.DONE', null, chore, choreQueue);
+      chore.$resolve$?.(value);
+      return value;
+    };
+    if (isPromise(returnValue)) {
+      chore.$promise$ = returnValue;
+      return returnValue.then(after);
+    }
+    return after(returnValue);
   }
 };
 
diff --git a/starters/apps/e2e/src/components/computed/computed.tsx b/starters/apps/e2e/src/components/computed/computed.tsx
@@ -26,7 +26,6 @@ export const ComputedBasic = component$(() => {
   const triple = useComputed$(() => plus3.value * 3);
   const sum = useComputed$(() => double.value + plus3.value + triple.value);
 
-  console.log("here");
   return (
     <div>
       <div class="result">count: {count.value}</div>
