@highstate/backend 0.9.18 → 0.9.19

package/src/business/worker.ts

@@ -1,398 +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 { RandomProvider } from "../common"
+import { randomBytes } from "node:crypto"
+import { PrismaClientKnownRequestError } from "@prisma/client/runtime/client"
+import { createProjectLogger } from "../common"
 import {
+  extractDigestFromImage,
   getWorkerIdentity,
-  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 random: RandomProvider,
+    private readonly pubsubManager: PubSubManager,
     private readonly logger: Logger,
   ) {}
 
   /**
-   * Updates the worker registrations for the given project and instance.
-   * It creates new registrations for each unit worker, updates existing ones,
-   * and deletes registrations that are no longer present.
+   * Updates unit worker registrations for an instance in a single transaction
+   * creates workers and worker versions as needed and cleans up removed registrations
    *
-   * @param projectId The ID of the project.
-   * @param instanceId The ID of the instance.
-   * @param existingRegistrationIds A mapping of unit worker names to existing registration IDs.
-   * @param unitWorkers The list of unit workers to register.
-   * @returns A new mapping of unit worker names to their registration IDs.
+   * @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,
-    instanceId: string,
-    existingRegistrationIds: Record<string, string>,
+    stateId: string,
     unitWorkers: UnitWorker[],
-  ): Promise<Record<string, string>> {
-    // we will lock on the worker images to ensure that
-    // concurrent units cannot create duplicate workers or leave dangling worker resources when removing them
-    // we will not lock on the registrations
-    // since it they are only updated within an operation
-    const lockKeys = unitWorkers.map(
-      unitWorker => ["worker-image", projectId, unitWorker.image] as const,
-    )
-
-    const newRegistrationIds: Record<string, string> = {}
-
-    return await this.lockManager.acquire(lockKeys, async () => {
-      const existingRegistrations = await this.stateManager
-        .getWorkerRegistrationRepository(projectId)
-        .getManyRecord(Object.values(existingRegistrationIds))
-
-      const nameToExistingRegistration = new Map<string, WorkerUnitRegistration>()
-      for (const [name, registrationId] of Object.entries(existingRegistrationIds)) {
-        const existingRegistration = existingRegistrations[registrationId]
-        if (existingRegistration) {
-          nameToExistingRegistration.set(name, existingRegistration)
-        }
-      }
-
-      const batch = this.stateManager.batch()
-
-      // create or update registrations for each unit worker
-      for (const unitWorker of unitWorkers) {
-        const workerId = await this.ensureWorkerCreated(projectId, unitWorker)
-        const existingRegistration = nameToExistingRegistration.get(unitWorker.name)
-
-        if (
-          existingRegistration &&
-          existingRegistration.workerId === workerId &&
-          JSON.stringify(existingRegistration.params) === JSON.stringify(unitWorker.params)
-        ) {
-          // no changes
-          newRegistrationIds[unitWorker.name] = existingRegistration.id
-          continue
-        }
-
-        const registration: WorkerUnitRegistration = {
-          id: existingRegistration?.id ?? this.random.uuidv7(),
-          meta: existingRegistration?.meta ?? {},
-          image: unitWorker.image,
-          params: unitWorker.params,
-          instanceId,
-          workerId,
-        }
+  ): Promise<void> {
+    // parse images first
+    const parsedWorkers = unitWorkers.map(w => {
+      const digest = extractDigestFromImage(w.image)
+      const identity = getWorkerIdentity(w.image)
 
-        const isNewRegistration = !existingRegistration
-        const isWorkerChange = existingRegistration?.workerId !== workerId
-        const isParamsChange =
-          existingRegistration &&
-          JSON.stringify(existingRegistration.params) !== JSON.stringify(unitWorker.params)
-
-        if (isNewRegistration) {
-          this.logger.info(
-            {
-              projectId,
-              registrationId: registration.id,
-              unitWorkerName: unitWorker.name,
-              workerId,
-              image: unitWorker.image,
-            },
-            `creating worker registration "%s" for unit worker "%s" in project "%s"`,
-            registration.id,
-            unitWorker.name,
-            projectId,
-          )
-        } else if (isWorkerChange) {
-          this.logger.info(
-            {
-              projectId,
-              registrationId: registration.id,
-              unitWorkerName: unitWorker.name,
-              oldWorkerId: existingRegistration?.workerId,
-              newWorkerId: workerId,
-              oldImage: existingRegistration?.image,
-              newImage: unitWorker.image,
-            },
-            `updating worker registration "%s" for unit worker "%s" in project "%s" (worker changed from "%s" to "%s")`,
-            registration.id,
-            unitWorker.name,
-            projectId,
-            existingRegistration?.workerId,
-            workerId,
-          )
-        } else if (isParamsChange) {
-          this.logger.info(
-            {
-              projectId,
-              registrationId: registration.id,
-              unitWorkerName: unitWorker.name,
-              workerId,
-            },
-            `updating worker registration "%s" for unit worker "%s" in project "%s" (params changed)`,
-            registration.id,
-            unitWorker.name,
-            projectId,
-          )
-        }
+      return { ...w, digest, identity }
+    })
 
-        await this.stateManager
-          .getWorkerRegistrationRepository(projectId)
-          .putItem(registration, batch)
+    const logger = createProjectLogger(this.logger, projectId)
 
-        // update the registration indexes for new and old workers
-        if (existingRegistration?.workerId !== workerId) {
-          await this.stateManager
-            .getWorkerRegistrationIndexRepository(projectId, workerId)
-            .indexRepository.put(registration.id, SAME_KEY)
+    const eventsToPublish: { workerVersionId: string; event: WorkerUnitRegistrationEvent }[] = []
 
-          if (existingRegistration?.workerId) {
-            await this.stateManager
-              .getWorkerRegistrationIndexRepository(projectId, existingRegistration.workerId)
-              .indexRepository.delete(registration.id)
+    // query all registrations for the instance
+    const existingRegistrations = await tx.workerUnitRegistration.findMany({
+      where: { stateId },
+      select: { stateId: true, name: true, params: true, workerVersionId: true },
+    })
 
-            // check if the old worker should be deleted
-            await this.deleteWorkerIfHasNoRegistrations(projectId, existingRegistration.workerId)
-          }
-        }
+    // the set of names we want to keep
+    const desiredNames = new Set(parsedWorkers.map(w => w.name))
 
-        newRegistrationIds[unitWorker.name] = registration.id
-      }
+    for (const worker of parsedWorkers) {
+      const workerRecord = await this.ensureWorker(tx, worker.identity)
+      const workerVersionRecord = await this.ensureWorkerVersion(tx, workerRecord, worker.digest)
 
-      // delete registrations that are no longer present
-      for (const [name, existingRegistration] of nameToExistingRegistration.entries()) {
-        if (unitWorkers.some(unitWorker => unitWorker.name === name)) {
-          continue
-        }
+      const existing = existingRegistrations.find(r => r.name === worker.name)
+      const stringifiedParams = JSON.stringify(worker.params)
 
-        this.logger.info(
-          {
-            projectId,
-            registrationId: existingRegistration.id,
-            unitWorkerName: name,
-            workerId: existingRegistration.workerId,
+      // 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,
           },
-          `deleting worker registration "%s" for unit worker "%s" in project "%s" (unit worker no longer present)`,
-          existingRegistration.id,
-          name,
-          projectId,
-        )
-
-        await this.stateManager
-          .getWorkerRegistrationRepository(projectId)
-          .delete(existingRegistration.id, batch)
-
-        await this.stateManager
-          .getWorkerRegistrationIndexRepository(projectId, existingRegistration.workerId)
-          .indexRepository.delete(existingRegistration.id)
-
-        // ensure the worker is deleted if it has no registrations left
-        await this.deleteWorkerIfHasNoRegistrations(projectId, existingRegistration.workerId)
-      }
-
-      await batch.write()
+        })
 
-      return newRegistrationIds
-    })
-  }
+        eventsToPublish.push({
+          workerVersionId: workerVersionRecord.id,
+          event: { type: "registered", instanceId: stateId, params: worker.params },
+        })
 
-  private async ensureWorkerCreated(projectId: string, unitWorker: UnitWorker): Promise<string> {
-    const identity = getWorkerIdentity(unitWorker.image)
-
-    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()
-
-      const existingWorker = workers.find(worker => worker.image === unitWorker.image)
-      if (existingWorker) {
-        this.logger.debug(
-          { projectId, workerId: existingWorker.id, image: unitWorker.image },
-          `worker with image "%s" already exists, reusing it`,
-          unitWorker.image,
-        )
-        return existingWorker.id
+        continue
       }
 
-      this.logger.info(
-        { projectId, image: unitWorker.image, identity },
-        `creating new worker for image "%s" with identity "%s" in project "%s"`,
-        unitWorker.image,
-        identity,
-        projectId,
-      )
-
-      let serviceAccountId: string | undefined
+      const paramsChanged = JSON.stringify(existing.params) !== stringifiedParams
+      const versionChanged = existing.workerVersionId !== workerVersionRecord.id
 
-      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,
-        )
+      // 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,
+          },
+        })
+
+        // deregister from old worker version if it changed
+        if (versionChanged) {
+          eventsToPublish.push({
+            workerVersionId: existing.workerVersionId,
+            event: { type: "deregistered", instanceId: stateId },
+          })
+        }
 
-        serviceAccountId = siblingWorker.serviceAccountId
+        eventsToPublish.push({
+          workerVersionId: workerVersionRecord.id,
+          event: { type: "registered", instanceId: stateId, params: worker.params },
+        })
       }
+    }
+
+    // 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 } },
+        })
+
+        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: this.random.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
+    // ensure all worker versions are started
+    void this.workerManager.syncWorkers(projectId)
 
-        this.logger.info(
-          { projectId, serviceAccountId, identity },
-          `creating service account "%s" for worker identity "%s" in project "%s"`,
-          serviceAccountId,
-          identity,
-          projectId,
-        )
+    logger.info(`updated worker registrations for instance state "%s"`, stateId)
+  }
 
-        await this.stateManager
-          .getServiceAccountRepository(projectId)
-          .putItem(serviceAccount, batch)
-      }
+  private async ensureWorker(tx: ProjectTransaction, identity: string): Promise<Worker> {
+    const existing = await tx.worker.findUnique({ where: { identity } })
+    if (existing) {
+      return existing
+    }
+
+    // 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 apiKey: ProjectApiKey = {
-        id: this.random.uuidv7(),
-        meta: {},
-        token: Buffer.from(this.random.bytes(32)).toString("hex"),
-        scopes: [
-          {
-            type: "service-account",
-            actions: ["full"],
-            serviceAccountIds: [serviceAccountId],
-          },
-        ],
-      }
+    return await tx.worker.create({
+      data: {
+        identity,
+        serviceAccountId: serviceAccount.id,
+      },
+    })
+  }
 
-      await this.stateManager.getApiKeyRepository(projectId).putItem(apiKey, batch)
+  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"),
+      },
+    })
 
-      const worker: Worker = {
-        id: this.random.uuidv7(),
-        meta: {},
-        status: "starting",
-        failedStartAttempts: 5,
-        identity,
-        image: unitWorker.image,
-        serviceAccountId,
+    return await tx.workerVersion.create({
+      data: {
+        workerId: worker.id,
+        digest,
+        meta: {
+          title: "Worker Version",
+          description: `Worker version with digest ${digest}`,
+        },
         apiKeyId: apiKey.id,
-      }
+      },
+    })
+  }
 
-      this.logger.info(
-        { projectId, workerId: worker.id, image: unitWorker.image, identity, serviceAccountId },
-        `creating worker "%s" for image "%s" in project "%s"`,
-        worker.id,
-        unitWorker.image,
-        projectId,
-      )
+  /**
+   * 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)
+    })
 
-      await this.stateManager.getWorkerRepository(projectId).putItem(worker, batch)
-      await batch.write()
+    // ensure all worker versions are started
+    void this.workerManager.syncWorkers(projectId)
+  }
 
-      void this.workerManager.startWorker(projectId, worker.id)
+  private async cleanupUnusedWorkerVersions(tx: ProjectTransaction): Promise<void> {
+    const unused = await tx.workerVersion.findMany({
+      where: {
+        unitRegistrations: { none: {} },
+      },
+      select: { id: true },
+    })
 
-      return worker.id
+    if (unused.length === 0) {
+      return
+    }
+
+    await tx.workerVersion.deleteMany({
+      where: { id: { in: unused.map(u => u.id) } },
     })
   }
 
-  private async deleteWorkerIfHasNoRegistrations(
+  /**
+   * 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,
-    workerId: string,
+    workerVersionId: string,
+    meta: CommonObjectMeta,
   ): Promise<void> {
-    await this.lockManager.acquire(["worker", workerId], async () => {
-      const registrations = await this.stateManager
-        .getWorkerRegistrationIndexRepository(projectId, workerId)
-        .getAllItems()
-
-      if (registrations.length > 0) {
-        // still has registrations, no need to delete
-        this.logger.debug(
-          { projectId, workerId, registrationCount: registrations.length },
-          `worker "%s" still has %d registrations, not deleting`,
-          workerId,
-          registrations.length,
-        )
-        return
+    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)
       }
 
-      const worker = await this.stateManager.getWorkerRepository(projectId).get(workerId)
+      throw error
+    }
 
-      if (!worker) {
-        this.logger.warn(
-          { projectId, workerId },
-          `worker "%s" not found in project "%s" while deleting`,
-          workerId,
-          projectId,
-        )
-        return
-      }
-
-      const batch = this.stateManager.batch()
-
-      this.logger.info(
-        { projectId, workerId, image: worker.image, identity: worker.identity },
-        `deleting worker "%s" with image "%s" in project "%s" (no registrations remaining)`,
-        workerId,
-        worker.image,
-        projectId,
-      )
-
-      await this.stateManager.getWorkerRepository(projectId).delete(workerId, batch)
-      await this.stateManager.getApiKeyRepository(projectId).delete(worker.apiKeyId, batch)
-
-      const workers = await this.stateManager.getWorkerRepository(projectId).getAllItems()
-      const hasSiblingWorker = workers.some(
-        siblingWorker =>
-          siblingWorker.identity === worker.identity && siblingWorker.id !== workerId,
-      )
-
-      if (!hasSiblingWorker) {
-        this.logger.info(
-          {
-            projectId,
-            workerId,
-            serviceAccountId: worker.serviceAccountId,
-            identity: worker.identity,
-          },
-          `deleting service account "%s" for worker "%s" in project "%s" (no sibling workers remaining)`,
-          worker.serviceAccountId,
-          workerId,
-          projectId,
-        )
-
-        await this.stateManager
-          .getServiceAccountRepository(projectId)
-          .delete(worker.serviceAccountId, batch)
-      } else {
-        this.logger.debug(
-          {
-            projectId,
-            workerId,
-            serviceAccountId: worker.serviceAccountId,
-            identity: worker.identity,
-          },
-          `not deleting service account "%s" for worker "%s" in project "%s" (has sibling workers)`,
-          worker.serviceAccountId,
-          workerId,
-          projectId,
-        )
-      }
-
-      // for now, we will keep the service account even if the last sibling worker is deleted
-      await batch.write()
+    logger.info(`updated worker version metadata for "%s"`, workerVersionId)
+  }
 
-      // stop the worker and clear logs after it
-      this.workerManager.stopWorker(projectId, workerId)
+  /**
+   * 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 await database.workerVersionLog.findMany({
+      where: {
+        workerVersionId,
+      },
+      orderBy: {
+        id: "asc",
+      },
     })
   }
 }