← 中文资源

简体中文版本

感知背压的批量写入器

把异步写入合并为有界批次,在高水位触发背压,串行化刷新并在关闭时排空。

阅读英文原文

功能说明

该示例把高频的小写入合并成批量操作,减少下游压力。add 将元素放入缓冲区并返回一个在该元素所在批次成功写入后解决的 Promise;当缓冲量达到 highWaterMark 时立即触发刷新。flush 串行执行,不允许多个刷新同时操作缓冲区;close 标记关闭并排空剩余数据。适用于批量插入数据库、日志聚合、对象存储上传等场景。

关键代码

缓冲区与等待者数组按相同顺序增长,每次从缓冲区切出最多 maxBatchSize 个元素时,也从等待者中切出相同数量,保证每个 Promise 对应自己的批次。刷新使用单个 flushing Promise 去重,drain 循环直到缓冲区为空,因此刷新期间新增的数据也会被后续批次带走。写入失败时只拒绝当前批次的等待者,保留剩余数据并抛出错误,调用方可再次 flush 重试。

使用注意

add 达到高水位后会主动刷新,但刷新失败不会通过 add 返回的 Promise 之外的渠道传播,代码内部用 catch 吞掉刷新 Promise 的拒绝以免未处理拒绝,真正的错误仍会传给各等待者。close 之后再 add 会立即拒绝;close 失败后缓冲区可能仍有数据。该实现不启动定时器,低于高水位的数据必须显式 flush 或 close 才会写出。

验证方式

用记录批次的假写入函数:连续 add 到高水位,断言 writeBatch 收到不超过 maxBatchSize 的批次且等待者全部解决;在写入函数中延迟并继续 add,验证刷新期间数据不会被丢失;让某批次拒绝,断言对应等待者拒绝且剩余数据仍可通过 flush 写出;close 后 add 拒绝;测试非法批次尺寸与高水位抛 RangeError。

代码示例

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;
      }
    }
  }
}