11package org .togetherjava .jshellapi .service ;
22
3- import org .apache .tomcat .util .http .fileupload .util .Closeable ;
43import org .slf4j .Logger ;
54import org .slf4j .LoggerFactory ;
65import org .springframework .lang .Nullable ;
1514import java .util .List ;
1615import java .util .Optional ;
1716
18- public class JShellService implements Closeable {
17+ public class JShellService {
1918 private static final Logger LOGGER = LoggerFactory .getLogger (JShellService .class );
2019 private final JShellSessionService sessionService ;
2120 private final String id ;
2221 private final BufferedWriter writer ;
2322 private final BufferedReader reader ;
24-
23+ private boolean markedAsDead ;
2524 private Instant lastTimeoutUpdate ;
2625 private final long timeout ;
2726 private final boolean renewable ;
@@ -63,21 +62,20 @@ public JShellService(DockerService dockerService, JShellSessionService sessionSe
6362 startupScriptSize = Integer .parseInt (reader .readLine ());
6463 } catch (Exception e ) {
6564 LOGGER .warn ("Unexpected error during creation." , e );
65+ markAsDead ();
6666 throw new DockerException ("Creation of the session failed." , e );
6767 }
6868 this .doingOperation = false ;
6969 }
7070
7171 public Optional <JShellResult > eval (String code ) throws DockerException {
72+ if (shouldDie ())
73+ throw new DockerException ("Session %s is already dead." .formatted (id ));
7274 synchronized (this ) {
7375 if (!tryStartOperation ()) {
7476 return Optional .empty ();
7577 }
7678 }
77- if (isClosed ()) {
78- close ();
79- return Optional .empty ();
80- }
8179 updateLastTimeout ();
8280 sessionService .scheduleEvalTimeoutValidation (id , evalTimeout + evalTimeoutValidationLeeway );
8381 if (!code .endsWith ("\n " ))
@@ -95,7 +93,7 @@ public Optional<JShellResult> eval(String code) throws DockerException {
9593 return Optional .of (readResult ());
9694 } catch (DockerException | IOException | NumberFormatException ex ) {
9795 LOGGER .warn ("Unexpected error." , ex );
98- close ();
96+ markAsDead ();
9997 throw new DockerException (ex );
10098 } finally {
10199 stopOperation ();
@@ -147,6 +145,8 @@ private JShellResult readResult() throws IOException, NumberFormatException, Doc
147145 }
148146
149147 public Optional <List <String >> snippets (boolean includeStartupScript ) throws DockerException {
148+ if (shouldDie ())
149+ throw new DockerException ("Session %s is already dead." .formatted (id ));
150150 synchronized (this ) {
151151 if (!tryStartOperation ()) {
152152 return Optional .empty ();
@@ -169,7 +169,7 @@ public Optional<List<String>> snippets(boolean includeStartupScript) throws Dock
169169 : snippets .subList (startupScriptSize , snippets .size ()));
170170 } catch (Exception ex ) {
171171 LOGGER .warn ("Unexpected error." , ex );
172- close ();
172+ markAsDead ();
173173 throw new DockerException (ex );
174174 } finally {
175175 stopOperation ();
@@ -186,46 +186,70 @@ public boolean isInvalidEvalTimeout() {
186186 .isBefore (Instant .now ());
187187 }
188188
189- public boolean shouldDie () {
190- return lastTimeoutUpdate .plusSeconds (timeout ).isBefore (Instant .now ());
189+ /**
190+ * Returns if this session should be killed in the next heartbeat of the session killer.
191+ *
192+ * @return true if this session should be killed in the next heartbeat of the session killer
193+ * false otherwise
194+ */
195+ public boolean isMarkedAsDead () {
196+ return this .markedAsDead ;
191197 }
192198
193- public void stop () throws DockerException {
199+ /**
200+ * Marks this session as dead and also tries to gracefully close it, so it can be killed in the
201+ * next heartbeat of the session killer.
202+ */
203+ public synchronized void markAsDead () {
204+ if (this .markedAsDead )
205+ return ;
206+ LOGGER .info ("Session {} marked as dead." , id );
207+ this .markedAsDead = true ;
208+
194209 try {
195210 writer .write ("exit" );
196211 writer .newLine ();
197212 writer .flush ();
198- } catch (IOException e ) {
199- throw new DockerException ( e );
213+ } catch (IOException ex ) {
214+ LOGGER . debug ( "Couldn't close session {} gracefully." , id , ex );
200215 }
201216 }
202217
218+ /**
219+ * Returns if this session should be killed. Returns true if either it is marked as dead, if the
220+ * timeout is reached or if the container is dead.
221+ *
222+ * @return true if this session should be killed, false otherwise
223+ */
224+ public boolean shouldDie () {
225+ return markedAsDead || lastTimeoutUpdate .plusSeconds (timeout ).isBefore (Instant .now ())
226+ || dockerService .isDead (containerName ());
227+ }
228+
203229 public String id () {
204230 return id ;
205231 }
206232
207- @ Override
208233 public void close () {
209- LOGGER .debug ("Close called for session {}." , id );
210234 try {
211- dockerService .killContainerByName (containerName ());
212- try {
213- writer .close ();
214- } finally {
215- reader .close ();
235+ writer .close ();
236+ } catch (Exception ex ) {
237+ LOGGER .error ("Unexpected error while closing the writer." , ex );
238+ }
239+ try {
240+ reader .close ();
241+ } catch (Exception ex ) {
242+ LOGGER .error ("Unexpected error while closing the reader." , ex );
243+ }
244+ try {
245+ if (!dockerService .isDead (containerName ())) {
246+ dockerService .killContainerByName (containerName ());
216247 }
217- } catch (IOException ex ) {
218- LOGGER .error ("Unexpected error while closing." , ex );
219- } finally {
220- sessionService .notifyDeath (id );
248+ } catch (Exception ex ) {
249+ LOGGER .error ("Unexpected error while destroying the container." , ex );
221250 }
222251 }
223252
224- @ Override
225- public boolean isClosed () {
226- return dockerService .isDead (containerName ());
227- }
228-
229253 private void updateLastTimeout () {
230254 if (renewable ) {
231255 lastTimeoutUpdate = Instant .now ();
@@ -243,7 +267,6 @@ private void checkContainerOK() throws DockerException {
243267 "Container of session " + id + " is dead because status was " + ok );
244268 }
245269 } catch (IOException ex ) {
246- close ();
247270 throw new DockerException (ex );
248271 }
249272 }
0 commit comments