[runtime][python] Release bridge-owned Pemja objects during cleanup

[joeyutong]

Linked issue: #942

Purpose of change

The embedded Python bridge owns several Pemja PyObject handles: the async thread pool and runner context in PythonActionExecutor, the Python resource context in PythonResourceAdapterImpl, and the Python Mem0 object in Mem0LongTermMemory. Each handle keeps a native Python reference, while Python objects can retain the surrounding Java task graph through Pemja proxies and JNI global references.

The existing close paths either performed only Python-level cleanup or did not close the PyObject at all. Repeated task failovers could therefore retain handles and old task graphs after the corresponding attempt had closed.

This change:

• runs Python-level cleanup before releasing each owned native handle;
• clears handle fields before releasing them, making repeated close calls safe against Pemja's unguarded native decRef;
• closes Mem0, the action executor, the Python resource adapter, the interpreter, and the environment in reverse creation order;
• continues closing sibling and upper-level resources after a cleanup failure while preserving the first exception and suppressing later ones; and
• clears the Python runner context's long-term-memory reference before cleanup, breaking the FlinkRunnerContext -> Mem0LongTermMemory -> FlinkRunnerContext cycle so native-handle release can collapse the full Java -> Python -> Java retention chain.

The PythonBridgeManager fallback also covers subtasks that initialized Mem0 but never created a Java RunnerContext, so cleanup does not depend on whether that subtask processed a Java action.

Tests

• Java lifecycle and failure-path tests: 55 passed across PythonActionExecutorTest, PythonResourceAdapterImplTest, Mem0LongTermMemoryTest, PythonBridgeManagerTest, and ActionExecutionOperatorTest.
• Python runner-context cleanup tests: 10 passed, including idempotent close and LTM cleanup failure paths.
• Ruff and Maven Spotless checks passed.
• Local 20-restart failover A/B for the original runner-context leak reduced post-Full-GC heap growth from 29.2 MiB to 1.7 MiB and removed the old runner-context/task graphs after termination.
• Focused HPROF A/B for the two sibling handles found 84 retained target PyObject handles after termination in each baseline; the native-close variants retained 0.

API

No user-facing API changes.

Documentation

• doc-needed
• doc-not-needed
• doc-included

Was this patch authored or co-authored using generative AI tooling?

• Yes
• No

Generated-by: Codex 0.147.0-alpha.6.5 (GPT-5)

[weiqingy]

Thanks for chasing this one down. I reverted close() locally and both new tests fail against the old body, so the regression coverage is real. A few questions inline.

[weiqingy]

nit: The copy-then-null reads as defensive style, but it looks load-bearing. PyObject.close() in pemja 0.5.7 is an unguarded decRef(tState, pyobject) with no null check and no double-close flag, so clearing the fields first is the only thing stopping a repeated close() from decrementing a second time on an already-released handle. Someone later tidying this into closePythonObject(CLOSE_ASYNC_THREAD_POOL, pythonAsyncThreadPool) would drop that quietly, and the only assertion that would notice is the three-line tail of releasesBothPythonObjectsWhenLogicalCleanupFails.

Would a short comment here save the next reader that trip? There is precedent right next door at ActionExecutionOperator.java:471 (// Must close before pythonInterpreter since cached resources may hold Python references.).

Something like this, if it helps:

// Clear the fields before releasing: PyObject.close() is an unguarded native decRef,
// so a repeated close() must not reach the same handle twice.

[joeyutong]

Good point. I added a short comment explaining why the fields are cleared before releasing the Pemja handles, so a repeated close() cannot reach the same handle twice.

[weiqingy]

Combining both failures and rethrowing is the right call. What I keep looking at is what happens to this exception one frame up:

// PythonBridgeManager.close(), lines 292-302
if (pythonActionExecutor != null) { pythonActionExecutor.close(); }
if (pythonInterpreter != null) { pythonInterpreter.close(); }
if (pythonEnvironmentManager != null) { pythonEnvironmentManager.close(); }

That is a plain sequence, so on exactly the failure path this PR is built for, pythonInterpreter.close() never runs, and that is the release that tears down the interpreter owning every handle still outstanding. ActionExecutionOperator.close() (lines 469-489) has the same shape across its five closes.

To be clear, this is pre-existing. The old close() threw on a failed interpreter.invoke too, so nothing has regressed here. But given the stated goal is "releases both handles even if one cleanup operation fails", how do you see that goal holding one frame up? Carrying the same firstOrSuppressed pattern into PythonBridgeManager.close() would make the guarantee end-to-end, though I may be missing a reason the interpreter is fine to leak on that path.

[joeyutong]

Great catch. The guarantee did not hold one frame up. I extended best-effort cleanup through both PythonBridgeManager and ActionExecutionOperator, preserving close order and suppressing later failures, with tests at both layers.

[weiqingy]

Both frames check out. This is about a third instance rather than those two.

ActionTaskContextManager.close() (runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java:319-330) still has the round-1 shape:

if (runnerContext != null) {
    try { runnerContext.close(); } finally { runnerContext = null; }
}
if (continuationActionExecutor != null) {
    continuationActionExecutor.close();
}

It is argument 2 of the operator's own new closeAll(...) list, so it sits inside the chain you just hardened. The throw path is this PR's own failure case: RunnerContextImpl.close() (context/RunnerContextImpl.java:340-344) calls this.ltm.close(), which reaches Mem0LongTermMemory.close() and a Pemja invoke.

The skipped close is not a no-op everywhere. The JDK 21 variant (runtime/src/main/java21/.../async/ContinuationActionExecutor.java:151-153) is asyncExecutor.shutdownNow(), while the JDK 11 one (:59) is empty. So on JDK 21 a failing LTM close leaks those threads.

Narrow exposure: JDK 21 only, and only when a close is already failing. Worth giving this frame the same treatment while you are in here?

One wrinkle if you do. ContinuationActionExecutor.close() does not declare throws, so it would need implements AutoCloseable or a lambda wrapper to go into the varargs list.

I checked the rest of the chain for the same shape and it is clean, so this looks like the last one rather than the first of many.

[weiqingy]

This is the shape the whole fix turns on: logical cleanup inside, native release on the way out. Two other Pemja handles in the same lifecycle still have the pre-PR shape.

Mem0LongTermMemory.close() (Mem0LongTermMemory.java:137-140) calls adapter.callMethod(pyMem0, "close", Map.of()) and never pyMem0.close(), which is the old PythonActionExecutor.close() exactly. It is live rather than dead code: RunnerContextImpl.java:339-345 calls ltm.close(), and the handle comes from PythonBridgeManager.java:237.

PythonResourceAdapterImpl.pythonResourceContext (PythonResourceAdapterImpl.java:86,99) is built as interpreter.invoke(GET_RESOURCE_CONTEXT, this), so it is a Java object handed into Python, the same JNI-global-ref pattern #942 describes. That class has no close() at all, and PythonBridgeManager.close() never touches the adapter.

I read the PR and #942 as deliberately scoped to the action executor, so I am not suggesting you widen this one. Is a follow-up issue the plan for the sibling handles, or is there something that already releases those two that I have missed?

[joeyutong]

Good catch. Both handles are live, and the A/B heap dumps confirmed the same retention shape, so I widened this PR instead of opening a follow-up:

• PythonResourceAdapterImpl now owns and closes pythonResourceContext; nulling the field first makes repeated close safe.
• Mem0LongTermMemory now runs Python-level cleanup and always closes pyMem0, including when logical cleanup fails; it is also idempotent.
• PythonBridgeManager now closes Mem0 and the resource adapter before the interpreter/environment, while continuing all closes after failures. This also covers Mem0-initialized subtasks that never created a Java RunnerContext.
• FlinkRunnerContext.close() clears __ltm before cleanup and still closes the resource cache if LTM cleanup fails. This breaks the Python context/Mem0 cycle so releasing the Pemja handles can collapse the full Java -> Python -> Java retention chain.
• Added lifecycle, failure-path, and repeated-close tests. In focused 20-restart HPROF A/B runs, each sibling baseline retained 84 target PyObject handles after termination; the native-close variants retained 0.

[weiqingy]

Thanks for addressing the comments. Both frames you named are fixed, and the new failure-path tests pin invariants that were only implicit before.

One instance of the aggregate-close question is still open, on the existing thread, plus two smaller things inline.

[weiqingy]

This exercises six closeables and verifies four. eventRouter and super::close are never injected or asserted.

I tried deleting each from the closeAll(...) list, and all 407 runtime-module tests stay green either way. Dropping super::close would silently stop disposing Flink operator state, which is the costlier of the two.

eventRouter looks cheap to cover: it is private final transient, and the setPrivateField helper this test already uses works on non-static final instance fields, so a throwing mock could join the InOrder. Any appetite for that?

super.close() I don't have a cheap way to assert, so I'd call that half a known gap rather than something to chase here.

[weiqingy]

nit: org.apache.flink.util.IOUtils.closeAll(AutoCloseable...) throws Exception already exists with the same signature, and ships in both flink-core 1.20.0 and 2.3.0, so this isn't a cross-version gap. It has the semantics this helper relies on: null-guards, skips null elements, aggregates through ExceptionUtils.firstOrSuppressed, rethrows. The repo already uses it at runtime/src/main/java/.../feedback/Checkpoints.java:22,63.

Moving the helper out of OperatorUtils makes sense, since that class does have a per-version variant under dist/flink-1.20/. Is there a reason to prefer a local helper over the Flink one?

If it did become IOUtils.closeAll, one knock-on: CloseableUtils is package-private in operator, so PythonActionExecutor.close() (PythonActionExecutor.java:200-214) still hand-rolls the same firstOrSuppressed accumulation and can't reach it. The public utility would let that collapse too.

[joeyutong]

You are right. I previously worried that the available utilities did not quite match the required aggregation semantics, and I was probably over-cautious when extracting the local helper. IOUtils.closeAll has exactly the behavior needed here. I replaced both local-helper call sites and the hand-rolled PythonActionExecutor path with it, removed CloseableUtils, and also used it for the new ActionTaskContextManager aggregation in 56795e1.

[weiqingy]

Thanks for addressing the comments. Nothing further from me.

[wenjin272]

If both ltm.close() and resource_cache.close() fail, Python's finally semantics make the resource-cache exception primary and leave the earlier LTM failure only in __context__. This contradicts the PR's stated behavior of preserving the first exception and suppressing later ones, and it differs from the Java IOUtils.closeAll paths. Could we explicitly aggregate these failures and add a test where both closes throw?

[joeyutong]

Good catch. I replaced the try/finally propagation with explicit aggregation: the long-term-memory failure remains the primary exception, while a later resource-cache failure is retained in its context. The new double-failure test asserts both identities and also verifies that repeated close() remains a no-op. Fixed in fc4c0fb.

[emecii]

Following up on @weiqingy's note in #987 that these two PRs disagree on mechanism rather than merely touching the same seven files. I'm the author of #987, so flagging my own interest up front.

I checked the Error case empirically against flink-core-2.3.0 rather than reading the source, and the concern holds. Three closeables that record whether close() ran, the first one throwing:

A (first throws OutOfMemoryError):   thrown=java.lang.OutOfMemoryError: boom
A closed flags -> a1=true a2=false a3=false
B (first throws IllegalStateException): thrown=java.lang.IllegalStateException: boom
B closed flags -> b1=true b2=true b3=true

The varargs overload delegates to closeAll(Iterable), which delegates to closeAll(Iterable, Class<T>) with suppressedException = Exception.class; that method rethrows anything not assignable to it before closing the rest. So closeAll continues past an Exception but stops dead on any non-Exception Throwable.

At this call site that means an Error out of longTermMemory.close() or pythonActionExecutor.close() leaves pythonInterpreter and pythonEnvironmentManager unclosed — the native Python state this PR exists to release, retained for the lifetime of the TaskManager JVM. The same shape appears at the other two sites:

• ActionExecutionOperator.close(): an Error from resourceCache skips contextManager, pythonBridge, eventRouter, durableExecManager, and the trailing super::close, so stateHandler.dispose() is skipped too.
• ActionTaskContextManager.close(): an Error from the runner context strands the continuation executor's thread pool.

Worth noting that ResourceCache.close() (ResourceCache.java:148 and :160) currently catches only Exception, so an Error out of a cached Resource.close() propagates unchanged and is a concrete way to reach the first of those.

Would you consider replacing closeAll with a catch (Throwable) ladder that aggregates via ExceptionUtils.firstOrSuppressed and rethrows via ExceptionUtils.rethrowException? It keeps the same first-failure-wins-with-later-ones-suppressed semantics, and rethrowException passes both Error and Exception through unwrapped, so callers still see the original type and instance. That is the shape KafkaActionStateStore.close() (#948) and ResourceCache.close() already use in this module.

On sequencing, since we overlap on seven files: #987 is the smaller change and only rewrites the three close() methods. If it lands first, this PR's rebase becomes additive — longTermMemory and pythonResourceAdapter slot into ladders that already exist — instead of a mechanism swap in one direction or the other. I'm equally happy to go the other way and rework #987 as a follow-up on top of this if you'd prefer not to reshuffle. Mainly I'd like the two not to land opposite decisions on the Error case. What works best for you?

[joeyutong]

Thanks for the detailed analysis. I agree that IOUtils.closeAll does not provide the required behavior for non-Exception Throwables. Since #987 is the smaller and more general fix, I prefer to let it land first and then rebase #944 on top of it, adding longTermMemory and pythonResourceAdapter to the resulting cleanup ladders. I also noticed that #944 introduces another IOUtils.closeAll inside PythonActionExecutor.close(); I will address that Error path as part of the rebase.