@highstate/backend 0.9.15 → 0.9.18

package/src/worker/manager.ts

@@ -0,0 +1,167 @@
+import type { Logger } from "pino"
+import type { WorkerBackend } from "./abstractions"
+import type { StateManager } from "../state"
+import type { ApiKeyService, ProjectUnlockService } from "../business"
+import type { LockManager } from "../lock"
+import { PassThrough } from "node:stream"
+import { z } from "zod"
+import c from "ansi-colors"
+import { v7 as uuidv7 } from "uuid"
+import { createAsyncBatcher } from "../shared"
+
+export const workerManagerConfig = z.object({
+  HIGHSTATE_WORKER_API_PATH: z.string().default("/var/run/highstate.sock"),
+})
+
+export class WorkerManager {
+  constructor(
+    readonly config: z.infer<typeof workerManagerConfig>,
+    private readonly workerBackend: WorkerBackend,
+    private readonly stateManager: StateManager,
+    private readonly projectUnlockService: ProjectUnlockService,
+    private readonly lockManager: LockManager,
+    private readonly apiKeyService: ApiKeyService,
+    private readonly logger: Logger,
+  ) {
+    this.projectUnlockService.registerUnlockTask(
+      //
+      "launch-active-workers",
+      projectId => this.launchActiveWorkers(projectId),
+    )
+  }
+
+  private readonly workerAbortControllers = new Map<string, AbortController>()
+
+  async startWorker(projectId: string, workerId: string, restart = false): Promise<void> {
+    if (this.workerAbortControllers.get(workerId)?.signal.aborted) {
+      // if the worker was aborted due to worker deletion, clear logs and exit
+      await this.stateManager.getWorkerLogRepository(projectId, workerId).clear()
+      return
+    }
+
+    await this.lockManager.acquire(["worker", workerId], async () => {
+      const worker = await this.stateManager.getWorkerRepository(projectId).get(workerId)
+      if (!worker) {
+        throw new Error(`Worker "${workerId}" not found in project "${projectId}"`)
+      }
+
+      if (!worker.apiKeyId) {
+        throw new Error(`Worker "${workerId}" does not have an API key associated`)
+      }
+
+      if (worker.failedStartAttempts === 5) {
+        this.logger.debug(
+          { projectId },
+          `skipping worker "%s" start attempt since no attempts left`,
+          worker.id,
+        )
+
+        return
+      }
+
+      worker.status = "starting"
+
+      if (restart) {
+        worker.failedStartAttempts += 1
+      }
+
+      await this.stateManager.getWorkerRepository(projectId).putItem(worker)
+
+      const apiKey = await this.apiKeyService.regenerateToken(projectId, worker.apiKeyId)
+      const stdout = new PassThrough()
+
+      await this.writeSystemMessage(
+        projectId,
+        worker.id,
+        `starting worker container, attempt ${worker.failedStartAttempts + 1}/5`,
+      )
+
+      const logBatcher = createAsyncBatcher<string>(async logs => {
+        await this.stateManager
+          .getWorkerLogRepository(projectId, worker.id)
+          .putMany(logs.map(log => [uuidv7(), log]))
+      })
+
+      stdout.on("data", chunk => {
+        const log = String(chunk)
+
+        logBatcher.call(log)
+      })
+
+      const abortController = new AbortController()
+      this.workerAbortControllers.set(workerId, abortController)
+
+      void this.workerBackend
+        .run({
+          projectId,
+          worker,
+          apiPath: this.config.HIGHSTATE_WORKER_API_PATH,
+          apiKey: apiKey.token,
+          stdout,
+          signal: abortController.signal,
+        })
+        // regardless the exit reason, we want to restart the worker if it has remaining attempts
+        .finally(() => void this.startWorker(projectId, worker.id, true))
+    })
+  }
+
+  async writeSystemMessage(projectId: string, workerId: string, message: string): Promise<void> {
+    const prefix = c.bold.magenta("[runtime]")
+
+    await this.stateManager
+      .getWorkerLogRepository(projectId, workerId)
+      .put(uuidv7(), `${prefix} ${message}`)
+  }
+
+  async setWorkerRunning(projectId: string, workerId: string): Promise<void> {
+    await this.lockManager.acquire(["worker", workerId], async () => {
+      const worker = await this.stateManager.getWorkerRepository(projectId).get(workerId)
+      if (!worker) {
+        throw new Error(`Worker "${workerId}" not found in project "${projectId}"`)
+      }
+
+      worker.status = "running"
+      worker.failedStartAttempts = 0
+
+      await this.stateManager.getWorkerRepository(projectId).putItem(worker)
+
+      this.logger.debug(
+        { projectId, workerId },
+        `worker "%s" is now running in project "%s"`,
+        workerId,
+        projectId,
+      )
+    })
+  }
+
+  stopWorker(projectId: string, workerId: string): void {
+    const abortController = this.workerAbortControllers.get(workerId)
+    if (!abortController) {
+      this.logger.warn(
+        { projectId, workerId },
+        `worker "%s" is not running in project "%s", cannot stop`,
+        workerId,
+        projectId,
+      )
+
+      return
+    }
+
+    abortController.abort()
+  }
+
+  private async launchActiveWorkers(projectId: string): Promise<void> {
+    const workers = await this.stateManager.getWorkerRepository(projectId).getAllItems()
+
+    this.logger.debug(
+      { projectId, workerCount: workers.length },
+      `launching %s active workers for project "%s"`,
+      workers.length,
+      projectId,
+    )
+
+    for (const worker of workers) {
+      void this.startWorker(projectId, worker.id)
+    }
+  }
+}