diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/artifact/ArtifactStagingService.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/artifact/ArtifactStagingService.java index 8b403f2f25f0..265665a11de3 100644 --- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/artifact/ArtifactStagingService.java +++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/artifact/ArtifactStagingService.java @@ -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; @@ -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(); @@ -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); } @@ -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 { @@ -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(); + responseObserver.onError( + Status.INTERNAL + .withDescription("Error staging artifacts: " + exn) + .withCause(exn) + .asException()); } break; @@ -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 responseObserver) { if (pendingResolves.isEmpty()) { diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/artifact/ArtifactStagingServiceTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/artifact/ArtifactStagingServiceTest.java index 5610e4f5bb52..abdce59458db 100644 --- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/artifact/ArtifactStagingServiceTest.java +++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/artifact/ArtifactStagingServiceTest.java @@ -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; @@ -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 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 response = @@ -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 expectedContents, List staged) { assertEquals(