Akka Stream自定义GraphStage异步回调最终调用未执行求助

看起来你踩中了Akka Streams自定义GraphStage处理异步任务时的几个经典陷阱——阻塞stage逻辑线程、错误处理上游完成事件，导致异步回调无法正常执行。咱们一步步拆解问题并修复：

核心问题分析 #

你的代码里有三个关键错误：

1. 阻塞式等待异步任务：onUpstreamFinish里的while(!canComplete)循环会完全占用GraphStageLogic的单线程，而异步回调也需要在这个线程上执行，导致回调根本没机会运行。这就是为什么最后一个元素的异步操作明明完成了，但called completedCb日志从未打印，直到stage超时结束后才抛出StreamDetachedException。
2. 简陋的任务跟踪方式：只用一个布尔值canComplete无法准确跟踪多个异步任务的状态，多元素场景下很容易出现状态混乱。
3. 错误使用invokeWithFeedback：这个方法是给回调需要向调用者返回结果的场景设计的，你的场景用普通的invoke就足够了。

修复方案 #

我们需要用异步方式跟踪待处理任务，而非阻塞，并正确利用Akka Streams的stage生命周期API：

1. 替换阻塞等待为任务计数器 #

用整数计数器跟踪正在处理的异步任务数量，替代canComplete布尔值：

private var pendingTasks = 0
private var upstreamFinished = false

2. 正确处理上游完成事件 #

去掉阻塞循环，改用holdUpstreamFinish()延迟stage完成，直到所有异步任务结束：

override def onUpstreamFinish(): Unit = {
  logger.info(s"Upstream finished")
  upstreamFinished = true
  // 没有待处理任务就直接完成stage，否则等待任务结束
  checkIfCanComplete()
}

3. 异步任务生命周期绑定计数器 #

启动异步任务时增加计数器，回调完成时减少计数器，并检查是否可以完成stage：

// 完成回调逻辑
val completedCb = createAsyncCallback[VSLazyItem](i => {
  logger.info(s"called completedCb")
  pendingTasks -= 1
  push(out, i)
  checkIfCanComplete() // 检查是否满足stage完成条件
})

// 失败回调逻辑
val failedCb = createAsyncCallback[Throwable](err => {
  logger.error("Called failedCallback: ", err)
  pendingTasks -= 1
  failStage(err) // 失败时直接终止stage
})

// 启动异步任务时更新计数器
pendingTasks += 1
bunchOfFuturesChainedTogetherGivingUploadResult()
  .flatMap(uploadResult => {
    logger.info(s"Uploaded to ${uploadResult.location}")
    completedCb.invoke(elem) // 改用invoke
    Future.successful(())
  }).recoverWith({
    case err: Throwable =>
      logger.error(s"Could not perform upload for any of shape $shapeNameAnyOf on item ${elem.itemId}: ", err)
      failedCb.invoke(err) // 改用invoke
      Future.failed(err)
  })

// 新增辅助方法：检查是否可以完成stage
private def checkIfCanComplete(): Unit = {
  if (upstreamFinished && pendingTasks == 0) {
    logger.info(s"Processing completed")
    completeStage()
  } else if (upstreamFinished) {
    // 告诉Akka不要立刻完成stage，等待我们手动触发
    holdUpstreamFinish()
  }
}

4. 优化下游拉取逻辑 #

在onPull时检查待处理任务数量，避免上游推送过多元素导致背压问题：

setHandler(out, new AbstractOutHandler {
  override def onPull(): Unit = {
    // 只有没有正在处理的异步任务时，才拉取上游的下一个元素
    if (pendingTasks == 0) {
      pull(in)
    }
  }

  override def onDownstreamFinish(): Unit = {
    logger.info("Downstream finished")
    // 下游完成时，清空待处理任务并终止stage
    pendingTasks = 0
    completeStage()
  }
})

完整修改后的核心逻辑片段 #

override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = new GraphStageLogic(shape) {
  private implicit val logger: org.slf4j.Logger = LoggerFactory.getLogger(getClass)
  private var pendingTasks = 0
  private var upstreamFinished = false

  private def checkIfCanComplete(): Unit = {
    if (upstreamFinished && pendingTasks == 0) {
      logger.info(s"Processing completed")
      completeStage()
    } else if (upstreamFinished) {
      holdUpstreamFinish()
    }
  }

  setHandler(in, new AbstractInHandler {
    override def onPush(): Unit = {
      val completedCb = createAsyncCallback[VSLazyItem](i => {
        logger.info(s"called completedCb")
        pendingTasks -= 1
        push(out, i)
        checkIfCanComplete()
      })

      val failedCb = createAsyncCallback[Throwable](err => {
        logger.error("Called failedCallback: ", err)
        pendingTasks -= 1
        failStage(err)
      })

      val elem = grab(in)
      val shapes = shapeNameAnyOf.map(shapeName => findShape(elem, shapeName)).collect({ case Some(s) => s })

      if (shapes.nonEmpty) {
        if (shapes.length > 1) {
          logger.warn(s"Got shapes multiple shapes $shapes for item ${elem.itemId}, using the first")
        }
        pendingTasks += 1
        bunchOfFuturesChainedTogetherGivingUploadResult()
          .flatMap(uploadResult => {
            logger.info(s"Uploaded to ${uploadResult.location}")
            completedCb.invoke(elem)
            Future.successful(())
          }).recoverWith({
            case err: Throwable =>
              logger.error(s"Could not perform upload for any of shape $shapeNameAnyOf on item ${elem.itemId}: ", err)
              failedCb.invoke(err)
              Future.failed(err)
          })
      } else {
        val actualShapeNames = elem.shapes.map(_.keySet)
        logger.error(s"No shapes could be found matching $shapeNameAnyOf on the given item (got $actualShapeNames)")
        push(out, elem)
      }
    }

    override def onUpstreamFinish(): Unit = {
      logger.info(s"Upstream finished")
      upstreamFinished = true
      checkIfCanComplete()
    }
  })

  setHandler(out, new AbstractOutHandler {
    override def onPull(): Unit = {
      if (pendingTasks == 0) {
        pull(in)
      }
    }

    override def onDownstreamFinish(): Unit = {
      logger.info("Downstream finished")
      pendingTasks = 0
      completeStage()
    }
  })
}

关键修复点说明 #

• 任务计数器：准确跟踪异步任务数量，避免布尔值的状态歧义。
• 非阻塞等待：去掉Thread.sleep和循环，改用回调触发的状态检查，让stage逻辑线程可以自由处理回调事件。
• 正确的回调调用：用invoke替代invokeWithFeedback，贴合你的使用场景。
• 背压处理：在onPull里检查待处理任务，避免上游过度推送元素导致内存压力。

这样修改后，你的测试场景中两个元素的异步回调都会正确执行，stage会等待所有异步任务完成后才结束，不会再出现StreamDetachedException。

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