feat: better scheduling strategy

README.md

@@ -41,9 +41,9 @@ The field `priority` can be can be `ImmediatePriority`, `NormalPriority`, or `Lo
 
 Now you can use jotai-scheduler to replace jotai in your project.
 
-- useAtom --> useAtomWithSchedule
-- useAtomValue --> useAtomValueWithSchedule
-- useSetAtom --> useSetAtomWithSchedule
+- `useAtom` --> `useAtomWithSchedule`
+- `useAtomValue` --> `useAtomValueWithSchedule`
+- `useSetAtom` --> `useSetAtomWithSchedule`
 
 ## Why we need jotai-scheduler?
 

package.json

@@ -1,6 +1,6 @@
 {
   "name": "jotai-scheduler",
-  "version": "0.0.1",
+  "version": "0.0.2",
   "author": "Lincoln",
   "repository": {
     "type": "git",

src/constants.ts

@@ -1,8 +1,3 @@
 export const ImmediatePriority = 1;
 export const NormalPriority = 2;
 export const LowPriority = 3;
-
-export type PriorityLevel =
-  | typeof ImmediatePriority
-  | typeof NormalPriority
-  | typeof LowPriority;

src/index.ts

@@ -1,4 +1,5 @@
-export * from './useAtomValueWithSchedule';
-export * from './useAtomWithSchedule';
-export * from './useSetAtomWithSchedule';
-export * from './constants';
+export { useAtomValueWithSchedule } from './useAtomValueWithSchedule';
+export { useAtomWithSchedule } from './useAtomWithSchedule';
+export { useSetAtomWithSchedule } from './useSetAtomWithSchedule';
+export { LowPriority, NormalPriority, ImmediatePriority } from './constants';
+export * from './types';

src/types.ts

@@ -0,0 +1,25 @@
+import { Atom, useStore } from 'jotai';
+import { ImmediatePriority, LowPriority, NormalPriority } from './constants';
+
+export type SetAtom<Args extends unknown[], Result> = (...args: Args) => Result;
+
+export type AnyAtom = Atom<unknown>;
+
+export type Store = ReturnType<typeof useStore>;
+
+export type Listener = () => void;
+
+export type PriorityLevel =
+  | typeof ImmediatePriority
+  | typeof NormalPriority
+  | typeof LowPriority;
+
+export type Options = Parameters<typeof useStore>[0] & {
+  delay?: number;
+  priority?: PriorityLevel;
+};
+
+export type Task = {
+  priority: PriorityLevel;
+  subscribe: Listener;
+};

src/useAtomValueWithSchedule.ts

@@ -4,45 +4,21 @@ import { useDebugValue, useEffect, useReducer } from 'react';
 import type { ReducerWithoutAction } from 'react';
 import type { Atom, ExtractAtomValue } from 'jotai';
 import { useStore } from 'jotai';
-import {
-  ImmediatePriority,
-  LowPriority,
-  NormalPriority,
-  PriorityLevel,
-} from './constants';
-import { isLastElement, isPromiseLike, use } from './utils';
-
-type Store = ReturnType<typeof useStore>;
-
-type AnyAtom = Atom<unknown>;
 
-type Options = Parameters<typeof useStore>[0] & {
-  delay?: number;
-  priority?: PriorityLevel;
-};
-
-const ImmediatePriorityReRenderMap = new Map<AnyAtom, Set<() => void>>();
-const NormalPriorityPriorityReRenderMap = new Map<AnyAtom, Set<() => void>>();
-const LowPriorityReRenderMap = new Map<AnyAtom, Set<() => void>>();
-const allReRenderMap = new Set<() => void>();
+import { NormalPriority } from './constants';
+import { isLastElement, isPromiseLike, use } from './utils';
+import { addTask, initiateWorkLoop } from './workLoop';
+import { AnyAtom, Options, Store, PriorityLevel, Listener } from './types';
 
-const priorityMap = new Map<
-  PriorityLevel,
-  | typeof ImmediatePriorityReRenderMap
-  | typeof NormalPriorityPriorityReRenderMap
-  | typeof LowPriorityReRenderMap
->([
-  [ImmediatePriority, ImmediatePriorityReRenderMap],
-  [NormalPriority, NormalPriorityPriorityReRenderMap],
-  [LowPriority, LowPriorityReRenderMap],
-]);
+const prioritySubscriptionsMap = new Map<Listener, PriorityLevel>();
+const atomListenersMap = new Map<AnyAtom, Set<Listener>>();
 
 export function useAtomValueWithSchedule<Value>(
   atom: Atom<Value>,
   options?: Options,
 ): Awaited<Value>;
 
-export function useAtomValueWithSchedule<AtomType extends Atom<unknown>>(
+export function useAtomValueWithSchedule<AtomType extends AnyAtom>(
   atom: AtomType,
   options?: Options,
 ): Awaited<ExtractAtomValue<AtomType>>;
@@ -81,6 +57,7 @@ export function useAtomValueWithSchedule<Value>(
 
   const delay = options?.delay;
   const priority = options?.priority ?? NormalPriority;
+
   useEffect(() => {
     const subscribe = () => {
       if (typeof delay === 'number') {
@@ -91,31 +68,29 @@ export function useAtomValueWithSchedule<Value>(
       rerender();
     };
 
-    const listener = priorityMap.get(priority)!.get(atom) ?? new Set();
-    allReRenderMap.add(subscribe);
-    listener.add(subscribe);
-    priorityMap.get(priority)!.set(atom, listener);
+    prioritySubscriptionsMap.set(subscribe, priority);
+    const listeners = atomListenersMap.get(atom) ?? new Set();
+    listeners.add(subscribe);
+    atomListenersMap.set(atom, listeners);
 
     const unsub = store.sub(atom, () => {
-      if (isLastElement(allReRenderMap, subscribe)) {
-        requestAnimationFrame(() => {
-          ImmediatePriorityReRenderMap.get(atom)?.forEach((l) => l());
-          requestAnimationFrame(() => {
-            NormalPriorityPriorityReRenderMap.get(atom)?.forEach((l) => l());
-            requestAnimationFrame(() => {
-              LowPriorityReRenderMap.get(atom)?.forEach((l) => l());
-            });
+      if (isLastElement(listeners, subscribe)) {
+        for (const listener of listeners) {
+          addTask({
+            subscribe: listener,
+            priority: prioritySubscriptionsMap.get(subscribe)!,
           });
-        });
+        }
+        initiateWorkLoop();
       }
     });
-    rerender();
+
     return () => {
       unsub();
-      listener.delete(rerender);
-      allReRenderMap.delete(rerender);
+      prioritySubscriptionsMap.delete(subscribe);
+      listeners.delete(subscribe);
     };
-  }, [store, atom, delay, priority]);
+  }, [atom, delay, priority, store]);
 
   useDebugValue(value);
   // TS doesn't allow using `use` always.

src/useAtomWithSchedule.ts

@@ -9,10 +9,7 @@ import type {
 } from 'jotai';
 import { useAtomValueWithSchedule } from './useAtomValueWithSchedule';
 import { useSetAtomWithSchedule } from './useSetAtomWithSchedule';
-
-type SetAtom<Args extends unknown[], Result> = (...args: Args) => Result;
-
-type Options = Parameters<typeof useAtomValueWithSchedule>[1];
+import { Options, SetAtom } from './types';
 
 export function useAtomWithSchedule<Value, Args extends unknown[], Result>(
   atom: WritableAtom<Value, Args, Result>,

src/useSetAtomWithSchedule.ts

@@ -1,9 +1,7 @@
 import { useCallback } from 'react';
 import type { ExtractAtomArgs, ExtractAtomResult, WritableAtom } from 'jotai';
 import { useStore } from 'jotai';
-
-type SetAtom<Args extends unknown[], Result> = (...args: Args) => Result;
-type Options = Parameters<typeof useStore>[0];
+import { Options, SetAtom } from './types';
 
 export function useSetAtomWithSchedule<Value, Args extends unknown[], Result>(
   atom: WritableAtom<Value, Args, Result>,

src/workLoop.ts

@@ -0,0 +1,48 @@
+import { Task } from './types';
+
+let isMessageLoopRunning: boolean = false;
+
+const taskQueue: Array<Task> = [];
+
+function workLoop(): boolean {
+  const highestPriority = Math.min(...taskQueue.map((task) => task.priority));
+  const highestPriorityList = taskQueue.filter(
+    (task) => task.priority === highestPriority,
+  );
+  highestPriorityList.forEach((task) => {
+    task.subscribe();
+    taskQueue.splice(taskQueue.indexOf(task), 1);
+  });
+  if (taskQueue.length > 0) {
+    return true;
+  }
+  return false;
+}
+
+function handleNextBatch() {
+  const hasMoreWork = workLoop();
+  if (hasMoreWork) {
+    enqueueWorkExecution();
+  } else {
+    isMessageLoopRunning = false;
+  }
+}
+
+const channel = new MessageChannel();
+const port = channel.port2;
+channel.port1.onmessage = handleNextBatch;
+
+export const enqueueWorkExecution = () => {
+  port.postMessage(null);
+};
+
+export const initiateWorkLoop = () => {
+  if (!isMessageLoopRunning) {
+    isMessageLoopRunning = true;
+    enqueueWorkExecution();
+  }
+};
+
+export const addTask = (newTask: Task) => {
+  taskQueue.push(newTask);
+};