[runtime] Fix lifecycle and metric bookkeeping gaps in parallel action execution

runtime/src/main/java/org/apache/flink/agents/runtime/lifecycle/ComponentExecutionListener.java

@@ -34,7 +34,10 @@
  * <p>Invariants a listener may rely on, and must not break:
  *
  * <ul>
- *   <li>The callback runs on the mailbox thread, so a listener needs no synchronization of its own.
+ *   <li>The callback runs on the thread that drives the execution being reported: the mailbox
+ *       thread on the default engine, but possibly an action worker thread when parallel action
+ *       execution is enabled. A listener must not rely on thread affinity and must synchronize its
+ *       own state when more than one thread can reach it.
  *   <li>An exception thrown by a listener is logged and swallowed, so reporting never fails the
  *       reporting component and never starves the remaining listeners.
  *   <li>The event carries the lifecycle status only; the reporting component is identified by the

runtime/src/main/java/org/apache/flink/agents/runtime/lifecycle/TaskLifecycleListener.java

@@ -23,7 +23,10 @@
 /**
  * Observes the per-record and per-action lifecycle of {@code ActionExecutionOperator}.
  *
- * <p>All callbacks run on the mailbox thread.
+ * <p>Callbacks run on the thread that drives the execution being reported, which is not fixed: the
+ * mailbox thread on the default engine, and either the mailbox thread or an action worker thread
+ * when parallel action execution is enabled. A listener must not rely on thread affinity and must
+ * synchronize its own state if more than one thread can reach it.
  *
  * <p>Event pairing semantics:
  *

runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java

@@ -1584,6 +1584,10 @@ private ActionTask pollNextActionTask() throws Exception {
     }
 
     private void tryResumeProcessActionTasks() throws Exception {
+        // Keys with at least one lifecycle-visible restored task, meaning a task other than the
+        // synthetic noop input action. The parallel restore path emits the matching record start
+        // for them, because completion emits one record finished per key.
+        Set<Object> keysWithVisibleTask = new LinkedHashSet<>();
         stateManager.forEachActionTaskKey(
                 getKeyedStateBackend(),
                 (key, state) -> {
@@ -1601,10 +1605,21 @@ private void tryResumeProcessActionTasks() throws Exception {
                     // its root caller still awaits it (the caller resumes only after the child
                     // quiesces), so the root is always retained here to drive the replay.
                     state.update(rootTasks);
+                    boolean lifecycleVisible = false;
                     for (ActionTask actionTask : rootTasks) {
+                        // The synthetic noop input action stays out of action metrics, as it does
+                        // on the dequeue path: counting it here would leave the pending-task gauge
+                        // nonzero once the restored task completes.
+                        if (isInternalNoopInputAction(actionTask.action)) {
+                            continue;
+                        }
                         builtInMetrics.restoreActionTask(
                                 actionTask.getTraceContext(),
                                 actionTask.hasExecutionStartedEventEmitted());
+                        lifecycleVisible = true;
+                    }
+                    if (lifecycleVisible) {
+                        keysWithVisibleTask.add(key);
                     }
                 });
 
@@ -1636,6 +1651,13 @@ private void tryResumeProcessActionTasks() throws Exception {
                     for (int i = 0; i < queued; i++) {
                         executionCoordinator.addTask(key);
                     }
+                    // Mirror the fresh-record path, which emits one record start per key: the
+                    // worker-completion path emits one record finished per key, so the resumed
+                    // round needs the matching start. A key whose restored tasks are all synthetic
+                    // noop stays silent, because its completion is silent too.
+                    if (keysWithVisibleTask.contains(key)) {
+                        notifyRecordStart(key);
+                    }
                 } else {
                     // The normal engine has no workers; the kick mail drives
                     // processActionTaskForKey.
@@ -1723,6 +1745,7 @@ public void execute() {
                 setCurrentKey(key);
                 ActionTask task = pollNextActionTask();
                 checkState(task != null, "Action task queue was empty.");
+                this.actionTask = task;
 
                 contextManager.createAndSetRunnerContext(
                         task,
@@ -1774,7 +1797,6 @@ public void execute() {
                 context.setActionTask(task);
                 context.setPriority(recordIndex, taskIndex);
 
-                this.actionTask = task;
                 this.sequenceNumber = seq;
                 this.replayCompletedAction = replay;
 
@@ -1848,6 +1870,14 @@ public boolean isDone() {
         @Override
         public void commit() throws Exception {
             if (failure != null) {
+                // This engine propagates an action failure through commit() rather than through
+                // processActionTaskForKey(), so the input run is closed here: the failure counter,
+                // the active-run gauge and the terminal latency samples all depend on it.
+                // actionTask is bound right after the poll, so it is non-null on a preparation
+                // failure too; guard anyway and always rethrow the original failure.
+                if (actionTask != null) {
+                    builtInMetrics.markInputRunFailed(actionTask.getTraceContext().getInputRunId());
+                }
                 throw new ActionTaskExecutionException("Failed to execute action task", failure);
             }
             checkState(result != null, "Action task completion result must not be null.");

runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorParallelMetricsTest.java

@@ -18,7 +18,9 @@
 package org.apache.flink.agents.runtime.operator;
 
 import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.api.EventType;
 import org.apache.flink.agents.api.InputEvent;
+import org.apache.flink.agents.api.agents.Agent;
 import org.apache.flink.agents.api.agents.AgentExecutionOptions;
 import org.apache.flink.agents.api.context.DurableCallable;
 import org.apache.flink.agents.api.context.RunnerContext;
@@ -27,6 +29,7 @@
 import org.apache.flink.agents.plan.JavaFunction;
 import org.apache.flink.agents.plan.actions.Action;
 import org.apache.flink.agents.runtime.async.ContinuationActionExecutor;
+import org.apache.flink.agents.runtime.lifecycle.TaskLifecycleListener;
 import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
 import org.apache.flink.api.java.functions.KeySelector;
@@ -44,6 +47,7 @@
 import java.util.concurrent.TimeUnit;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.assertj.core.api.Assumptions.assumeThat;
 
 /** Component-level tests for metrics bookkeeping on the JDK<21 parallel execution path. */
@@ -232,4 +236,78 @@ private static AgentPlan getAgentPlan() throws Exception {
             return new AgentPlan(actions, new LinkedHashMap<>(), parallelConfig());
         }
     }
+
+    @Test
+    @Timeout(20)
+    void failedParallelActionMarksInputRunFailed() throws Exception {
+        assumeThat(ContinuationActionExecutor.isContinuationSupported()).isFalse();
+
+        try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+                newHarness(new AgentPlan(new FailingAgent(), parallelConfig()))) {
+            testHarness.open();
+            ActionExecutionOperator<Long, Object> operator = operator(testHarness);
+
+            assertThatThrownBy(
+                            () -> {
+                                testHarness.processElement(new StreamRecord<>(0L));
+                                operator.waitInFlightEventsFinished();
+                            })
+                    .hasStackTraceContaining("parallel action boom");
+
+            assertThat(metricGroup(operator).getCounter("numOfInputRunsFailed").getCount())
+                    .as("a failure propagated through commit() must close the input run")
+                    .isEqualTo(1L);
+        }
+    }
+
+    /** Agent whose only action always fails, exercising the parallel failure path. */
+    public static class FailingAgent extends Agent {
+
+        @org.apache.flink.agents.api.annotation.Action(EventType.InputEvent)
+        public static void handleInput(Event event, RunnerContext context) {
+            throw new IllegalStateException("parallel action boom");
+        }
+    }
+
+    @Test
+    @Timeout(20)
+    void preparationFailureKeepsCauseAndMarksInputRunFailed() throws Exception {
+        assumeThat(ContinuationActionExecutor.isContinuationSupported()).isFalse();
+
+        try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+                newHarness(new AgentPlan(new PlainAgent(), parallelConfig()))) {
+            testHarness.open();
+            ActionExecutionOperator<Long, Object> operator = operator(testHarness);
+
+            // A preparation-time listener failure must surface as the original exception,
+            // not as a NullPointerException from commit() dereferencing a null actionTask,
+            // and must still close the input run so the failure is counted.
+            operator.addTaskLifecycleListener(
+                    new TaskLifecycleListener() {
+                        @Override
+                        public void onActionPrepared(ActionTask task) {
+                            throw new IllegalStateException("prepare boom");
+                        }
+                    });
+
+            assertThatThrownBy(
+                            () -> {
+                                testHarness.processElement(new StreamRecord<>(0L));
+                                operator.waitInFlightEventsFinished();
+                            })
+                    .hasStackTraceContaining("prepare boom");
+
+            assertThat(metricGroup(operator).getCounter("numOfInputRunsFailed").getCount())
+                    .as(
+                            "a preparation failure propagated through commit() must close the input run")
+                    .isEqualTo(1L);
+        }
+    }
+
+    /** Agent whose action does nothing, so the only failure comes from the listener. */
+    public static class PlainAgent extends Agent {
+
+        @org.apache.flink.agents.api.annotation.Action(EventType.InputEvent)
+        public static void handleInput(Event event, RunnerContext context) {}
+    }
 }