This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new 7b399377 [hotfix][runtime][python] Close every component when an
earlier close fails (#987)
7b399377 is described below
commit 7b3993774802e96e3b26c56e0ed2cbd3ff9ef306
Author: Yifan Chen <[email protected]>
AuthorDate: Thu Aug 27 20:08:32 2026 -0700
[hotfix][runtime][python] Close every component when an earlier close fails
(#987)
Generated-by: Claude Code 2.1.227 (Claude Opus 5)
Co-Authored-By: Claude Opus 5 <[email protected]>
---
python/flink_agents/runtime/resource_cache.py | 58 ++++++-
.../runtime/tests/test_resource_cache_close.py | 130 ++++++++++++++
.../apache/flink/agents/runtime/ResourceCache.java | 31 ++--
.../runtime/operator/ActionExecutionOperator.java | 44 +++--
.../runtime/operator/ActionTaskContextManager.java | 21 ++-
.../runtime/operator/PythonBridgeManager.java | 30 +++-
.../runtime/python/utils/PythonActionExecutor.java | 37 ++--
.../flink/agents/runtime/skill/SkillManager.java | 22 +--
.../flink/agents/runtime/ResourceCacheTest.java | 193 +++++++++++++++++++++
.../operator/ActionExecutionOperatorTest.java | 175 +++++++++++++++++++
.../operator/ActionTaskContextManagerTest.java | 61 +++++++
.../runtime/operator/PythonBridgeManagerTest.java | 109 ++++++++++++
.../python/utils/PythonActionExecutorTest.java | 85 +++++++++
.../agents/runtime/skill/SkillManagerTest.java | 45 ++++-
14 files changed, 979 insertions(+), 62 deletions(-)
diff --git a/python/flink_agents/runtime/resource_cache.py
b/python/flink_agents/runtime/resource_cache.py
index 3f98a8b7..9c96636b 100644
--- a/python/flink_agents/runtime/resource_cache.py
+++ b/python/flink_agents/runtime/resource_cache.py
@@ -15,6 +15,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#################################################################################
+import logging
+from collections.abc import Callable
from typing import Any, Dict
from flink_agents.api.resource import Resource, ResourceType
@@ -24,6 +26,38 @@ from flink_agents.plan.resource_provider import
JavaResourceProvider, ResourcePr
from flink_agents.plan.tools.function_tool import FunctionTool
from flink_agents.runtime.resource_context import ResourceContextImpl
+_LOG = logging.getLogger(__name__)
+
+
+def _failure_of(close: Callable[[], None]) -> Exception | None:
+ """Run ``close``, returning any failure instead of raising it.
+
+ Keeps the caller's cleanup loop free of a ``try`` block so one bad
component
+ cannot end the iteration.
+ """
+ try:
+ close()
+ except Exception as e:
+ return e
+ return None
+
+
+def _first_or_logged(
+ failure: Exception | None, previous: Exception | None, what: str
+) -> Exception | None:
+ """Keep the first failure and log any later one.
+
+ The Python analogue of Flink's ``ExceptionUtils.firstOrSuppressed``. Later
+ failures are logged rather than attached, because ``ExceptionGroup``
requires
+ 3.11 and this package supports 3.10.
+ """
+ if failure is None:
+ return previous
+ if previous is None:
+ return failure
+ _LOG.warning("Suppressed failure closing %s.", what, exc_info=failure)
+ return previous
+
class ResourceCache:
"""Lazily resolves and caches Resource instances from ResourceProviders.
@@ -100,9 +134,29 @@ class ResourceCache:
cached ``SkillManager`` (releasing materialized skill temp dirs). This
is what releases skill resources on operator close, including Flink
failover when the JVM stays up.
+
+ Every resource is closed even when an earlier one fails, so a single
bad
+ resource cannot strand the ones behind it, the cache clear, or the
+ resource context. The first failure is re-raised; later ones are
logged.
+
+ This mirrors the Java ``ResourceCache.close()`` contract. Two
differences
+ are deliberate. Java catches ``Throwable`` to keep a non-``Exception``
+ failure from stopping the loop; the Python equivalent is ``Exception``,
+ since it already covers what Java calls ``Error`` (``MemoryError`` and
+ friends) while ``BaseException`` would also swallow
``KeyboardInterrupt``
+ and ``SystemExit``, which cleanup must not do. And Java attaches later
+ failures with ``addSuppressed``; ``ExceptionGroup`` needs 3.11 and this
+ package supports 3.10, so later failures are logged instead of
attached.
"""
+ first_failure: Exception | None = None
for typed in self._cache.values():
for resource in typed.values():
- resource.close()
+ first_failure = _first_or_logged(
+ _failure_of(resource.close), first_failure, "resource"
+ )
self._cache.clear()
- self._resource_context.close()
+ first_failure = _first_or_logged(
+ _failure_of(self._resource_context.close), first_failure,
"resource context"
+ )
+ if first_failure is not None:
+ raise first_failure
diff --git a/python/flink_agents/runtime/tests/test_resource_cache_close.py
b/python/flink_agents/runtime/tests/test_resource_cache_close.py
new file mode 100644
index 00000000..61332e2c
--- /dev/null
+++ b/python/flink_agents/runtime/tests/test_resource_cache_close.py
@@ -0,0 +1,130 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#################################################################################
+"""Close-path contract tests for the Python ResourceCache.
+
+These mirror the Java ``ResourceCacheTest`` close tests: a failing resource
must
+not strand the resources behind it, the cache clear, or the resource context.
+"""
+
+import pytest
+
+from flink_agents.api.resource import Resource, ResourceType
+from flink_agents.runtime.resource_cache import ResourceCache
+
+
+class RecordingResource(Resource):
+ """Records whether ``close()`` ran, and optionally fails it."""
+
+ closed: bool = False
+ failure: Exception | None = None
+
+ @classmethod
+ def resource_type(cls) -> ResourceType:
+ return ResourceType.TOOL
+
+ def close(self) -> None:
+ """Record the call, then fail if this resource was configured to."""
+ self.closed = True
+ if self.failure is not None:
+ raise self.failure
+
+
+class RecordingContext:
+ """Stands in for the resource context so its close is observable."""
+
+ def __init__(self) -> None:
+ self.closed = False
+
+ def close(self) -> None:
+ """Record that the cache reached the resource context."""
+ self.closed = True
+
+
+def _cache_with(*resources: RecordingResource) -> tuple[ResourceCache,
RecordingContext]:
+ cache = ResourceCache({})
+ context = RecordingContext()
+ cache._resource_context = context
+ for i, resource in enumerate(resources):
+ cache._cache.setdefault(ResourceType.TOOL, {})[f"r{i}"] = resource
+ return cache, context
+
+
+def test_close_closes_every_resource_when_an_earlier_one_fails() -> None:
+ """A failing resource must not strand the ones behind it.
+
+ The cache clear and the resource context close are on the same straight
+ line behind the failure, so both are asserted rather than assumed.
+ """
+ failure = RuntimeError("resource close failed")
+ failing = RecordingResource(failure=failure)
+ surviving = RecordingResource()
+ cache, context = _cache_with(failing, surviving)
+
+ with pytest.raises(RuntimeError) as excinfo:
+ cache.close()
+
+ # The failure reaches the caller unchanged in identity, not wrapped.
+ assert excinfo.value is failure
+ assert failing.closed
+ assert surviving.closed
+ assert cache._cache == {}
+ assert context.closed
+
+
+def test_close_reports_first_failure_when_several_fail() -> None:
+ """The first failure is the one raised; later ones do not replace it."""
+ first = RuntimeError("first")
+ second = RuntimeError("second")
+ cache, context = _cache_with(
+ RecordingResource(failure=first), RecordingResource(failure=second)
+ )
+
+ with pytest.raises(RuntimeError) as excinfo:
+ cache.close()
+
+ assert excinfo.value is first
+ assert context.closed
+
+
+def test_close_reports_resource_context_failure() -> None:
+ """A resource context failure surfaces when no resource failed before
it."""
+ cache, context = _cache_with(RecordingResource())
+ failure = RuntimeError("resource context close failed")
+
+ def failing_close() -> None:
+ context.closed = True
+ raise failure
+
+ context.close = failing_close # type: ignore[method-assign]
+
+ with pytest.raises(RuntimeError) as excinfo:
+ cache.close()
+
+ assert excinfo.value is failure
+
+
+def test_close_returns_normally_when_nothing_fails() -> None:
+ """The healthy path stays a plain no-raise close."""
+ surviving = RecordingResource()
+ cache, context = _cache_with(surviving)
+
+ cache.close()
+
+ assert surviving.closed
+ assert cache._cache == {}
+ assert context.closed
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/ResourceCache.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/ResourceCache.java
index e3584ac2..cd84da8c 100644
--- a/runtime/src/main/java/org/apache/flink/agents/runtime/ResourceCache.java
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/ResourceCache.java
@@ -25,6 +25,7 @@ import
org.apache.flink.agents.plan.resourceprovider.PythonResourceProvider;
import org.apache.flink.agents.plan.resourceprovider.ResourceProvider;
import org.apache.flink.agents.plan.tools.FunctionTool;
import org.apache.flink.agents.runtime.resource.ResourceContextImpl;
+import org.apache.flink.util.ExceptionUtils;
import java.util.HashMap;
import java.util.Map;
@@ -158,32 +159,32 @@ public class ResourceCache implements AutoCloseable {
@Override
public void close() throws Exception {
- Exception firstException = null;
+ // Close every cached resource, then the resource context, even when
an earlier close
+ // fails. The first failure is rethrown with the later ones suppressed.
+ //
+ // The ladders catch Throwable, not Exception:
ActionExecutionOperator.close() closes this
+ // cache before the Python interpreter because cached resources may
hold Python references,
+ // so a non-Exception Throwable escaping here would leave the
remaining resources open
+ // while the interpreter behind them is torn down anyway.
ExceptionUtils.rethrowException
+ // passes Error and Exception through unchanged, so the caller still
sees the original.
+ Throwable firstFailure = null;
for (Map<String, Resource> resources : cache.values()) {
for (Resource resource : resources.values()) {
try {
resource.close();
- } catch (Exception e) {
- if (firstException == null) {
- firstException = e;
- } else {
- firstException.addSuppressed(e);
- }
+ } catch (Throwable t) {
+ firstFailure = ExceptionUtils.firstOrSuppressed(t,
firstFailure);
}
}
}
cache.clear();
try {
resourceContext.close();
- } catch (Exception e) {
- if (firstException == null) {
- firstException = e;
- } else {
- firstException.addSuppressed(e);
- }
+ } catch (Throwable t) {
+ firstFailure = ExceptionUtils.firstOrSuppressed(t, firstFailure);
}
- if (firstException != null) {
- throw firstException;
+ if (firstFailure != null) {
+ ExceptionUtils.rethrowException(firstFailure);
}
}
}
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
index 18931ee9..f82d8047 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
@@ -591,24 +591,38 @@ public class ActionExecutionOperator<IN, OUT> extends
AbstractStreamOperator<OUT
@Override
public void close() throws Exception {
- // Must close before pythonInterpreter since cached resources may hold
Python references.
- if (resourceCache != null) {
- resourceCache.close();
- }
- if (contextManager != null) {
- contextManager.close();
- }
- if (pythonBridge != null) {
- pythonBridge.close();
- }
- if (eventLogWriter != null) {
- eventLogWriter.close();
+ // Close every component even when an earlier one fails, so a failing
close cannot leak
+ // the components behind it or skip super.close(). The first failure
is rethrown with
+ // the later ones suppressed. Order is preserved: the resource cache
must close before
+ // pythonInterpreter since cached resources may hold Python references.
+ //
+ // The ladder catches Throwable, not Exception, and IOUtils.closeAll
is deliberately not
+ // used: both stop at the first non-Exception Throwable without
closing what follows,
+ // which is the very leak this method has to avoid.
+ Throwable firstFailure = null;
+ for (AutoCloseable closeable :
+ new AutoCloseable[] {
+ resourceCache, contextManager, pythonBridge,
eventLogWriter, durableExecManager
+ }) {
+ if (closeable == null) {
+ continue;
+ }
+ try {
+ closeable.close();
+ } catch (Throwable t) {
+ firstFailure = ExceptionUtils.firstOrSuppressed(t,
firstFailure);
+ }
}
- if (durableExecManager != null) {
- durableExecManager.close();
+
+ try {
+ super.close();
+ } catch (Throwable t) {
+ firstFailure = ExceptionUtils.firstOrSuppressed(t, firstFailure);
}
- super.close();
+ if (firstFailure != null) {
+ ExceptionUtils.rethrowException(firstFailure);
+ }
}
@Override
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java
index 0fd0caa6..6da5592e 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java
@@ -35,6 +35,7 @@ import
org.apache.flink.agents.runtime.python.context.PythonRunnerContextImpl;
import org.apache.flink.agents.runtime.trace.ExecutionEventSink;
import org.apache.flink.agents.runtime.trace.ReportedExecutionKey;
import org.apache.flink.api.common.state.MapState;
+import org.apache.flink.util.ExceptionUtils;
import javax.annotation.Nullable;
@@ -350,15 +351,33 @@ class ActionTaskContextManager implements AutoCloseable {
@Override
public void close() throws Exception {
+ // Close the continuation executor even when the runner context fails
to close. The first
+ // failure is rethrown with the later one suppressed.
+ //
+ // The ladder catches Throwable, not Exception, so a non-Exception
Throwable from the
+ // runner context cannot strand the executor's thread pool. Neither
type implements
+ // AutoCloseable, so the aggregation is spelled out rather than
delegated. Both rungs go
+ // through firstOrSuppressed even though the first one cannot yet have
a previous failure,
+ // so that a close inserted above it later suppresses rather than
overwrites.
+ Throwable firstFailure = null;
if (runnerContext != null) {
try {
runnerContext.close();
+ } catch (Throwable t) {
+ firstFailure = ExceptionUtils.firstOrSuppressed(t,
firstFailure);
} finally {
runnerContext = null;
}
}
if (continuationActionExecutor != null) {
- continuationActionExecutor.close();
+ try {
+ continuationActionExecutor.close();
+ } catch (Throwable t) {
+ firstFailure = ExceptionUtils.firstOrSuppressed(t,
firstFailure);
+ }
+ }
+ if (firstFailure != null) {
+ ExceptionUtils.rethrowException(firstFailure);
}
}
}
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java
index ae3bedb3..3c4127ac 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java
@@ -36,6 +36,7 @@ import
org.apache.flink.agents.runtime.python.utils.PythonResourceAdapterImpl;
import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.JobID;
import org.apache.flink.python.env.PythonDependencyInfo;
+import org.apache.flink.util.ExceptionUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import pemja.core.PythonInterpreter;
@@ -315,14 +316,29 @@ class PythonBridgeManager implements AutoCloseable {
@Override
public void close() throws Exception {
- if (pythonActionExecutor != null) {
- pythonActionExecutor.close();
- }
- if (pythonInterpreter != null) {
- pythonInterpreter.close();
+ // Close every component even when an earlier one fails, so a failing
action executor
+ // cannot leak the interpreter or the environment manager. The first
failure is
+ // rethrown with the later ones suppressed.
+ //
+ // The ladder catches Throwable, not Exception, and IOUtils.closeAll
is deliberately not
+ // used: both stop at the first non-Exception Throwable without
closing what follows, and
+ // what follows here is the native Python state.
+ Throwable firstFailure = null;
+ for (AutoCloseable closeable :
+ new AutoCloseable[] {
+ pythonActionExecutor, pythonInterpreter,
pythonEnvironmentManager
+ }) {
+ if (closeable == null) {
+ continue;
+ }
+ try {
+ closeable.close();
+ } catch (Throwable t) {
+ firstFailure = ExceptionUtils.firstOrSuppressed(t,
firstFailure);
+ }
}
- if (pythonEnvironmentManager != null) {
- pythonEnvironmentManager.close();
+ if (firstFailure != null) {
+ ExceptionUtils.rethrowException(firstFailure);
}
}
}
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java
index d3740277..39f73f42 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java
@@ -25,6 +25,7 @@ import org.apache.flink.agents.plan.AgentPlan;
import org.apache.flink.agents.plan.PythonFunction;
import org.apache.flink.agents.runtime.python.context.PythonRunnerContextImpl;
import org.apache.flink.types.Row;
+import org.apache.flink.util.ExceptionUtils;
import pemja.core.PythonInterpreter;
import pemja.core.object.PyObject;
@@ -34,7 +35,7 @@ import java.util.concurrent.atomic.AtomicLong;
import static org.apache.flink.util.Preconditions.checkState;
/** Execute the corresponding Python action in the agent. */
-public class PythonActionExecutor {
+public class PythonActionExecutor implements AutoCloseable {
private static final String PYTHON_IMPORTS =
"from flink_agents.plan import function\n"
@@ -201,20 +202,36 @@ public class PythonActionExecutor {
return (boolean) ((Object[]) invokeResult)[0];
}
+ @Override
public void close() throws Exception {
- if (interpreter != null) {
- if (pythonAsyncThreadPool != null) {
+ // The two Python-side cleanups are independent, so attempt both even
when the first
+ // fails. Skipping the runner-context cleanup leaves that context's
long-term memory and
+ // resource cache unreleased, and PythonBridgeManager closes the
interpreter right behind
+ // us, so there is no later chance to run it. The first failure is
rethrown with the later
+ // one suppressed, matching the ladders in the managers above.
+ if (interpreter == null) {
+ return;
+ }
+ Throwable firstFailure = null;
+ if (pythonAsyncThreadPool != null) {
+ try {
interpreter.invoke(CLOSE_ASYNC_THREAD_POOL,
pythonAsyncThreadPool);
+ } catch (Throwable t) {
+ firstFailure = ExceptionUtils.firstOrSuppressed(t,
firstFailure);
}
-
- if (pythonRunnerContext != null) {
- try {
- interpreter.invoke(CLOSE_FLINK_RUNNER_CONTEXT,
pythonRunnerContext);
- } finally {
- pythonRunnerContext = null;
- }
+ }
+ if (pythonRunnerContext != null) {
+ try {
+ interpreter.invoke(CLOSE_FLINK_RUNNER_CONTEXT,
pythonRunnerContext);
+ } catch (Throwable t) {
+ firstFailure = ExceptionUtils.firstOrSuppressed(t,
firstFailure);
+ } finally {
+ pythonRunnerContext = null;
}
}
+ if (firstFailure != null) {
+ ExceptionUtils.rethrowException(firstFailure);
+ }
}
/** Failed to execute Python action. */
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/skill/SkillManager.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/skill/SkillManager.java
index 9111edf6..0acee277 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/skill/SkillManager.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/skill/SkillManager.java
@@ -20,6 +20,7 @@ package org.apache.flink.agents.runtime.skill;
import org.apache.flink.agents.api.skills.SkillSourceSpec;
import org.apache.flink.agents.api.skills.Skills;
+import org.apache.flink.util.ExceptionUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -211,7 +212,8 @@ public class SkillManager implements AutoCloseable {
* <p>Mirrors {@code ResourceCache.close()}: the first failure is rethrown
after every repo has
* been attempted, with subsequent failures attached as suppressed
exceptions. This surfaces
* real shutdown bugs (locked files, permission denied, disk full) instead
of silently
- * swallowing them.
+ * swallowing them. As there, a non-{@code Exception} {@code Throwable}
does not stop the
+ * remaining repos, and reaches the caller unwrapped.
*/
@Override
public void close() throws Exception {
@@ -225,20 +227,20 @@ public class SkillManager implements AutoCloseable {
// contributes multiple skills.
Set<SkillRepository> unique = Collections.newSetFromMap(new
IdentityHashMap<>());
unique.addAll(openedRepos);
- Exception firstException = null;
+ // Catches Throwable, not Exception: a repo failing with an Error
would otherwise skip the
+ // repos behind it, and ResourceContextImpl.close() clears its manager
reference in a
+ // finally, so nothing can retry the ones that were skipped. This is
the nested branch of
+ // the close-all guarantee that ResourceCache.close() makes one level
up.
+ Throwable firstFailure = null;
for (SkillRepository repo : unique) {
try {
repo.close();
- } catch (Exception e) {
- if (firstException == null) {
- firstException = e;
- } else {
- firstException.addSuppressed(e);
- }
+ } catch (Throwable t) {
+ firstFailure = ExceptionUtils.firstOrSuppressed(t,
firstFailure);
}
}
- if (firstException != null) {
- throw firstException;
+ if (firstFailure != null) {
+ ExceptionUtils.rethrowException(firstFailure);
}
}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/ResourceCacheTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/ResourceCacheTest.java
index b0e5a685..b51f7e0c 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/ResourceCacheTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/ResourceCacheTest.java
@@ -34,18 +34,31 @@ import org.apache.flink.agents.api.resource.ResourceType;
import org.apache.flink.agents.api.resource.SerializableResource;
import org.apache.flink.agents.api.resource.python.PythonResourceAdapter;
import org.apache.flink.agents.api.resource.python.PythonResourceWrapper;
+import org.apache.flink.agents.api.skills.SkillSourceSpec;
+import org.apache.flink.agents.api.skills.Skills;
import org.apache.flink.agents.api.vectorstores.Document;
import org.apache.flink.agents.api.vectorstores.VectorStoreQuery;
import org.apache.flink.agents.api.vectorstores.VectorStoreQueryResult;
import org.apache.flink.agents.plan.AgentPlan;
+import org.apache.flink.agents.runtime.resource.ResourceContextImpl;
+import org.apache.flink.agents.runtime.skill.AgentSkill;
+import org.apache.flink.agents.runtime.skill.SkillManager;
+import org.apache.flink.agents.runtime.skill.SkillRepository;
+import org.apache.flink.agents.runtime.skill.SkillSourceRegistry;
import org.junit.jupiter.api.Test;
import pemja.core.object.PyObject;
+import java.lang.reflect.Field;
+import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.assertj.core.api.Assertions.catchThrowable;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
/** Tests for {@link ResourceCache}. */
public class ResourceCacheTest {
@@ -274,4 +287,184 @@ public class ResourceCacheTest {
Resource myToolAgain = cache.getResource("myTool", ResourceType.TOOL);
assertThat(myTool).isSameAs(myToolAgain);
}
+
+ /**
+ * A cached resource failing with a non-{@code Exception} {@code
Throwable} must not strand the
+ * remaining resources, the cache clear, or the resource context. {@code
+ * ActionExecutionOperator.close()} closes this cache before the Python
interpreter precisely
+ * because cached resources may hold Python references, so leaving
resources open here while the
+ * interpreter behind them is torn down would break that ordering.
+ */
+ @Test
+ public void closeClosesEveryResourceWhenAnEarlierResourceThrowsError()
throws Exception {
+ ResourceCache cache = new ResourceCache(new HashMap<>());
+ OutOfMemoryError failure = new OutOfMemoryError("resource close
failed");
+ RecordingResource failing = new RecordingResource(ResourceType.TOOL,
failure);
+ RecordingResource surviving = new
RecordingResource(ResourceType.CHAT_MODEL, null);
+ cache.put("failing", ResourceType.TOOL, failing);
+ cache.put("surviving", ResourceType.CHAT_MODEL, surviving);
+ // Stands in for the lazily-cached skill manager, so that
resourceContext.close() running
+ // is observable rather than merely assumed from its position in the
method.
+ SkillManager skillManager = mock(SkillManager.class);
+ setSkillManager(cache.getResourceContext(), skillManager);
+
+ // The Error reaches the caller unchanged rather than wrapped in an
Exception, and with
+ // nothing attached to it.
+ assertThatThrownBy(cache::close)
+ .isSameAs(failure)
+ .satisfies(thrown ->
assertThat(thrown.getSuppressed()).isEmpty());
+
+ assertThat(failing.closed).isTrue();
+ assertThat(surviving.closed).isTrue();
+ // Neither the cache clear nor the resource context is skipped by the
Error.
+ assertThat(cachedResources(cache)).isEmpty();
+ verify(skillManager).close();
+ }
+
+ /** The first failure is rethrown and any later one is attached as
suppressed, never dropped. */
+ @Test
+ public void closeReportsFirstResourceFailureWithLaterOnesSuppressed()
throws Exception {
+ ResourceCache cache = new ResourceCache(new HashMap<>());
+ RecordingResource first =
+ new RecordingResource(ResourceType.TOOL, new
IllegalStateException("first"));
+ RecordingResource second =
+ new RecordingResource(ResourceType.TOOL, new
IllegalStateException("second"));
+ cache.put("first", ResourceType.TOOL, first);
+ cache.put("second", ResourceType.TOOL, second);
+
+ Throwable thrown = catchThrowable(cache::close);
+
+ // Iteration order over the cache is unspecified, so pin the
aggregation rather than which
+ // of the two lands first: one is thrown and the other is suppressed
on it.
+ assertThat(thrown).isInstanceOf(IllegalStateException.class);
+ assertThat(thrown.getSuppressed()).hasSize(1);
+ assertThat(new String[] {thrown.getMessage(),
thrown.getSuppressed()[0].getMessage()})
+ .containsExactlyInAnyOrder("first", "second");
+ assertThat(first.closed).isTrue();
+ assertThat(second.closed).isTrue();
+ }
+
+ /**
+ * The close-all guarantee has to reach the nested skill repositories, not
stop at {@code
+ * ResourceContextImpl}. Exercised through the real production path —
{@code
+ * ResourceCache.close()} → {@code ResourceContextImpl.close()} → {@code
SkillManager.close()} →
+ * the repos — because {@code ResourceContextImpl} clears its manager
reference in a {@code
+ * finally}, so a repo skipped here can never be retried and leaks its
temp directory.
+ */
+ @Test
+ public void closeClosesEverySkillRepositoryWhenAnEarlierRepoThrowsError()
throws Exception {
+ // Both repos fail, so the assertions do not depend on the de-dup
set's iteration order:
+ // a handler narrowed to Exception anywhere along the chain stops at
whichever runs first
+ // and leaves the other unclosed, which fails here either way round.
+ Error firstBoom = new Error("repo close failed");
+ Error secondBoom = new Error("other repo close failed");
+ RecordingRepo failing = new RecordingRepo("alpha", firstBoom);
+ RecordingRepo surviving = new RecordingRepo("beta", secondBoom);
+ AtomicInteger seq = new AtomicInteger();
+ List<RecordingRepo> ordered = List.of(failing, surviving);
+ SkillSourceRegistry.register(
+ "test-resource-cache-close-error",
+ (params, cl) -> ordered.get(seq.getAndIncrement()));
+ Skills skills =
+ new Skills(
+ List.of(
+ new
SkillSourceSpec("test-resource-cache-close-error", Map.of()),
+ new
SkillSourceSpec("test-resource-cache-close-error", Map.of())));
+
+ ResourceCache cache = new ResourceCache(new HashMap<>());
+ cache.put(Skills.SKILLS_CONFIG, ResourceType.SKILLS, skills);
+ // Force the lazily-cached SkillManager to exist, so close() has repos
to release.
+ cache.getResourceContext().getSkillDirs(List.of("alpha"));
+
+ // The Error reaches the caller unwrapped, through both intervening
close() methods.
+ Throwable thrown = catchThrowable(cache::close);
+
+ assertThat(thrown).isInstanceOf(Error.class);
+ assertThat(failing.closed).isTrue();
+ assertThat(surviving.closed).isTrue();
+ assertThat(thrown.getSuppressed()).hasSize(1);
+ assertThat(List.of(thrown, thrown.getSuppressed()[0]))
+ .containsExactlyInAnyOrder(firstBoom, secondBoom);
+ }
+
+ /** A skill repository that records its close and can be made to fail it.
*/
+ private static final class RecordingRepo implements SkillRepository {
+ private final AgentSkill skill;
+ private final Throwable failure;
+ private boolean closed = false;
+
+ private RecordingRepo(String skillName, Throwable failure) {
+ this.skill = new AgentSkill(skillName, "fake", "body", null, null,
null);
+ this.failure = failure;
+ }
+
+ @Override
+ public AgentSkill getSkill(String name) {
+ return name.equals(skill.getName()) ? skill : null;
+ }
+
+ @Override
+ public List<AgentSkill> getSkills() {
+ return List.of(skill);
+ }
+
+ @Override
+ public Map<String, String> getResources(String name) {
+ return Map.of();
+ }
+
+ @Override
+ public void close() {
+ closed = true;
+ if (failure instanceof Error) {
+ throw (Error) failure;
+ }
+ if (failure instanceof RuntimeException) {
+ throw (RuntimeException) failure;
+ }
+ }
+ }
+
+ private static void setSkillManager(ResourceContextImpl context,
SkillManager skillManager)
+ throws Exception {
+ Field field =
ResourceContextImpl.class.getDeclaredField("skillManager");
+ field.setAccessible(true);
+ field.set(context, skillManager);
+ }
+
+ @SuppressWarnings("unchecked")
+ private static Map<ResourceType, Map<String, Resource>>
cachedResources(ResourceCache cache)
+ throws Exception {
+ Field field = ResourceCache.class.getDeclaredField("cache");
+ field.setAccessible(true);
+ return (Map<ResourceType, Map<String, Resource>>) field.get(cache);
+ }
+
+ /** Records whether {@code close()} ran, and optionally fails it. */
+ private static final class RecordingResource extends Resource {
+ private final ResourceType type;
+ private final Throwable failure;
+ private boolean closed = false;
+
+ private RecordingResource(ResourceType type, Throwable failure) {
+ this.type = type;
+ this.failure = failure;
+ }
+
+ @Override
+ public ResourceType getResourceType() {
+ return type;
+ }
+
+ @Override
+ public void close() throws Exception {
+ closed = true;
+ if (failure instanceof Error) {
+ throw (Error) failure;
+ }
+ if (failure instanceof Exception) {
+ throw (Exception) failure;
+ }
+ }
+ }
}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java
index 9fcab04d..b92f887e 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java
@@ -48,11 +48,13 @@ import org.apache.flink.agents.plan.actions.ToolCallAction;
import
org.apache.flink.agents.plan.resourceprovider.JavaSerializableResourceProvider;
import org.apache.flink.agents.plan.resourceprovider.ResourceProvider;
import org.apache.flink.agents.plan.tools.FunctionTool;
+import org.apache.flink.agents.runtime.ResourceCache;
import org.apache.flink.agents.runtime.actionstate.ActionState;
import org.apache.flink.agents.runtime.actionstate.ActionStateSerde;
import org.apache.flink.agents.runtime.actionstate.ActionStateUtil;
import org.apache.flink.agents.runtime.actionstate.CallResult;
import org.apache.flink.agents.runtime.actionstate.InMemoryActionStateStore;
+import org.apache.flink.agents.runtime.eventlog.EventLogWriter;
import org.apache.flink.agents.runtime.eventlog.FileEventLogger;
import org.apache.flink.agents.runtime.eventlog.Slf4jEventLogger;
import org.apache.flink.agents.runtime.memory.Mem0LongTermMemory;
@@ -60,6 +62,8 @@ import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
import org.apache.flink.runtime.state.KeyGroupRangeAssignment;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.StreamOperatorStateHandler;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.runtime.tasks.mailbox.TaskMailbox;
@@ -69,6 +73,7 @@ import org.apache.flink.util.ExceptionUtils;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
+import org.mockito.InOrder;
import java.io.IOException;
import java.io.Serializable;
@@ -90,6 +95,9 @@ import java.util.stream.Collectors;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.assertj.core.api.Assertions.catchThrowable;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.inOrder;
+import static org.mockito.Mockito.mock;
/** Tests for {@link ActionExecutionOperator}. */
public class ActionExecutionOperatorTest {
@@ -615,6 +623,173 @@ public class ActionExecutionOperatorTest {
ltmField.set(operator, ltm);
}
+ private static void replaceOperatorField(
+ ActionExecutionOperator<?, ?> operator, String name, Object value)
throws Exception {
+ Field field = ActionExecutionOperator.class.getDeclaredField(name);
+ field.setAccessible(true);
+ field.set(operator, value);
+ }
+
+ /**
+ * Swaps the state handler {@link AbstractStreamOperator} inherits,
returning the previous one.
+ *
+ * <p>{@code super.close()} compiles to {@code stateHandler.dispose()} and
binds statically, so
+ * a subclass cannot intercept the call. Replacing the inherited handler
is what makes the super
+ * call observable, and what lets it be made to fail.
+ */
+ private static StreamOperatorStateHandler replaceStateHandler(
+ ActionExecutionOperator<?, ?> operator, StreamOperatorStateHandler
handler)
+ throws Exception {
+ Field field =
AbstractStreamOperator.class.getDeclaredField("stateHandler");
+ field.setAccessible(true);
+ StreamOperatorStateHandler previous = (StreamOperatorStateHandler)
field.get(operator);
+ field.set(operator, handler);
+ return previous;
+ }
+
+ private static KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
openCloseTestHarness()
+ throws Exception {
+ KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness
=
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new
ActionExecutionOperatorFactory(TestAgent.getAgentPlan(false), true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class));
+ testHarness.open();
+ return testHarness;
+ }
+
+ /** The operator's five closeable components, stubbed so each close is
observable. */
+ private static final class CloseComponents {
+ private final ResourceCache resourceCache = mock(ResourceCache.class);
+ private final ActionTaskContextManager contextManager =
+ mock(ActionTaskContextManager.class);
+ private final PythonBridgeManager pythonBridge =
mock(PythonBridgeManager.class);
+ private final EventLogWriter eventLogWriter =
mock(EventLogWriter.class);
+ private final DurableExecutionManager durableExecManager =
+ mock(DurableExecutionManager.class);
+
+ private void installInto(ActionExecutionOperator<?, ?> operator)
throws Exception {
+ replaceOperatorField(operator, "resourceCache", resourceCache);
+ replaceOperatorField(operator, "contextManager", contextManager);
+ replaceOperatorField(operator, "pythonBridge", pythonBridge);
+ replaceOperatorField(operator, "eventLogWriter", eventLogWriter);
+ replaceOperatorField(operator, "durableExecManager",
durableExecManager);
+ }
+
+ /** Detaches the mocks so the harness teardown does not re-trigger the
failure. */
+ private void detachFrom(ActionExecutionOperator<?, ?> operator) throws
Exception {
+ replaceOperatorField(operator, "resourceCache", null);
+ replaceOperatorField(operator, "contextManager", null);
+ replaceOperatorField(operator, "pythonBridge", null);
+ replaceOperatorField(operator, "eventLogWriter", null);
+ replaceOperatorField(operator, "durableExecManager", null);
+ }
+
+ /**
+ * Verifies every component was released, in the documented order,
with {@code
+ * super.close()} last.
+ *
+ * <p>Order is load-bearing rather than incidental: {@code
resourceCache} must close before
+ * {@code pythonBridge} because cached resources may hold Python
references, and {@code
+ * super.close()} disposes the state backends the components run
against.
+ */
+ private void verifyClosedInOrder(StreamOperatorStateHandler
stateHandler) throws Exception {
+ InOrder inOrder =
+ inOrder(
+ resourceCache,
+ contextManager,
+ pythonBridge,
+ eventLogWriter,
+ durableExecManager,
+ stateHandler);
+ inOrder.verify(resourceCache).close();
+ inOrder.verify(contextManager).close();
+ inOrder.verify(pythonBridge).close();
+ inOrder.verify(eventLogWriter).close();
+ inOrder.verify(durableExecManager).close();
+ inOrder.verify(stateHandler).dispose();
+ }
+ }
+
+ /**
+ * A failing component must not strand the ones behind it. This matters
most for {@code
+ * resourceCache}, which closes first and aggregates its own failures, and
for {@code
+ * pythonBridge}, which releases the embedded Python interpreter.
+ *
+ * <p>Also pins that {@code super.close()} still runs. {@link
AbstractStreamOperator#close()}
+ * disposes the state handler, so skipping it strands the state backends.
+ */
+ @Test
+ void closeClosesEveryComponentWhenAnEarlierCloseFails() throws Exception {
+ KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness
=
+ openCloseTestHarness();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>)
testHarness.getOperator();
+
+ CloseComponents components = new CloseComponents();
+ doThrow(new IllegalStateException("resource cache close failed"))
+ .when(components.resourceCache)
+ .close();
+ components.installInto(operator);
+ StreamOperatorStateHandler stateHandler =
mock(StreamOperatorStateHandler.class);
+ StreamOperatorStateHandler realStateHandler =
replaceStateHandler(operator, stateHandler);
+
+ try {
+ assertThatThrownBy(operator::close)
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessage("resource cache close failed")
+ // Contract 3: with super.close() healthy, nothing is
attached to the failure.
+ .satisfies(thrown ->
assertThat(thrown.getSuppressed()).isEmpty());
+
+ components.verifyClosedInOrder(stateHandler);
+ } finally {
+ components.detachFrom(operator);
+ replaceStateHandler(operator, realStateHandler);
+ testHarness.close();
+ }
+ }
+
+ /**
+ * A {@code super.close()} failure must aggregate with the component
failures rather than
+ * replace them: the earlier component failure still reaches the caller,
with the super failure
+ * attached to it as suppressed.
+ */
+ @Test
+ void closeAggregatesSuperCloseFailureWithComponentFailure() throws
Exception {
+ KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness
=
+ openCloseTestHarness();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>)
testHarness.getOperator();
+
+ CloseComponents components = new CloseComponents();
+ doThrow(new IllegalStateException("resource cache close failed"))
+ .when(components.resourceCache)
+ .close();
+ components.installInto(operator);
+ StreamOperatorStateHandler stateHandler =
mock(StreamOperatorStateHandler.class);
+ doThrow(new IllegalStateException("state handler dispose failed"))
+ .when(stateHandler)
+ .dispose();
+ StreamOperatorStateHandler realStateHandler =
replaceStateHandler(operator, stateHandler);
+
+ try {
+ assertThatThrownBy(operator::close)
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessage("resource cache close failed")
+ .satisfies(
+ thrown ->
+ assertThat(thrown.getSuppressed())
+ .extracting(Throwable::getMessage)
+ .containsExactly("state handler
dispose failed"));
+
+ components.verifyClosedInOrder(stateHandler);
+ } finally {
+ components.detachFrom(operator);
+ replaceStateHandler(operator, realStateHandler);
+ testHarness.close();
+ }
+ }
+
/** Java-side stand-in for the Python-backed LTM wrapper used to observe
the failure path. */
private static final class RecordingMem0LongTermMemory extends
Mem0LongTermMemory {
private final List<String> recordedKeys = new ArrayList<>();
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java
index 468a3ed7..a1ea1b0c 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java
@@ -25,6 +25,7 @@ import org.apache.flink.agents.plan.actions.Action;
import org.apache.flink.agents.runtime.ResourceCache;
import org.apache.flink.agents.runtime.actionstate.ActionState;
import org.apache.flink.agents.runtime.actionstate.InMemoryActionStateStore;
+import org.apache.flink.agents.runtime.async.ContinuationActionExecutor;
import org.apache.flink.agents.runtime.async.ContinuationContext;
import org.apache.flink.agents.runtime.context.JavaRunnerContextImpl;
import org.apache.flink.agents.runtime.context.RunnerContextImpl;
@@ -40,6 +41,7 @@ import org.apache.flink.core.memory.DataInputDeserializer;
import org.apache.flink.core.memory.DataOutputSerializer;
import org.junit.jupiter.api.Test;
+import java.lang.reflect.Field;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
@@ -50,6 +52,7 @@ import static
org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
+import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.verify;
@@ -57,6 +60,64 @@ import static org.mockito.Mockito.verify;
/** Contract tests for {@link ActionTaskContextManager}. */
class ActionTaskContextManagerTest {
+ /**
+ * A failing runner context must not strand the continuation executor,
which owns the async
+ * thread pool and would otherwise keep non-daemon threads alive after the
operator closes.
+ */
+ @Test
+ void closeClosesContinuationExecutorWhenRunnerContextFails() throws
Exception {
+ ActionTaskContextManager mgr = new ActionTaskContextManager(1);
+ RunnerContextImpl failingContext = mock(RunnerContextImpl.class);
+ ContinuationActionExecutor continuationExecutor =
mock(ContinuationActionExecutor.class);
+ doThrow(new IllegalStateException("runner context close failed"))
+ .when(failingContext)
+ .close();
+
+ setField(mgr, "runnerContext", failingContext);
+ setField(mgr, "continuationActionExecutor", continuationExecutor);
+
+ assertThatThrownBy(mgr::close)
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessage("runner context close failed")
+ // Contract 3: a lone failure arrives with nothing attached to
it.
+ .satisfies(thrown ->
assertThat(thrown.getSuppressed()).isEmpty());
+
+ verify(continuationExecutor).close();
+ }
+
+ /** The first failure is rethrown and the later one is attached as
suppressed, never dropped. */
+ @Test
+ void closeReportsFirstFailureWithLaterOneSuppressed() throws Exception {
+ ActionTaskContextManager mgr = new ActionTaskContextManager(1);
+ RunnerContextImpl failingContext = mock(RunnerContextImpl.class);
+ ContinuationActionExecutor failingExecutor =
mock(ContinuationActionExecutor.class);
+ doThrow(new IllegalStateException("runner context close failed"))
+ .when(failingContext)
+ .close();
+ doThrow(new IllegalStateException("continuation executor close
failed"))
+ .when(failingExecutor)
+ .close();
+
+ setField(mgr, "runnerContext", failingContext);
+ setField(mgr, "continuationActionExecutor", failingExecutor);
+
+ assertThatThrownBy(mgr::close)
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessage("runner context close failed")
+ .satisfies(
+ thrown ->
+ assertThat(thrown.getSuppressed())
+ .extracting(Throwable::getMessage)
+ .containsExactly("continuation
executor close failed"));
+ }
+
+ private static void setField(ActionTaskContextManager mgr, String name,
Object value)
+ throws Exception {
+ Field field = ActionTaskContextManager.class.getDeclaredField(name);
+ field.setAccessible(true);
+ field.set(mgr, value);
+ }
+
@Test
void perTaskMapsAreIsolatedAcrossPutGetRemove() throws Exception {
try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) {
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/PythonBridgeManagerTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/PythonBridgeManagerTest.java
index 17dce19f..b4daf548 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/PythonBridgeManagerTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/PythonBridgeManagerTest.java
@@ -20,14 +20,24 @@ package org.apache.flink.agents.runtime.operator;
import org.apache.flink.agents.api.InputEvent;
import org.apache.flink.agents.plan.AgentPlan;
import org.apache.flink.agents.plan.actions.Action;
+import org.apache.flink.agents.runtime.env.PythonEnvironmentManager;
+import org.apache.flink.agents.runtime.python.utils.PythonActionExecutor;
import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.JobID;
import org.junit.jupiter.api.Test;
+import org.mockito.InOrder;
+import pemja.core.PythonInterpreter;
+import java.lang.reflect.Field;
import java.util.List;
import java.util.Map;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.inOrder;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
/** Contract tests for {@link PythonBridgeManager}. */
class PythonBridgeManagerTest {
@@ -59,4 +69,103 @@ class PythonBridgeManagerTest {
assertThat(bridge.getPythonRunnerContext()).isNull();
}
}
+
+ /**
+ * A failing action executor must not strand the interpreter or the
environment manager: both
+ * hold native Python state that leaks for the lifetime of the TaskManager
if never closed.
+ *
+ * <p>Also pins the close order documented on the class, which is
load-bearing rather than
+ * incidental: {@link PythonActionExecutor#close()} calls back into the
interpreter, so it has
+ * to run before the interpreter is closed.
+ */
+ @Test
+ void closeReleasesInterpreterAndEnvironmentWhenActionExecutorFails()
throws Exception {
+ PythonBridgeManager bridge = new PythonBridgeManager();
+ PythonActionExecutor actionExecutor = mock(PythonActionExecutor.class);
+ PythonInterpreter interpreter = mock(PythonInterpreter.class);
+ PythonEnvironmentManager environmentManager =
mock(PythonEnvironmentManager.class);
+ doThrow(new IllegalStateException("action executor close failed"))
+ .when(actionExecutor)
+ .close();
+
+ setField(bridge, "pythonActionExecutor", actionExecutor);
+ setField(bridge, "pythonInterpreter", interpreter);
+ setField(bridge, "pythonEnvironmentManager", environmentManager);
+
+ assertThatThrownBy(bridge::close)
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessage("action executor close failed")
+ // Contract 3: a lone failure arrives with nothing attached to
it.
+ .satisfies(thrown ->
assertThat(thrown.getSuppressed()).isEmpty());
+
+ InOrder inOrder = inOrder(actionExecutor, interpreter,
environmentManager);
+ inOrder.verify(actionExecutor).close();
+ inOrder.verify(interpreter).close();
+ inOrder.verify(environmentManager).close();
+ }
+
+ /** The first failure is rethrown and any later one is attached as
suppressed, never dropped. */
+ @Test
+ void closeReportsFirstFailureWithLaterOnesSuppressed() throws Exception {
+ PythonBridgeManager bridge = new PythonBridgeManager();
+ PythonActionExecutor actionExecutor = mock(PythonActionExecutor.class);
+ PythonInterpreter interpreter = mock(PythonInterpreter.class);
+ PythonEnvironmentManager environmentManager =
mock(PythonEnvironmentManager.class);
+ doThrow(new IllegalStateException("action executor close failed"))
+ .when(actionExecutor)
+ .close();
+ doThrow(new IllegalStateException("environment manager close failed"))
+ .when(environmentManager)
+ .close();
+
+ setField(bridge, "pythonActionExecutor", actionExecutor);
+ setField(bridge, "pythonInterpreter", interpreter);
+ setField(bridge, "pythonEnvironmentManager", environmentManager);
+
+ assertThatThrownBy(bridge::close)
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessage("action executor close failed")
+ .satisfies(
+ thrown ->
+ assertThat(thrown.getSuppressed())
+ .extracting(Throwable::getMessage)
+ .containsExactly("environment manager
close failed"));
+
+ verify(interpreter).close();
+ }
+
+ /**
+ * Pins the handler to {@code Throwable}. A non-{@code Exception} failure
from the action
+ * executor must still release the native Python state; a {@code catch
(Exception)} ladder or
+ * {@code IOUtils.closeAll} would stop here and leak it.
+ */
+ @Test
+ void closeReleasesInterpreterAndEnvironmentWhenActionExecutorThrowsError()
throws Exception {
+ PythonBridgeManager bridge = new PythonBridgeManager();
+ PythonActionExecutor actionExecutor = mock(PythonActionExecutor.class);
+ PythonInterpreter interpreter = mock(PythonInterpreter.class);
+ PythonEnvironmentManager environmentManager =
mock(PythonEnvironmentManager.class);
+ OutOfMemoryError failure = new OutOfMemoryError("action executor close
failed");
+ doThrow(failure).when(actionExecutor).close();
+
+ setField(bridge, "pythonActionExecutor", actionExecutor);
+ setField(bridge, "pythonInterpreter", interpreter);
+ setField(bridge, "pythonEnvironmentManager", environmentManager);
+
+ // The Error reaches the caller unchanged rather than wrapped in an
Exception, and with
+ // nothing attached to it.
+ assertThatThrownBy(bridge::close)
+ .isSameAs(failure)
+ .satisfies(thrown ->
assertThat(thrown.getSuppressed()).isEmpty());
+
+ verify(interpreter).close();
+ verify(environmentManager).close();
+ }
+
+ private static void setField(PythonBridgeManager bridge, String name,
Object value)
+ throws Exception {
+ Field field = PythonBridgeManager.class.getDeclaredField(name);
+ field.setAccessible(true);
+ field.set(bridge, value);
+ }
}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutorTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutorTest.java
index baf584fa..5169269e 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutorTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutorTest.java
@@ -20,6 +20,9 @@ package org.apache.flink.agents.runtime.python.utils;
import org.apache.flink.types.Row;
import org.junit.jupiter.api.Test;
import pemja.core.PythonInterpreter;
+import pemja.core.object.PyObject;
+
+import java.lang.reflect.Field;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -87,6 +90,88 @@ class PythonActionExecutorTest {
.hasMessage("bad pickle");
}
+ /**
+ * A failing thread-pool shutdown must not skip the runner-context
cleanup. That cleanup
+ * releases the Python context's long-term memory and resource cache, and
{@code
+ * PythonBridgeManager} closes the interpreter immediately behind this
call, so a skipped
+ * cleanup never runs at all.
+ */
+ @Test
+ void closeCleansRunnerContextWhenThreadPoolShutdownFails() throws
Exception {
+ PythonInterpreter interpreter = mock(PythonInterpreter.class);
+ PythonActionExecutor executor = newExecutor(interpreter);
+ PyObject threadPool = mock(PyObject.class);
+ PyObject runnerContext = mock(PyObject.class);
+ setField(executor, "pythonAsyncThreadPool", threadPool);
+ setField(executor, "pythonRunnerContext", runnerContext);
+
when(interpreter.invoke("flink_runner_context.close_async_thread_pool",
threadPool))
+ .thenThrow(new RuntimeException("thread pool shutdown
failed"));
+
+ assertThatThrownBy(executor::close)
+ .isInstanceOf(RuntimeException.class)
+ .hasMessage("thread pool shutdown failed")
+ // A lone failure arrives with nothing attached to it.
+ .satisfies(thrown ->
assertThat(thrown.getSuppressed()).isEmpty());
+
+ verify(interpreter)
+ .invoke("flink_runner_context.close_flink_runner_context",
runnerContext);
+ // The handle is released even on the failing path, so a repeated
close cannot double-free.
+ assertThat(executor.getPythonRunnerContext()).isNull();
+ }
+
+ /** The first failure is rethrown and the later one is attached as
suppressed, never dropped. */
+ @Test
+ void closeReportsFirstFailureWithLaterOneSuppressed() throws Exception {
+ PythonInterpreter interpreter = mock(PythonInterpreter.class);
+ PythonActionExecutor executor = newExecutor(interpreter);
+ PyObject threadPool = mock(PyObject.class);
+ PyObject runnerContext = mock(PyObject.class);
+ setField(executor, "pythonAsyncThreadPool", threadPool);
+ setField(executor, "pythonRunnerContext", runnerContext);
+
when(interpreter.invoke("flink_runner_context.close_async_thread_pool",
threadPool))
+ .thenThrow(new RuntimeException("thread pool shutdown
failed"));
+
when(interpreter.invoke("flink_runner_context.close_flink_runner_context",
runnerContext))
+ .thenThrow(new RuntimeException("runner context cleanup
failed"));
+
+ assertThatThrownBy(executor::close)
+ .isInstanceOf(RuntimeException.class)
+ .hasMessage("thread pool shutdown failed")
+ .satisfies(
+ thrown ->
+ assertThat(thrown.getSuppressed())
+ .extracting(Throwable::getMessage)
+ .containsExactly("runner context
cleanup failed"));
+ }
+
+ /**
+ * Pins the handler to {@code Throwable}: a non-{@code Exception} failure
from the thread-pool
+ * shutdown must still release the runner context, and must reach the
caller unwrapped.
+ */
+ @Test
+ void closeCleansRunnerContextWhenThreadPoolShutdownThrowsError() throws
Exception {
+ PythonInterpreter interpreter = mock(PythonInterpreter.class);
+ PythonActionExecutor executor = newExecutor(interpreter);
+ PyObject threadPool = mock(PyObject.class);
+ PyObject runnerContext = mock(PyObject.class);
+ setField(executor, "pythonAsyncThreadPool", threadPool);
+ setField(executor, "pythonRunnerContext", runnerContext);
+ OutOfMemoryError failure = new OutOfMemoryError("thread pool shutdown
failed");
+
when(interpreter.invoke("flink_runner_context.close_async_thread_pool",
threadPool))
+ .thenThrow(failure);
+
+ assertThatThrownBy(executor::close).isSameAs(failure);
+
+ verify(interpreter)
+ .invoke("flink_runner_context.close_flink_runner_context",
runnerContext);
+ }
+
+ private static void setField(PythonActionExecutor executor, String name,
Object value)
+ throws Exception {
+ Field field = PythonActionExecutor.class.getDeclaredField(name);
+ field.setAccessible(true);
+ field.set(executor, value);
+ }
+
private static PythonActionExecutor newExecutor(PythonInterpreter
interpreter)
throws Exception {
return new PythonActionExecutor(interpreter, null, null, null,
"test-job");
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/skill/SkillManagerTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/skill/SkillManagerTest.java
index daf8660f..5f72d4a8 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/skill/SkillManagerTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/skill/SkillManagerTest.java
@@ -455,8 +455,9 @@ class SkillManagerTest {
// fails its registration with one Error and then its close() with
another: a load guard
// narrowed to Exception would skip the cleanup entirely, and a
cleanup guard narrowed to
// Exception would let the close() Error escape and replace the
registration failure.
- // One source keeps the assertions deterministic — closeRepos()
catches only Exception per
- // repo, so an Error from any repo's close() ends the iteration over
the remaining ones.
+ // One source keeps the assertions deterministic: with a single repo
there is no iteration
+ // order to depend on. closeRepos() now continues past an Error too —
see
+ // closeAttemptsEveryRepoWhenAnEarlierRepoThrowsError.
Error registrationBoom = new Error("registration-error");
Error closeBoom = new Error("close-error");
FakeRepo repo = new FakeRepo("alpha", closeBoom, registrationBoom);
@@ -511,6 +512,46 @@ class SkillManagerTest {
assertTrue(all.contains(good2Boom), "good2-close must surface");
}
+ @Test
+ void closeAttemptsEveryRepoWhenAnEarlierRepoThrowsError() throws Exception
{
+ // An Error from one repo's close() must not strand the repos behind
it. Nothing can retry
+ // them: ResourceContextImpl.close() clears its SkillManager reference
in a finally, so a
+ // skipped repo leaks its materialized temp directory for the life of
the JVM.
+ //
+ // Both repos fail so the assertions do not depend on the de-dup set's
iteration order,
+ // which is unspecified. With a handler narrowed to Exception,
whichever repo runs first
+ // ends the iteration and the other is left unclosed — that fails here
either way round.
+ Error firstBoom = new Error("close-error-1");
+ Error secondBoom = new Error("close-error-2");
+ FakeRepo first = new FakeRepo("alpha", firstBoom);
+ FakeRepo second = new FakeRepo("beta", secondBoom);
+
+ AtomicInteger seq = new AtomicInteger();
+ List<FakeRepo> ordered = List.of(first, second);
+ SkillSourceRegistry.register(
+ "test-close-error", (params, cl) ->
ordered.get(seq.getAndIncrement()));
+
+ Skills config =
+ new Skills(
+ List.of(
+ new SkillSourceSpec("test-close-error",
Map.of()),
+ new SkillSourceSpec("test-close-error",
Map.of())));
+
+ SkillManager manager = new SkillManager(config);
+ // The Error reaches the caller unwrapped rather than boxed in an
Exception.
+ Error thrown = assertThrows(Error.class, manager::close);
+
+ assertTrue(first.closed.get(), "the first repo must be attempted");
+ assertTrue(second.closed.get(), "the repo behind the Error must still
be closed");
+ List<Throwable> all = new ArrayList<>();
+ all.add(thrown);
+ for (Throwable s : thrown.getSuppressed()) {
+ all.add(s);
+ }
+ assertTrue(all.contains(firstBoom), "the first Error must surface");
+ assertTrue(all.contains(secondBoom), "the second Error must surface");
+ }
+
@Test
void closeReleasesRepoDisplacedByDuplicateSkillName() throws Exception {
// Two sources both contribute a skill named "dup". The skill-name →
repo map keeps only
