Skip to content

Commit

Permalink
[enhancement/improveclose] Improved close feature (#49)
Browse files Browse the repository at this point in the history
* [enhancement/improveclose] Improved close feature

* [enhancement/improveclose] Improved code inside close to have less nesting

* [enhancement/improveclose] Moved "Session died" log up a call

* [enhancement/improveclose] Improved error logs in JShellService#close()

* [enhancement/improveclose] Sonar
  • Loading branch information
Alathreon authored Aug 30, 2024
1 parent daddfe9 commit b359b3c
Show file tree
Hide file tree
Showing 6 changed files with 85 additions and 80 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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"))
Expand All @@ -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();
Expand Down Expand Up @@ -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();
Expand All @@ -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();
Expand All @@ -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();
Expand All @@ -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);
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand All @@ -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)
Expand Down Expand Up @@ -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);
}
Expand Down
6 changes: 4 additions & 2 deletions JShellAPI/src/main/resources/application.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ jshellapi:
regularSessionTimeoutSeconds: 1800
oneTimeSessionTimeoutSeconds: 30
evalTimeoutSeconds: 15
evalTimeoutValidationLeeway: 10
evalTimeoutValidationLeeway: 5
sysOutCharLimit: 1024
maxAliveSessions: 10

Expand All @@ -14,7 +14,7 @@ jshellapi:
dockerCPUSetCPUs: 0

# Internal config
schedulerSessionKillScanRateSeconds: 60
schedulerSessionKillScanRateSeconds: 10

# Docker service config
dockerResponseTimeout: 60
Expand All @@ -29,3 +29,5 @@ logging:
org:
springframework:
web: DEBUG
togetherjava:
jshellapi: DEBUG
20 changes: 8 additions & 12 deletions JShellWrapper/src/test/java/JShellWrapperTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -49,12 +49,10 @@ void testHelloWorld() {

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

0 comments on commit b359b3c

Please sign in to comment.