Backpressure-Aware Batch Writer With Drain

export interface BatchWriterOptions<T> {
  writeBatch: (items: readonly T[]) => Promise<void>;
  maxBatchSize: number;
  highWaterMark?: number;
}

interface Waiter {
  resolve: () => void;
  reject: (error: unknown) => void;
}

export class BatchWriter<T> {
  private readonly writeBatch: (items: readonly T[]) => Promise<void>;
  private readonly maxBatchSize: number;
  private readonly highWaterMark: number;
  private readonly buffer: T[] = [];
  private readonly waiters: Waiter[] = [];
  private flushing: Promise<void> | null = null;
  private closed = false;

  constructor(options: BatchWriterOptions<T>) {
    if (typeof options.writeBatch !== "function") {
      throw new TypeError("writeBatch must be a function");
    }
    if (!Number.isInteger(options.maxBatchSize) || options.maxBatchSize < 1) {
      throw new RangeError("maxBatchSize must be a positive integer");
    }
    const highWaterMark = options.highWaterMark ?? options.maxBatchSize;
    if (!Number.isInteger(highWaterMark) || highWaterMark < 1) {
      throw new RangeError("highWaterMark must be a positive integer");
    }
    this.writeBatch = options.writeBatch;
    this.maxBatchSize = options.maxBatchSize;
    this.highWaterMark = highWaterMark;
  }

  get pending(): number {
    return this.buffer.length;
  }

  add(item: T): Promise<void> {
    if (this.closed) return Promise.reject(new Error("writer is closed"));
    const waiter = new Promise<void>((resolve, reject) => {
      this.waiters.push({ resolve, reject });
    });
    this.buffer.push(item);
    if (this.buffer.length >= this.highWaterMark) {
      void this.flush().catch(() => {});
    }
    return waiter;
  }

  flush(): Promise<void> {
    if (this.flushing) return this.flushing;
    this.flushing = this.drain().finally(() => {
      this.flushing = null;
    });
    return this.flushing;
  }

  async close(): Promise<void> {
    this.closed = true;
    await this.flush();
  }

  private async drain(): Promise<void> {
    while (this.buffer.length > 0) {
      const batch = this.buffer.splice(0, this.maxBatchSize);
      const waiters = this.waiters.splice(0, batch.length);
      try {
        await this.writeBatch(batch);
        for (const waiter of waiters) waiter.resolve();
      } catch (error) {
        for (const waiter of waiters) waiter.reject(error);
        throw error;
      }
    }
  }
}