Skip to content

Fail fast in artifact staging when storing an artifact fails #39367

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Eliaaazzz wants to merge 1 commit into
base: master
Choose a base branch
from
+102 −10
Open

Fail fast in artifact staging when storing an artifact fails #39367

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 34 additions & 10 deletions ...on/src/main/java/org/apache/beam/runners/fnexecution/artifact/ArtifactStagingService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import org.apache.beam.model.jobmanagement.v1.ArtifactApi;
import org.apache.beam.model.jobmanagement.v1.ArtifactStagingServiceGrpc;
import org.apache.beam.model.pipeline.v1.RunnerApi;
Expand Down Expand Up @@ -223,15 +224,19 @@ public OverflowingSemaphore(int totalPermits) {
}

synchronized void aquire(int permits) throws Exception {
while (usedPermits >= totalPermits) {
if (exception != null) {
throw exception;
}
while (exception == null && usedPermits >= totalPermits) {
this.wait();
}
checkException();
usedPermits += permits;
}

synchronized void checkException() throws Exception {
if (exception != null) {
throw exception;
}
}

synchronized void release(int permits) {
usedPermits -= permits;
this.notifyAll();
Expand Down Expand Up @@ -282,13 +287,16 @@ public RunnerApi.ArtifactInformation call() throws IOException {
.setTypeUrn(dest.getTypeUrn())
.setTypePayload(dest.getTypePayload())
.build();
} catch (IOException | InterruptedException exn) {
} catch (Exception exn) {
// As this thread will no longer be draining the queue, we don't want to get stuck writing
// to it.
// to it. This must happen for unchecked exceptions as well: getDestination can throw e.g.
// InvalidPathException, and leaving the error unset would block the producer forever.
totalPendingBytes.setException(exn);
LOG.error("Exception staging artifacts", exn);
if (exn instanceof IOException) {
throw (IOException) exn;
} else if (exn instanceof RuntimeException) {
throw (RuntimeException) exn;
} else {
throw new RuntimeException(exn);
Comment on lines 296 to 301

@gemini-code-assist gemini-code-assist Bot Jul 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

When catching Exception (which includes InterruptedException), the interrupted status of the thread is cleared. If the exception is an InterruptedException, we should restore the interrupted status by calling Thread.currentThread().interrupt() before rethrowing or wrapping it, so that calling code or the thread pool can handle the interruption properly.

        if (exn instanceof IOException) {
          throw (IOException) exn;
        } else if (exn instanceof InterruptedException) {
          Thread.currentThread().interrupt();
          throw new RuntimeException(exn);
        } else if (exn instanceof RuntimeException) {
          throw (RuntimeException) exn;
        } else {
          throw new RuntimeException(exn);
        }

}
Expand Down Expand Up @@ -409,10 +417,10 @@ public synchronized void onNext(ArtifactApi.ArtifactResponseWrapper responseWrap
ByteString chunk = responseWrapper.getGetArtifactResponse().getData();
if (chunk.size() > 0) { // Make sure we don't accidentally send the EOF value.
totalPendingBytes.aquire(chunk.size());
currentOutput.put(chunk);
putChunk(chunk);
}
if (responseWrapper.getIsLast()) {
currentOutput.put(ByteString.EMPTY); // The EOF value.
putChunk(ByteString.EMPTY); // The EOF value.
if (pendingGets.isEmpty()) {
resolveNextEnvironment(responseObserver);
} else {
Expand All @@ -421,8 +429,16 @@ public synchronized void onNext(ArtifactApi.ArtifactResponseWrapper responseWrap
}
}
} catch (Exception exn) {
LOG.error("Error submitting.", exn);
onError(exn);
// The write of a previous chunk failed; surface the failure to the client rather
// than leaving the stream unterminated, which would make the client block forever.
LOG.error("Error staging artifacts", exn);
state = State.ERROR;
stagingExecutor.shutdownNow();
Comment on lines +435 to +436

@gemini-code-assist gemini-code-assist Bot Jul 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

Calling stagingExecutor.shutdownNow() will shut down the shared executor service of the ArtifactStagingService instance. Since ArtifactStagingService is typically a long-running gRPC service (e.g., in a JobServer), shutting down this executor will permanently prevent any future artifact staging requests from succeeding.

Since the StoreArtifact task has already failed and terminated (which is what set the exception on totalPendingBytes and triggered this catch block), there is no need to shut down or cancel anything in the executor. We can safely remove this call.

Suggested change
state = State.ERROR;
stagingExecutor.shutdownNow();
state = State.ERROR;

responseObserver.onError(
Status.INTERNAL
.withDescription("Error staging artifacts: " + exn)
.withCause(exn)
.asException());
}
break;

Expand All @@ -433,6 +449,14 @@ public synchronized void onNext(ArtifactApi.ArtifactResponseWrapper responseWrap
}
}

private void putChunk(ByteString chunk) throws Exception {
// The queue is drained by a StoreArtifact task. Don't block forever on a plain put if
// that task has died with an error; poll so its exception can interrupt the wait.
while (!currentOutput.offer(chunk, 1, TimeUnit.SECONDS)) {
totalPendingBytes.checkException();
}
}

private void resolveNextEnvironment(
StreamObserver<ArtifactApi.ArtifactRequestWrapper> responseObserver) {
if (pendingResolves.isEmpty()) {
Expand Down
68 changes: 68 additions & 0 deletions ...rc/test/java/org/apache/beam/runners/fnexecution/artifact/ArtifactStagingServiceTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,10 @@
package org.apache.beam.runners.fnexecution.artifact;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;

import java.nio.file.InvalidPathException;
import java.nio.file.Path;
import java.util.Iterator;
import java.util.List;
Expand Down Expand Up @@ -135,6 +138,22 @@ public static RunnerApi.ArtifactInformation unresolvedArtifact(String contents)
}
}

/** Streams each artifact one byte per chunk, to exercise the service's chunk buffering. */
private static class OneBytePerChunkArtifactRetrievalService
extends FakeArtifactRetrievalService {
@Override
public void getArtifact(
ArtifactApi.GetArtifactRequest request,
StreamObserver<ArtifactApi.GetArtifactResponse> responseObserver) {
ByteString data = request.getArtifact().getTypePayload();
for (int i = 0; i < data.size(); i++) {
responseObserver.onNext(
ArtifactApi.GetArtifactResponse.newBuilder().setData(data.substring(i, i + 1)).build());
}
responseObserver.onCompleted();
}
}

private String getArtifact(RunnerApi.ArtifactInformation artifact) {
ByteString all = ByteString.EMPTY;
Iterator<ArtifactApi.GetArtifactResponse> response =
Expand Down Expand Up @@ -166,6 +185,55 @@ public void testStageArtifacts() throws InterruptedException, ExecutionException
checkArtifacts(contentsList, staged.get("env2"));
}

@SuppressWarnings("InlineMeInliner") // inline `Strings.repeat()` - Java 11+ API only
@Test(timeout = 60_000)
public void testDestinationFailureFailsOfferInsteadOfHanging() throws Exception {
// Resolving the destination of a staged artifact can throw an unchecked exception, e.g.
// InvalidPathException on Windows where the generated filename may contain characters that
// are illegal in paths (https://github.com/apache/beam/issues/39364). This must fail the
// offering client instead of stalling the transfer forever.
ArtifactStagingService failingStagingService =
new ArtifactStagingService(
new ArtifactStagingService.ArtifactDestinationProvider() {
@Override
public ArtifactStagingService.ArtifactDestination getDestination(
String stagingToken, String name) {
throw new InvalidPathException(name, "Illegal char simulated");
}

@Override
public void removeStagedArtifacts(String stagingToken) {}
});
grpcCleanup.register(
InProcessServerBuilder.forName("failing-server")
.directExecutor()
.addService(failingStagingService)
.build()
.start());
ManagedChannel failingChannel =
grpcCleanup.register(InProcessChannelBuilder.forName("failing-server").build());
ArtifactStagingServiceGrpc.ArtifactStagingServiceStub failingStub =
ArtifactStagingServiceGrpc.newStub(failingChannel);

// More chunks than the service buffers per artifact, so staging cannot run to completion
// before the destination failure is observed.
String contents = Strings.repeat("x", 300);
failingStagingService.registerJob(
"failingToken",
ImmutableMap.of(
"env1", ImmutableList.of(FakeArtifactRetrievalService.resolvedArtifact(contents))));

ExecutionException exn =
assertThrows(
ExecutionException.class,
() ->
ArtifactStagingService.offer(
new OneBytePerChunkArtifactRetrievalService(), failingStub, "failingToken"));
assertTrue(
"Expected the destination failure, got: " + exn.getCause(),
exn.getCause().getMessage().contains("Illegal char simulated"));
}

private void checkArtifacts(
List<String> expectedContents, List<RunnerApi.ArtifactInformation> staged) {
assertEquals(
Expand Down
Loading