Skip to content

[Prism] Wait for the job service before starting loopback workers #40773

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
bvolpato wants to merge 1 commit into
base: master
Choose a base branch
from
+273 −35
Open

[Prism] Wait for the job service before starting loopback workers #40773

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
3 changes: 3 additions & 0 deletions runners/prism/java/build.gradle
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
45 changes: 26 additions & 19 deletions runners/prism/java/src/main/java/org/apache/beam/runners/prism/PrismExecutor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
}
}

Expand Down
117 changes: 101 additions & 16 deletions runners/prism/java/src/main/java/org/apache/beam/runners/prism/PrismRunner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -46,6 +56,7 @@
public class PrismRunner extends PipelineRunner<PipelineResult> {

private static final Logger LOG = LoggerFactory.getLogger(PrismRunner.class);
private static final int MAX_START_ATTEMPTS = 3;

private final PrismPipelineOptions prismPipelineOptions;

Expand Down Expand Up @@ -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());
Expand All @@ -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) {
Expand All @@ -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();
}
Expand Down
35 changes: 35 additions & 0 deletions runners/prism/java/src/test/java/org/apache/beam/runners/prism/PrismExecutorTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,13 +20,18 @@
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;
import java.io.IOException;
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;
Expand All @@ -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();
Expand Down
Loading
Loading