@highstate/backend 0.9.16 → 0.9.18

package/src/business/worker.ts

@@ -1,12 +1,9 @@
 import type { Logger } from "pino"
 import type { LockManager } from "../lock"
 import type { WorkerManager } from "../worker"
-import type { InstanceStateService } from "./instance-state"
-import { randomBytes } from "node:crypto"
-import { v7 as uuidv7 } from "uuid"
+import type { RandomProvider } from "../common"
 import {
   getWorkerIdentity,
-  type InstanceState,
   type ProjectApiKey,
   type ServiceAccount,
   type UnitWorker,
@@ -20,66 +17,186 @@ export class WorkerService {
     private readonly stateManager: StateManager,
     private readonly workerManager: WorkerManager,
     private readonly lockManager: LockManager,
-    private readonly instanceStateService: InstanceStateService,
+    private readonly random: RandomProvider,
     private readonly logger: Logger,
   ) {}
 
-  async handleUnitWorker(
+  /**
+   * 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.
+   *
+   * @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.
+   */
+  async updateUnitRegistrations(
     projectId: string,
-    state: InstanceState,
-    unitWorker: UnitWorker,
-  ): Promise<void> {
-    await this.lockManager.acquire(["worker-image", projectId, unitWorker.image], async () => {
+    instanceId: string,
+    existingRegistrationIds: Record<string, 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)
-        .getManyItems(state.extra?.workerRegistrationIds ?? [])
+        .getManyRecord(Object.values(existingRegistrationIds))
 
-      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
+      const nameToExistingRegistration = new Map<string, WorkerUnitRegistration>()
+      for (const [name, registrationId] of Object.entries(existingRegistrationIds)) {
+        const existingRegistration = existingRegistrations[registrationId]
+        if (existingRegistration) {
+          nameToExistingRegistration.set(name, existingRegistration)
+        }
       }
 
-      let workerId = existingReg?.workerId
-      if (!workerId) {
-        workerId = await this.ensureWorkerCreated(projectId, unitWorker)
-      }
+      const batch = this.stateManager.batch()
 
-      const registration: WorkerUnitRegistration = {
-        id: existingReg?.id ?? uuidv7(),
-        meta: existingReg?.meta ?? {},
-        image: unitWorker.image,
-        params: unitWorker.params,
-        instanceId: state.id,
-        workerId,
-      }
+      // 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)
 
-      const batch = this.stateManager.batch()
+        if (
+          existingRegistration &&
+          existingRegistration.workerId === workerId &&
+          JSON.stringify(existingRegistration.params) === JSON.stringify(unitWorker.params)
+        ) {
+          // no changes
+          newRegistrationIds[unitWorker.name] = existingRegistration.id
+          continue
+        }
 
-      await this.stateManager
-        .getWorkerRegistrationRepository(projectId)
-        .putItem(registration, batch)
+        const registration: WorkerUnitRegistration = {
+          id: existingRegistration?.id ?? this.random.uuidv7(),
+          meta: existingRegistration?.meta ?? {},
+          image: unitWorker.image,
+          params: unitWorker.params,
+          instanceId,
+          workerId,
+        }
+
+        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,
+          )
+        }
 
-      if (!existingReg) {
-        // update the index if this is a new registration
         await this.stateManager
-          .getWorkerRegistrationIndexRepository(projectId, workerId)
-          .indexRepository.put(registration.id, SAME_KEY, batch)
-
-        // update the instance state with the new registration ID
-        await this.instanceStateService.patchInstanceState(projectId, {
-          id: state.id,
-          extra: {
-            workerRegistrationIds: [...(state.extra?.workerRegistrationIds ?? []), registration.id],
+          .getWorkerRegistrationRepository(projectId)
+          .putItem(registration, batch)
+
+        // 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)
+
+          if (existingRegistration?.workerId) {
+            await this.stateManager
+              .getWorkerRegistrationIndexRepository(projectId, existingRegistration.workerId)
+              .indexRepository.delete(registration.id)
+
+            // check if the old worker should be deleted
+            await this.deleteWorkerIfHasNoRegistrations(projectId, existingRegistration.workerId)
+          }
+        }
+
+        newRegistrationIds[unitWorker.name] = registration.id
+      }
+
+      // delete registrations that are no longer present
+      for (const [name, existingRegistration] of nameToExistingRegistration.entries()) {
+        if (unitWorkers.some(unitWorker => unitWorker.name === name)) {
+          continue
+        }
+
+        this.logger.info(
+          {
+            projectId,
+            registrationId: existingRegistration.id,
+            unitWorkerName: name,
+            workerId: existingRegistration.workerId,
           },
-        })
+          `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
     })
   }
 
@@ -92,10 +209,22 @@ export class WorkerService {
 
       const existingWorker = workers.find(worker => worker.image === unitWorker.image)
       if (existingWorker) {
-        this.logger.debug(`worker with image "%s" already exists, reusing it`, unitWorker.image)
+        this.logger.debug(
+          { projectId, workerId: existingWorker.id, image: unitWorker.image },
+          `worker with image "%s" already exists, reusing it`,
+          unitWorker.image,
+        )
         return existingWorker.id
       }
 
+      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 siblingWorker = workers.find(worker => worker.identity === identity)
@@ -114,20 +243,29 @@ export class WorkerService {
       // the meta of the account will be updated by the worker itself
       if (!serviceAccountId) {
         const serviceAccount: ServiceAccount = {
-          id: uuidv7(),
+          id: this.random.uuidv7(),
           meta: {},
         }
 
         serviceAccountId = serviceAccount.id
+
+        this.logger.info(
+          { projectId, serviceAccountId, identity },
+          `creating service account "%s" for worker identity "%s" in project "%s"`,
+          serviceAccountId,
+          identity,
+          projectId,
+        )
+
         await this.stateManager
           .getServiceAccountRepository(projectId)
           .putItem(serviceAccount, batch)
       }
 
       const apiKey: ProjectApiKey = {
-        id: uuidv7(),
+        id: this.random.uuidv7(),
         meta: {},
-        token: randomBytes(32).toString("hex"),
+        token: Buffer.from(this.random.bytes(32)).toString("hex"),
         scopes: [
           {
             type: "service-account",
@@ -140,7 +278,7 @@ export class WorkerService {
       await this.stateManager.getApiKeyRepository(projectId).putItem(apiKey, batch)
 
       const worker: Worker = {
-        id: uuidv7(),
+        id: this.random.uuidv7(),
         meta: {},
         status: "starting",
         failedStartAttempts: 5,
@@ -150,12 +288,111 @@ export class WorkerService {
         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,
+      )
+
       await this.stateManager.getWorkerRepository(projectId).putItem(worker, batch)
       await batch.write()
 
-      this.workerManager.startWorker(projectId, worker)
+      void this.workerManager.startWorker(projectId, worker.id)
 
       return worker.id
     })
   }
+
+  private async deleteWorkerIfHasNoRegistrations(
+    projectId: string,
+    workerId: string,
+  ): 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 worker = await this.stateManager.getWorkerRepository(projectId).get(workerId)
+
+      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()
+
+      // stop the worker and clear logs after it
+      this.workerManager.stopWorker(projectId, workerId)
+    })
+  }
 }

package/src/common/clock.ts

@@ -0,0 +1,18 @@
+export interface ClockProvider {
+  /**
+   * Returns the current time in milliseconds since the Unix epoch.
+   */
+  now(): number
+}
+
+export class SystemClockProvider implements ClockProvider {
+  now(): number {
+    return Date.now()
+  }
+}
+
+export class ReproducibleClockProvider implements ClockProvider {
+  now(): number {
+    return 0
+  }
+}

package/src/common/index.ts

@@ -2,3 +2,6 @@ export * from "./utils"
 export * from "./local"
 export * from "./tree"
 export * from "./performance"
+export * from "./test"
+export * from "./clock"
+export * from "./random"

package/src/common/random.ts

@@ -0,0 +1,68 @@
+import { randomBytes, randomUUID } from "node:crypto"
+import { v4 as uuidv4, v7 as uuidv7 } from "uuid"
+
+export interface RandomProvider {
+  /**
+   * Generates a random string of the specified length.
+   *
+   * @param length The length of the random string to generate.
+   */
+  bytes(length: number): Uint8Array
+
+  /**
+   * Generates a random UUID (version 4).
+   */
+  uuidv4(): string
+
+  /**
+   * Generates a random UUID (version 7).
+   */
+  uuidv7(): string
+}
+
+export class CryptoRandomProvider implements RandomProvider {
+  bytes(length: number): Uint8Array {
+    return randomBytes(length)
+  }
+
+  uuidv4(): string {
+    return randomUUID()
+  }
+
+  uuidv7(): string {
+    return uuidv7()
+  }
+}
+
+export class ReproducibleRandomProvider implements RandomProvider {
+  private readonly seed: Uint8Array
+  private counter = 0
+
+  constructor(seed: Uint8Array) {
+    this.seed = seed
+  }
+
+  bytes(length: number): Uint8Array {
+    const buffer = new Uint8Array(length)
+
+    for (let i = 0; i < length; i++) {
+      // use a simple counter to generate reproducible random bytes
+      buffer[i] = (this.seed[i % this.seed.length] + this.counter) % 256
+      this.counter += 1
+    }
+
+    return buffer
+  }
+
+  uuidv4(): string {
+    return uuidv4({ random: this.bytes(16) })
+  }
+
+  uuidv7(): string {
+    return uuidv7({
+      random: this.bytes(16),
+      seq: this.counter++,
+      msecs: this.counter++,
+    })
+  }
+}

package/src/common/test/index.ts

@@ -0,0 +1,2 @@
+export * from "./render"
+export * from "./tracer"

package/src/common/test/render.ts

@@ -0,0 +1,98 @@
+import { camelCaseToHumanReadable, registerKnownAbbreviations } from "@highstate/contract"
+import * as md from "ts-markdown-builder"
+
+type RenderTraceEntryField = {
+  value: unknown
+  raw?: boolean
+  alwaysRender?: boolean
+}
+
+type RenderTraceEntryOptions = {
+  icon: string
+  title: string
+  code?: string
+  codeBlock?: string
+  body?: string
+  fields?: Record<string, RenderTraceEntryField>
+}
+
+export function codeBlock(code: string): string {
+  return `\`\`\`js\n${code}\n\`\`\``
+}
+
+function isPrimitive(value: unknown): boolean {
+  return (
+    typeof value === "string" ||
+    typeof value === "number" ||
+    typeof value === "boolean" ||
+    value === null ||
+    value === undefined
+  )
+}
+
+export function renderMdValue(value: unknown): string {
+  if (typeof value === "string") {
+    return md.code(`"${value}"`)
+  }
+
+  if (isPrimitive(value)) {
+    return md.code(String(value))
+  }
+
+  if (value instanceof Error) {
+    return codeBlock(value.stack ?? value.message)
+  }
+
+  return codeBlock(JSON.stringify(value, null, 2))
+}
+
+registerKnownAbbreviations(["ID"])
+
+export function renderTraceEntry(options: RenderTraceEntryOptions): string {
+  const blocks: string[] = [`${options.icon} ${options.title}`]
+
+  if (options.code) {
+    blocks.push(md.code(options.code))
+  }
+
+  if (options.codeBlock) {
+    blocks.push(codeBlock(options.codeBlock))
+  }
+
+  if (options.fields) {
+    for (const [key, field] of Object.entries(options.fields)) {
+      const prefix = md.bold(camelCaseToHumanReadable(key) + ":")
+
+      if (field.raw && typeof field.value === "string") {
+        blocks.push(`${prefix} ${field.value}`)
+        continue
+      }
+
+      if (!field.alwaysRender && field.value === undefined) {
+        continue
+      }
+
+      if (Array.isArray(field.value) && field.value.length > 0) {
+        blocks.push(prefix, ...field.value.map(item => renderMdValue(item)))
+      } else if (Array.isArray(field.value)) {
+        blocks.push(`${prefix} ${md.code("[]")}`)
+      } else if (
+        typeof field.value === "object" &&
+        field.value !== null &&
+        Object.keys(field.value).length === 0
+      ) {
+        blocks.push(`${prefix} ${md.code("{}")}`)
+      } else if (isPrimitive(field.value)) {
+        blocks.push(`${prefix} ${renderMdValue(field.value)}`)
+      } else {
+        blocks.push(prefix, renderMdValue(field.value))
+      }
+    }
+  }
+
+  if (options.body) {
+    blocks.push(options.body)
+  }
+
+  return md.joinBlocks(blocks)
+}