Apache Beam写入MongoDB时连接数无法控制的问题求助

问题背景 #

基于Apache Beam v2.43开发流处理管道，部署在DataFlow上，通过自定义DoFn以批量Upsert模式将数据写入MongoDB。已限制worker数量为15，且每个worker的Mongo连接池最大连接数设为10，预期总连接数最多150，但当PubSub出现输入峰值时，MongoDB侧连接数突破20K，直接导致数据库被压垮。

核心问题分析 #

1. MongoClient实例重复创建：当前DoFn通过构造函数生成唯一instanceId，并在@Setup阶段为每个DoFn实例创建独立的MongoClient。而Beam的DoFn会被实例化多次（如bundle并行处理场景），导致每个实例都带10连接的池，总连接数远超预期。
2. batch容器非线程安全：ProcessElement中使用的ArrayList并非线程安全，若Beam多线程并行调用该方法，会引发batch并发修改问题，可能导致频繁flush或数据丢失，间接加剧连接占用。
3. 连接回收机制失效：虽有@Teardown关闭客户端，但DoFn实例若被重复创建且未正确销毁，会导致大量闲置连接堆积。

解决思路 #

1. 单worker内共享唯一MongoClient实例 #

MongoClient本身是线程安全的，应保证每个worker进程内仅存在一个实例，共享同一连接池。修改代码如下：

static class BatchUpsertFn extends DoFn<KV<Document, Document>, CustomBulkWriteError> {

  private static volatile MongoClient mongoClient;
  private List<WriteModel<Document>> batch;

  @Setup
  public void createMongoClient() {
    // 双重检查锁确保单worker内唯一实例
    if (mongoClient == null) {
      synchronized (BatchUpsertFn.class) {
        if (mongoClient == null) {
          mongoClient = MongoClients.create(MongoClientSettings.builder()
              .applyConnectionString(new ConnectionString("myUri"))
              .applyToConnectionPoolSettings(builder -> builder
                  .maxConnectionIdleTime(30, SECONDS)
                  .maxConnectionLifeTime(30, SECONDS)
                  .maxSize(10)
                  .minSize(1)
                  .waitQueueTimeout(5, SECONDS)
              )
              .applyToSocketSettings(builder -> builder
                  .connectTimeout(10, SECONDS)
                  .readTimeout(20, SECONDS))
              .build());
        }
      }
    }
  }

  @StartBundle
  public void startBundle() {
    batch = new ArrayList<>();
  }

  @ProcessElement
  public void processElement(ProcessContext ctx) {
    synchronized (this) {
      batch.add(
            new UpdateManyModel<>(
                ctx.element().getKey(),
                ctx.element().getValue(),
                new UpdateOptions().upsert(true)));

      if (batch.size() >= 1024L) {
        try {
          flush();
        } catch (MongoBulkWriteException pException) {
          pException.getWriteErrors()
              .forEach(pBulkWriteError ->
                  ctx.output(new CustomBulkWriteError(pBulkWriteError)));
        }
      }
    }
  }

  @FinishBundle
  public void finishBundle(FinishBundleContext ctx) {
    try {
      flush();
    } catch (MongoBulkWriteException pException) {
      pException.getWriteErrors()
          .forEach(pBulkWriteError ->
              ctx.output(new CustomBulkWriteError(pBulkWriteError),
                  Instant.now(),
                  GlobalWindow.INSTANCE));
    }
  }

  private void flush() throws MongoBulkWriteException {
    if (batch.isEmpty()) {
      return;
    }
    MongoDatabase mongoDatabase = mongoClient.getDatabase("myDatabase");
    MongoCollection<Document> mongoCollection = mongoDatabase.getCollection("myCollection");
    try {
      mongoCollection.bulkWrite(batch, new BulkWriteOptions().ordered(false));
      batch.clear();
    } catch (MongoBulkWriteException pException) {
      batch.clear();
      throw pException;
    }
  }

  @Teardown
  public void closeMongoClient() {
    if (mongoClient != null) {
      mongoClient.close();
      mongoClient = null;
    }
  }
}

2. 优化连接池参数 #

• 缩短maxConnectionIdleTime和maxConnectionLifeTime至30秒，加快闲置连接回收
• 添加waitQueueTimeout，当连接池满时，超时直接抛出异常，避免连接请求堆积

3. 改用Beam官方MongoDB连接器 #

替换自定义DoFn为官方beam-sdks-java-io-mongodb连接器，官方实现已优化连接池与并发控制，无需手动管理客户端：

MongoDbIO.write()
    .withUri("myUri")
    .withDatabase("myDatabase")
    .withCollection("myCollection")
    .withBulkWriteOptions(new BulkWriteOptions().ordered(false))
    .withWriteFn(doc -> new UpdateManyModel<>(doc.getKey(), doc.getValue(), new UpdateOptions().upsert(true)))

4. 限制DataFlow worker规模 #

结合MongoDB最大连接数计算合理的worker上限：MongoDB最大连接数 ÷ 单worker连接池大小，避免worker扩容后连接数超出数据库承载能力。同时确保worker资源配置充足，减少因重启导致的连接泄漏。

内容的提问来源于stack exchange，提问作者Thibault Coudert