dianfu commented on code in PR #28716:
URL: https://github.com/apache/flink/pull/28716#discussion_r4355845587
##########
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:
finishBundle() calls RemoteBundle.close(), which flushes buffered input and
waits for Python processing and output delivery, in addition to releasing
bundle resources. During cancellation, completing the remaining computation and
collecting its results are unnecessary (which may take very long time if it's
slow, e.g. calling external model services), but releasing the bundle's
registrations and client reference is still required.
--
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]
