Add active evals (#25)

* Add active evals

* Convert type of stats to date or string

* Fix generic overloads for run

* Lower evaluate timeout

* Add comment for createWriteStream

packages/protocol/src/node/evaluate.ts

@@ -1,9 +1,14 @@
 import * as vm from "vm";
-import { NewEvalMessage, TypedValue, EvalFailedMessage, EvalDoneMessage, ServerMessage } from "../proto";
+import { NewEvalMessage, TypedValue, EvalFailedMessage, EvalDoneMessage, ServerMessage, EvalEventMessage, ClientMessage } from "../proto";
 import { SendableConnection } from "../common/connection";
+import { EventEmitter } from "events";
+
+export interface ActiveEvaluation {
+	onEvent(msg: EvalEventMessage): void;
+}
 
 declare var __non_webpack_require__: typeof require;
-export const evaluate = async (connection: SendableConnection, message: NewEvalMessage): Promise<void> => {
+export const evaluate = (connection: SendableConnection, message: NewEvalMessage, onDispose: () => void): ActiveEvaluation | void => {
 	const argStr: string[] = [];
 	message.getArgsList().forEach((value) => {
 		argStr.push(value);
@@ -50,18 +55,55 @@ export const evaluate = async (connection: SendableConnection, message: NewEvalM
 		serverMsg.setEvalFailed(evalFailed);
 		connection.send(serverMsg.serializeBinary());
 	};
+	let eventEmitter: EventEmitter | undefined;
 	try {
-		const value = vm.runInNewContext(`(${message.getFunction()})(${argStr.join(",")})`, {
+		if (message.getActive()) {
+			eventEmitter = new EventEmitter();
+		}
+
+		const value = vm.runInNewContext(`(${message.getFunction()})(${eventEmitter ? `eventEmitter, ` : ""}${argStr.join(",")})`, {
+			eventEmitter: eventEmitter ? {
+				on: (event: string, cb: (...args: any[]) => void): void => {
+					eventEmitter!.on(event, cb);
+				},
+				emit: (event: string, ...args: any[]): void => {
+					const eventMsg = new EvalEventMessage();
+					eventMsg.setEvent(event);
+					eventMsg.setArgsList(args.filter(a => a).map(a => JSON.stringify(a)));
+					eventMsg.setId(message.getId());
+					const serverMsg = new ServerMessage();
+					serverMsg.setEvalEvent(eventMsg);
+					connection.send(serverMsg.serializeBinary());
+				},
+			} : undefined,
 			Buffer,
 			require: typeof __non_webpack_require__ !== "undefined" ? __non_webpack_require__ : require,
 			_require: typeof __non_webpack_require__ !== "undefined" ? __non_webpack_require__ : require,
 			tslib_1: require("tslib"), // TODO: is there a better way to do this?
 			setTimeout,
 		}, {
-			timeout: message.getTimeout() || 30000,
-		});
-		sendResp(await value);
+				timeout: message.getTimeout() || 15000,
+			});
+		if (eventEmitter) {
+			// Is an active evaluation and should NOT be ended
+			eventEmitter.on("close", () => onDispose());
+			eventEmitter.on("error", () => onDispose());
+		} else {
+			if ((value as Promise<void>).then) {
+				// Is promise
+				(value as Promise<void>).then(r => sendResp(r)).catch(ex => sendErr(EvalFailedMessage.Reason.EXCEPTION, ex.toString()));
+			} else {
+				sendResp(value);
+			}
+			onDispose();
+		}
 	} catch (ex) {
 		sendErr(EvalFailedMessage.Reason.EXCEPTION, ex.toString());
 	}
+
+	return eventEmitter ? {
+		onEvent: (eventMsg: EvalEventMessage): void => {
+			eventEmitter!.emit(eventMsg.getEvent(), ...eventMsg.getArgsList().map(a => JSON.parse(a)));
+		},
+	} : undefined;
 };

packages/protocol/src/node/server.ts

@@ -6,7 +6,7 @@ import { promisify } from "util";
 import { TextDecoder } from "text-encoding";
 import { logger, field } from "@coder/logger";
 import { ClientMessage, WorkingInitMessage, ServerMessage, NewSessionMessage, WriteToSessionMessage } from "../proto";
-import { evaluate } from "./evaluate";
+import { evaluate, ActiveEvaluation } from "./evaluate";
 import { ReadWriteConnection } from "../common/connection";
 import { Process, handleNewSession, handleNewConnection, handleNewServer } from "./command";
 import * as net from "net";
@@ -23,6 +23,7 @@ export class Server {
 	private readonly sessions: Map<number, Process> = new Map();
 	private readonly connections: Map<number, net.Socket> = new Map();
 	private readonly servers: Map<number, net.Server> = new Map();
+	private readonly evals: Map<number, ActiveEvaluation> = new Map();
 
 	private connectionId: number = Number.MAX_SAFE_INTEGER;
 
@@ -114,7 +115,18 @@ export class Server {
 				field("args", evalMessage.getArgsList()),
 				field("function", evalMessage.getFunction()),
 			]);
-			evaluate(this.connection, evalMessage);
+			const resp = evaluate(this.connection, evalMessage, () => {
+				this.evals.delete(evalMessage.getId());
+			});
+			if (resp) {
+				this.evals.set(evalMessage.getId(), resp);
+			}
+		} else if (message.hasEvalEvent()) {
+			const e = this.evals.get(message.getEvalEvent()!.getId());
+			if (!e) {
+				return;
+			}
+			e.onEvent(message.getEvalEvent()!);
 		} else if (message.hasNewSession()) {
 			const sessionMessage = message.getNewSession()!;
 			logger.debug("NewSession", field("id", sessionMessage.getId()));