From f92183bcfa8c0959c0eb68b0f3efb24087a87e08 Mon Sep 17 00:00:00 2001
From: Rahim Alwer <rahim.alwer@gmail.com>
Date: Wed, 6 Jul 2022 18:22:56 +1000
Subject: [PATCH] feat: new export `getScheduler`

---
 README.md               | 21 +++++++++++++++++++++
 src/observables.ts      |  9 ++++++++-
 src/scheduler.ts        | 12 +++++++-----
 tests/syncFlush.test.ts | 29 +++++++++++++++++++++++++++++
 4 files changed, 65 insertions(+), 6 deletions(-)
 create mode 100644 tests/syncFlush.test.ts

diff --git a/README.md b/README.md
index 20d0f2d..446cd41 100644
--- a/README.md
+++ b/README.md
@@ -109,6 +109,7 @@ $: yarn add @maverick-js/observables
 - [`onDispose`](#ondispose)
 - [`isObservable`](#isobservable)
 - [`isSubject`](#issubject)
+- [`getScheduler`](#getscheduler)
 
 ## `$root`
 
@@ -391,6 +392,26 @@ isSubject($computed(() => 10));
 isSubject($readonly($observable(10)));
 ```
 
+### `getScheduler`
+
+Returns the global scheduler which can be used to queue additional tasks or synchronously flush
+the queue.
+
+```js
+const scheduler = getScheduler();
+
+// Queue task
+scheduler.enqueue(() => {
+  // ...
+});
+
+// Flush queue synchronously
+scheduler.syncFlush();
+```
+
+> **Note**
+> See the ['Scheduler'](#scheduler) section for more information.
+
 ## Debugging
 
 The `$observable`, `$computed`, and `$effect` functions accept a debugging ID (string) as part
diff --git a/src/observables.ts b/src/observables.ts
index f3adf3a..3e07ade 100644
--- a/src/observables.ts
+++ b/src/observables.ts
@@ -1,4 +1,4 @@
-import { createScheduler } from './scheduler';
+import { createScheduler, type Scheduler } from './scheduler';
 
 export type Observable<T> = {
   id?: string;
@@ -377,6 +377,13 @@ export function isSubject<T>(fn: MaybeObservable<T>): fn is ObservableSubject<T>
   return isObservable(fn) && !!(fn as ObservableSubject<T>).set;
 }
 
+/**
+ * Returns the global scheduler.
+ */
+export function getScheduler(): Scheduler {
+  return _scheduler;
+}
+
 type Node = {
   id?: string;
   (): any;
diff --git a/src/scheduler.ts b/src/scheduler.ts
index b5829c0..a84f860 100644
--- a/src/scheduler.ts
+++ b/src/scheduler.ts
@@ -5,6 +5,7 @@ export type Scheduler = {
   enqueue: (task: ScheduledTask) => void;
   served: (task: ScheduledTask) => boolean;
   flush: () => void;
+  syncFlush: () => void;
   tick: Promise<void>;
 };
 
@@ -29,14 +30,14 @@ export type Scheduler = {
  * ```
  */
 export function createScheduler(onFlush?: SchedulerFlushed): Scheduler {
-  const processed = new Set<ScheduledTask>();
+  const served = new Set<ScheduledTask>();
   const queue = new Set<ScheduledTask>();
   const microtask = Promise.resolve();
   const queueTask = typeof queueMicrotask !== 'undefined' ? queueMicrotask : microtask.then;
 
   function enqueue(task: ScheduledTask) {
     // `processed` is only populated during a flush.
-    if (!processed.has(task)) queue.add(task);
+    if (!served.has(task)) queue.add(task);
     scheduleFlush();
   }
 
@@ -52,20 +53,21 @@ export function createScheduler(onFlush?: SchedulerFlushed): Scheduler {
     do {
       for (const task of queue) {
         task();
-        processed.add(task);
+        served.add(task);
         queue.delete(task);
       }
     } while (queue.size > 0);
 
-    processed.clear();
+    served.clear();
     flushing = false;
     onFlush?.();
   }
 
   return {
     enqueue,
-    served: (task) => processed.has(task),
+    served: (task) => served.has(task),
     flush: scheduleFlush,
+    syncFlush: flush,
     tick: microtask,
   };
 }
diff --git a/tests/syncFlush.test.ts b/tests/syncFlush.test.ts
new file mode 100644
index 0000000..78af45b
--- /dev/null
+++ b/tests/syncFlush.test.ts
@@ -0,0 +1,29 @@
+import { $effect, $observable, getScheduler } from '../src';
+
+it('should flush queue synchronously', async () => {
+  const scheduler = getScheduler();
+
+  const $a = $observable(0);
+  const $b = $observable(0);
+
+  const effect = vi.fn();
+  $effect(() => {
+    $a(), $b(), effect();
+  });
+
+  $a.set(2);
+  $b.set(3);
+
+  scheduler.syncFlush();
+  expect($a()).toBe(2);
+  expect($b()).toBe(3);
+  expect(effect).toHaveBeenCalledTimes(2);
+
+  $a.set(4);
+  $b.set(6);
+
+  scheduler.syncFlush();
+  expect($a()).toBe(4);
+  expect($b()).toBe(6);
+  expect(effect).toHaveBeenCalledTimes(3);
+});