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 403261d5 [runtime][python] Release bridge-owned Pemja objects (#944)
403261d5 is described below
commit 403261d56b3e8f81d3fd973c92e8c8c9cd037090
Author: Yu Tang <[email protected]>
AuthorDate: Mon Aug 31 14:50:56 2026 +0800
[runtime][python] Release bridge-owned Pemja objects (#944)
Co-Authored-By: Claude Code <[email protected]>
---
.../flink_agents/runtime/flink_runner_context.py | 27 +++--
.../tests/test_flink_runner_context_close.py | 89 +++++++++++++++
.../agents/runtime/memory/Mem0LongTermMemory.java | 13 ++-
.../runtime/operator/PythonBridgeManager.java | 18 ++--
.../runtime/python/utils/PythonActionExecutor.java | 39 ++++---
.../python/utils/PythonResourceAdapterImpl.java | 12 ++-
.../runtime/memory/Mem0LongTermMemoryTest.java | 18 +++-
.../runtime/operator/PythonBridgeManagerTest.java | 46 ++++++++
.../python/utils/PythonActionExecutorTest.java | 119 +++++++++++++++++++++
.../utils/PythonResourceAdapterImplTest.java | 14 +++
10 files changed, 362 insertions(+), 33 deletions(-)
diff --git a/python/flink_agents/runtime/flink_runner_context.py
b/python/flink_agents/runtime/flink_runner_context.py
index 0d6e08a0..284b3c6e 100644
--- a/python/flink_agents/runtime/flink_runner_context.py
+++ b/python/flink_agents/runtime/flink_runner_context.py
@@ -60,7 +60,11 @@ from
flink_agents.runtime.memory.internal_base_long_term_memory import (
from flink_agents.runtime.memory.mem0.mem0_long_term_memory import (
Mem0LongTermMemory,
)
-from flink_agents.runtime.resource_cache import ResourceCache
+from flink_agents.runtime.resource_cache import (
+ ResourceCache,
+ _failure_of,
+ _first_or_logged,
+)
logger = logging.getLogger(__name__)
@@ -1210,14 +1214,21 @@ class FlinkRunnerContext(RunnerContext,
ExecutionReporter):
@override
def close(self) -> None:
- if self.long_term_memory is not None:
- self.long_term_memory.close()
+ ltm = self.__ltm
+ self.__ltm = None
- if self.__resource_cache is not None:
- try:
- self.__resource_cache.close()
- finally:
- self.__resource_cache = None
+ first_failure = _failure_of(ltm.close) if ltm is not None else None
+
+ resource_cache = self.__resource_cache
+ self.__resource_cache = None
+ first_failure = _first_or_logged(
+ _failure_of(resource_cache.close) if resource_cache is not None
else None,
+ first_failure,
+ "runner context resource cache",
+ )
+
+ if first_failure is not None:
+ raise first_failure
def create_flink_runner_context(
diff --git
a/python/flink_agents/runtime/tests/test_flink_runner_context_close.py
b/python/flink_agents/runtime/tests/test_flink_runner_context_close.py
new file mode 100644
index 00000000..24ee5572
--- /dev/null
+++ b/python/flink_agents/runtime/tests/test_flink_runner_context_close.py
@@ -0,0 +1,89 @@
+################################################################################
+# 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.
+#################################################################################
+from unittest.mock import MagicMock
+
+import pytest
+
+from flink_agents.runtime.flink_runner_context import FlinkRunnerContext
+
+
+def _create_context() -> tuple[FlinkRunnerContext, MagicMock, MagicMock]:
+ ctx = FlinkRunnerContext.__new__(FlinkRunnerContext)
+ ltm = MagicMock()
+ resource_cache = MagicMock()
+ ctx._FlinkRunnerContext__ltm = ltm
+ ctx._FlinkRunnerContext__resource_cache = resource_cache
+ return ctx, ltm, resource_cache
+
+
+def test_close_releases_long_term_memory_and_resource_cache_once() -> None:
+ ctx, ltm, resource_cache = _create_context()
+
+ ctx.close()
+ ctx.close()
+
+ ltm.close.assert_called_once_with()
+ resource_cache.close.assert_called_once_with()
+ assert ctx.long_term_memory is None
+
+
+def test_close_clears_long_term_memory_before_logical_cleanup() -> None:
+ ctx, ltm, resource_cache = _create_context()
+ ltm.close.side_effect = RuntimeError("logical close failed")
+
+ with pytest.raises(RuntimeError, match="logical close failed"):
+ ctx.close()
+
+ assert ctx.long_term_memory is None
+ resource_cache.close.assert_called_once_with()
+ ctx.close()
+ ltm.close.assert_called_once_with()
+ resource_cache.close.assert_called_once_with()
+
+
+def test_close_preserves_first_failure_when_both_cleanups_fail(
+ caplog: pytest.LogCaptureFixture,
+) -> None:
+ ctx, ltm, resource_cache = _create_context()
+ ltm_failure = RuntimeError("logical close failed")
+ resource_cache_failure = RuntimeError("resource cache close failed")
+ ltm.close.side_effect = ltm_failure
+ resource_cache.close.side_effect = resource_cache_failure
+
+ with pytest.raises(RuntimeError, match="logical close failed") as exc_info:
+ ctx.close()
+
+ assert exc_info.value is ltm_failure
+ assert "Suppressed failure closing runner context resource cache." in
caplog.text
+ assert "resource cache close failed" in caplog.text
+
+ ctx.close()
+ ltm.close.assert_called_once_with()
+ resource_cache.close.assert_called_once_with()
+
+
+def test_close_does_not_demote_system_exit_behind_an_earlier_failure() -> None:
+ ctx, ltm, resource_cache = _create_context()
+ ltm.close.side_effect = RuntimeError("logical close failed")
+ system_exit = SystemExit("resource cache close interrupted")
+ resource_cache.close.side_effect = system_exit
+
+ with pytest.raises(SystemExit) as exc_info:
+ ctx.close()
+
+ assert exc_info.value is system_exit
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemory.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemory.java
index 11d95ac8..f414422c 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemory.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemory.java
@@ -44,7 +44,7 @@ public class Mem0LongTermMemory implements
InteranlBaseLongTermMemory {
private static final String MEM0_ITEMS_TO_JAVA =
"python_java_utils.mem0_items_to_java";
private final PythonResourceAdapter adapter;
- private final PyObject pyMem0;
+ private PyObject pyMem0;
public Mem0LongTermMemory(PythonResourceAdapter adapter, PyObject pyMem0) {
this.adapter = adapter;
@@ -168,8 +168,15 @@ public class Mem0LongTermMemory implements
InteranlBaseLongTermMemory {
}
@Override
- public void close() {
- adapter.callMethod(pyMem0, "close", Map.of());
+ public void close() throws Exception {
+ // Clear first because Pemja's close performs an unguarded native
decRef.
+ PyObject memory = pyMem0;
+ pyMem0 = null;
+ if (memory != null) {
+ try (memory) {
+ adapter.callMethod(memory, "close", Map.of());
+ }
+ }
}
private Object buildPyMemorySet(MemorySet memorySet) {
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 3c4127ac..0383faee 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
@@ -58,17 +58,19 @@ import static
org.apache.flink.agents.plan.actions.Utils.supportAsync;
* <ul>
* <li>The {@link PythonEnvironmentManager} that prepares dependencies and
the Pemja runtime.
* <li>The {@link PythonInterpreter} obtained from that environment.
- * <li>The {@link PythonActionExecutor} (when the plan contains Python
actions).
+ * <li>The {@link PythonActionExecutor} (when the plan contains Python
actions or Mem0).
* <li>The {@link PythonRunnerContextImpl} consumed by Python actions.
* <li>The Java/Python resource adapters that bridge resource lookups across
languages.
+ * <li>The Java wrapper around Python Mem0 long-term memory (when
configured).
* </ul>
*
* <p>Lifecycle: instantiated by the operator's {@code open()} (lazy — not in
the operator
* constructor), then immediately initialized via {@link #open} in the same
call. {@link #open} is a
- * no-op when the agent plan contains no Python actions and no Python
resources — in that case all
- * accessors return {@code null} and {@link #isInitialized()} returns {@code
false}. {@link
- * #close()} closes the owned resources in the reverse order of creation:
{@code
- * pythonActionExecutor} → {@code pythonInterpreter} → {@code
pythonEnvironmentManager}.
+ * no-op when the agent plan contains no Python actions, Python resources, or
Mem0 configuration —
+ * in that case all accessors return {@code null} and {@link #isInitialized()}
returns {@code
+ * false}. {@link #close()} closes the owned resources in the reverse order of
creation: {@code
+ * longTermMemory} → {@code pythonActionExecutor} → {@code
pythonResourceAdapter} → {@code
+ * pythonInterpreter} → {@code pythonEnvironmentManager}.
*
* <p>Design constraint: package-private; no manager-to-manager held
references. Other managers
* receive what they need (e.g. the Python runner context, the action
executor) via method
@@ -326,7 +328,11 @@ class PythonBridgeManager implements AutoCloseable {
Throwable firstFailure = null;
for (AutoCloseable closeable :
new AutoCloseable[] {
- pythonActionExecutor, pythonInterpreter,
pythonEnvironmentManager
+ longTermMemory,
+ pythonActionExecutor,
+ pythonResourceAdapter,
+ pythonInterpreter,
+ pythonEnvironmentManager
}) {
if (closeable == null) {
continue;
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 39f73f42..f39958e0 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
@@ -212,28 +212,39 @@ public class PythonActionExecutor implements
AutoCloseable {
if (interpreter == null) {
return;
}
+
+ // Clear the fields before releasing: PyObject.close() performs an
unguarded native decRef,
+ // so a repeated close() must not reach the same handle twice.
+ PyObject asyncThreadPool = pythonAsyncThreadPool;
+ PyObject runnerContext = pythonRunnerContext;
+ pythonAsyncThreadPool = null;
+ pythonRunnerContext = null;
+
Throwable firstFailure = null;
- if (pythonAsyncThreadPool != null) {
- try {
- interpreter.invoke(CLOSE_ASYNC_THREAD_POOL,
pythonAsyncThreadPool);
- } catch (Throwable t) {
- firstFailure = ExceptionUtils.firstOrSuppressed(t,
firstFailure);
- }
+ try {
+ closePythonObject(CLOSE_ASYNC_THREAD_POOL, asyncThreadPool);
+ } catch (Throwable t) {
+ firstFailure = ExceptionUtils.firstOrSuppressed(t, firstFailure);
}
- if (pythonRunnerContext != null) {
- try {
- interpreter.invoke(CLOSE_FLINK_RUNNER_CONTEXT,
pythonRunnerContext);
- } catch (Throwable t) {
- firstFailure = ExceptionUtils.firstOrSuppressed(t,
firstFailure);
- } finally {
- pythonRunnerContext = null;
- }
+ try {
+ closePythonObject(CLOSE_FLINK_RUNNER_CONTEXT, runnerContext);
+ } catch (Throwable t) {
+ firstFailure = ExceptionUtils.firstOrSuppressed(t, firstFailure);
}
+
if (firstFailure != null) {
ExceptionUtils.rethrowException(firstFailure);
}
}
+ private void closePythonObject(String closeFunction, PyObject
pythonObject) throws Exception {
+ if (pythonObject != null) {
+ try (pythonObject) {
+ interpreter.invoke(closeFunction, pythonObject);
+ }
+ }
+ }
+
/** Failed to execute Python action. */
public static class PythonActionExecutionException extends Exception {
public PythonActionExecutionException(String message, Throwable cause)
{
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImpl.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImpl.java
index 5cf4aec0..31c89713 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImpl.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImpl.java
@@ -38,7 +38,7 @@ import java.util.HashMap;
import java.util.List;
import java.util.Map;
-public class PythonResourceAdapterImpl implements PythonResourceAdapter {
+public class PythonResourceAdapterImpl implements PythonResourceAdapter,
AutoCloseable {
static final String PYTHON_IMPORTS = "from flink_agents.runtime import
python_java_utils";
@@ -99,6 +99,16 @@ public class PythonResourceAdapterImpl implements
PythonResourceAdapter {
pythonResourceContext = (PyObject)
interpreter.invoke(GET_RESOURCE_CONTEXT, this);
}
+ @Override
+ public void close() throws Exception {
+ // Clear first because Pemja's close performs an unguarded native
decRef.
+ PyObject resourceContext = pythonResourceContext;
+ pythonResourceContext = null;
+ if (resourceContext != null) {
+ resourceContext.close();
+ }
+ }
+
public Object getResource(String resourceName, String resourceType) {
Resource resource;
try {
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemoryTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemoryTest.java
index f2575be4..72498d77 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemoryTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemoryTest.java
@@ -32,8 +32,10 @@ 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.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
@@ -183,11 +185,12 @@ public class Mem0LongTermMemoryTest {
}
@Test
- void testSwitchContextAndCloseForward() {
+ void testSwitchContextAndCloseForward() throws Exception {
ltm.configureObservation(true, false, true);
ltm.switchContext("k1", "observation-1", true);
ltm.drainObservationRecordsJson("k1", "observation-1");
ltm.close();
+ ltm.close();
verify(mockAdapter)
.callMethod(
@@ -219,5 +222,18 @@ public class Mem0LongTermMemoryTest {
eq("drain_ltm_observation_records"),
eq(Map.of("key", "k1", "observation_id",
"observation-1")));
verify(mockAdapter).callMethod(eq(mockPyMem0), eq("close"),
eq(Map.of()));
+ verify(mockPyMem0).close();
+ }
+
+ @Test
+ void testCloseReleasesPythonObjectWhenLogicalCleanupFails() throws
Exception {
+ RuntimeException failure = new RuntimeException("logical close
failed");
+ doThrow(failure).when(mockAdapter).callMethod(mockPyMem0, "close",
Map.of());
+
+ assertThatThrownBy(ltm::close).isSameAs(failure);
+ ltm.close();
+
+ verify(mockAdapter).callMethod(mockPyMem0, "close", Map.of());
+ verify(mockPyMem0).close();
}
}
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 b4daf548..fa25afff 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
@@ -21,7 +21,9 @@ 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.memory.Mem0LongTermMemory;
import org.apache.flink.agents.runtime.python.utils.PythonActionExecutor;
+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.junit.jupiter.api.Test;
@@ -42,6 +44,50 @@ import static org.mockito.Mockito.verify;
/** Contract tests for {@link PythonBridgeManager}. */
class PythonBridgeManagerTest {
+ @Test
+ void closeAttemptsAllResourcesAndSuppressesLaterFailures() throws
Exception {
+ PythonBridgeManager bridge = new PythonBridgeManager();
+ Mem0LongTermMemory longTermMemory = mock(Mem0LongTermMemory.class);
+ PythonActionExecutor actionExecutor = mock(PythonActionExecutor.class);
+ PythonResourceAdapterImpl resourceAdapter =
mock(PythonResourceAdapterImpl.class);
+ PythonInterpreter interpreter = mock(PythonInterpreter.class);
+ PythonEnvironmentManager environmentManager =
mock(PythonEnvironmentManager.class);
+ RuntimeException actionExecutorFailure =
+ new RuntimeException("action executor close failed");
+ RuntimeException interpreterFailure = new
RuntimeException("interpreter close failed");
+ RuntimeException environmentFailure = new
RuntimeException("environment close failed");
+
+ doThrow(actionExecutorFailure).when(actionExecutor).close();
+ RuntimeException resourceAdapterFailure =
+ new RuntimeException("resource adapter close failed");
+ doThrow(resourceAdapterFailure).when(resourceAdapter).close();
+ doThrow(interpreterFailure).when(interpreter).close();
+ doThrow(environmentFailure).when(environmentManager).close();
+ setField(bridge, "longTermMemory", longTermMemory);
+ setField(bridge, "pythonActionExecutor", actionExecutor);
+ setField(bridge, "pythonResourceAdapter", resourceAdapter);
+ setField(bridge, "pythonInterpreter", interpreter);
+ setField(bridge, "pythonEnvironmentManager", environmentManager);
+
+ assertThatThrownBy(bridge::close)
+ .isSameAs(actionExecutorFailure)
+ .hasSuppressedException(resourceAdapterFailure)
+ .hasSuppressedException(interpreterFailure)
+ .hasSuppressedException(environmentFailure);
+ InOrder closeOrder =
+ inOrder(
+ longTermMemory,
+ actionExecutor,
+ resourceAdapter,
+ interpreter,
+ environmentManager);
+ closeOrder.verify(longTermMemory).close();
+ closeOrder.verify(actionExecutor).close();
+ closeOrder.verify(resourceAdapter).close();
+ closeOrder.verify(interpreter).close();
+ closeOrder.verify(environmentManager).close();
+ }
+
@Test
void openIsNoOpWhenPlanHasNeitherPythonActionsNorResources() throws
Exception {
// Java-only plan: one Java action, no resources.
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 5169269e..5bdac5dc 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
@@ -17,21 +17,41 @@
*/
package org.apache.flink.agents.runtime.python.utils;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.flink.agents.api.agents.AgentExecutionOptions;
+import org.apache.flink.agents.plan.AgentPlan;
+import org.apache.flink.agents.runtime.python.context.PythonRunnerContextImpl;
import org.apache.flink.types.Row;
import org.junit.jupiter.api.Test;
+import org.mockito.InOrder;
import pemja.core.PythonInterpreter;
import pemja.core.object.PyObject;
import java.lang.reflect.Field;
+import java.util.HashMap;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.Mockito.clearInvocations;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.inOrder;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoInteractions;
import static org.mockito.Mockito.when;
+/** Tests for {@link PythonActionExecutor}. */
class PythonActionExecutorTest {
+ private static final String CREATE_ASYNC_THREAD_POOL =
+ "flink_runner_context.create_async_thread_pool";
+ private static final String CLOSE_ASYNC_THREAD_POOL =
+ "flink_runner_context.close_async_thread_pool";
+ private static final String CREATE_FLINK_RUNNER_CONTEXT =
+ "flink_runner_context.create_flink_runner_context";
+ private static final String CLOSE_FLINK_RUNNER_CONTEXT =
+ "flink_runner_context.close_flink_runner_context";
+
@Test
void resolvesPickledPythonKeyTextFromPyFlinkKeyRow() throws Exception {
PythonInterpreter interpreter = mock(PythonInterpreter.class);
@@ -172,8 +192,107 @@ class PythonActionExecutorTest {
field.set(executor, value);
}
+ @Test
+ void releasesPythonObjectsAfterLogicalCleanup() throws Exception {
+ TestFixture fixture = createOpenedExecutor();
+ clearInvocations(fixture.interpreter, fixture.asyncThreadPool,
fixture.runnerContextObject);
+
+ fixture.executor.close();
+
+ InOrder closeOrder =
+ inOrder(fixture.interpreter, fixture.asyncThreadPool,
fixture.runnerContextObject);
+ closeOrder
+ .verify(fixture.interpreter)
+ .invoke(CLOSE_ASYNC_THREAD_POOL, fixture.asyncThreadPool);
+ closeOrder.verify(fixture.asyncThreadPool).close();
+ closeOrder
+ .verify(fixture.interpreter)
+ .invoke(CLOSE_FLINK_RUNNER_CONTEXT,
fixture.runnerContextObject);
+ closeOrder.verify(fixture.runnerContextObject).close();
+ }
+
+ @Test
+ void releasesBothPythonObjectsWhenLogicalCleanupFails() throws Exception {
+ TestFixture fixture = createOpenedExecutor();
+ RuntimeException asyncFailure = new RuntimeException("async cleanup
failed");
+ RuntimeException contextFailure = new RuntimeException("context
cleanup failed");
+ doThrow(asyncFailure)
+ .when(fixture.interpreter)
+ .invoke(CLOSE_ASYNC_THREAD_POOL, fixture.asyncThreadPool);
+ doThrow(contextFailure)
+ .when(fixture.interpreter)
+ .invoke(CLOSE_FLINK_RUNNER_CONTEXT,
fixture.runnerContextObject);
+
+ assertThatThrownBy(fixture.executor::close)
+ .isSameAs(asyncFailure)
+ .hasSuppressedException(contextFailure);
+ InOrder closeOrder =
+ inOrder(fixture.interpreter, fixture.asyncThreadPool,
fixture.runnerContextObject);
+ closeOrder
+ .verify(fixture.interpreter)
+ .invoke(CLOSE_ASYNC_THREAD_POOL, fixture.asyncThreadPool);
+ closeOrder.verify(fixture.asyncThreadPool).close();
+ closeOrder
+ .verify(fixture.interpreter)
+ .invoke(CLOSE_FLINK_RUNNER_CONTEXT,
fixture.runnerContextObject);
+ closeOrder.verify(fixture.runnerContextObject).close();
+
+ clearInvocations(fixture.interpreter, fixture.asyncThreadPool,
fixture.runnerContextObject);
+ fixture.executor.close();
+ verifyNoInteractions(
+ fixture.interpreter, fixture.asyncThreadPool,
fixture.runnerContextObject);
+ }
+
private static PythonActionExecutor newExecutor(PythonInterpreter
interpreter)
throws Exception {
return new PythonActionExecutor(interpreter, null, null, null,
"test-job");
}
+
+ private static TestFixture createOpenedExecutor() throws Exception {
+ PythonInterpreter interpreter = mock(PythonInterpreter.class);
+ PythonRunnerContextImpl runnerContext =
mock(PythonRunnerContextImpl.class);
+ JavaResourceAdapter resourceAdapter = mock(JavaResourceAdapter.class);
+ PyObject asyncThreadPool = mock(PyObject.class);
+ PyObject runnerContextObject = mock(PyObject.class);
+ AgentPlan plan = new AgentPlan(new HashMap<>(), new HashMap<>());
+ String planJson = new ObjectMapper().writeValueAsString(plan);
+ String jobIdentifier = "job-1";
+
+ when(interpreter.invoke(
+ CREATE_ASYNC_THREAD_POOL,
+
plan.getConfig().get(AgentExecutionOptions.NUM_ASYNC_THREADS)))
+ .thenReturn(asyncThreadPool);
+ when(interpreter.invoke(
+ CREATE_FLINK_RUNNER_CONTEXT,
+ runnerContext,
+ planJson,
+ asyncThreadPool,
+ resourceAdapter,
+ jobIdentifier))
+ .thenReturn(runnerContextObject);
+
+ PythonActionExecutor executor =
+ new PythonActionExecutor(
+ interpreter, plan, resourceAdapter, runnerContext,
jobIdentifier);
+ executor.open();
+ return new TestFixture(interpreter, asyncThreadPool,
runnerContextObject, executor);
+ }
+
+ private static final class TestFixture {
+ private final PythonInterpreter interpreter;
+ private final PyObject asyncThreadPool;
+ private final PyObject runnerContextObject;
+ private final PythonActionExecutor executor;
+
+ private TestFixture(
+ PythonInterpreter interpreter,
+ PyObject asyncThreadPool,
+ PyObject runnerContextObject,
+ PythonActionExecutor executor) {
+ this.interpreter = interpreter;
+ this.asyncThreadPool = asyncThreadPool;
+ this.runnerContextObject = runnerContextObject;
+ this.executor = executor;
+ }
+ }
}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImplTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImplTest.java
index e46e3ea6..f372821a 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImplTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImplTest.java
@@ -92,6 +92,20 @@ public class PythonResourceAdapterImplTest {
.invoke(PythonResourceAdapterImpl.GET_RESOURCE_CONTEXT,
pythonResourceAdapter);
}
+ @Test
+ void testCloseReleasesPythonResourceContextOnce() throws Exception {
+ PyObject pythonResourceContext = mock(PyObject.class);
+ when(mockInterpreter.invoke(
+ PythonResourceAdapterImpl.GET_RESOURCE_CONTEXT,
pythonResourceAdapter))
+ .thenReturn(pythonResourceContext);
+ pythonResourceAdapter.open();
+
+ pythonResourceAdapter.close();
+ pythonResourceAdapter.close();
+
+ verify(pythonResourceContext).close();
+ }
+
@Test
void testGetResourceWithPythonResourceWrapper() throws Exception {
String resourceName = "test_resource";
