Skip to content

Commit f0b95f4

Browse files
committed
api/async exec: cleanup api, more testing
1 parent 0475484 commit f0b95f4

File tree

10 files changed

+287
-145
lines changed

10 files changed

+287
-145
lines changed

src/packages/backend/execute-code.test.ts

+101-15
Original file line numberDiff line numberDiff line change
@@ -90,27 +90,113 @@ describe("test timeout", () => {
9090
});
9191
});
9292

93+
describe("test env", () => {
94+
it("allows to specify environment variables", async () => {
95+
const { stdout, stderr } = await executeCode({
96+
command: "sh",
97+
args: ["-c", "echo $FOO;"],
98+
err_on_exit: false,
99+
bash: false,
100+
env: { FOO: "bar" },
101+
});
102+
expect(stdout).toBe("bar\n");
103+
expect(stderr).toBe("");
104+
});
105+
});
106+
93107
describe("async", () => {
94-
it("use ID to get async execution result", async () => {
108+
it("use ID to get async result", async () => {
95109
const c = await executeCode({
96110
command: "sh",
97-
args: ["-c", "sleep .5; echo foo;"],
111+
args: ["-c", "echo foo; sleep .5; echo bar; sleep .5; echo baz;"],
98112
bash: false,
99113
timeout: 10,
100-
async_mode: true,
114+
async_call: true,
101115
});
102-
expect(c.async_status).toEqual("running");
103-
expect(c.async_start).toBeGreaterThan(1);
104-
const id = c.async_id;
105-
expect(typeof id).toEqual("string");
106-
if (typeof id === "string") {
107-
await new Promise((done) => setTimeout(done, 1000));
108-
const status = await executeCode({ async_get: id });
109-
expect(status.async_status).toEqual("completed");
110-
expect(status.stdout).toEqual("foo\n");
111-
expect(status.elapsed_s).toBeGreaterThan(0.1);
112-
expect(status.elapsed_s).toBeLessThan(3);
113-
expect(status.async_start).toBeGreaterThan(1);
116+
expect(c.type).toEqual("async");
117+
if (c.type !== "async") return;
118+
const { status, start, job_id } = c;
119+
expect(status).toEqual("running");
120+
expect(start).toBeGreaterThan(1);
121+
expect(typeof job_id).toEqual("string");
122+
if (typeof job_id !== "string") return;
123+
await new Promise((done) => setTimeout(done, 250));
124+
{
125+
const s = await executeCode({ async_get: job_id });
126+
expect(s.type).toEqual("async");
127+
if (s.type !== "async") return;
128+
expect(s.status).toEqual("running");
129+
// partial stdout result
130+
expect(s.stdout).toEqual("foo\n");
131+
expect(s.elapsed_s).toBeUndefined();
132+
expect(s.start).toBeGreaterThan(1);
133+
expect(s.exit_code).toEqual(0);
134+
}
135+
136+
await new Promise((done) => setTimeout(done, 900));
137+
{
138+
const s = await executeCode({ async_get: job_id });
139+
expect(s.type).toEqual("async");
140+
if (s.type !== "async") return;
141+
expect(s.status).toEqual("completed");
142+
expect(s.stdout).toEqual("foo\nbar\nbaz\n");
143+
expect(s.elapsed_s).toBeGreaterThan(0.1);
144+
expect(s.elapsed_s).toBeLessThan(3);
145+
expect(s.start).toBeGreaterThan(1);
146+
expect(s.stderr).toEqual("");
147+
expect(s.exit_code).toEqual(0);
114148
}
115149
});
150+
151+
it("with an error", async () => {
152+
const c = await executeCode({
153+
command: ">&2 echo baz; exit 3",
154+
bash: true,
155+
async_call: true,
156+
});
157+
expect(c.type).toEqual("async");
158+
if (c.type !== "async") return;
159+
const { job_id } = c;
160+
expect(typeof job_id).toEqual("string");
161+
if (typeof job_id !== "string") return;
162+
await new Promise((done) => setTimeout(done, 250));
163+
const s = await executeCode({ async_get: job_id });
164+
expect(s.type).toEqual("async");
165+
if (s.type !== "async") return;
166+
expect(s.status).toEqual("error");
167+
expect(s.stdout).toEqual("");
168+
expect(s.stderr).toEqual("baz\n");
169+
// any error is code 1 it seems?
170+
expect(s.exit_code).toEqual(1);
171+
});
172+
173+
it("trigger a timeout", async () => {
174+
const c = await executeCode({
175+
command: "sh",
176+
args: ["-c", "echo foo; sleep 1; echo bar;"],
177+
bash: false,
178+
timeout: 0.1,
179+
async_call: true,
180+
});
181+
expect(c.type).toEqual("async");
182+
if (c.type !== "async") return;
183+
const { status, start, job_id } = c;
184+
expect(status).toEqual("running");
185+
expect(start).toBeGreaterThan(1);
186+
expect(typeof job_id).toEqual("string");
187+
if (typeof job_id !== "string") return;
188+
await new Promise((done) => setTimeout(done, 250));
189+
const s = await executeCode({ async_get: job_id });
190+
expect(s.type).toEqual("async");
191+
if (s.type !== "async") return;
192+
expect(s.status).toEqual("error");
193+
expect(s.stdout).toEqual("");
194+
expect(s.elapsed_s).toBeGreaterThan(0.01);
195+
expect(s.elapsed_s).toBeLessThan(3);
196+
expect(s.start).toBeGreaterThan(1);
197+
expect(s.stderr).toEqual(
198+
"killed command 'sh -c echo foo; sleep 1; echo bar;'",
199+
);
200+
expect(s.exit_code).toEqual(1);
201+
});
116202
});

src/packages/backend/execute-code.ts

+79-59
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@ import { to_json, trunc, uuid, walltime } from "@cocalc/util/misc";
2424
import { envForSpawn } from "./misc";
2525

2626
import {
27+
ExecuteCodeOutputAsync,
28+
ExecuteCodeOutputBlocking,
2729
isExecuteCodeOptionsAsyncGet,
2830
type ExecuteCodeFunctionWithCallback,
2931
type ExecuteCodeOptions,
@@ -34,7 +36,7 @@ import {
3436

3537
const log = getLogger("execute-code");
3638

37-
const asyncCache = new LRU<string, ExecuteCodeOutput>({
39+
const asyncCache = new LRU<string, ExecuteCodeOutputAsync>({
3840
max: 100,
3941
ttl: 1000 * 60 * 60,
4042
ttlAutopurge: true,
@@ -64,6 +66,12 @@ export const execute_code: ExecuteCodeFunctionWithCallback = aggregate(
6466
},
6567
);
6668

69+
async function clean_up_tmp(tempDir: string | undefined) {
70+
if (tempDir) {
71+
await rm(tempDir, { force: true, recursive: true });
72+
}
73+
}
74+
6775
// actual implementation, without the aggregate wrapper
6876
async function executeCodeNoAggregate(
6977
opts: ExecuteCodeOptions | ExecuteCodeOptionsAsyncGet,
@@ -77,16 +85,15 @@ async function executeCodeNoAggregate(
7785
}
7886
}
7987

80-
if (opts.args == null) opts.args = [];
81-
if (opts.timeout == null) opts.timeout = 10;
82-
if (opts.ulimit_timeout == null) opts.ulimit_timeout = true;
83-
if (opts.err_on_exit == null) opts.err_on_exit = true;
84-
if (opts.verbose == null) opts.verbose = true;
88+
opts.args ??= [];
89+
opts.timeout ??= 10;
90+
opts.ulimit_timeout ??= true;
91+
opts.err_on_exit ??= true;
92+
opts.verbose ??= true;
8593

8694
if (opts.verbose) {
8795
log.debug(`input: ${opts.command} ${opts.args?.join(" ")}`);
8896
}
89-
9097
const s = opts.command.split(/\s+/g); // split on whitespace
9198
if (opts.args?.length === 0 && s.length > 1) {
9299
opts.bash = true;
@@ -141,49 +148,59 @@ async function executeCodeNoAggregate(
141148
await chmod(tempPath, 0o700);
142149
}
143150

144-
if (opts.async_mode) {
151+
if (opts.async_call) {
145152
// we return an ID, the caller can then use it to query the status
146153
opts.max_output ??= 1024 * 1024; // we limit how much we keep in memory, to avoid problems;
147154
opts.timeout ??= 10 * 60;
148-
const id = uuid();
155+
const job_id = uuid();
149156
const start = new Date();
150-
const started: ExecuteCodeOutput = {
157+
const started: ExecuteCodeOutputAsync = {
158+
type: "async",
151159
stdout: `Process started running at ${start.toISOString()}`,
152160
stderr: "",
153161
exit_code: 0,
154-
async_start: start.getTime(),
155-
async_id: id,
156-
async_status: "running",
162+
start: start.getTime(),
163+
job_id,
164+
status: "running",
157165
};
158-
asyncCache.set(id, started);
159-
160-
doSpawn({ ...opts, origCommand, async_id: id }, (err, result) => {
161-
const started = asyncCache.get(id)?.async_start ?? 0;
162-
const info: Partial<ExecuteCodeOutput> = {
163-
elapsed_s: (Date.now() - started) / 1000,
164-
async_start: start.getTime(),
165-
async_status: "error",
166-
};
167-
if (err) {
168-
asyncCache.set(id, {
169-
stdout: "",
170-
stderr: `${err}`,
171-
exit_code: 1,
172-
...info,
173-
});
174-
} else if (result != null) {
175-
asyncCache.set(id, {
176-
...result,
177-
...info,
178-
...{ async_status: "completed" },
179-
});
180-
} else {
181-
asyncCache.set(id, {
182-
stdout: "",
183-
stderr: `No result`,
184-
exit_code: 1,
185-
...info,
186-
});
166+
asyncCache.set(job_id, started);
167+
168+
doSpawn({ ...opts, origCommand, job_id }, async (err, result) => {
169+
try {
170+
const started = asyncCache.get(job_id)?.start ?? 0;
171+
const info: Omit<
172+
ExecuteCodeOutputAsync,
173+
"stdout" | "stderr" | "exit_code"
174+
> = {
175+
job_id,
176+
type: "async",
177+
elapsed_s: (Date.now() - started) / 1000,
178+
start: start.getTime(),
179+
status: "error",
180+
};
181+
if (err) {
182+
asyncCache.set(job_id, {
183+
stdout: "",
184+
stderr: `${err}`,
185+
exit_code: 1,
186+
...info,
187+
});
188+
} else if (result != null) {
189+
asyncCache.set(job_id, {
190+
...result,
191+
...info,
192+
...{ status: "completed" },
193+
});
194+
} else {
195+
asyncCache.set(job_id, {
196+
stdout: "",
197+
stderr: `No result`,
198+
exit_code: 1,
199+
...info,
200+
});
201+
}
202+
} finally {
203+
await clean_up_tmp(tempDir);
187204
}
188205
});
189206

@@ -193,28 +210,26 @@ async function executeCodeNoAggregate(
193210
return await callback(doSpawn, { ...opts, origCommand });
194211
}
195212
} finally {
196-
// clean up
197-
if (tempDir) {
198-
await rm(tempDir, { force: true, recursive: true });
199-
}
213+
// do not delete the tempDir in async mode!
214+
if (!opts.async_call) await clean_up_tmp(tempDir);
200215
}
201216
}
202217

203218
function update_async(
204-
async_id: string | undefined,
219+
job_id: string | undefined,
205220
stream: "stdout" | "stderr",
206221
data: string,
207222
) {
208-
if (!async_id) return;
209-
const obj = asyncCache.get(async_id);
223+
if (!job_id) return;
224+
const obj = asyncCache.get(job_id);
210225
if (obj != null) {
211226
obj[stream] = data;
212227
}
213228
}
214229

215230
function doSpawn(
216231
opts,
217-
cb: (err: string | undefined, result?: ExecuteCodeOutput) => void,
232+
cb: (err: string | undefined, result?: ExecuteCodeOutputBlocking) => void,
218233
) {
219234
const start_time = walltime();
220235

@@ -278,7 +293,7 @@ function doSpawn(
278293
} else {
279294
stdout += data;
280295
}
281-
update_async(opts.async_id, "stdout", stdout);
296+
update_async(opts.job_id, "stdout", stdout);
282297
});
283298

284299
r.stderr.on("data", (data) => {
@@ -290,7 +305,7 @@ function doSpawn(
290305
} else {
291306
stderr += data;
292307
}
293-
update_async(opts.async_id, "stderr", stderr);
308+
update_async(opts.job_id, "stderr", stderr);
294309
});
295310

296311
let stderr_is_done = false;
@@ -363,12 +378,17 @@ function doSpawn(
363378
const x = opts.origCommand
364379
? opts.origCommand
365380
: `'${opts.command}' (args=${opts.args?.join(" ")})`;
366-
cb(
367-
`command '${x}' exited with nonzero code ${exit_code} -- stderr='${trunc(
368-
stderr,
369-
1024,
370-
)}'`,
371-
);
381+
if (opts.job_id) {
382+
cb(stderr);
383+
} else {
384+
// sync behavor, like it was before
385+
cb(
386+
`command '${x}' exited with nonzero code ${exit_code} -- stderr='${trunc(
387+
stderr,
388+
1024,
389+
)}'`,
390+
);
391+
}
372392
} else if (!ran_code) {
373393
// regardless of opts.err_on_exit !
374394
const x = opts.origCommand
@@ -390,7 +410,7 @@ function doSpawn(
390410
// if exit-code not set, may have been SIGKILL so we set it to 1
391411
exit_code = 1;
392412
}
393-
cb(undefined, { stdout, stderr, exit_code });
413+
cb(undefined, { type: "blocking", stdout, stderr, exit_code });
394414
}
395415
};
396416

0 commit comments

Comments
 (0)