diff --git a/engine/src/main/java/com/odysseusinc/arachne/executionengine/execution/r/DockerOverseer.java b/engine/src/main/java/com/odysseusinc/arachne/executionengine/execution/r/DockerOverseer.java index 91102b4e..dffe7e78 100644 --- a/engine/src/main/java/com/odysseusinc/arachne/executionengine/execution/r/DockerOverseer.java +++ b/engine/src/main/java/com/odysseusinc/arachne/executionengine/execution/r/DockerOverseer.java @@ -4,7 +4,6 @@ import com.github.dockerjava.api.DockerClient; import com.github.dockerjava.api.async.ResultCallback; -import com.github.dockerjava.api.command.LogContainerCmd; import com.github.dockerjava.api.command.WaitContainerResultCallback; import com.github.dockerjava.api.exception.DockerException; import com.github.dockerjava.api.exception.NotFoundException; @@ -27,7 +26,10 @@ @Slf4j @Getter public class DockerOverseer extends AbstractOverseer { - private final ScheduledExecutorService executor = new ScheduledThreadPoolExecutor(1) {{ + /** + * Two threads: one blocks on awaitStatusCode for the whole run, the other runs the periodic log flush. + */ + private final ScheduledExecutorService executor = new ScheduledThreadPoolExecutor(2) {{ setRemoveOnCancelPolicy(true); }}; private final CompletableFuture init; @@ -37,74 +39,91 @@ public class DockerOverseer extends AbstractOverseer { */ private volatile int pos; private final DockerClient client; + private volatile boolean aborting = false; public DockerOverseer( long id, DockerClient client, Instant started, int timeoutSec, StringBuffer stdout, CompletableFuture init, int updateInterval, BiConsumer callback, String image, int killTimeoutSec ) { - super(id, callback, started, image, killTimeoutSec, stdout, init.handle((containerId, throwable) -> { - if (throwable != null) { - String out = stdout.append("\r\n").append(ExceptionUtils.getStackTrace(throwable)).toString(); - return new ExecutionOutcome(Stage.INITIALIZE, throwable.getMessage(), out); - } else { - LogContainerCmd cmd = client.logContainerCmd(containerId).withStdOut(true).withStdErr(true).withFollowStream(true); - cmd.exec(logAdapter(id, stdout)); - Integer exitCode = client.waitContainerCmd(containerId).exec(new WaitContainerResultCallback()).awaitStatusCode(timeoutSec, TimeUnit.SECONDS); - log.info("Execution [{}] Rscript exit code {}", id, exitCode); - String out = stdout.toString(); - return (exitCode == 0) - ? new ExecutionOutcome(Stage.COMPLETED, null, out) - : new ExecutionOutcome(Stage.EXECUTE, "Exit code " + exitCode, out); - } - })); + super(id, callback, started, image, killTimeoutSec, stdout, new CompletableFuture<>()); pos = stdout.length(); this.client = client; - init.thenAccept(containerId -> - executor.scheduleWithFixedDelay(this::writeLogs, updateInterval, updateInterval, TimeUnit.MILLISECONDS) - ); this.init = init; - } - @Override - public CompletableFuture abort() { - if (!init.isDone()) { - init.cancel(true); - } + // Shut the executor down exactly once, when the run ends (no leak). + outcome.whenComplete((r, t) -> executor.shutdown()); - return init.handle((containerId, throwable) -> { - if (containerId == null) { - if (throwable != null) { - log.error("Error during initialization: {}", throwable.getMessage(), throwable); - stdout.append("\r\n\"Initialization failed: ").append(throwable.getMessage()); - return new ExecutionOutcome(Stage.INITIALIZE, "Initialization failed: " + throwable.getMessage(), stdout.toString()); - } else { - stdout.append("\r\nContainer initialization was canceled"); - return new ExecutionOutcome(Stage.ABORTED, null, stdout.toString()); - } + // Only REGISTER the container-startup handler here; it never blocks, so the shared pool thread is freed instantly. + init.whenComplete((containerId, throwable) -> { + if (throwable != null) { + onInitFailed(throwable); } else { - if (result.isDone()) { - return result.join(); - } else { - try { - client.stopContainerCmd(containerId).exec(); - stdout.append("\r\nDocker container aborted successfully"); - return new ExecutionOutcome(Stage.ABORTED, null, stdout.toString()); - } catch (NotFoundException e) { - log.info("Container not found or already stopped: " + e.getMessage()); - return result.join(); - } catch (DockerException e) { - log.error("Error stopping the Docker container: {}", e.getMessage()); - stdout.append("\r\n\"Error aborting Docker container:").append(e.getMessage()); - return new ExecutionOutcome(Stage.ABORT, "Error aborting Docker container: " + e.getMessage(), stdout.toString()); - } - } + onContainerStarted(containerId, timeoutSec, updateInterval); } }); + } + + private void onInitFailed(Throwable throwable) { + if (aborting) { + outcome.complete(new ExecutionOutcome(Stage.ABORTED, null, stdout.toString())); + } else { + String out = stdout.append("\r\n").append(ExceptionUtils.getStackTrace(throwable)).toString(); + outcome.complete(new ExecutionOutcome(Stage.INITIALIZE, throwable.getMessage(), out)); + } + } + + private void onContainerStarted(String containerId, int timeoutSec, int updateInterval) { + // Flush logs to the callback periodically (one pool thread). + executor.scheduleWithFixedDelay(this::writeLogs, updateInterval, updateInterval, TimeUnit.MILLISECONDS); + // Block on the wait on our OWN pool (another thread), not on the shared analysisTaskExecutor. + executor.execute(() -> awaitAndComplete(containerId, timeoutSec)); + } + + private void awaitAndComplete(String containerId, int timeoutSec) { + try { + client.logContainerCmd(containerId).withStdOut(true).withStdErr(true).withFollowStream(true).exec(logAdapter(id, stdout)); + Integer exitCode = client.waitContainerCmd(containerId).exec(new WaitContainerResultCallback()).awaitStatusCode(timeoutSec, TimeUnit.SECONDS); + completeFromExit(exitCode); + } catch (Exception e) { + log.error("Execution [{}] error waiting for container", id, e); + writeLogs(); + outcome.complete(aborting + ? new ExecutionOutcome(Stage.ABORTED, null, stdout.toString()) + : new ExecutionOutcome(Stage.EXECUTE, e.getMessage(), stdout.toString())); + } + } + private void completeFromExit(int exitCode) { + writeLogs(); + log.info("Execution [{}] Rscript exit code {}", id, exitCode); + String out = stdout.toString(); + outcome.complete(exitCode == 0 + ? new ExecutionOutcome(Stage.COMPLETED, null, out) + : aborting + ? new ExecutionOutcome(Stage.ABORTED, null, out) + : new ExecutionOutcome(Stage.EXECUTE, "Exit code " + exitCode, out)); + } + + @Override + public CompletableFuture abort() { + aborting = true; + init.thenAccept(this::stopContainer); + return CompletableFuture.completedFuture(new ExecutionOutcome(Stage.ABORT, null, stdout.toString())); + } + + private void stopContainer(String containerId) { + try { + client.stopContainerCmd(containerId).withTimeout(0).exec(); + log.info("Execution [{}] stop command sent to Docker container", id); + } catch (NotFoundException e) { + log.info("Execution [{}] container not found or already stopped: {}", id, e.getMessage()); + } catch (DockerException e) { + log.error("Execution [{}] error stopping Docker container: {}", id, e.getMessage()); + } } private static ResultCallback.Adapter logAdapter(long id, StringBuffer stdout) { - return new ResultCallback.Adapter() { + return new ResultCallback.Adapter<>() { @Override public void onNext(Frame item) { super.onNext(item); @@ -116,7 +135,7 @@ public void onNext(Frame item) { @Override public void onError(Throwable throwable) { if (!(throwable instanceof NotFoundException)) { - log.error("Execution [{}] error: {}", id, throwable); + log.error("Execution [{}] error: {}", id, throwable.getMessage(), throwable); stdout.append("Execution error: ").append(throwable.getMessage()); super.onError(throwable); } diff --git a/engine/src/main/java/com/odysseusinc/arachne/executionengine/execution/r/TarballROverseer.java b/engine/src/main/java/com/odysseusinc/arachne/executionengine/execution/r/TarballROverseer.java index 95ca0f2e..705a7035 100644 --- a/engine/src/main/java/com/odysseusinc/arachne/executionengine/execution/r/TarballROverseer.java +++ b/engine/src/main/java/com/odysseusinc/arachne/executionengine/execution/r/TarballROverseer.java @@ -25,6 +25,7 @@ public class TarballROverseer extends AbstractOverseer { private final Process process; private final BufferedReader reader; private final ScheduledFuture logFlush; + private volatile boolean aborting = false; /** * Creates a new process overseer. @@ -62,12 +63,14 @@ public AnalysisRequestTypeDTO getType() { @Override public CompletableFuture abort() { + aborting = true; if (process.isAlive()) { log.info("Overseer [{}] processing abort request", id); if (terminate()) { outcome.complete(new ExecutionOutcome(Stage.ABORTED, null, stdout.toString())); } else { callback.accept(Stage.ABORT, "Timed out waiting for termination"); + outcome.complete(new ExecutionOutcome(Stage.ABORT, "Timed out waiting for termination", stdout.toString())); } } else { log.info("Overseer [{}] received abort, but process exited already", id); @@ -108,10 +111,12 @@ private void writeLogs(String stage) { private void complete(int exitValue) { executor.shutdown(); - ExecutionOutcome outcome = (exitValue == 0) + ExecutionOutcome executionOutcome = exitValue == 0 ? new ExecutionOutcome(Stage.COMPLETED, null, stdout.toString()) - : new ExecutionOutcome(Stage.EXECUTE, "Exit code " + exitValue, stdout.toString()); - this.outcome.complete(outcome); + : aborting + ? new ExecutionOutcome(Stage.ABORTED, null, stdout.toString()) + : new ExecutionOutcome(Stage.EXECUTE, "Exit code " + exitValue, stdout.toString()); + this.outcome.complete(executionOutcome); } private boolean terminate() { diff --git a/engine/src/test/java/com/odysseusinc/arachne/executionengine/execution/r/DockerOverseerIntegrationTest.java b/engine/src/test/java/com/odysseusinc/arachne/executionengine/execution/r/DockerOverseerIntegrationTest.java new file mode 100644 index 00000000..fe457958 --- /dev/null +++ b/engine/src/test/java/com/odysseusinc/arachne/executionengine/execution/r/DockerOverseerIntegrationTest.java @@ -0,0 +1,129 @@ +package com.odysseusinc.arachne.executionengine.execution.r; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.github.dockerjava.api.DockerClient; +import com.github.dockerjava.api.command.CreateContainerResponse; +import com.github.dockerjava.api.command.PullImageResultCallback; +import com.github.dockerjava.api.exception.NotFoundException; +import com.github.dockerjava.api.model.HostConfig; +import com.github.dockerjava.core.DefaultDockerClientConfig; +import com.github.dockerjava.core.DockerClientImpl; +import com.github.dockerjava.httpclient5.ApacheDockerHttpClient; +import com.github.dockerjava.transport.DockerHttpClient; +import com.odysseusinc.arachne.execution_engine_common.api.v1.dto.ExecutionOutcome; +import com.odysseusinc.arachne.execution_engine_common.api.v1.dto.Stage; +import java.time.Instant; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assumptions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +/** + * Integration test for {@link DockerOverseer}. Requires a running Docker daemon; the whole class is + * skipped (assumption aborted) when Docker is not reachable. Each test starts a real short-lived + * container running actual R code and asserts the resulting {@link ExecutionOutcome} status. + */ +class DockerOverseerIntegrationTest { + + private static final String IMAGE = "r-base:4.4.1"; + private static final int TIMEOUT_SEC = 120; + private static final int UPDATE_INTERVAL_MS = 500; + private static final int KILL_TIMEOUT_SEC = 10; + + private static DockerClient client; + private final List startedContainers = new ArrayList<>(); + + @BeforeAll + static void setUp() { + boolean available; + try { + DefaultDockerClientConfig config = DefaultDockerClientConfig.createDefaultConfigBuilder().build(); + DockerHttpClient httpClient = new ApacheDockerHttpClient.Builder() + .dockerHost(config.getDockerHost()) + .sslConfig(config.getSSLConfig()) + .maxConnections(10) + .build(); + client = DockerClientImpl.getInstance(config, httpClient); + client.pingCmd().exec(); + available = true; + } catch (Throwable t) { + available = false; + } + Assumptions.assumeTrue(available, "Docker daemon is not available, skipping DockerOverseer integration test"); + ensureImage(); + } + + private static void ensureImage() { + try { + client.inspectImageCmd(IMAGE).exec(); + } catch (NotFoundException notPresent) { + try { + client.pullImageCmd(IMAGE).exec(new PullImageResultCallback()).awaitCompletion(3, TimeUnit.MINUTES); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + Assumptions.assumeTrue(false, "Interrupted while pulling " + IMAGE); + } + } + } + + @AfterEach + void cleanUp() { + for (String id : startedContainers) { + try { + client.removeContainerCmd(id).withForce(true).exec(); + } catch (RuntimeException alreadyGone) { + // container used --rm and was auto-removed on exit, or never created + } + } + startedContainers.clear(); + } + + @Test + void completedOnSuccess() throws Exception { + ExecutionOutcome outcome = runR("x <- 1:10; cat('sum:', sum(x), '\\n')").getResult().get(60, TimeUnit.SECONDS); + + assertEquals(Stage.COMPLETED, outcome.getStage()); + assertNull(outcome.getError()); + assertTrue(outcome.getStdout().contains("sum: 55"), "stdout should contain R computation result"); + } + + @Test + void executeStageOnNonZeroExit() throws Exception { + ExecutionOutcome outcome = runR("quit(status=3, save='no')").getResult().get(60, TimeUnit.SECONDS); + + assertEquals(Stage.EXECUTE, outcome.getStage()); + assertEquals("Exit code 3", outcome.getError()); + } + + @Test + void abortedOnCancel() throws Exception { + DockerOverseer overseer = runR("Sys.sleep(120)"); + + ExecutionOutcome ack = overseer.abort().get(5, TimeUnit.SECONDS); + assertEquals(Stage.ABORT, ack.getStage(), "abort() should acknowledge immediately with ABORT"); + + ExecutionOutcome outcome = overseer.getResult().get(30, TimeUnit.SECONDS); + assertEquals(Stage.ABORTED, outcome.getStage()); + } + + private DockerOverseer runR(String rExpr) { + CreateContainerResponse container = client.createContainerCmd(IMAGE) + .withHostConfig(HostConfig.newHostConfig().withAutoRemove(true)) + .withCmd("Rscript", "-e", rExpr) + .exec(); + startedContainers.add(container.getId()); + client.startContainerCmd(container.getId()).exec(); + return new DockerOverseer( + 1L, client, Instant.now(), TIMEOUT_SEC, new StringBuffer(), + CompletableFuture.completedFuture(container.getId()), UPDATE_INTERVAL_MS, + (stage, log) -> { }, IMAGE, KILL_TIMEOUT_SEC + ); + } +}