From 8bdc964cd4fa6503c735d44cfb14a9e33201b15b Mon Sep 17 00:00:00 2001 From: bvolpato Date: Sun, 25 Oct 2026 01:30:29 -0400 Subject: [PATCH] [Prism] Wait for the job service before starting loopback workers --- runners/prism/java/build.gradle | 3 + .../beam/runners/prism/PrismExecutor.java | 45 ++++--- .../beam/runners/prism/PrismRunner.java | 117 +++++++++++++++--- .../beam/runners/prism/PrismExecutorTest.java | 35 ++++++ .../beam/runners/prism/PrismRunnerTest.java | 108 ++++++++++++++++ 5 files changed, 273 insertions(+), 35 deletions(-) diff --git a/runners/prism/java/build.gradle b/runners/prism/java/build.gradle index 8bd585b3ae0c..27a54f356401 100644 --- a/runners/prism/java/build.gradle +++ b/runners/prism/java/build.gradle @@ -28,11 +28,14 @@ description = "Apache Beam :: Runners :: Prism :: Java" ext.summary = "Support for executing a pipeline on Prism." dependencies { + implementation project(path: ":model:job-management", configuration: "shadow") + implementation project(path: ":model:pipeline", configuration: "shadow") implementation project(path: ":sdks:java:core", configuration: "shadow") implementation project(":runners:portability:java") implementation library.java.joda_time implementation library.java.slf4j_api + implementation library.java.vendored_grpc_1_69_0 implementation library.java.vendored_guava_32_1_2_jre compileOnly library.java.hamcrest diff --git a/runners/prism/java/src/main/java/org/apache/beam/runners/prism/PrismExecutor.java b/runners/prism/java/src/main/java/org/apache/beam/runners/prism/PrismExecutor.java index 0e13b526b0fc..77313a09d558 100644 --- a/runners/prism/java/src/main/java/org/apache/beam/runners/prism/PrismExecutor.java +++ b/runners/prism/java/src/main/java/org/apache/beam/runners/prism/PrismExecutor.java @@ -72,26 +72,33 @@ static Builder builder() { /** Stops the execution of the {@link Process}, created as a result of {@link #execute}. */ void stop() { LOG.info("Stopping Prism..."); - if (future != null) { - future.cancel(true); - } - executorService.shutdown(); - try { - boolean ignored = executorService.awaitTermination(5000L, TimeUnit.MILLISECONDS); - } catch (InterruptedException ignored) { - // Ignore InterruptedException during shutdown. - } - if (process == null) { - return; - } - if (!process.isAlive()) { - return; - } - process.destroy(); + boolean interrupted = false; try { - process.waitFor(); - } catch (InterruptedException ignored) { - // Ignore InterruptedException during shutdown. + if (future != null) { + future.cancel(true); + } + executorService.shutdown(); + try { + boolean ignored = executorService.awaitTermination(5000L, TimeUnit.MILLISECONDS); + } catch (InterruptedException ignored) { + interrupted = true; + } + if (process == null) { + return; + } + if (!process.isAlive()) { + return; + } + process.destroy(); + try { + process.waitFor(); + } catch (InterruptedException ignored) { + interrupted = true; + } + } finally { + if (interrupted) { + Thread.currentThread().interrupt(); + } } } diff --git a/runners/prism/java/src/main/java/org/apache/beam/runners/prism/PrismRunner.java b/runners/prism/java/src/main/java/org/apache/beam/runners/prism/PrismRunner.java index ac1e68237faf..037060bfdea0 100644 --- a/runners/prism/java/src/main/java/org/apache/beam/runners/prism/PrismRunner.java +++ b/runners/prism/java/src/main/java/org/apache/beam/runners/prism/PrismRunner.java @@ -18,18 +18,28 @@ package org.apache.beam.runners.prism; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; import java.io.IOException; +import java.io.InterruptedIOException; +import java.net.InetSocketAddress; import java.net.ServerSocket; +import java.net.Socket; import java.util.Arrays; +import java.util.concurrent.TimeUnit; +import org.apache.beam.model.jobmanagement.v1.JobApi.GetJobsRequest; +import org.apache.beam.model.jobmanagement.v1.JobServiceGrpc; +import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor; import org.apache.beam.runners.portability.PortableRunner; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.PipelineRunner; +import org.apache.beam.sdk.fn.channel.ManagedChannelFactory; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PortablePipelineOptions; import org.apache.beam.sdk.util.construction.Environments; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ManagedChannel; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Status; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.StatusRuntimeException; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -46,6 +56,7 @@ public class PrismRunner extends PipelineRunner { private static final Logger LOG = LoggerFactory.getLogger(PrismRunner.class); + private static final int MAX_START_ATTEMPTS = 3; private final PrismPipelineOptions prismPipelineOptions; @@ -92,9 +103,9 @@ public PipelineResult run(Pipeline pipeline) { } PrismExecutor startPrism() throws IOException { + checkInterrupted(); PrismLocator locator = new PrismLocator(prismPipelineOptions); - int port = findAvailablePort(); - String portFlag = String.format(PrismExecutor.JOB_PORT_FLAG_TEMPLATE, port); + String command = locator.resolve(); String serveHttpFlag = String.format( PrismExecutor.SERVE_HTTP_FLAG_TEMPLATE, prismPipelineOptions.getEnableWebUI()); @@ -104,18 +115,92 @@ PrismExecutor startPrism() throws IOException { String logLevelFlag = String.format( PrismExecutor.LOG_LEVEL_FLAG_TEMPLATE, prismPipelineOptions.getPrismLogLevel()); - String endpoint = "localhost:" + port; - prismPipelineOptions.setJobEndpoint(endpoint); - String command = locator.resolve(); - PrismExecutor executor = - PrismExecutor.builder() - .setCommand(command) - .setArguments( - Arrays.asList(portFlag, serveHttpFlag, idleShutdownTimeoutFlag, logLevelFlag)) - .build(); - executor.execute(); - checkState(executor.isAlive()); - return executor; + for (int attempt = 1; attempt <= MAX_START_ATTEMPTS; attempt++) { + checkInterrupted(); + int port = findAvailablePort(); + String endpoint = "localhost:" + port; + prismPipelineOptions.setJobEndpoint(endpoint); + PrismExecutor executor = + PrismExecutor.builder() + .setCommand(command) + .setArguments( + Arrays.asList( + String.format(PrismExecutor.JOB_PORT_FLAG_TEMPLATE, port), + serveHttpFlag, + idleShutdownTimeoutFlag, + logLevelFlag)) + .build(); + boolean ready = false; + try { + executor.execute(); + ready = waitForJobService(executor, port, endpoint); + if (ready) { + return executor; + } + } finally { + if (!ready) { + executor.stop(); + } + } + LOG.warn("Prism did not start its JobService at {} (attempt {})", endpoint, attempt); + } + throw new IOException( + "Prism failed to start its JobService after " + MAX_START_ATTEMPTS + " attempts"); + } + + private boolean waitForJobService(PrismExecutor executor, int port, String endpoint) + throws IOException { + long deadline = + System.nanoTime() + TimeUnit.SECONDS.toNanos(prismPipelineOptions.getJobServerTimeout()); + ManagedChannel channel = + ManagedChannelFactory.createDefault() + .forDescriptor(ApiServiceDescriptor.newBuilder().setUrl(endpoint).build()); + try { + while (executor.isAlive() && System.nanoTime() < deadline) { + checkInterrupted(); + // Wait for the listener before connecting gRPC to avoid its reconnect backoff on startup. + try (Socket socket = new Socket()) { + socket.connect( + new InetSocketAddress("localhost", port), + (int) + Math.max( + 1, + Math.min(100, TimeUnit.NANOSECONDS.toMillis(deadline - System.nanoTime())))); + long remaining = deadline - System.nanoTime(); + if (remaining <= 0) { + return false; + } + JobServiceGrpc.newBlockingStub(channel) + .withDeadlineAfter( + Math.min(remaining, TimeUnit.SECONDS.toNanos(1)), TimeUnit.NANOSECONDS) + .getJobs(GetJobsRequest.getDefaultInstance()); + return executor.isAlive(); + } catch (StatusRuntimeException e) { + checkInterrupted(); + if (e.getStatus().getCode() != Status.Code.UNAVAILABLE + && e.getStatus().getCode() != Status.Code.DEADLINE_EXCEEDED) { + return false; + } + } catch (IOException ignored) { + // The child process has not bound its port yet. + } + try { + Thread.sleep(10); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new InterruptedIOException("Interrupted while waiting for Prism's JobService"); + } + } + return false; + } finally { + channel.shutdownNow(); + } + } + + private static void checkInterrupted() throws InterruptedIOException { + if (Thread.currentThread().isInterrupted()) { + throw new InterruptedIOException("Interrupted while starting Prism's JobService"); + } } private static void assignDefaultsIfNeeded(PrismPipelineOptions prismPipelineOptions) { @@ -124,7 +209,7 @@ private static void assignDefaultsIfNeeded(PrismPipelineOptions prismPipelineOpt } } - private static int findAvailablePort() throws IOException { + int findAvailablePort() throws IOException { try (ServerSocket socket = new ServerSocket(0)) { return socket.getLocalPort(); } diff --git a/runners/prism/java/src/test/java/org/apache/beam/runners/prism/PrismExecutorTest.java b/runners/prism/java/src/test/java/org/apache/beam/runners/prism/PrismExecutorTest.java index a81e3e24ee69..a780f77fcfa1 100644 --- a/runners/prism/java/src/test/java/org/apache/beam/runners/prism/PrismExecutorTest.java +++ b/runners/prism/java/src/test/java/org/apache/beam/runners/prism/PrismExecutorTest.java @@ -20,6 +20,10 @@ import static com.google.common.truth.Truth.assertThat; import static org.apache.beam.runners.prism.PrismPipelineOptions.JOB_PORT_FLAG_NAME; import static org.apache.beam.runners.prism.PrismRunnerTest.getLocalPrismBuildOrIgnoreTest; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; import java.io.ByteArrayOutputStream; import java.io.File; @@ -27,6 +31,7 @@ import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.util.Collections; +import java.util.concurrent.ExecutorService; import java.util.stream.Collectors; import java.util.stream.Stream; import org.junit.Rule; @@ -43,6 +48,36 @@ public class PrismExecutorTest { @Rule public TestName testName = new TestName(); + @Test + public void stopPreservesInterruptionWhileWaitingForExecutor() throws InterruptedException { + PrismExecutor executor = PrismExecutor.builder().setCommand("unused").build(); + ExecutorService executorService = mock(ExecutorService.class); + when(executorService.awaitTermination(anyLong(), any())).thenThrow(new InterruptedException()); + executor.executorService.shutdown(); + executor.executorService = executorService; + try { + executor.stop(); + assertThat(Thread.currentThread().isInterrupted()).isTrue(); + } finally { + Thread.interrupted(); + } + } + + @Test + public void stopPreservesInterruptionWhileWaitingForProcess() throws InterruptedException { + PrismExecutor executor = PrismExecutor.builder().setCommand("unused").build(); + Process process = mock(Process.class); + when(process.isAlive()).thenReturn(true); + when(process.waitFor()).thenThrow(new InterruptedException()); + executor.process = process; + try { + executor.stop(); + assertThat(Thread.currentThread().isInterrupted()).isTrue(); + } finally { + Thread.interrupted(); + } + } + @Test public void executeThenStop() throws IOException { PrismExecutor executor = underTest().build(); diff --git a/runners/prism/java/src/test/java/org/apache/beam/runners/prism/PrismRunnerTest.java b/runners/prism/java/src/test/java/org/apache/beam/runners/prism/PrismRunnerTest.java index e4d239275988..ff93f7fcbac6 100644 --- a/runners/prism/java/src/test/java/org/apache/beam/runners/prism/PrismRunnerTest.java +++ b/runners/prism/java/src/test/java/org/apache/beam/runners/prism/PrismRunnerTest.java @@ -21,9 +21,19 @@ import static org.junit.Assert.assertThrows; import static org.junit.Assume.assumeTrue; +import java.io.IOException; +import java.io.InterruptedIOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; +import org.apache.beam.model.jobmanagement.v1.JobApi.GetJobsRequest; +import org.apache.beam.model.jobmanagement.v1.JobServiceGrpc; +import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor; +import org.apache.beam.sdk.fn.channel.ManagedChannelFactory; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.testing.PAssert; import org.apache.beam.sdk.testing.TestPipeline; @@ -36,11 +46,15 @@ import org.apache.beam.sdk.util.construction.Environments; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ManagedChannel; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.Server; +import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.ServerBuilder; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; import org.joda.time.Duration; import org.joda.time.Instant; import org.junit.Rule; import org.junit.Test; +import org.junit.rules.TemporaryFolder; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -49,10 +63,104 @@ public class PrismRunnerTest { @Rule public TestPipeline pipeline = TestPipeline.fromOptions(options()); + @Rule public TemporaryFolder temporaryFolder = new TemporaryFolder(); // See build.gradle for test task configuration. private static final String PRISM_BUILD_TARGET_PROPERTY_NAME = "prism.buildTarget"; + @Test + public void startPrismRetriesAfterPortConflict() throws Exception { + Server otherService = ServerBuilder.forPort(0).build().start(); + TestPrismPipelineOptions options = options(); + options.setJobServerTimeout(5); + AtomicBoolean firstAttempt = new AtomicBoolean(true); + PrismRunner runner = + new PrismRunner(options) { + @Override + int findAvailablePort() throws IOException { + return firstAttempt.getAndSet(false) + ? otherService.getPort() + : super.findAvailablePort(); + } + }; + try { + PrismExecutor executor = runner.startPrism(); + try { + assertThat(options.getJobEndpoint()).isNotEqualTo("localhost:" + otherService.getPort()); + ManagedChannel channel = + ManagedChannelFactory.createDefault() + .forDescriptor( + ApiServiceDescriptor.newBuilder().setUrl(options.getJobEndpoint()).build()); + try { + JobServiceGrpc.newBlockingStub(channel) + .withDeadlineAfter(5, java.util.concurrent.TimeUnit.SECONDS) + .getJobs(GetJobsRequest.getDefaultInstance()); + } finally { + channel.shutdownNow(); + } + } finally { + executor.stop(); + } + } finally { + otherService.shutdownNow(); + } + } + + @Test + public void startPrismRejectsLiveProcessWithoutJobService() throws Exception { + assumeTrue("requires a POSIX shell", Files.isExecutable(Path.of("/bin/sh"))); + Path command = temporaryFolder.newFile("not-prism").toPath(); + Path pids = Path.of(command + ".pids"); + Files.write( + command, + Arrays.asList("#!/bin/sh", "echo $$ >> \"$0.pids\"", "exec sleep 30"), + StandardCharsets.UTF_8); + assertThat(command.toFile().setExecutable(true)).isTrue(); + TestPrismPipelineOptions options = options(); + options.setPrismLocation(command.toString()); + options.setJobServerTimeout(1); + Server otherService = ServerBuilder.forPort(0).build().start(); + AtomicBoolean firstAttempt = new AtomicBoolean(true); + PrismRunner runner = + new PrismRunner(options) { + @Override + int findAvailablePort() throws IOException { + return firstAttempt.getAndSet(false) + ? otherService.getPort() + : super.findAvailablePort(); + } + }; + try { + assertThrows(IOException.class, () -> runner.startPrism().stop()); + if (Files.exists(pids)) { + for (String pid : Files.readAllLines(pids)) { + assertThat( + ProcessHandle.of(Long.parseLong(pid)).map(ProcessHandle::isAlive).orElse(false)) + .isFalse(); + } + } + } finally { + otherService.shutdownNow(); + if (Files.exists(pids)) { + for (String pid : Files.readAllLines(pids)) { + ProcessHandle.of(Long.parseLong(pid)).ifPresent(ProcessHandle::destroyForcibly); + } + } + } + } + + @Test + public void startPrismPreservesInterruption() { + PrismRunner runner = new PrismRunner(options()); + Thread.currentThread().interrupt(); + try { + assertThrows(InterruptedIOException.class, runner::startPrism); + assertThat(Thread.currentThread().isInterrupted()).isTrue(); + } finally { + Thread.interrupted(); + } + } + @Test public void givenJobEndpointSet_TestPrismRunner_validateThrows() { TestPrismPipelineOptions options =