Flink Agent：ActionTask 与可续跑状态机 (Coroutine/Continuation)

Flink Agent：ActionTask 与可续跑状态机 (Coroutine/Continuation)

本解析聚焦于 ActionTask.java 及其子类（JavaActionTask、PythonActionTask、PythonGeneratorActionTask）。这部分代码解决的工程问题是：当一个用户编写的 Action 函数内部包含了网络阻塞调用时，如何将这个原本必须"一气呵成"的函数，硬生生"劈成两半"，让它在前半段执行完后挂起让出线程，并在网络返回后从断点处继续执行后半段？

我们将严格遵循"演进式推导"来拆解这段极其微妙的设计。

步骤 1：寻找"第一性原理"（The Naive Approach）

• 核心业务问题：用户写了一个 Action，里面调了 LLM。我们需要执行它。
• 最朴素的硬编码实现：直接调用用户的函数。

// 最朴素的伪代码
public void executeAction(Event event, RunnerContext ctx) {
    // 用户代码：
    // String prompt = buildPrompt(event);
    // String result = llm.chat(prompt); // <--- 这里会死等 5 秒！
    // ctx.sendEvent(new OutputEvent(result));
    
    action.invoke(event, ctx); // 框架直接调用
}

步骤 2：识别痛点与第一次演进

• 致命痛点 ：Flink 主线程绝对不能被 llm.chat() 这种网络请求卡死。
• 第一次抽象（强制用户拆分回调） ：
  既然不能卡死，最传统的做法是逼迫用户不要写同步代码，而是把代码拆成两截，用 Callback（回调）来写。

// 第一次演进后的伪代码：回调地狱 (Callback Hell)
public void executeAction(Event event, RunnerContext ctx) {
    // 用户必须这么写：
    // String prompt = buildPrompt(event);
    // llm.chatAsync(prompt, new Callback() {
    //     public void onSuccess(String result) {
    //         ctx.sendEvent(new OutputEvent(result));
    //     }
    // });
    
    action.invokeAsync(event, ctx); 
    // 方法立刻返回，但代码被割裂了
}

步骤 3：识别新痛点与第二次演进（Iterative Evolution）

• 新的棘手问题（心智负担与状态丢失） ：
  1. 心智负担：如果 Agent 有个循环（比如 ReAct 循环调工具），用 Callback 写循环会变成恐怖的"回调地狱"，代码根本没法维护。
  2. 状态丢失 ：
     • 疑问：为什么 Flink 保持不了 Callback 状态？
     • 解答：问题不在于 "Callback 语法本身绝对不能序列化"，而在于它没有被框架提升为一个可恢复的运行时单元。Flink 的 Checkpoint 能稳定保存的是明确挂在状态里的数据对象，而不是某个外部线程池里已经注册好的回调关系。即使这个匿名类在语法层面可序列化，框架也仍然缺少这些恢复信息：这个回调属于哪个 Action、当前执行到第几段、恢复后该由谁重新调度。机器宕机后，Flink 只能恢复状态，不能自动恢复"那次已经发出的异步请求与它的回调绑定关系"。所以这段处理逻辑依然会断掉。
• 第二次抽象（引入"协程/可续跑状态机"与 ActionTask） ：
  我们希望用户依然像写同步代码一样写 （比如用 Python 的 await，或者 Java 21 的同步写法），由框架底层 负责把函数挂起，并把"当前正在执行哪个 Action、由哪个 Event 触发、后续该继续哪一步"提升成一个可被调度的 ActionTask。
  这里要特别注意：ActionTask 保存的是可恢复的执行语义，不是完整的语言调用栈镜像。

// 第二次演进：框架底层的 ActionTask 封装（对用户完全透明）
class ActionTask {
    public ActionTaskResult invoke() {
        // 注意：这段代码是 Flink 引擎底层的代码，不是用户写的！
        // 用户依然像单机脚本一样写同步代码（如 await llm.chat()）。
        // 引擎底层的执行器（比如 Python解释器 或 Java虚拟线程执行器）会去执行用户的代码。
        ExecutionResult res = executeCoroutine(); 
        // 一旦用户的代码碰到了 await 或者发生网络 IO，底层的 executeCoroutine() 会自动被挂起并立即返回。
        
        if (res.isFinished()) {
            return new ActionTaskResult(finished=true, nextTask=null);
        } else {
            // 核心设计：用户代码被挂起了（没跑完）。
            // 框架不会直接序列化整个调用栈，而是生成一个"下一次继续执行的 Task"。
            // 这个 Task 只携带足够恢复调度语义的信息。
            ActionTask generatedTask = createNextTaskFromRuntimeState();
            return new ActionTaskResult(finished=false, nextTask=generatedTask);
        }
    }
}

这里的 createNextTaskFromRuntimeState() 不是一个真实存在的源码函数，而是上面伪代码里故意抽象出来的一步。它的作用是：把"语言运行时里的挂起现场"翻译成"Flink 调度层下次还能继续处理的任务"。

// 伪代码：把运行时挂起现场转换成下一个可调度的 Task
ActionTask createNextTaskFromRuntimeState() {
    if (language == JAVA) {
        // Java 侧并不会 new 一个完全不同的任务对象，
        // 而是复用当前 JavaActionTask，把它自己作为"下次继续调度的任务"
        saveContinuationContext(this, continuationContext);
        return this;
    }

    if (language == PYTHON) {
        // Python 侧第一次执行后，会拿到一个 awaitableRef，
        // 再包装成 PythonGeneratorActionTask，后续专门负责 send(None) 续跑
        saveAwaitableRef(awaitableRef);
        return new PythonGeneratorActionTask(action, event, key, awaitableRef);
    }

    throw new IllegalStateException("unknown runtime");
}

• How：

  1. 先判断当前是 Java 续跑还是 Python 续跑；
  2. 把语言运行时里的挂起现场单独保存起来；
  3. 生成一个"下次要继续执行的 Task"交还给算子；
  4. 算子把这个 Task 放进 actionTasksKState，等以后再调度。

• Why：

  • 这一步必须拆开做，因为 Flink 能调度的是 Task，不能直接调度语言内部的 continuation / awaitable。
  • 同时，Task 也不能把所有语言运行时对象都直接塞进状态里，否则会碰到不可序列化的问题。
  • 所以这里实际上做了两次转换：
    1. 挂起现场保存：保存 continuationContext 或 awaitableRef；
    2. 调度对象生成：返回一个 generatedTask 给 Flink。

• 一个最小例子：

  • 假设一次库存查询要 3 秒。
  • 第 1 次执行 CheckInventoryAction 时，Java 侧已经发出了 RPC，但结果还没回来。
  • 此时 createNextTaskFromRuntimeState() 做的不是"把整个 Java 栈塞进 Task"，而是：
    • 把 ContinuationContext 暂存在算子内存；
    • 返回当前这个 JavaActionTask 作为 generatedTask。
  • 这样下次调度时，Flink 知道"继续跑谁"；而 Java 执行器知道"从哪里继续跑"。

步骤 4：映射到真实源码（Mapping to Reality）

翻开源码，这种"多阶段续跑"的状态机设计被实现成了统一的运行时模型：

1. ActionTask.java#L33-L132 (基类设计)：

   • 它的注释写得很清楚："一个 Action 被拆分为多个代码块... 如果还有代码块没执行完，你可以通过 getGeneratedActionTask() 拿到下一个 Task 并继续执行。"
   • 疑问：用户需要注意这个 generatedActionTask 吗？它到底是谁在调用？
   • 解答 ：用户绝对不需要注意它！ ActionTask 及其 invoke() 方法完全是 Flink 算子内部的底层组件。
     • 谁在生成它？ 是底层的语言桥接层。比如 Python 解释器发现代码被 await 挂起时，底层的 PythonActionTask 会自动 new 一个 PythonGeneratorActionTask（这就叫 generatedActionTask）返回出去。
     • 谁在调用它？ 是外层的管家 ActionExecutionOperator。算子调用 invoke() 后，如果发现返回了 generatedActionTask，算子就知道"哦，这段代码没跑完"，算子就会把这个新生成的 Task 塞进 Flink 的状态队列 actionTasksKState 里 。等下次网络回调唤醒时，算子会从队列里把这个 generatedActionTask 拿出来，再次调用它的 invoke()，从而实现断点续跑。
   • 这种设计让 Flink 的 ActionExecutionOperator 可以拥有统一的调度模型："只要底层语言引擎返回了 generatedTask，我就存状态；等唤醒了，我再调你的 invoke"。
   • 还要注意一个关键细节：ActionTask 里真正不能序列化的 runnerContext 被明确声明成了 transient，不会随着 Task 一起进状态。ActionTask.java#L50-L55 这说明框架保存的是"任务外壳与调度语义"，不是把整个执行环境原样塞进 Checkpoint。

2. PythonActionTask.java#L46-L72 (Python 协程的实现)：

   • 当调用 Python 的 @action 函数时，如果里面有 await ctx.durable_execute_async()，Python 解释器遇到 yield 会立刻挂起并返回一个协程对象 (awaitable)。
   • PythonActionExecutor 拿到这个协程后，并不是简单存入某个 Java Map，而是通过 interpreter.set(...) 把它挂到 Python 解释器上下文里，再生成一个字符串引用 pythonAwaitableRef 返回给上层。PythonActionExecutor.java#L133-L143
   • 接着在 PythonActionTask 中，pythonAwaitableRef 被交给 PythonRunnerContextImpl，然后立即切换成一个新的 PythonGeneratorActionTask 去承接后续续跑。PythonActionTask.java#L61-L67

3. PythonGeneratorActionTask.java#L46-L62 (Python 续跑机制)：

   • Flink 会把这个 GeneratorTask 放回队列。当信箱再次调度到它时，它会调用 callPythonAwaitable，底层等于对 Python 协程执行一次 send(None)。function.py#L339-L370
   • 如果协程还没结束，call_python_awaitable 会返回 finished=false；如果结束了，则返回 finished=true 和最终值。也就是说，Java 侧并不知道 Python 栈长什么样，它只是在反复驱动一个 awaitable 直到结束。

4. ContinuationActionExecutor.java#L99-L134 (Java 21 Continuation 实现)：

   • How ：
     1. 用户代码调用 ctx.durableExecuteAsync(callable)；
     2. JavaRunnerContextImpl 将 callable 包装为 Supplier，交给 ContinuationActionExecutor.executeAsync(...)；JavaRunnerContextImpl.java#L62-L105
     3. 执行器把 supplier.get() 提交到后台线程池拿到 future；
     4. 当前线程进入
        while (!future.isDone()) { Continuation.yield(SCOPE); }，在未完成时主动让出；ContinuationActionExecutor.java#L127-L148
     5. 当后台线程把结果或异常写回 ContinuationContext 后，下一次调度再继续执行该 continuation。
   • Why ：yield 不是阻塞，而是把当前 Java 执行流"挂起"。这让算子线程可以继续处理同 Task 中其他 Key 和 Barrier。
   • 恢复与状态 ：算子并不把 Java 调用栈塞进 Flink 状态；而是把"这次还要继续执行的 Task"放进 actionTasksKState，并在内存映射里维护 generatedTask -> ContinuationContext 的对应关系。ActionExecutionOperator.java#L565-L570
   • 具象例子 （外部调用耗时 3 秒）：
     • t=0：提交 callable 到线程池；当前线程在 yield() 处挂起；JavaActionTask 返回 finished=false 和 generatedTask=this；
     • t=1：算子继续干别的事；
     • t=3：后台线程把结果写回 ContinuationContext；该 JavaActionTask 再次被调度；continuation 从 yield() 后继续执行，拿到结果并落到用户后续代码。

步骤 5：批判性总结（Critical Trade-offs）

• 优势 ：
  • 开发体验稳定 ：开发者可以继续用同步的心智（Python await 或 Java 直接调）写复杂的循环逻辑，而不需要自己维护 Callback 链。
  • 契合 Flink 调度模型 ：通过把"当前该继续执行哪一个动作实例"抽象成 ActionTask，框架可以把 Task 本身放进 ListState，并由算子统一调度。ActionExecutionOperator.java#L580-L610
• 代价与局限（函数调用栈序列化的难点） ：
  • 这套机制并不完美。最大的痛点在于Java 和 Python 无法原生序列化函数调用栈。
  • 为什么不能序列化调用栈？
    1. 内存指针与本地资源绑定：在 Java/Python 中，一个运行中的函数栈包含局部变量、寄存器状态，甚至可能间接依赖底层 C/C++ 资源（如 Socket 文件描述符）。这些东西和当前进程、当前机器强绑定，不能被当作普通业务对象那样稳定搬运到另一台机器恢复。
    2. JVM/Python 解释器的黑盒 ：Java 的 Continuation（虚拟线程底层）和 Python 的 coroutine 对象，在语言规范层面上都没有实现 Serializable 接口。虚拟机厂商不允许你把一个活着的执行上下文序列化到硬盘上。
  • 导致的脆弱性（语言运行时上下文的丢失） ：
    进入 Checkpoint 的并不是完整调用栈，而是：
    • ActionTask 这类可序列化的任务外壳；
    • ActionState 里的 CallResult 等细粒度 durable 记录；
    • 以及部分可恢复的业务状态。
      反过来，下面这些东西是算子进程内的瞬时对象：
    • Python 的 awaitableRef 映射；
    • Java 的 ContinuationContext；
    • runnerContext 本身。
      因此，机器宕机后，即使 ActionTask 被恢复出来，语言运行时内部那个"活着的挂起对象"仍然可能已经丢失。
  • 框架的补救措施（必要的恢复退化） ：
    • Python 侧最直接：如果恢复后发现 pythonAwaitableRef 丢了，PythonGeneratorActionTask 会退化成重新从头执行 PythonActionTask。PythonGeneratorActionTask.java#L39-L57
    • Java 侧更温和：恢复后如果 ContinuationContext 不在内存里，会新建一个上下文；已经完成过的 durable 调用则依靠 CallResult 跳过重放。ActionExecutionOperator.java#L903-L918 、 CallResult.java#L25-L34
    • 所以更准确的说法是：框架能保证"调度语义"和"已完成 durable 调用结果"可恢复，但不能保证语言层挂起对象本身总能原样续上。
• 更优解的探讨 ：
  • 这种"手动把控制流切片并存入状态"的做法，是当前主流语言接入外部非确定性长调用时比较现实的方案。如果未来语言运行时能原生提供稳定、可序列化的 Continuation，那么 ActionTask 这一层有可能继续变薄；但在当前 Java / Python 条件下，这层抽象仍然是必要的。