[enhancement/improveclose] Improved close feature

JShellAPI/src/main/java/org/togetherjava/jshellapi/rest/JShellController.java

@@ -64,7 +64,7 @@ public List<String> snippets(@PathVariable String id,
     }
 
     @DeleteMapping("/{id}")
-    public void delete(@PathVariable String id) throws DockerException {
+    public void delete(@PathVariable String id) {
         validateId(id);
         if (!service.hasSession(id))
             throw new ResponseStatusException(HttpStatus.NOT_FOUND, "Id " + id + " not found");

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

@@ -133,9 +133,10 @@ public void onNext(Frame object) {
     public void killContainerByName(String name) {
         LOGGER.debug("Fetching container to kill {}.", name);
         List<Container> containers = client.listContainersCmd().withNameFilter(Set.of(name)).exec();
-        LOGGER.debug("Number of containers to kill: {} for name {}.", containers.size(), name);
-        if (containers.size() != 1) {
-            LOGGER.error("There is more than 1 container for name {}.", name);
+        if (containers.size() == 1) {
+            LOGGER.debug("Found 1 container for name {}.", name);
+        } else {
+            LOGGER.error("Expected 1 container but found {} for name {}.", containers.size(), name);
         }
         for (Container container : containers) {
             client.killContainerCmd(container.getId()).exec();

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

@@ -1,6 +1,5 @@
 package org.togetherjava.jshellapi.service;
 
-import org.apache.tomcat.util.http.fileupload.util.Closeable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.lang.Nullable;
@@ -15,13 +14,13 @@
 import java.util.List;
 import java.util.Optional;
 
-public class JShellService implements Closeable {
+public class JShellService {
     private static final Logger LOGGER = LoggerFactory.getLogger(JShellService.class);
     private final JShellSessionService sessionService;
     private final String id;
     private final BufferedWriter writer;
     private final BufferedReader reader;
-
+    private boolean markedAsDead;
     private Instant lastTimeoutUpdate;
     private final long timeout;
     private final boolean renewable;
@@ -63,21 +62,20 @@ public JShellService(DockerService dockerService, JShellSessionService sessionSe
             startupScriptSize = Integer.parseInt(reader.readLine());
         } catch (Exception e) {
             LOGGER.warn("Unexpected error during creation.", e);
+            markAsDead();
             throw new DockerException("Creation of the session failed.", e);
         }
         this.doingOperation = false;
     }
 
     public Optional<JShellResult> eval(String code) throws DockerException {
+        if (shouldDie())
+            throw new DockerException("Session %s is already dead.".formatted(id));
         synchronized (this) {
             if (!tryStartOperation()) {
                 return Optional.empty();
             }
         }
-        if (isClosed()) {
-            close();
-            return Optional.empty();
-        }
         updateLastTimeout();
         sessionService.scheduleEvalTimeoutValidation(id, evalTimeout + evalTimeoutValidationLeeway);
         if (!code.endsWith("\n"))
@@ -95,7 +93,7 @@ public Optional<JShellResult> eval(String code) throws DockerException {
             return Optional.of(readResult());
         } catch (DockerException | IOException | NumberFormatException ex) {
             LOGGER.warn("Unexpected error.", ex);
-            close();
+            markAsDead();
             throw new DockerException(ex);
         } finally {
             stopOperation();
@@ -147,6 +145,8 @@ private JShellResult readResult() throws IOException, NumberFormatException, Doc
     }
 
     public Optional<List<String>> snippets(boolean includeStartupScript) throws DockerException {
+        if (shouldDie())
+            throw new DockerException("Session %s is already dead.".formatted(id));
         synchronized (this) {
             if (!tryStartOperation()) {
                 return Optional.empty();
@@ -169,7 +169,7 @@ public Optional<List<String>> snippets(boolean includeStartupScript) throws Dock
                     : snippets.subList(startupScriptSize, snippets.size()));
         } catch (Exception ex) {
             LOGGER.warn("Unexpected error.", ex);
-            close();
+            markAsDead();
             throw new DockerException(ex);
         } finally {
             stopOperation();
@@ -186,46 +186,70 @@ public boolean isInvalidEvalTimeout() {
                     .isBefore(Instant.now());
     }
 
-    public boolean shouldDie() {
-        return lastTimeoutUpdate.plusSeconds(timeout).isBefore(Instant.now());
+    /**
+     * Returns if this session should be killed in the next heartbeat of the session killer.
+     * 
+     * @return true if this session should be killed in the next heartbeat of the session killer
+     *         false otherwise
+     */
+    public boolean isMarkedAsDead() {
+        return this.markedAsDead;
     }
 
-    public void stop() throws DockerException {
+    /**
+     * Marks this session as dead and also tries to gracefully close it, so it can be killed in the
+     * next heartbeat of the session killer.
+     */
+    public synchronized void markAsDead() {
+        if (this.markedAsDead)
+            return;
+        LOGGER.info("Session {} marked as dead.", id);
+        this.markedAsDead = true;
+
         try {
             writer.write("exit");
             writer.newLine();
             writer.flush();
-        } catch (IOException e) {
-            throw new DockerException(e);
+        } catch (IOException ex) {
+            LOGGER.debug("Couldn't close session {} gracefully.", id, ex);
         }
     }
 
+    /**
+     * Returns if this session should be killed. Returns true if either it is marked as dead, if the
+     * timeout is reached or if the container is dead.
+     * 
+     * @return true if this session should be killed, false otherwise
+     */
+    public boolean shouldDie() {
+        return markedAsDead || lastTimeoutUpdate.plusSeconds(timeout).isBefore(Instant.now())
+                || dockerService.isDead(containerName());
+    }
+
     public String id() {
         return id;
     }
 
-    @Override
     public void close() {
-        LOGGER.debug("Close called for session {}.", id);
         try {
-            dockerService.killContainerByName(containerName());
-            try {
-                writer.close();
-            } finally {
-                reader.close();
+            writer.close();
+        } catch (Exception ex) {
+            LOGGER.error("Unexpected error while closing the writer.", ex);
+        }
+        try {
+            reader.close();
+        } catch (Exception ex) {
+            LOGGER.error("Unexpected error while closing the reader.", ex);
+        }
+        try {
+            if (!dockerService.isDead(containerName())) {
+                dockerService.killContainerByName(containerName());
             }
-        } catch (IOException ex) {
-            LOGGER.error("Unexpected error while closing.", ex);
-        } finally {
-            sessionService.notifyDeath(id);
+        } catch (Exception ex) {
+            LOGGER.error("Unexpected error while destroying the container.", ex);
         }
     }
 
-    @Override
-    public boolean isClosed() {
-        return dockerService.isDead(containerName());
-    }
-
     private void updateLastTimeout() {
         if (renewable) {
             lastTimeoutUpdate = Instant.now();
@@ -243,7 +267,6 @@ private void checkContainerOK() throws DockerException {
                         "Container of session " + id + " is dead because status was " + ok);
             }
         } catch (IOException ex) {
-            close();
             throw new DockerException(ex);
         }
     }

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

@@ -28,40 +28,28 @@ public class JShellSessionService {
     private void initScheduler() {
         scheduler = Executors.newSingleThreadScheduledExecutor();
         scheduler.scheduleAtFixedRate(() -> {
-            LOGGER.info("Scheduler heartbeat: started.");
-            jshellSessions.keySet()
-                .stream()
-                .filter(id -> jshellSessions.get(id).isClosed())
-                .forEach(this::notifyDeath);
             List<String> toDie = jshellSessions.keySet()
                 .stream()
                 .filter(id -> jshellSessions.get(id).shouldDie())
                 .toList();
-            LOGGER.info("Scheduler heartbeat: sessions ready to die: {}", toDie);
+            LOGGER.info("Scheduler heartbeat, sessions ready to die: {}", toDie);
             for (String id : toDie) {
-                try {
-                    deleteSession(id);
-                } catch (DockerException ex) {
-                    LOGGER.error("Unexpected error when deleting session.", ex);
+                JShellService service = jshellSessions.get(id);
+                if (service.isMarkedAsDead()) {
+                    try {
+                        jshellSessions.remove(id).close();
+                        LOGGER.info("Session {} died.", id);
+                    } catch (Exception ex) {
+                        LOGGER.error("Unexpected exception for session {}", id, ex);
+                    }
+                } else {
+                    service.markAsDead();
                 }
             }
         }, config.schedulerSessionKillScanRateSeconds(),
                 config.schedulerSessionKillScanRateSeconds(), TimeUnit.SECONDS);
     }
 
-    void notifyDeath(String id) {
-        JShellService shellService = jshellSessions.remove(id);
-        if (shellService == null) {
-            LOGGER.debug("Notify death on already removed session {}.", id);
-            return;
-        }
-        if (!shellService.isClosed()) {
-            LOGGER.error("JShell Service isn't dead when it should for id {}.", id);
-            return;
-        }
-        LOGGER.info("Session {} died.", id);
-    }
-
     public boolean hasSession(String id) {
         return jshellSessions.containsKey(id);
     }
@@ -85,13 +73,8 @@ public JShellService oneTimeSession(@Nullable StartupScriptId startupScriptId)
                 true, config));
     }
 
-    public void deleteSession(String id) throws DockerException {
-        JShellService service = jshellSessions.remove(id);
-        try {
-            service.stop();
-        } finally {
-            scheduler.schedule(service::close, 500, TimeUnit.MILLISECONDS);
-        }
+    public void deleteSession(String id) {
+        jshellSessions.get(id).markAsDead();
     }
 
     private synchronized JShellService createSession(SessionInfo sessionInfo)
@@ -128,7 +111,7 @@ public void scheduleEvalTimeoutValidation(String id, long timeSeconds) {
             if (service == null)
                 return;
             if (service.isInvalidEvalTimeout()) {
-                service.close();
+                deleteSession(id);
             }
         }, timeSeconds, TimeUnit.SECONDS);
     }

JShellAPI/src/main/resources/application.yaml

@@ -4,7 +4,7 @@ jshellapi:
   regularSessionTimeoutSeconds: 1800
   oneTimeSessionTimeoutSeconds: 30
   evalTimeoutSeconds: 15
-  evalTimeoutValidationLeeway: 10
+  evalTimeoutValidationLeeway: 5
   sysOutCharLimit: 1024
   maxAliveSessions: 10
 
@@ -14,7 +14,7 @@ jshellapi:
   dockerCPUSetCPUs: 0
 
   # Internal config
-  schedulerSessionKillScanRateSeconds: 60
+  schedulerSessionKillScanRateSeconds: 10
 
   # Docker service config
   dockerResponseTimeout: 60
@@ -29,3 +29,5 @@ logging:
     org:
       springframework:
         web: DEBUG
+      togetherjava:
+        jshellapi: DEBUG

JShellWrapper/src/test/java/JShellWrapperTest.java

@@ -49,12 +49,10 @@ void testHelloWorld() {
 
     @Test
     void testExpressionResult() {
-        evalTest(
-                """
-            eval
-            1
-            "Hello world!\"""",
-                """
+        evalTest("""
+                eval
+                1
+                "Hello world!\"""", """
                 OK
                 0
                 OK
@@ -67,12 +65,10 @@ void testExpressionResult() {
 
                 false
                 """);
-        evalTest(
-                """
-            eval
-            1
-            2+2""",
-                """
+        evalTest("""
+                eval
+                1
+                2+2""", """
                 OK
                 0
                 OK