Skip to content

Commit e8d86ef

Browse files
Persist input/output streams for containers and fix test case issues
- Added functionality to persist input and output streams of newly created or cached Docker containers using a ConcurrentHashMap. - Ensured that container streams are properly managed and retrievable for future requests, improving container reuse. - Fixed test case issues related to stream management and container lifecycle, ensuring tests pass consistently.
1 parent 2758d1a commit e8d86ef

File tree

2 files changed

+32
-45
lines changed

2 files changed

+32
-45
lines changed

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

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,13 @@
33
import java.io.BufferedReader;
44
import java.io.BufferedWriter;
55

6-
public record ContainerState(boolean isCached, String containerId, BufferedReader containerOutput,
6+
/**
7+
* Data record for the state of a container.
8+
*
9+
* @param containerId The id of the container.
10+
* @param containerOutput The output of the container.
11+
* @param containerInput The input of the container.
12+
*/
13+
public record ContainerState(String containerId, BufferedReader containerOutput,
714
BufferedWriter containerInput) {
815
}

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

Lines changed: 24 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ public class DockerService implements DisposableBean {
3131
private final DockerClient client;
3232
private final Config config;
3333
private final ExecutorService executor = Executors.newSingleThreadExecutor();
34-
private final ConcurrentHashMap<StartupScriptId, String> cachedContainers =
34+
private final ConcurrentHashMap<StartupScriptId, ContainerState> cachedContainers =
3535
new ConcurrentHashMap<>();
3636
private final StartupScriptsService startupScriptsService;
3737

@@ -93,12 +93,10 @@ private boolean isImagePresentLocally() {
9393
* Pulls the Docker image.
9494
*/
9595
private void pullImage() throws InterruptedException {
96-
if (!isImagePresentLocally()) {
97-
client.pullImageCmd(jshellWrapperBaseImageName)
98-
.withTag("master")
99-
.exec(new PullImageResultCallback())
100-
.awaitCompletion(5, TimeUnit.MINUTES);
101-
}
96+
client.pullImageCmd(jshellWrapperBaseImageName)
97+
.withTag("master")
98+
.exec(new PullImageResultCallback())
99+
.awaitCompletion(5, TimeUnit.MINUTES);
102100
}
103101

104102
/**
@@ -107,7 +105,7 @@ private void pullImage() throws InterruptedException {
107105
* @param name The name of the container to create.
108106
* @return The ID of the created container.
109107
*/
110-
public String createContainer(String name) {
108+
private String createContainer(String name) {
111109
HostConfig hostConfig = HostConfig.newHostConfig()
112110
.withAutoRemove(true)
113111
.withInit(true)
@@ -146,13 +144,13 @@ public ContainerState initializeContainer(String name, StartupScriptId startupSc
146144
if (startupScriptId == null || cachedContainers.isEmpty()
147145
|| !cachedContainers.containsKey(startupScriptId)) {
148146
String containerId = createContainer(name);
149-
return setupContainerWithScript(containerId, false, startupScriptId);
147+
return setupContainerWithScript(containerId, startupScriptId);
150148
}
151-
String containerId = cachedContainers.get(startupScriptId);
149+
ContainerState containerState = cachedContainers.get(startupScriptId);
152150
executor.submit(() -> initializeCachedContainer(startupScriptId));
153151

154-
client.renameContainerCmd(containerId).withName(name).exec();
155-
return setupContainerWithScript(containerId, true, startupScriptId);
152+
client.renameContainerCmd(containerState.containerId()).withName(name).exec();
153+
return containerState;
156154
}
157155

158156
/**
@@ -165,58 +163,44 @@ private void initializeCachedContainer(StartupScriptId startupScriptId) {
165163
String id = createContainer(containerName);
166164
startContainer(id);
167165

168-
try (PipedInputStream containerInput = new PipedInputStream();
169-
BufferedWriter writer = new BufferedWriter(
170-
new OutputStreamWriter(new PipedOutputStream(containerInput)))) {
171-
InputStream containerOutput = attachToContainer(id, containerInput, true);
172-
173-
writer.write(Utils.sanitizeStartupScript(startupScriptsService.get(startupScriptId)));
174-
writer.newLine();
175-
writer.flush();
176-
containerOutput.close();
177-
178-
cachedContainers.put(startupScriptId, id);
166+
try {
167+
ContainerState containerState = setupContainerWithScript(id, startupScriptId);
168+
cachedContainers.put(startupScriptId, containerState);
179169
} catch (IOException e) {
180170
killContainerByName(containerName);
181171
throw new RuntimeException(e);
182172
}
183173
}
184174

185175
/**
186-
*
187176
* @param containerId The id of the container
188-
* @param isCached Indicator if the container is cached or new
189177
* @param startupScriptId The startup script id of the session
190178
* @return ContainerState of the spawned container.
191179
* @throws IOException if an I/O error occurs
192180
*/
193-
private ContainerState setupContainerWithScript(String containerId, boolean isCached,
181+
private ContainerState setupContainerWithScript(String containerId,
194182
StartupScriptId startupScriptId) throws IOException {
195-
if (!isCached) {
196-
startContainer(containerId);
197-
}
183+
startContainer(containerId);
198184
PipedInputStream containerInput = new PipedInputStream();
199185
BufferedWriter writer =
200186
new BufferedWriter(new OutputStreamWriter(new PipedOutputStream(containerInput)));
201187

202-
InputStream containerOutput = attachToContainer(containerId, containerInput, false);
188+
InputStream containerOutput = attachToContainer(containerId, containerInput);
203189
BufferedReader reader = new BufferedReader(new InputStreamReader(containerOutput));
204190

205-
if (!isCached) {
206-
writer.write(Utils.sanitizeStartupScript(startupScriptsService.get(startupScriptId)));
207-
writer.newLine();
208-
writer.flush();
209-
}
191+
writer.write(Utils.sanitizeStartupScript(startupScriptsService.get(startupScriptId)));
192+
writer.newLine();
193+
writer.flush();
210194

211-
return new ContainerState(isCached, containerId, reader, writer);
195+
return new ContainerState(containerId, reader, writer);
212196
}
213197

214198
/**
215199
* Creates a new container
216200
*
217201
* @param containerId the ID of the container to start
218202
*/
219-
public void startContainer(String containerId) {
203+
private void startContainer(String containerId) {
220204
if (!isContainerRunning(containerId)) {
221205
client.startContainerCmd(containerId).exec();
222206
}
@@ -228,12 +212,11 @@ public void startContainer(String containerId) {
228212
*
229213
* @param containerId The ID of the running container to attach to.
230214
* @param containerInput The input stream (containerInput) to send to the container.
231-
* @param isCached Indicator if the container is cached to prevent writing to output stream.
232215
* @return InputStream to read the container's stdout
233216
* @throws IOException if an I/O error occurs
234217
*/
235-
public InputStream attachToContainer(String containerId, InputStream containerInput,
236-
boolean isCached) throws IOException {
218+
private InputStream attachToContainer(String containerId, InputStream containerInput)
219+
throws IOException {
237220
PipedInputStream pipeIn = new PipedInputStream();
238221
PipedOutputStream pipeOut = new PipedOutputStream(pipeIn);
239222

@@ -250,9 +233,7 @@ public void onNext(Frame object) {
250233
String payloadString =
251234
new String(object.getPayload(), StandardCharsets.UTF_8);
252235
if (object.getStreamType() == StreamType.STDOUT) {
253-
if (!isCached) {
254-
pipeOut.write(object.getPayload()); // Write stdout data to pipeOut
255-
}
236+
pipeOut.write(object.getPayload()); // Write stdout data to pipeOut
256237
} else {
257238
LOGGER.warn("Received STDERR from container {}: {}", containerId,
258239
payloadString);
@@ -262,7 +243,6 @@ public void onNext(Frame object) {
262243
}
263244
}
264245
});
265-
266246
return pipeIn;
267247
}
268248

0 commit comments

Comments
 (0)