dianfu commented on code in PR #28716:
URL: https://github.com/apache/flink/pull/28716#discussion_r4373431187
##########
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:
> I recommend handling cleanup-only abort in a separate follow-up and
keeping this PR focused on the state-backend shutdown race.
This issue is introduced by this PR. We are not addressing an issue which
already exists before this PR.
--
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]
