bowenli86 commented on code in PR #28716:
URL: https://github.com/apache/flink/pull/28716#discussion_r4365239669
##########
flink-python/src/main/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunner.java:
##########
@@ -442,29 +514,309 @@ public Tuple3<String, byte[], Integer> takeResult()
throws Exception {
@Override
public void flush() throws Exception {
- if (bundleStarted) {
- try {
- finishBundle();
- } finally {
+ final BundleCloseOperation closeOperation;
+ final boolean ownsBundleClose;
+ synchronized (this) {
+ if (cancelRequested) {
+ return;
+ }
+
+ if (bundleStarted) {
+ closeOperation = new BundleCloseOperation(remoteBundle);
+ bundleCloseOperation = closeOperation;
bundleStarted = false;
+ clearBundleReferences();
+ ownsBundleClose = true;
+ } else {
+ closeOperation = bundleCloseOperation;
+ ownsBundleClose = false;
}
}
+
+ if (closeOperation == null) {
+ return;
+ }
+
+ if (ownsBundleClose) {
+ finishBundle(closeOperation);
+ } else {
+ closeOperation.await();
+ }
}
/** Interrupts the progress of takeResult. */
public void notifyNoMoreResults() {
resultBuffer.add(Tuple2.of(null, new byte[0]));
}
- private void finishBundle() {
+ private void finishBundle(BundleCloseOperation closeOperation) {
+ RuntimeException failure = null;
try {
- remoteBundle.close();
+ closeOperation.remoteBundle.close();
} catch (Throwable t) {
- throw new RuntimeException("Failed to close remote bundle", t);
+ failure = new RuntimeException("Failed to close remote bundle", t);
+ throw failure;
} finally {
- remoteBundle = null;
- mainInputReceiver = null;
- timerInputReceiver = null;
+ closeOperation.complete(failure);
+ if (isCancelRequested()) {
+ try {
+ closeResources();
+ } catch (Exception e) {
+ LOG.warn("Failed to clean up Python resources after
cancellation.", e);
+ }
+ }
+ }
+ }
+
+ private synchronized boolean startClose() {
+ if (closed) {
+ return false;
+ }
+ closed = true;
+ closeCompletion = new CompletableFuture<>();
+ return true;
+ }
+
+ /**
+ * Starts cancellation without waiting for a concurrent bundle close.
+ *
+ * <p>If no flush thread owns the active bundle yet, cancellation detaches
it and becomes
+ * responsible for finishing it asynchronously.
+ */
+ @Nullable
+ private synchronized BundleCloseClaim requestCancel() {
+ cancelRequested = true;
+ closed = true;
+ if (closeCompletion == null) {
+ closeCompletion = new CompletableFuture<>();
+ }
+
+ if (bundleStarted) {
+ final BundleCloseOperation closeOperation = new
BundleCloseOperation(remoteBundle);
+ bundleCloseOperation = closeOperation;
+ bundleStarted = false;
+ clearBundleReferences();
+ return new BundleCloseClaim(closeOperation, true);
+ }
+ if (bundleCloseOperation != null && !bundleCloseOperation.isDone()) {
+ return new BundleCloseClaim(bundleCloseOperation, false);
+ }
+ clearBundleReferences();
+ return null;
+ }
+
+ private synchronized boolean isCancelRequested() {
+ return cancelRequested;
+ }
+
+ private void awaitCloseCompletion() throws Exception {
+ final CompletableFuture<Void> currentCloseCompletion;
+ synchronized (this) {
+ currentCloseCompletion = closeCompletion;
+ }
+ if (currentCloseCompletion != null) {
+ awaitCompletion(currentCloseCompletion);
+ }
+ }
+
+ private void completeClose() {
+ final CompletableFuture<Void> currentCloseCompletion;
+ synchronized (this) {
+ currentCloseCompletion = closeCompletion;
+ }
+ if (currentCloseCompletion != null) {
+ currentCloseCompletion.complete(null);
+ }
+ }
+
+ private void closeResources() throws Exception {
+ final JobBundleFactory ownedJobBundleFactory;
+ final OpaqueMemoryResource<PythonSharedResources>
currentSharedResources;
+ final CompletableFuture<Void> closeCompletion;
+ final boolean ownsClose;
+ synchronized (this) {
+ if (resourceCloseCompletion == null) {
+ closeCompletion = new CompletableFuture<>();
+ resourceCloseCompletion = closeCompletion;
+ ownedJobBundleFactory = jobBundleFactory;
+ currentSharedResources = sharedResources;
+ jobBundleFactory = null;
+ sharedResources = null;
+ ownsClose = true;
+ } else {
+ closeCompletion = resourceCloseCompletion;
+ ownedJobBundleFactory = null;
+ currentSharedResources = null;
+ ownsClose = false;
+ }
+ }
+
+ if (!ownsClose) {
+ // Disposal holds the slot's SharedResources lock. Even an
interrupted cancellation
+ // must wait for that lock to be released before task cleanup can
finish.
+ try {
+ closeCompletion.join();
+ } catch (CompletionException e) {
+ ExceptionUtils.rethrowException(e.getCause());
+ }
+ return;
+ }
+
+ try {
+ if (currentSharedResources != null) {
+ try {
+ currentSharedResources.close();
+ } finally {
+ clearBundleReferences();
+ }
+ } else {
+ try {
+ if (ownedJobBundleFactory != null) {
+ ownedJobBundleFactory.close();
+ }
+ } finally {
+ try {
+ environmentManager.close();
+ } finally {
+ clearBundleReferences();
+ }
+ }
+ }
+ closeCompletion.complete(null);
+ } catch (Throwable t) {
+ closeCompletion.completeExceptionally(t);
+ throw t;
+ }
+ }
+
+ private void closeBundleAfterCancellation(BundleCloseOperation
closeOperation) {
+ final ExecutorService executor;
+ synchronized (this) {
+ if (cancellationBundleCloseExecutor == null) {
+ cancellationBundleCloseExecutor =
+ Executors.newSingleThreadExecutor(
+ new
ExecutorThreadFactory("beam-python-bundle-cancellation"));
+ }
+ executor = cancellationBundleCloseExecutor;
+ }
+ executor.execute(
+ () -> {
+ try {
+ finishBundle(closeOperation);
Review Comment:
Thanks, I agree that cancellation should have a cleanup-only path. Beam
2.75's `RemoteBundle` doesn't expose an abort operation that skips
flushing/processing while releasing the bundle's registrations and client
reference. Implementing that needs broader Beam integration changes, including
handling late output and cancellation racing an existing close.
I recommend handling cleanup-only abort in a separate follow-up and keeping
this PR focused on the state-backend shutdown race. The current cancellation
path gates state access before returning, but background bundle cleanup can
still wait on Python and retain bundle resources, potentially indefinitely for
a hung call. This PR therefore doesn't guarantee prompt per-bundle resource
reclamation.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
