@highstate/backend 0.9.16 → 0.9.19

package/src/business/worker.ts

@@ -1,161 +1,297 @@
+import type { CommonObjectMeta, UnitWorker } from "@highstate/contract"
 import type { Logger } from "pino"
-import type { LockManager } from "../lock"
+import type {
+  DatabaseManager,
+  ProjectTransaction,
+  Worker,
+  WorkerVersion,
+  WorkerVersionLog,
+} from "../database"
+import type { PubSubManager } from "../pubsub"
 import type { WorkerManager } from "../worker"
-import type { InstanceStateService } from "./instance-state"
 import { randomBytes } from "node:crypto"
-import { v7 as uuidv7 } from "uuid"
+import { PrismaClientKnownRequestError } from "@prisma/client/runtime/client"
+import { createProjectLogger } from "../common"
 import {
+  extractDigestFromImage,
   getWorkerIdentity,
-  type InstanceState,
-  type ProjectApiKey,
-  type ServiceAccount,
-  type UnitWorker,
-  type Worker,
-  type WorkerUnitRegistration,
+  type WorkerUnitRegistrationEvent,
 } from "../shared"
-import { SAME_KEY, type StateManager } from "../state"
+import { WorkerVersionNotFoundError } from "../shared/models/errors"
 
 export class WorkerService {
   constructor(
-    private readonly stateManager: StateManager,
+    private readonly database: DatabaseManager,
     private readonly workerManager: WorkerManager,
-    private readonly lockManager: LockManager,
-    private readonly instanceStateService: InstanceStateService,
+    private readonly pubsubManager: PubSubManager,
     private readonly logger: Logger,
   ) {}
 
-  async handleUnitWorker(
+  /**
+   * Updates unit worker registrations for an instance in a single transaction
+   * creates workers and worker versions as needed and cleans up removed registrations
+   *
+   * @param tx The transaction to use for database operations.
+   * @param projectId The ID of the project containing the instance.
+   * @param stateId The ID of the instance state to update registrations for.
+   * @param unitWorkers The list workers provided by the unit.
+   */
+  async updateUnitRegistrations(
+    tx: ProjectTransaction,
     projectId: string,
-    state: InstanceState,
-    unitWorker: UnitWorker,
+    stateId: string,
+    unitWorkers: UnitWorker[],
   ): Promise<void> {
-    await this.lockManager.acquire(["worker-image", projectId, unitWorker.image], async () => {
-      const existingRegistrations = await this.stateManager
-        .getWorkerRegistrationRepository(projectId)
-        .getManyItems(state.extra?.workerRegistrationIds ?? [])
-
-      const existingReg = existingRegistrations.find(reg => reg.image === unitWorker.image)
-      if (existingReg && JSON.stringify(existingReg.params) === JSON.stringify(unitWorker.params)) {
-        this.logger.debug(
-          `instance "%s" already registered with worker "%s" with the same parameters`,
-          state.id,
-          unitWorker.image,
-        )
-        return
-      }
+    // parse images first
+    const parsedWorkers = unitWorkers.map(w => {
+      const digest = extractDigestFromImage(w.image)
+      const identity = getWorkerIdentity(w.image)
 
-      let workerId = existingReg?.workerId
-      if (!workerId) {
-        workerId = await this.ensureWorkerCreated(projectId, unitWorker)
-      }
+      return { ...w, digest, identity }
+    })
 
-      const registration: WorkerUnitRegistration = {
-        id: existingReg?.id ?? uuidv7(),
-        meta: existingReg?.meta ?? {},
-        image: unitWorker.image,
-        params: unitWorker.params,
-        instanceId: state.id,
-        workerId,
-      }
+    const logger = createProjectLogger(this.logger, projectId)
+
+    const eventsToPublish: { workerVersionId: string; event: WorkerUnitRegistrationEvent }[] = []
+
+    // query all registrations for the instance
+    const existingRegistrations = await tx.workerUnitRegistration.findMany({
+      where: { stateId },
+      select: { stateId: true, name: true, params: true, workerVersionId: true },
+    })
 
-      const batch = this.stateManager.batch()
+    // the set of names we want to keep
+    const desiredNames = new Set(parsedWorkers.map(w => w.name))
 
-      await this.stateManager
-        .getWorkerRegistrationRepository(projectId)
-        .putItem(registration, batch)
+    for (const worker of parsedWorkers) {
+      const workerRecord = await this.ensureWorker(tx, worker.identity)
+      const workerVersionRecord = await this.ensureWorkerVersion(tx, workerRecord, worker.digest)
 
-      if (!existingReg) {
-        // update the index if this is a new registration
-        await this.stateManager
-          .getWorkerRegistrationIndexRepository(projectId, workerId)
-          .indexRepository.put(registration.id, SAME_KEY, batch)
+      const existing = existingRegistrations.find(r => r.name === worker.name)
+      const stringifiedParams = JSON.stringify(worker.params)
 
-        // update the instance state with the new registration ID
-        await this.instanceStateService.patchInstanceState(projectId, {
-          id: state.id,
-          extra: {
-            workerRegistrationIds: [...(state.extra?.workerRegistrationIds ?? []), registration.id],
+      // create a new registration if it doesn't exist
+      if (!existing) {
+        await tx.workerUnitRegistration.create({
+          data: {
+            stateId,
+            name: worker.name,
+            params: worker.params,
+            workerVersionId: workerVersionRecord.id,
           },
         })
+
+        eventsToPublish.push({
+          workerVersionId: workerVersionRecord.id,
+          event: { type: "registered", instanceId: stateId, params: worker.params },
+        })
+
+        continue
       }
 
-      await batch.write()
-    })
-  }
+      const paramsChanged = JSON.stringify(existing.params) !== stringifiedParams
+      const versionChanged = existing.workerVersionId !== workerVersionRecord.id
 
-  private async ensureWorkerCreated(projectId: string, unitWorker: UnitWorker): Promise<string> {
-    const identity = getWorkerIdentity(unitWorker.image)
+      // update existing registration if params or version changed
+      if (paramsChanged || versionChanged) {
+        await tx.workerUnitRegistration.update({
+          where: { stateId_name: { stateId, name: worker.name } },
+          data: {
+            params: worker.params,
+            workerVersionId: workerVersionRecord.id,
+          },
+        })
 
-    return await this.lockManager.acquire(["worker-image", projectId, identity], async () => {
-      // it is not critical to fetch all workers since their count is not expected to be high
-      const workers = await this.stateManager.getWorkerRepository(projectId).getAllItems()
+        // deregister from old worker version if it changed
+        if (versionChanged) {
+          eventsToPublish.push({
+            workerVersionId: existing.workerVersionId,
+            event: { type: "deregistered", instanceId: stateId },
+          })
+        }
 
-      const existingWorker = workers.find(worker => worker.image === unitWorker.image)
-      if (existingWorker) {
-        this.logger.debug(`worker with image "%s" already exists, reusing it`, unitWorker.image)
-        return existingWorker.id
+        eventsToPublish.push({
+          workerVersionId: workerVersionRecord.id,
+          event: { type: "registered", instanceId: stateId, params: worker.params },
+        })
       }
+    }
 
-      let serviceAccountId: string | undefined
-
-      const siblingWorker = workers.find(worker => worker.identity === identity)
-      if (siblingWorker) {
-        this.logger.debug(
-          `sibling worker with identity "%s" already exists, using its service account`,
-          identity,
-        )
+    // remove registrations that are no longer desired
+    for (const registration of existingRegistrations) {
+      if (!desiredNames.has(registration.name)) {
+        await tx.workerUnitRegistration.delete({
+          where: { stateId_name: { stateId, name: registration.name } },
+        })
 
-        serviceAccountId = siblingWorker.serviceAccountId
+        eventsToPublish.push({
+          workerVersionId: registration.workerVersionId,
+          event: { type: "deregistered", instanceId: stateId },
+        })
       }
+    }
 
-      const batch = this.stateManager.batch()
+    await this.cleanupUnusedWorkerVersions(tx)
 
-      // create an empty service account if it is the first worker with this identity
-      // the meta of the account will be updated by the worker itself
-      if (!serviceAccountId) {
-        const serviceAccount: ServiceAccount = {
-          id: uuidv7(),
-          meta: {},
-        }
+    // publish events after transaction commits
+    for (const { workerVersionId, event } of eventsToPublish) {
+      void this.pubsubManager.publish(
+        ["worker-unit-registration", projectId, workerVersionId],
+        event,
+      )
+    }
 
-        serviceAccountId = serviceAccount.id
-        await this.stateManager
-          .getServiceAccountRepository(projectId)
-          .putItem(serviceAccount, batch)
-      }
+    // ensure all worker versions are started
+    void this.workerManager.syncWorkers(projectId)
 
-      const apiKey: ProjectApiKey = {
-        id: uuidv7(),
-        meta: {},
-        token: randomBytes(32).toString("hex"),
-        scopes: [
-          {
-            type: "service-account",
-            actions: ["full"],
-            serviceAccountIds: [serviceAccountId],
-          },
-        ],
-      }
+    logger.info(`updated worker registrations for instance state "%s"`, stateId)
+  }
+
+  private async ensureWorker(tx: ProjectTransaction, identity: string): Promise<Worker> {
+    const existing = await tx.worker.findUnique({ where: { identity } })
+    if (existing) {
+      return existing
+    }
 
-      await this.stateManager.getApiKeyRepository(projectId).putItem(apiKey, batch)
+    // create a new service account for the worker
+    const serviceAccount = await tx.serviceAccount.create({
+      select: { id: true },
+      data: {
+        meta: {
+          // this generic meta should be replaced by the worker when it starts
+          title: "Worker Service Account",
+          description: `Automatically created for worker "${identity}".`,
+        },
+      },
+    })
 
-      const worker: Worker = {
-        id: uuidv7(),
-        meta: {},
-        status: "starting",
-        failedStartAttempts: 5,
+    return await tx.worker.create({
+      data: {
         identity,
-        image: unitWorker.image,
-        serviceAccountId,
+        serviceAccountId: serviceAccount.id,
+      },
+    })
+  }
+
+  private async ensureWorkerVersion(
+    tx: ProjectTransaction,
+    worker: Worker,
+    digest: string,
+  ): Promise<WorkerVersion> {
+    const existing = await tx.workerVersion.findUnique({ where: { digest } })
+    if (existing) {
+      return existing
+    }
+
+    // create an API key for the worker granting full access to its service account
+    const apiKey = await tx.apiKey.create({
+      data: {
+        meta: {
+          title: `Worker API Key for "${worker.identity}"`,
+          description: `Automatically created for worker "${worker.identity}" with digest "${digest}".`,
+        },
+        serviceAccountId: worker.serviceAccountId,
+        token: randomBytes(32).toString("hex"),
+      },
+    })
+
+    return await tx.workerVersion.create({
+      data: {
+        workerId: worker.id,
+        digest,
+        meta: {
+          title: "Worker Version",
+          description: `Worker version with digest ${digest}`,
+        },
         apiKeyId: apiKey.id,
+      },
+    })
+  }
+
+  /**
+   * Performs cleanup of unused worker versions and syncs workers for a project.
+   * This method should be called after operations that may leave unused workers.
+   *
+   * @param projectId The ID of the project to cleanup and sync workers for.
+   */
+  async cleanupWorkerUsageAndSync(projectId: string): Promise<void> {
+    const database = await this.database.forProject(projectId)
+
+    await database.$transaction(async tx => {
+      await this.cleanupUnusedWorkerVersions(tx)
+    })
+
+    // ensure all worker versions are started
+    void this.workerManager.syncWorkers(projectId)
+  }
+
+  private async cleanupUnusedWorkerVersions(tx: ProjectTransaction): Promise<void> {
+    const unused = await tx.workerVersion.findMany({
+      where: {
+        unitRegistrations: { none: {} },
+      },
+      select: { id: true },
+    })
+
+    if (unused.length === 0) {
+      return
+    }
+
+    await tx.workerVersion.deleteMany({
+      where: { id: { in: unused.map(u => u.id) } },
+    })
+  }
+
+  /**
+   * Updates the metadata for a worker version.
+   *
+   * @param projectId The ID of the project.
+   * @param workerVersionId The ID of the worker version to update.
+   * @param meta The new metadata to set.
+   */
+  async updateWorkerVersionMeta(
+    projectId: string,
+    workerVersionId: string,
+    meta: CommonObjectMeta,
+  ): Promise<void> {
+    const database = await this.database.forProject(projectId)
+    const logger = createProjectLogger(this.logger, projectId)
+
+    try {
+      await database.workerVersion.update({
+        where: { id: workerVersionId },
+        data: { meta },
+      })
+    } catch (error) {
+      if (error instanceof PrismaClientKnownRequestError && error.code === "P2025") {
+        throw new WorkerVersionNotFoundError(projectId, workerVersionId)
       }
 
-      await this.stateManager.getWorkerRepository(projectId).putItem(worker, batch)
-      await batch.write()
+      throw error
+    }
 
-      this.workerManager.startWorker(projectId, worker)
+    logger.info(`updated worker version metadata for "%s"`, workerVersionId)
+  }
+
+  /**
+   * Gets logs for a worker version.
+   *
+   * @param projectId The ID of the project.
+   * @param workerVersionId The ID of the worker version to get logs for.
+   */
+  async getWorkerVersionLogs(
+    projectId: string,
+    workerVersionId: string,
+  ): Promise<WorkerVersionLog[]> {
+    const database = await this.database.forProject(projectId)
 
-      return worker.id
+    return await database.workerVersionLog.findMany({
+      where: {
+        workerVersionId,
+      },
+      orderBy: {
+        id: "asc",
+      },
     })
   }
 }

package/src/common/codebase.ts

@@ -0,0 +1,65 @@
+import type { Logger } from "pino"
+import { mkdir } from "node:fs/promises"
+import { dirname, relative } from "node:path"
+import { resolvePackageJSON } from "pkg-types"
+import z from "zod"
+
+export const codebaseConfig = z.object({
+  HIGHSTATE_CODEBASE_PATH: z.string().optional(),
+})
+
+let codebasePath: Promise<string> | undefined
+let codebaseHighstatePath: Promise<string> | undefined
+
+async function _getCodebasePath(
+  config: z.infer<typeof codebaseConfig>,
+  logger: Logger,
+): Promise<string> {
+  if (config.HIGHSTATE_CODEBASE_PATH) {
+    return config.HIGHSTATE_CODEBASE_PATH
+  }
+
+  const packageJson = await resolvePackageJSON()
+  const path = dirname(packageJson)
+
+  if (path !== process.cwd()) {
+    const relativePath = relative(process.cwd(), path)
+    logger.info(`detected "%s" as codebase path`, relativePath)
+  }
+
+  return path
+}
+
+export async function getCodebasePath(
+  config: z.infer<typeof codebaseConfig>,
+  logger: Logger,
+): Promise<string> {
+  if (!codebasePath) {
+    codebasePath = _getCodebasePath(config, logger)
+  }
+
+  return codebasePath
+}
+
+async function _getCodebaseHighstatePath(
+  config: z.infer<typeof codebaseConfig>,
+  logger: Logger,
+): Promise<string> {
+  const path = await getCodebasePath(config, logger)
+
+  const highstatePath = `${path}/.highstate`
+  await mkdir(highstatePath, { recursive: true })
+
+  return highstatePath
+}
+
+export async function getCodebaseHighstatePath(
+  config: z.infer<typeof codebaseConfig>,
+  logger: Logger,
+): Promise<string> {
+  if (!codebaseHighstatePath) {
+    codebaseHighstatePath = _getCodebaseHighstatePath(config, logger)
+  }
+
+  return codebaseHighstatePath
+}

package/src/common/index.ts

@@ -1,4 +1,5 @@
-export * from "./utils"
+export * from "./codebase"
 export * from "./local"
+export * from "./logger"
 export * from "./tree"
-export * from "./performance"
+export * from "./utils"

package/src/common/logger.ts

@@ -0,0 +1,5 @@
+import type { Logger } from "pino"
+
+export function createProjectLogger(logger: Logger, projectId: string): Logger {
+  return logger.child({ projectId }, { msgPrefix: `[project:${projectId}] ` })
+}

package/src/common/utils.ts

@@ -25,8 +25,9 @@ export async function runWithRetryOnError<T>(
 }
 
 export class AbortError extends Error {
-  constructor(options?: ErrorOptions) {
-    super("Operation aborted", options)
+  constructor(message: string, options?: ErrorOptions) {
+    super(message, options)
+    this.name = "AbortError"
   }
 }
 
@@ -58,7 +59,7 @@ export function isAbortErrorLike(error: unknown): boolean {
 
 export function tryWrapAbortErrorLike(error: unknown): unknown {
   if (isAbortErrorLike(error)) {
-    return new AbortError({ cause: error })
+    return new AbortError("Operation aborted", { cause: error })
   }
 
   return error

package/src/config.ts

@@ -1,13 +1,13 @@
 import { z } from "zod"
-import { libraryBackendConfig } from "./library"
-import { projectBackendConfig } from "./project"
-import { terminalBackendConfig } from "./terminal"
-import { runnerBackendConfig } from "./runner"
-import { stateBackendConfig, stateManagerConfig } from "./state"
 import { artifactBackendConfig } from "./artifact"
-import { pubSubBackendConfig } from "./pubsub"
+import { projectUnlockServiceConfig } from "./business"
+import { codebaseConfig } from "./common"
+import { databaseConfig } from "./database"
+import { libraryBackendConfig } from "./library"
 import { lockBackendConfig } from "./lock"
-import { hotStateBackendConfig } from "./hotstate"
+import { pubSubBackendConfig } from "./pubsub"
+import { runnerBackendConfig } from "./runner"
+import { terminalBackendConfig } from "./terminal"
 import { workerBackendConfig, workerManagerConfig } from "./worker"
 
 const loggerConfig = z.object({
@@ -15,13 +15,12 @@ const loggerConfig = z.object({
 })
 
 const configSchema = z.object({
-  ...hotStateBackendConfig.shape,
+  ...codebaseConfig.shape,
+  ...databaseConfig.shape,
   ...pubSubBackendConfig.shape,
   ...lockBackendConfig.shape,
   ...libraryBackendConfig.shape,
-  ...projectBackendConfig.shape,
-  ...stateBackendConfig.shape,
-  ...stateManagerConfig.shape,
+  ...projectUnlockServiceConfig.shape,
   ...runnerBackendConfig.shape,
   ...terminalBackendConfig.shape,
   ...workerBackendConfig.shape,
@@ -40,5 +39,9 @@ export async function loadConfig(
     await import("dotenv/config")
   }
 
-  return configSchema.parse(env)
+  try {
+    return configSchema.parse(env)
+  } catch (error) {
+    throw new Error("Failed to parse backend configuration", { cause: error })
+  }
 }

package/src/database/_generated/backend/postgresql/client.ts

@@ -0,0 +1,72 @@
+
+/* !!! This is code generated by Prisma. Do not edit directly. !!! */
+/* eslint-disable */
+// @ts-nocheck 
+/**
+ * This file should be your main import to use Prisma. Through it you get access to all the models, enums, and input types.
+ *
+ * 🟢 You can import this file directly.
+ */
+
+import * as process from 'node:process'
+import * as path from 'node:path'
+import { fileURLToPath } from 'node:url'
+const __dirname = path.dirname(fileURLToPath(import.meta.url))
+
+import * as runtime from "@prisma/client/runtime/client"
+import * as $Enums from "./enums.ts"
+import * as $Class from "./internal/class.ts"
+import * as Prisma from "./internal/prismaNamespace.ts"
+
+export * as $Enums from './enums.ts'
+/**
+ * ## Prisma Client
+ * 
+ * Type-safe database client for TypeScript
+ * @example
+ * ```
+ * const prisma = new PrismaClient()
+ * // Fetch zero or more UserWorkspaseLayouts
+ * const userWorkspaseLayouts = await prisma.userWorkspaseLayout.findMany()
+ * ```
+ * 
+ * Read more in our [docs](https://www.prisma.io/docs/reference/tools-and-interfaces/prisma-client).
+ */
+export const PrismaClient = $Class.getPrismaClientClass(__dirname)
+export type PrismaClient<LogOpts extends Prisma.LogLevel = never, OmitOpts extends Prisma.PrismaClientOptions["omit"] = Prisma.PrismaClientOptions["omit"], ExtArgs extends runtime.Types.Extensions.InternalArgs = runtime.Types.Extensions.DefaultArgs> = $Class.PrismaClient<LogOpts, OmitOpts, ExtArgs>
+export { Prisma }
+
+
+
+/**
+ * Model UserWorkspaseLayout
+ * 
+ */
+export type UserWorkspaseLayout = Prisma.UserWorkspaseLayoutModel
+/**
+ * Model Library
+ * 
+ */
+export type Library = Prisma.LibraryModel
+/**
+ * Model Project
+ * 
+ */
+export type Project = Prisma.ProjectModel
+/**
+ * Model ProjectSpace
+ * 
+ */
+export type ProjectSpace = Prisma.ProjectSpaceModel
+/**
+ * Model ProjectModelStorage
+ * 
+ */
+export type ProjectModelStorage = Prisma.ProjectModelStorageModel
+/**
+ * Model PulumiBackend
+ * 
+ */
+export type PulumiBackend = Prisma.PulumiBackendModel
+
+