diff --git a/src/exec/gateway.ts b/src/exec/gateway.ts new file mode 100644 index 0000000..dc0c849 --- /dev/null +++ b/src/exec/gateway.ts @@ -0,0 +1,66 @@ +import { runtimeStore } from "../runtime/store.js"; +import { readVsockResponse } from "../runtime/protocol.js"; +import { getVmSocket } from "../runtime/transport.js"; +import { getSession, createSession, touchSession } from "./session.js"; +import { gatewayLogger } from "../utils/logger.js"; +import crypto from "crypto"; +import type { Vm } from "../types/types.js"; + +import { createVm } from "../runtime/vm-manager.js"; +import { Deque } from "../runtime/deque.js"; + +export async function ensureSession(sessionId: string): Promise { + let session = getSession(sessionId); + if (session?.state === "active") { + const fn = runtimeStore.functions.get(sessionId); + if (fn && fn.vms.length > 0) return fn.vms[0]!; + } + + session = session || createSession(sessionId); + + let fn = runtimeStore.functions.get(sessionId); + if (!fn) { + fn = { + functionId: sessionId, + queue: new Deque(), + vms: [], + readyVms: new Set(), + weight: 1, + inflightCount: 0, + deficit: 0, + pendingCreations: 0, + }; + runtimeStore.functions.set(sessionId, fn); + } + + if (fn.vms.length === 0) { + await createVm(sessionId, fn, "__exec__"); + } + + session.state = "active"; + return fn.vms[0]!; +} + +export async function sendSessionMessage( + sessionId: string, + message: Record, + onStream?: (chunk: any) => void, + timeout: number = 60000, +): Promise { + const vm = await ensureSession(sessionId); + touchSession(sessionId); + + const id = message.id || crypto.randomUUID(); + const fullMessage = { ...message, id }; + + const socket = await getVmSocket(vm); + socket.write(JSON.stringify(fullMessage) + "\n"); + + gatewayLogger.debug( + { sessionId, messageType: message.type, messageId: id }, + "message sent to VM" + ); + + const result = await readVsockResponse(socket, timeout, onStream); + return { ...result, messageId: id }; +} diff --git a/src/runtime/vm-manager.ts b/src/runtime/vm-manager.ts index 172b3e7..5dca8dc 100644 --- a/src/runtime/vm-manager.ts +++ b/src/runtime/vm-manager.ts @@ -11,6 +11,7 @@ import type { RuntimeFunction, Vm } from "../types/types.js"; export async function createVm( functionId: string, fn: RuntimeFunction, + snapshotId?: string, ): Promise { const instanceId = crypto.randomBytes(4).toString("hex"); const apiSock = `/tmp/firecracker-${functionId}-${instanceId}.socket`; @@ -38,7 +39,7 @@ export async function createVm( await waitForFirecrackerApiSocket(apiSock); const client = createFcClient(apiSock); - await restoreVm(client, functionId, vsock); + await restoreVm(client, snapshotId || functionId, vsock); const vm: Vm = { id: instanceId, diff --git a/src/utils/logger.ts b/src/utils/logger.ts index cdc7d0c..7cf2fa6 100644 --- a/src/utils/logger.ts +++ b/src/utils/logger.ts @@ -33,7 +33,8 @@ export const vmManagerLogger = logger.child({ module: "vm-manager" }); export const transportLogger = logger.child({ module: "transport" }); export const protocolLogger = logger.child({ module: "protocol" }); export const cleanupLogger = logger.child({ module: "cleanup" }); -export const sessionLogger = logger.child({module: "session"}) +export const sessionLogger = logger.child({ module: "session" }) +export const gatewayLogger = logger.child({module: "gateway"}) export const httpLoggerOptions: Options = { logger: logger.child({ module: "http" }),