util: convert aggregate from js to ts and tweak execute-code to run async #7666

Author: haraldschilly

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

@@ -1,3 +1,4 @@
+import { ExecuteCodeOptionsAsyncGet } from "@cocalc/util/types/execute-code";
 import { executeCode } from "./execute-code";
 
 describe("hello world", () => {
@@ -88,4 +89,29 @@ describe("test timeout", () => {
       expect(err).toContain("killed command");
     }
   });
+
+  it("returns ID for async execution", async () => {
+    const retID = await executeCode({
+      command: "sh",
+      args: ["-c", "sleep .1; echo foo;"],
+      bash: false,
+      timeout: 10,
+      async_exec: true,
+    });
+    const id = retID["async_id"];
+    expect(typeof id).toEqual("string");
+    if (typeof id === "string") {
+      await new Promise((done) => setTimeout(done, 1000));
+      const status = await executeCode({ async_get: id });
+      console.log("status", status)
+      expect(status.stdout).toEqual("foo\n");
+      expect(status.elapsed_s).toBeGreaterThan(0.1)
+    }
+  });
 });
+
+export function isExecuteCodeOptionsAsyncGet(
+  opts: unknown,
+): opts is ExecuteCodeOptionsAsyncGet {
+  return typeof (opts as any)?.async_get === "string";
+}

src/packages/backend/execute-code.ts

@@ -6,7 +6,12 @@
 // Execute code in a subprocess.
 
 import { callback } from "awaiting";
-import { spawn } from "node:child_process";
+import LRU from "lru-cache";
+import {
+  ChildProcessWithoutNullStreams,
+  spawn,
+  SpawnOptionsWithoutStdio,
+} from "node:child_process";
 import { chmod, mkdtemp, rm, writeFile } from "node:fs/promises";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
@@ -15,21 +20,32 @@ import shellEscape from "shell-escape";
 import getLogger from "@cocalc/backend/logger";
 import { aggregate } from "@cocalc/util/aggregate";
 import { callback_opts } from "@cocalc/util/async-utils";
-import { to_json, trunc, walltime } from "@cocalc/util/misc";
+import { to_json, trunc, uuid, walltime } from "@cocalc/util/misc";
 import { envForSpawn } from "./misc";
 
 import type {
   ExecuteCodeFunctionWithCallback,
   ExecuteCodeOptions,
+  ExecuteCodeOptionsAsyncGet,
   ExecuteCodeOptionsWithCallback,
   ExecuteCodeOutput,
 } from "@cocalc/util/types/execute-code";
+import { isExecuteCodeOptionsAsyncGet } from "./execute-code.test";
 
 const log = getLogger("execute-code");
 
+const asyncCache = new LRU<string, ExecuteCodeOutput>({
+  max: 100,
+  ttl: 1000 * 60 * 60,
+  ttlAutopurge: true,
+  allowStale: true,
+  updateAgeOnGet: true,
+  updateAgeOnHas: true,
+});
+
 // Async/await interface to executing code.
 export async function executeCode(
-  opts: ExecuteCodeOptions,
+  opts: ExecuteCodeOptions | ExecuteCodeOptionsAsyncGet,
 ): Promise<ExecuteCodeOutput> {
   return await callback_opts(execute_code)(opts);
 }
@@ -50,8 +66,17 @@ export const execute_code: ExecuteCodeFunctionWithCallback = aggregate(
 
 // actual implementation, without the aggregate wrapper
 async function executeCodeNoAggregate(
-  opts: ExecuteCodeOptions,
+  opts: ExecuteCodeOptions | ExecuteCodeOptionsAsyncGet,
 ): Promise<ExecuteCodeOutput> {
+  if (isExecuteCodeOptionsAsyncGet(opts)) {
+    const cached = asyncCache.get(opts.async_get);
+    if (cached != null) {
+      return cached;
+    } else {
+      throw new Error(`Status or result of '${opts.async_get}' not found.`);
+    }
+  }
+
   if (opts.args == null) opts.args = [];
   if (opts.timeout == null) opts.timeout = 10;
   if (opts.ulimit_timeout == null) opts.ulimit_timeout = true;
@@ -115,7 +140,48 @@ async function executeCodeNoAggregate(
       await writeFile(tempPath, cmd);
       await chmod(tempPath, 0o700);
     }
-    return await callback(doSpawn, { ...opts, origCommand });
+
+    if (opts.async_exec) {
+      // we return an ID, the caller can then use it to query the status
+      const async_limit = 1024 * 1024; // we limit how much we keep in memory, to avoid problems
+      opts.max_output = Math.min(async_limit, opts.max_output ?? async_limit);
+      const id = uuid();
+      const start = new Date();
+      const started = {
+        stdout: `Process started running at ${start.toISOString()}`,
+        stderr: "",
+        exit_code: start.getTime(),
+        async_id: id,
+      };
+      asyncCache.set(id, started);
+
+      doSpawn({ ...opts, origCommand }, (err, result) => {
+        const started = asyncCache.get(id)?.exit_code ?? 0;
+        const info = { elapsed_s: (Date.now() - started) / 1000 };
+        if (err) {
+          asyncCache.set(id, {
+            stdout: "",
+            stderr: `${err}`,
+            exit_code: 1,
+            ...info,
+          });
+        } else if (result != null) {
+          asyncCache.set(id, { ...result, ...info });
+        } else {
+          asyncCache.set(id, {
+            stdout: "",
+            stderr: `No result`,
+            exit_code: 1,
+            ...info,
+          });
+        }
+      });
+
+      return started;
+    } else {
+      // This is the blocking variant
+      return await callback(doSpawn, { ...opts, origCommand });
+    }
   } finally {
     // clean up
     if (tempDir) {
@@ -124,7 +190,10 @@ async function executeCodeNoAggregate(
   }
 }
 
-function doSpawn(opts, cb) {
+function doSpawn(
+  opts,
+  cb: (err: string | undefined, result?: ExecuteCodeOutput) => void,
+) {
   const start_time = walltime();
 
   if (opts.verbose) {
@@ -138,7 +207,7 @@ function doSpawn(opts, cb) {
       "seconds",
     );
   }
-  const spawnOptions = {
+  const spawnOptions: SpawnOptionsWithoutStdio = {
     detached: true, // so we can kill the entire process group if it times out
     cwd: opts.path,
     ...(opts.uid ? { uid: opts.uid } : undefined),
@@ -150,8 +219,8 @@ function doSpawn(opts, cb) {
     },
   };
 
-  let r,
-    ran_code = false;
+  let r: ChildProcessWithoutNullStreams;
+  let ran_code = false;
   try {
     r = spawn(opts.command, opts.args, spawnOptions);
     if (r.stdout == null || r.stderr == null) {
@@ -215,7 +284,7 @@ function doSpawn(opts, cb) {
   });
 
   r.on("exit", (code) => {
-    exit_code = code;
+    exit_code = code != null ? code : undefined;
     finish();
   });
 
@@ -317,8 +386,10 @@ function doSpawn(opts, cb) {
         );
       }
       try {
-        killed = true;
-        process.kill(-r.pid, "SIGKILL"); // this should kill process group
+        killed = true; // we set the kill flag in any case – i.e. process will no longer exist
+        if (r.pid != null) {
+          process.kill(-r.pid, "SIGKILL"); // this should kill process group
+        }
       } catch (err) {
         // Exceptions can happen, which left uncaught messes up calling code big time.
         if (opts.verbose) {

src/packages/next/lib/api/schema/exec.ts

@@ -10,14 +10,14 @@ export const ExecInputSchema = z
   .object({
     project_id: ProjectIdSchema,
     compute_server_id: ComputeServerIdSchema.describe(
-      `If provided, the desired shell command will be run on the compute server whose id 
+      `If provided, the desired shell command will be run on the compute server whose id
        is specified in this field (if available).`,
     ).optional(),
     filesystem: z
       .boolean()
       .optional()
       .describe(
-        `If \`true\`, this shell command runs in the fileserver container on the compute 
+        `If \`true\`, this shell command runs in the fileserver container on the compute
          server; otherwise, it runs on the main compute container.`,
       ),
     path: z
@@ -46,7 +46,7 @@ export const ExecInputSchema = z
       .boolean()
       .optional()
       .describe(
-        `If \`true\`, this command runs in a \`bash\` shell. To do so, the provided shell 
+        `If \`true\`, this command runs in a \`bash\` shell. To do so, the provided shell
          command is written to a file and then executed via the \`bash\` command.`,
       ),
     aggregate: z
@@ -57,24 +57,32 @@ export const ExecInputSchema = z
       ])
       .optional()
       .describe(
-        `If provided, this shell command is aggregated as in 
+        `If provided, this shell command is aggregated as in
          \`src/packages/backend/aggregate.js\`. This parameter allows one to specify
-         multiple callbacks to be executed against the output of the same command 
+         multiple callbacks to be executed against the output of the same command
          (given identical arguments) within a 60-second window.`,
       ),
     err_on_exit: z
       .boolean()
       .optional()
       .describe(
-        `When \`true\`, this call will throw an error whenever the provided shell command 
-       exits with a non-zero exit code.`,
+        `When \`true\`, this call will throw an error whenever the provided shell command
+         exits with a non-zero exit code.`,
       ),
     env: z
       .record(z.string(), z.string())
       .optional()
       .describe(
         "Environment variables to be passed to the shell command upon execution.",
       ),
+    async_exec: z.boolean().optional()
+      .describe(`If \`true\`, the execution happens asynchroneously.
+        This means it the API call does not block and returns an ID (\`async_id\`).
+        Later, use that ID in a call to \`async_get\` to eventually get the result`),
+    async_get: z.string().optional()
+      .describe(`For a given \`async_id\` returned by \`async\`,
+        retun the status, or the result as if it is called synchroneously.
+        Results are only cached temporarily!`),
   })
   .describe("Perform arbitrary shell commands in a compute server or project.");
 

src/packages/next/pages/api/v2/exec.ts

@@ -38,6 +38,8 @@ async function get(req) {
     aggregate,
     err_on_exit,
     env,
+    async_exec,
+    async_get,
   } = getParams(req);
 
   if (!(await isCollaborator({ account_id, project_id }))) {
@@ -61,6 +63,8 @@ async function get(req) {
       aggregate,
       err_on_exit,
       env,
+      async_exec,
+      async_get,
     },
   });
   // event and id don't make sense for http post api

src/packages/util/aggregate.d.ts

@@ -63,16 +63,17 @@ Where options is an object.
 
 */
 
-const { copy_without, field_cmp } = require("./misc");
 const json_stable = require("json-stable-stringify");
 
+import { copy_without, field_cmp } from "./misc";
+
 // To avoid using up too much memory, results are cached at most this long
 // (so long as function is called periodically to clear the cache... if not,
 // no point in clearing, since won't grow much.)
 const DONE_CACHE_TIMEOUT_MS = 60000;
 
 function clear_old(done) {
-  const now = new Date();
+  const now = Date.now();
   for (let key in done) {
     const s = done[key];
     if (now - s.time >= DONE_CACHE_TIMEOUT_MS) {
@@ -93,7 +94,7 @@ function leq(a, b) {
   return a <= b;
 }
 
-exports.aggregate = function (options, f) {
+export function aggregate(options, f?: any) {
   if (f == null) {
     f = options;
     options = undefined;
@@ -120,7 +121,7 @@ exports.aggregate = function (options, f) {
 
   function aggregate_call_f(opts) {
     // Key is a string that determines the inputs to f **that matter**.
-    const key = json_stable(copy_without(opts, omitted_fields));
+    const key: string = json_stable(copy_without(opts, omitted_fields));
     // Check state
     const current = state[key];
     const recent = done[key];
@@ -190,4 +191,4 @@ exports.aggregate = function (options, f) {
       aggregate_call_f(opts);
     }
   };
-};
+}

src/packages/util/aggregate.js renamed to src/packages/util/aggregate.ts

@@ -2,6 +2,8 @@ export interface ExecuteCodeOutput {
   stdout: string;
   stderr: string;
   exit_code: number;
+  async_id?: string;
+  elapsed_s?: number; // how long it took, async execution
 }
 
 export interface ExecuteCodeOptions {
@@ -20,12 +22,17 @@ export interface ExecuteCodeOptions {
   env?: object; // if given, added to exec environment
   aggregate?: string | number; // if given, aggregates multiple calls with same sequence number into one -- see @cocalc/util/aggregate; typically make this a timestamp for compiling code (e.g., latex).
   verbose?: boolean; // default true -- impacts amount of logging
+  async_exec?: boolean; // default false -- if true, return an ID and execute it asynchroneously
+}
+
+export interface ExecuteCodeOptionsAsyncGet {
+  async_get: string; // if set, everything else is ignored and the status/output of the async call is returned
 }
 
 export interface ExecuteCodeOptionsWithCallback extends ExecuteCodeOptions {
   cb?: (err: undefined | Error, output?: ExecuteCodeOutput) => void;
 }
 
 export type ExecuteCodeFunctionWithCallback = (
-  opts: ExecuteCodeOptionsWithCallback
+  opts: ExecuteCodeOptionsWithCallback,
 ) => void;