feat: Implement pre-initialized Docker container pool to improve /eval

JShellAPI/src/main/java/org/togetherjava/jshellapi/dto/ContainerState.java

@@ -0,0 +1,15 @@
+package org.togetherjava.jshellapi.dto;
+
+import java.io.BufferedReader;
+import java.io.BufferedWriter;
+
+/**
+ * Data record for the state of a container.
+ * 
+ * @param containerId The id of the container.
+ * @param containerOutput The output of the container.
+ * @param containerInput The input of the container.
+ */
+public record ContainerState(String containerId, BufferedReader containerOutput,
+        BufferedWriter containerInput) {
+}

JShellAPI/src/main/java/org/togetherjava/jshellapi/service/DockerService.java

@@ -2,6 +2,7 @@
 
 import com.github.dockerjava.api.DockerClient;
 import com.github.dockerjava.api.async.ResultCallback;
+import com.github.dockerjava.api.command.InspectContainerResponse;
 import com.github.dockerjava.api.command.PullImageResultCallback;
 import com.github.dockerjava.api.model.*;
 import com.github.dockerjava.core.DefaultDockerClientConfig;
@@ -14,12 +15,13 @@
 import org.springframework.stereotype.Service;
 
 import org.togetherjava.jshellapi.Config;
+import org.togetherjava.jshellapi.dto.ContainerState;
 
 import java.io.*;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.*;
-import java.util.concurrent.TimeUnit;
+import java.util.concurrent.*;
 
 @Service
 public class DockerService implements DisposableBean {
@@ -28,10 +30,17 @@ public class DockerService implements DisposableBean {
     private static final UUID WORKER_UNIQUE_ID = UUID.randomUUID();
 
     private final DockerClient client;
+    private final Config config;
+    private final ExecutorService executor = Executors.newSingleThreadExecutor();
+    private final ConcurrentHashMap<StartupScriptId, ContainerState> cachedContainers =
+            new ConcurrentHashMap<>();
+    private final StartupScriptsService startupScriptsService;
 
     private final String jshellWrapperBaseImageName;
 
-    public DockerService(Config config) {
+    public DockerService(Config config, StartupScriptsService startupScriptsService)
+            throws InterruptedException {
+        this.startupScriptsService = startupScriptsService;
         DefaultDockerClientConfig clientConfig =
                 DefaultDockerClientConfig.createDefaultConfigBuilder().build();
         ApacheDockerHttpClient httpClient =
@@ -41,11 +50,16 @@ public DockerService(Config config) {
                     .connectionTimeout(Duration.ofSeconds(config.dockerConnectionTimeout()))
                     .build();
         this.client = DockerClientImpl.getInstance(clientConfig, httpClient);
+        this.config = config;
 
         this.jshellWrapperBaseImageName =
                 config.jshellWrapperImageName().split(Config.JSHELL_WRAPPER_IMAGE_NAME_TAG)[0];
 
+        if (!isImagePresentLocally()) {
+            pullImage();
+        }
         cleanupLeftovers(WORKER_UNIQUE_ID);
+        executor.submit(() -> initializeCachedContainer(StartupScriptId.CUSTOM_DEFAULT));
     }
 
     private void cleanupLeftovers(UUID currentId) {
@@ -62,48 +76,162 @@ private void cleanupLeftovers(UUID currentId) {
         }
     }
 
-    public String spawnContainer(long maxMemoryMegs, long cpus, @Nullable String cpuSetCpus,
-            String name, Duration evalTimeout, long sysoutLimit) throws InterruptedException {
-
-        boolean presentLocally = client.listImagesCmd()
+    /**
+     * Checks if the Docker image with the given name and tag is present locally.
+     *
+     * @return true if the image is present, false otherwise.
+     */
+    private boolean isImagePresentLocally() {
+        return client.listImagesCmd()
             .withFilter("reference", List.of(jshellWrapperBaseImageName))
             .exec()
             .stream()
             .flatMap(it -> Arrays.stream(it.getRepoTags()))
             .anyMatch(it -> it.endsWith(Config.JSHELL_WRAPPER_IMAGE_NAME_TAG));
+    }
 
-        if (!presentLocally) {
-            client.pullImageCmd(jshellWrapperBaseImageName)
-                .withTag("master")
-                .exec(new PullImageResultCallback())
-                .awaitCompletion(5, TimeUnit.MINUTES);
-        }
+    /**
+     * Pulls the Docker image.
+     */
+    private void pullImage() throws InterruptedException {
+        client.pullImageCmd(jshellWrapperBaseImageName)
+            .withTag("master")
+            .exec(new PullImageResultCallback())
+            .awaitCompletion(5, TimeUnit.MINUTES);
+    }
+
+    /**
+     * Creates a Docker container with the given name.
+     *
+     * @param name The name of the container to create.
+     * @return The ID of the created container.
+     */
+    private String createContainer(String name) {
+        LOGGER.debug("Creating container '{}'", name);
+        HostConfig hostConfig = HostConfig.newHostConfig()
+            .withAutoRemove(true)
+            .withInit(true)
+            .withCapDrop(Capability.ALL)
+            .withNetworkMode("none")
+            .withPidsLimit(2000L)
+            .withReadonlyRootfs(true)
+            .withMemory((long) config.dockerMaxRamMegaBytes() * 1024 * 1024)
+            .withCpuCount((long) Math.ceil(config.dockerCPUsUsage()))
+            .withCpusetCpus(config.dockerCPUSetCPUs());
 
         return client
             .createContainerCmd(jshellWrapperBaseImageName + Config.JSHELL_WRAPPER_IMAGE_NAME_TAG)
-            .withHostConfig(HostConfig.newHostConfig()
-                .withAutoRemove(true)
-                .withInit(true)
-                .withCapDrop(Capability.ALL)
-                .withNetworkMode("none")
-                .withPidsLimit(2000L)
-                .withReadonlyRootfs(true)
-                .withMemory(maxMemoryMegs * 1024 * 1024)
-                .withCpuCount(cpus)
-                .withCpusetCpus(cpuSetCpus))
+            .withHostConfig(hostConfig)
             .withStdinOpen(true)
             .withAttachStdin(true)
             .withAttachStderr(true)
             .withAttachStdout(true)
-            .withEnv("evalTimeoutSeconds=" + evalTimeout.toSeconds(),
-                    "sysOutCharLimit=" + sysoutLimit)
+            .withEnv("evalTimeoutSeconds=" + config.evalTimeoutSeconds(),
+                    "sysOutCharLimit=" + config.sysOutCharLimit())
             .withLabels(Map.of(WORKER_LABEL, WORKER_UNIQUE_ID.toString()))
             .withName(name)
             .exec()
             .getId();
     }
 
-    public InputStream startAndAttachToContainer(String containerId, InputStream stdin)
+    /**
+     * Spawns a new Docker container with specified configurations.
+     *
+     * @param name Name of the container.
+     * @param startupScriptId Script to initialize the container with.
+     * @throws IOException if an I/O error occurs.
+     * @return The ContainerState of the newly created container.
+     */
+    public ContainerState initializeContainer(String name,
+            @Nullable StartupScriptId startupScriptId) throws IOException {
+        LOGGER.info("Initializing container '{}' with Startup script ID: {}", name,
+                startupScriptId);
+        if (startupScriptId == null || cachedContainers.isEmpty()
+                || !cachedContainers.containsKey(startupScriptId)) {
+            String containerId = createContainer(name);
+            return setupContainerWithScript(containerId, startupScriptId);
+        }
+        ContainerState containerState = cachedContainers.get(startupScriptId);
+        executor.submit(() -> initializeCachedContainer(startupScriptId));
+
+        client.renameContainerCmd(containerState.containerId()).withName(name).exec();
+        return containerState;
+    }
+
+    /**
+     * Initializes a new cached docker container with specified configurations.
+     *
+     * @param startupScriptId Script to initialize the container with.
+     */
+    private void initializeCachedContainer(StartupScriptId startupScriptId) {
+        LOGGER.info("Initializing cached container with Startup script ID: {}", startupScriptId);
+        String containerName = newCachedContainerName();
+        String id = createContainer(containerName);
+        startContainer(id);
+
+        try {
+            ContainerState containerState = setupContainerWithScript(id, startupScriptId);
+            cachedContainers.put(startupScriptId, containerState);
+        } catch (IOException e) {
+            LOGGER.error("Could not initialize container {}", id, e);
+            killContainerByName(containerName);
+            throw new RuntimeException(e);
+        }
+    }
+
+    /**
+     * Setup container with startup script and also initializes input and output streams for the
+     * container.
+     *
+     * @param containerId The id of the container
+     * @param startupScriptId The startup script id of the session
+     * @return ContainerState of the spawned container.
+     * @throws IOException if an I/O error occurs
+     */
+    private ContainerState setupContainerWithScript(String containerId,
+            StartupScriptId startupScriptId) throws IOException {
+        LOGGER.info("Setting up container with id {} with Startup script ID: {}", containerId,
+                startupScriptId);
+        startContainer(containerId);
+        PipedInputStream containerInput = new PipedInputStream();
+        BufferedWriter writer =
+                new BufferedWriter(new OutputStreamWriter(new PipedOutputStream(containerInput)));
+
+        InputStream containerOutput = attachToContainer(containerId, containerInput);
+        BufferedReader reader = new BufferedReader(new InputStreamReader(containerOutput));
+
+        writer.write(Utils.sanitizeStartupScript(startupScriptsService.get(startupScriptId)));
+        writer.newLine();
+        writer.flush();
+
+        return new ContainerState(containerId, reader, writer);
+    }
+
+    /**
+     * Creates a new container
+     *
+     * @param containerId the ID of the container to start
+     */
+    private void startContainer(String containerId) {
+        boolean isRunning = isContainerRunning(containerId);
+        if (isRunning) {
+            LOGGER.debug("Container {} is already running.", containerId);
+            return;
+        }
+        LOGGER.debug("Container {} is not running. Starting it now.", containerId);
+        client.startContainerCmd(containerId).exec();
+    }
+
+    /**
+     * Attaches to a running Docker container's input (stdin) and output streams (stdout, stderr).
+     * Logs any output from stderr and returns an InputStream to read stdout.
+     *
+     * @param containerId The ID of the running container to attach to.
+     * @param containerInput The input stream (containerInput) to send to the container.
+     * @return InputStream to read the container's stdout
+     * @throws IOException if an I/O error occurs
+     */
+    private InputStream attachToContainer(String containerId, InputStream containerInput)
             throws IOException {
         PipedInputStream pipeIn = new PipedInputStream();
         PipedOutputStream pipeOut = new PipedOutputStream(pipeIn);
@@ -113,7 +241,7 @@ public InputStream startAndAttachToContainer(String containerId, InputStream std
             .withFollowStream(true)
             .withStdOut(true)
             .withStdErr(true)
-            .withStdIn(stdin)
+            .withStdIn(containerInput)
             .exec(new ResultCallback.Adapter<>() {
                 @Override
                 public void onNext(Frame object) {
@@ -131,11 +259,24 @@ public void onNext(Frame object) {
                     }
                 }
             });
-
-        client.startContainerCmd(containerId).exec();
         return pipeIn;
     }
 
+    /**
+     * Checks if the Docker container with the given ID is currently running.
+     *
+     * @param containerId the ID of the container to check
+     * @return true if the container is running, false otherwise
+     */
+    public boolean isContainerRunning(String containerId) {
+        InspectContainerResponse containerResponse = client.inspectContainerCmd(containerId).exec();
+        return Boolean.TRUE.equals(containerResponse.getState().getRunning());
+    }
+
+    private String newCachedContainerName() {
+        return "cached_session_" + UUID.randomUUID();
+    }
+
     public void killContainerByName(String name) {
         LOGGER.debug("Fetching container to kill {}.", name);
         List<Container> containers = client.listContainersCmd().withNameFilter(Set.of(name)).exec();
@@ -156,6 +297,7 @@ public boolean isDead(String containerName) {
     @Override
     public void destroy() throws Exception {
         LOGGER.info("destroy() called. Destroying all containers...");
+        executor.shutdown();
         cleanupLeftovers(UUID.randomUUID());
         client.close();
     }

JShellAPI/src/main/java/org/togetherjava/jshellapi/service/JShellService.java

@@ -2,13 +2,12 @@
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
-import org.springframework.lang.Nullable;
 
+import org.togetherjava.jshellapi.Config;
 import org.togetherjava.jshellapi.dto.*;
 import org.togetherjava.jshellapi.exceptions.DockerException;
 
 import java.io.*;
-import java.time.Duration;
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
@@ -31,33 +30,26 @@ public class JShellService {
     private final int startupScriptSize;
 
     public JShellService(DockerService dockerService, JShellSessionService sessionService,
-            String id, long timeout, boolean renewable, long evalTimeout,
-            long evalTimeoutValidationLeeway, int sysOutCharLimit, int maxMemory, double cpus,
-            @Nullable String cpuSetCpus, String startupScript) throws DockerException {
+            SessionInfo sessionInfo, Config config) throws DockerException {
         this.dockerService = dockerService;
         this.sessionService = sessionService;
-        this.id = id;
-        this.timeout = timeout;
-        this.renewable = renewable;
-        this.evalTimeout = evalTimeout;
-        this.evalTimeoutValidationLeeway = evalTimeoutValidationLeeway;
+        this.id = sessionInfo.id();
+        this.timeout = config.dockerConnectionTimeout();
+        this.renewable = sessionInfo.renewable();
+        this.evalTimeout = sessionInfo.evalTimeout();
+        this.evalTimeoutValidationLeeway = sessionInfo.evalTimeoutValidationLeeway();
         this.lastTimeoutUpdate = Instant.now();
+
         if (!dockerService.isDead(containerName())) {
             LOGGER.warn("Tried to create an existing container {}.", containerName());
             throw new DockerException("The session isn't completely destroyed, try again later.");
         }
+
         try {
-            String containerId = dockerService.spawnContainer(maxMemory, (long) Math.ceil(cpus),
-                    cpuSetCpus, containerName(), Duration.ofSeconds(evalTimeout), sysOutCharLimit);
-            PipedInputStream containerInput = new PipedInputStream();
-            this.writer = new BufferedWriter(
-                    new OutputStreamWriter(new PipedOutputStream(containerInput)));
-            InputStream containerOutput =
-                    dockerService.startAndAttachToContainer(containerId, containerInput);
-            reader = new BufferedReader(new InputStreamReader(containerOutput));
-            writer.write(sanitize(startupScript));
-            writer.newLine();
-            writer.flush();
+            ContainerState containerState = dockerService.initializeContainer(containerName(),
+                    sessionInfo.startupScriptId());
+            this.writer = containerState.containerInput();
+            this.reader = containerState.containerOutput();
             checkContainerOK();
             startupScriptSize = Integer.parseInt(reader.readLine());
         } catch (Exception e) {
@@ -127,7 +119,7 @@ private JShellResult readResult() throws IOException, NumberFormatException, Doc
                     int errorCount = Integer.parseInt(reader.readLine());
                     List<String> errors = new ArrayList<>();
                     for (int i = 0; i < errorCount; i++) {
-                        errors.add(desanitize(reader.readLine()));
+                        errors.add(Utils.deSanitizeStartupScript((reader.readLine())));
                     }
                     yield new JShellEvalAbortionCause.CompileTimeErrorAbortionCause(errors);
                 }
@@ -140,7 +132,7 @@ private JShellResult readResult() throws IOException, NumberFormatException, Doc
             abortion = new JShellEvalAbortion(causeSource, remainingSource, abortionCause);
         }
         boolean stdoutOverflow = Boolean.parseBoolean(reader.readLine());
-        String stdout = desanitize(reader.readLine());
+        String stdout = Utils.deSanitizeStartupScript(reader.readLine());
         return new JShellResult(snippetResults, abortion, stdoutOverflow, stdout);
     }
 
@@ -282,14 +274,6 @@ private void stopOperation() {
         doingOperation = false;
     }
 
-    private static String sanitize(String s) {
-        return s.replace("\\", "\\\\").replace("\n", "\\n");
-    }
-
-    private static String desanitize(String text) {
-        return text.replace("\\n", "\n").replace("\\\\", "\\");
-    }
-
     private static String cleanCode(String code) {
         return code.translateEscapes();
     }