From f2aadee00039328150c740145f6b6215c3f84e64 Mon Sep 17 00:00:00 2001 From: Lincoln <778157949@qq.com> Date: Fri, 23 Aug 2024 09:34:14 +0000 Subject: [PATCH] feat: better scheduling strategy --- README.md | 6 +-- package.json | 2 +- src/constants.ts | 5 --- src/index.ts | 9 +++-- src/types.ts | 25 ++++++++++++ src/useAtomValueWithSchedule.ts | 71 +++++++++++---------------------- src/useAtomWithSchedule.ts | 5 +-- src/useSetAtomWithSchedule.ts | 4 +- src/workLoop.ts | 48 ++++++++++++++++++++++ 9 files changed, 107 insertions(+), 68 deletions(-) create mode 100644 src/types.ts create mode 100644 src/workLoop.ts diff --git a/README.md b/README.md index dc6bf74..1cbb149 100644 --- a/README.md +++ b/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? diff --git a/package.json b/package.json index b4e3803..6ee4430 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "jotai-scheduler", - "version": "0.0.1", + "version": "0.0.2", "author": "Lincoln", "repository": { "type": "git", diff --git a/src/constants.ts b/src/constants.ts index 0197716..9bbab58 100644 --- a/src/constants.ts +++ b/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; diff --git a/src/index.ts b/src/index.ts index a08fe9d..f37f996 100644 --- a/src/index.ts +++ b/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'; diff --git a/src/types.ts b/src/types.ts new file mode 100644 index 0000000..0ecb13b --- /dev/null +++ b/src/types.ts @@ -0,0 +1,25 @@ +import { Atom, useStore } from 'jotai'; +import { ImmediatePriority, LowPriority, NormalPriority } from './constants'; + +export type SetAtom = (...args: Args) => Result; + +export type AnyAtom = Atom; + +export type Store = ReturnType; + +export type Listener = () => void; + +export type PriorityLevel = + | typeof ImmediatePriority + | typeof NormalPriority + | typeof LowPriority; + +export type Options = Parameters[0] & { + delay?: number; + priority?: PriorityLevel; +}; + +export type Task = { + priority: PriorityLevel; + subscribe: Listener; +}; diff --git a/src/useAtomValueWithSchedule.ts b/src/useAtomValueWithSchedule.ts index 9ef61e4..3f98516 100644 --- a/src/useAtomValueWithSchedule.ts +++ b/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; - -type AnyAtom = Atom; -type Options = Parameters[0] & { - delay?: number; - priority?: PriorityLevel; -}; - -const ImmediatePriorityReRenderMap = new Map void>>(); -const NormalPriorityPriorityReRenderMap = new Map void>>(); -const LowPriorityReRenderMap = new Map 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(); +const atomListenersMap = new Map>(); export function useAtomValueWithSchedule( atom: Atom, options?: Options, ): Awaited; -export function useAtomValueWithSchedule>( +export function useAtomValueWithSchedule( atom: AtomType, options?: Options, ): Awaited>; @@ -81,6 +57,7 @@ export function useAtomValueWithSchedule( const delay = options?.delay; const priority = options?.priority ?? NormalPriority; + useEffect(() => { const subscribe = () => { if (typeof delay === 'number') { @@ -91,31 +68,29 @@ export function useAtomValueWithSchedule( 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. diff --git a/src/useAtomWithSchedule.ts b/src/useAtomWithSchedule.ts index 8a3dbb9..b149c76 100644 --- a/src/useAtomWithSchedule.ts +++ b/src/useAtomWithSchedule.ts @@ -9,10 +9,7 @@ import type { } from 'jotai'; import { useAtomValueWithSchedule } from './useAtomValueWithSchedule'; import { useSetAtomWithSchedule } from './useSetAtomWithSchedule'; - -type SetAtom = (...args: Args) => Result; - -type Options = Parameters[1]; +import { Options, SetAtom } from './types'; export function useAtomWithSchedule( atom: WritableAtom, diff --git a/src/useSetAtomWithSchedule.ts b/src/useSetAtomWithSchedule.ts index 726185b..f7646ab 100644 --- a/src/useSetAtomWithSchedule.ts +++ b/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: Args) => Result; -type Options = Parameters[0]; +import { Options, SetAtom } from './types'; export function useSetAtomWithSchedule( atom: WritableAtom, diff --git a/src/workLoop.ts b/src/workLoop.ts new file mode 100644 index 0000000..8991801 --- /dev/null +++ b/src/workLoop.ts @@ -0,0 +1,48 @@ +import { Task } from './types'; + +let isMessageLoopRunning: boolean = false; + +const taskQueue: Array = []; + +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); +};