Using a linked list instead of an array for subscribers queue

Author: divdavem

src/index.spec.ts

@@ -2848,6 +2848,28 @@ describe('stores', () => {
       a.next(4);
       expect(values).toEqual([0, 6]);
     });
+
+    it('should work to unsubscribe inside batch', () => {
+      const valuesA: number[] = [];
+      const valuesB: number[] = [];
+      const a = writable(0);
+      const b = computed(() => a() + 1);
+      const unsubscribeA = a.subscribe((v) => {
+        valuesA.push(v);
+      });
+      const unsubscribeB = b.subscribe((v) => {
+        valuesB.push(v);
+      });
+      expect(valuesA).toEqual([0]);
+      expect(valuesB).toEqual([1]);
+      batch(() => {
+        a.set(5);
+        unsubscribeB();
+      });
+      expect(valuesA).toEqual([0, 5]);
+      expect(valuesB).toEqual([1]);
+      unsubscribeA();
+    });
   });
 
   describe('Listeners and batch timing', () => {

src/internal/batch.ts

@@ -1,7 +1,10 @@
+import { createQueue } from './linkedQueue';
 import type { SubscribeConsumer } from './subscribeConsumer';
 
-export const subscribersQueue: SubscribeConsumer<any, any>[] = [];
 let willProcessQueue = false;
+const { add, remove, shift } = createQueue<SubscribeConsumer<any, any>>();
+
+export { add as addToQueue, remove as removeFromQueue };
 
 /**
  * Batches multiple changes to stores while calling the provided function,
@@ -53,8 +56,8 @@ export const batch = <T>(fn: () => T): T => {
     res = fn();
   } finally {
     if (needsProcessQueue) {
-      while (subscribersQueue.length > 0) {
-        const consumer = subscribersQueue.shift()!;
+      let consumer = shift();
+      while (consumer) {
         try {
           consumer.notify();
         } catch (e) {
@@ -65,6 +68,7 @@ export const batch = <T>(fn: () => T): T => {
             error = e;
           }
         }
+        consumer = shift();
       }
       willProcessQueue = false;
     }

src/internal/linkedQueue.ts

@@ -0,0 +1,63 @@
+export interface QueueItem<T> {
+  next: T | null;
+  prev: T | null;
+}
+
+export const createQueue = <T extends QueueItem<T>>(): {
+  add: (item: T) => boolean;
+  remove: (item: T) => boolean;
+  shift: () => T | null;
+} => {
+  let first: T | null = null;
+  let last: T | null = null;
+
+  const remove = (item: T): boolean => {
+    const { prev, next } = item;
+    if (prev || next || first === item) {
+      item.prev = null;
+      item.next = null;
+      if (prev) {
+        prev.next = next;
+      } else {
+        first = next;
+      }
+      if (next) {
+        next.prev = prev;
+      } else {
+        last = prev;
+      }
+      return true;
+    }
+    return false;
+  };
+
+  const add = (item: T): boolean => {
+    if (item === last) {
+      // already the last in the queue, nothing to do !
+      return false;
+    }
+    const existing = remove(item);
+    item.prev = last;
+    if (last) {
+      last.next = item;
+    } else {
+      first = item;
+    }
+    last = item;
+    return !existing;
+  };
+
+  const shift = (): T | null => {
+    const item = first;
+    if (item) {
+      remove(item);
+    }
+    return item;
+  };
+
+  return {
+    add,
+    remove,
+    shift,
+  };
+};

src/internal/subscribeConsumer.ts

@@ -1,5 +1,6 @@
 import type { Subscriber, SubscriberObject } from '../types';
-import { subscribersQueue } from './batch';
+import { addToQueue, removeFromQueue } from './batch';
+import type { QueueItem } from './linkedQueue';
 import { updateLinkProducerValue, type BaseLink, type Consumer, type RawStore } from './store';
 
 export const noop = (): void => {};
@@ -21,10 +22,14 @@ const toSubscriberObject = <T>(subscriber: Subscriber<T>): SubscriberObject<T> =
   resume: bind(subscriber, 'resume'),
 });
 
-export class SubscribeConsumer<T, Link extends BaseLink<T>> implements Consumer {
+export class SubscribeConsumer<T, Link extends BaseLink<T>>
+  implements Consumer, QueueItem<SubscribeConsumer<any, any>>
+{
   private readonly link: Link;
   private subscriber: SubscriberObject<T>;
-  dirtyCount = 1;
+  prev: SubscribeConsumer<any, any> | null = null;
+  next: SubscribeConsumer<any, any> | null = null;
+
   constructor(producer: RawStore<T, Link>, subscriber: Subscriber<T>) {
     this.subscriber = toSubscriberObject(subscriber);
     this.link = producer.registerConsumer(producer.newLink(this));
@@ -34,31 +39,27 @@ export class SubscribeConsumer<T, Link extends BaseLink<T>> implements Consumer
   unsubscribe(): void {
     if (this.subscriber !== noopSubscriber) {
       this.subscriber = noopSubscriber;
+      removeFromQueue(this);
       this.link.producer.unregisterConsumer(this.link);
     }
   }
 
   markDirty(): void {
-    this.dirtyCount++;
-    subscribersQueue.push(this);
-    if (this.dirtyCount === 1) {
+    if (addToQueue(this)) {
       this.subscriber.pause();
     }
   }
 
   notify(first = false): void {
-    this.dirtyCount--;
-    if (this.dirtyCount === 0 && this.subscriber !== noopSubscriber) {
-      const link = this.link;
-      const producer = link.producer;
-      updateLinkProducerValue(link);
-      if (producer.isLinkUpToDate(link) && !first) {
-        this.subscriber.resume();
-      } else {
-        // note that the following line can throw
-        const value = producer.updateLink(link);
-        this.subscriber.next(value);
-      }
+    const link = this.link;
+    const producer = link.producer;
+    updateLinkProducerValue(link);
+    if (producer.isLinkUpToDate(link) && !first) {
+      this.subscriber.resume();
+    } else {
+      // note that the following line can throw
+      const value = producer.updateLink(link);
+      this.subscriber.next(value);
     }
   }
 }