@highstate/backend 0.9.16 → 0.9.18

package/src/state/test.ts

@@ -0,0 +1,457 @@
+import type { Fixtures } from "@vitest/runner"
+import type { StateBatch, StateSnapshot, CollectionQuery } from "./abstractions"
+import type { EncryptionBackend, KeyObfuscationBackend } from "./encryption"
+import type { StateCollection } from "./abstractions"
+import type { CollectionQueryResult } from "../shared"
+import { z } from "zod"
+import * as md from "ts-markdown-builder"
+import {
+  type TraceEntry,
+  type TestTracer,
+  type TestBaseFixtures,
+  linkTraceEntry,
+  renderMdValue,
+  renderTraceEntry,
+} from "../common"
+import { createTestHotStateManager, HotStateManager } from "../hotstate"
+import { MemoryProjectUnlockBackend, type ProjectUnlockBackend } from "../unlock"
+import { StateManager } from "./manager"
+import { PassThroughEncryptionBackend, PassThroughKeyObfuscationBackend } from "./encryption"
+import { MemoryStateBackend, MemoryStateBatch, MemoryStateSnapshot } from "./memory"
+import { StateRepository } from "./repository"
+
+/**
+ * Formats multiple key-value pairs as individual blocks.
+ */
+function formatKeyValuePairs(entries: [string, unknown][]): string {
+  return md.joinBlocks(
+    entries.flatMap(([key, value]) => {
+      if (value === undefined) {
+        return `${md.bold("Key:")} ${md.code(`"${key}"`)}`
+      }
+
+      return [
+        //
+        `${md.bold("Key:")} ${md.code(`"${key}"`)}`,
+        md.bold("Value:"),
+        renderMdValue(value),
+      ]
+    }),
+  )
+}
+
+class StateReadEntry implements TraceEntry {
+  constructor(
+    readonly id: number,
+    private readonly operation: string,
+    private readonly collection: string,
+    private readonly keys: string[] | undefined,
+    private readonly result: unknown,
+    private readonly snapshot?: StateSnapshot,
+  ) {
+    if (snapshot instanceof MemoryStateSnapshot) {
+      snapshot.usageEntries.push(this)
+    }
+  }
+
+  render(): string {
+    const fields: Record<string, unknown> = {
+      collection: this.collection,
+    }
+
+    if (this.keys) {
+      fields.keys = this.keys
+    }
+
+    fields.result = this.result
+
+    return renderTraceEntry({
+      icon: "📖",
+      title: this.operation,
+      fields: {
+        collection: {
+          value: md.code(this.collection),
+          raw: true,
+        },
+        keys: {
+          value: this.keys,
+        },
+        result: {
+          value: this.result,
+          alwaysRender: true,
+        },
+        snapshotId: {
+          value:
+            this.snapshot instanceof MemoryStateSnapshot
+              ? `${md.code(this.snapshot.id.toString())}, captured at ${linkTraceEntry(this.snapshot.traceEntry)}`
+              : undefined,
+          raw: true,
+        },
+      },
+    })
+  }
+}
+
+class StateWriteEntry implements TraceEntry {
+  constructor(
+    readonly id: number,
+    private readonly operation: string,
+    private readonly collection: string,
+    private readonly entries: [string, unknown][],
+    private readonly batch?: StateBatch,
+  ) {
+    if (batch instanceof MemoryStateBatch) {
+      batch.opEntries.push(this)
+    }
+  }
+
+  render(): string {
+    const icon = this.operation.startsWith("delete") || this.operation === "clear" ? "🗑️" : "💾"
+
+    const batchInfo =
+      this.batch instanceof MemoryStateBatch
+        ? this.batch.traceEntry
+          ? `${md.code(this.batch.id.toString())}, written at ${linkTraceEntry(this.batch.traceEntry)}`
+          : md.code(this.batch.id.toString())
+        : undefined
+
+    return renderTraceEntry({
+      icon,
+      title: this.operation,
+      fields: {
+        collection: {
+          value: md.code(this.collection),
+          raw: true,
+        },
+        batchId: {
+          value: batchInfo,
+          raw: true,
+        },
+      },
+      body: formatKeyValuePairs(this.entries),
+    })
+  }
+}
+
+class TestStateRepository<TSchema extends z.ZodType> {
+  private readonly collectionName: Promise<string>
+
+  constructor(
+    collection: StateCollection,
+    private readonly repository: StateRepository<TSchema>,
+    private readonly tracer: TestTracer,
+  ) {
+    this.collectionName = collection.namespace
+  }
+
+  async get(key: string, snapshot?: StateSnapshot): Promise<z.infer<TSchema> | undefined> {
+    const result = await this.repository.get(key, snapshot)
+
+    const entry = new StateReadEntry(
+      this.tracer.nextEntryId(),
+      "get",
+      await this.collectionName,
+      [key],
+      result,
+      snapshot,
+    )
+    this.tracer.addEntry(entry)
+
+    return result
+  }
+
+  async getMany(keys: string[], snapshot?: StateSnapshot): Promise<[string, z.infer<TSchema>][]> {
+    const result = await this.repository.getMany(keys, snapshot)
+
+    const entry = new StateReadEntry(
+      this.tracer.nextEntryId(),
+      "getMany",
+      await this.collectionName,
+      keys,
+      result,
+      snapshot,
+    )
+    this.tracer.addEntry(entry)
+
+    return result
+  }
+
+  async getManyRecord(
+    keys: string[],
+    snapshot?: StateSnapshot,
+  ): Promise<Record<string, z.infer<TSchema>>> {
+    const result = await this.repository.getManyRecord(keys, snapshot)
+
+    const entry = new StateReadEntry(
+      this.tracer.nextEntryId(),
+      "getManyRecord",
+      await this.collectionName,
+      keys,
+      result,
+      snapshot,
+    )
+    this.tracer.addEntry(entry)
+
+    return result
+  }
+
+  async getManyItems(ids: string[], snapshot?: StateSnapshot): Promise<z.infer<TSchema>[]> {
+    const result = await this.repository.getManyItems(ids, snapshot)
+
+    const entry = new StateReadEntry(
+      this.tracer.nextEntryId(),
+      "getManyItems",
+      await this.collectionName,
+      ids,
+      result,
+      snapshot,
+    )
+    this.tracer.addEntry(entry)
+
+    return result
+  }
+
+  async getAll(snapshot?: StateSnapshot): Promise<Array<[string, z.infer<TSchema>]>> {
+    const result = await this.repository.getAll(snapshot)
+
+    const entry = new StateReadEntry(
+      this.tracer.nextEntryId(),
+      "getAll",
+      await this.collectionName,
+      undefined,
+      result,
+      snapshot,
+    )
+    this.tracer.addEntry(entry)
+
+    return result
+  }
+
+  async getAllRecord(snapshot?: StateSnapshot): Promise<Record<string, z.infer<TSchema>>> {
+    const result = await this.repository.getAllRecord(snapshot)
+
+    const entry = new StateReadEntry(
+      this.tracer.nextEntryId(),
+      "getAllRecord",
+      await this.collectionName,
+      undefined,
+      result,
+      snapshot,
+    )
+    this.tracer.addEntry(entry)
+
+    return result
+  }
+
+  async getAllItems(snapshot?: StateSnapshot): Promise<z.infer<TSchema>[]> {
+    const result = await this.repository.getAllItems(snapshot)
+
+    const entry = new StateReadEntry(
+      this.tracer.nextEntryId(),
+      "getAllItems",
+      await this.collectionName,
+      undefined,
+      result,
+      snapshot,
+    )
+    this.tracer.addEntry(entry)
+
+    return result
+  }
+
+  async put(key: string, value: z.infer<TSchema>, batch?: StateBatch): Promise<void> {
+    await this.repository.put(key, value, batch)
+
+    const entry = new StateWriteEntry(
+      this.tracer.nextEntryId(),
+      "put",
+      await this.collectionName,
+      [[key, value]],
+      batch,
+    )
+    this.tracer.addEntry(entry)
+  }
+
+  async putMany(entries: [string, z.infer<TSchema>][], batch?: StateBatch): Promise<void> {
+    await this.repository.putMany(entries, batch)
+
+    const entry = new StateWriteEntry(
+      this.tracer.nextEntryId(),
+      "putMany",
+      await this.collectionName,
+      entries,
+      batch,
+    )
+    this.tracer.addEntry(entry)
+  }
+
+  async putItem(item: z.infer<TSchema> & { id: string }, batch?: StateBatch): Promise<void> {
+    await this.repository.putItem(item, batch)
+
+    const entry = new StateWriteEntry(
+      this.tracer.nextEntryId(),
+      "putItem",
+      await this.collectionName,
+      [[item.id, item]],
+      batch,
+    )
+    this.tracer.addEntry(entry)
+  }
+
+  async putManyItems(
+    items: (z.infer<TSchema> & { id: string })[],
+    batch?: StateBatch,
+  ): Promise<void> {
+    await this.repository.putManyItems(items, batch)
+
+    const entries = items.map(item => [item.id, item] as [string, unknown])
+    const entry = new StateWriteEntry(
+      this.tracer.nextEntryId(),
+      "putManyItems",
+      await this.collectionName,
+      entries,
+      batch,
+    )
+    this.tracer.addEntry(entry)
+  }
+
+  async delete(key: string, batch?: StateBatch): Promise<void> {
+    await this.repository.delete(key, batch)
+
+    const entry = new StateWriteEntry(
+      this.tracer.nextEntryId(),
+      "delete",
+      await this.collectionName,
+      [[key, undefined]],
+      batch,
+    )
+    this.tracer.addEntry(entry)
+  }
+
+  async deleteMany(keys: string[], batch?: StateBatch): Promise<void> {
+    await this.repository.deleteMany(keys, batch)
+
+    const entry = new StateWriteEntry(
+      this.tracer.nextEntryId(),
+      "deleteMany",
+      await this.collectionName,
+      keys.map(key => [key, undefined]),
+      batch,
+    )
+    this.tracer.addEntry(entry)
+  }
+
+  async clear(): Promise<void> {
+    await this.repository.clear()
+
+    const entry = new StateWriteEntry(
+      this.tracer.nextEntryId(),
+      "clear",
+      await this.collectionName,
+      [],
+    )
+
+    this.tracer.addEntry(entry)
+  }
+
+  async query(
+    query: CollectionQuery,
+    snapshot?: StateSnapshot,
+  ): Promise<CollectionQueryResult<[string, z.infer<TSchema>]>> {
+    const result = await this.repository.query(query, snapshot)
+
+    const entry = new StateReadEntry(
+      this.tracer.nextEntryId(),
+      "query",
+      await this.collectionName,
+      undefined,
+      result,
+      snapshot,
+    )
+    this.tracer.addEntry(entry)
+
+    return result
+  }
+
+  async queryItems(
+    query: CollectionQuery,
+    snapshot?: StateSnapshot,
+  ): Promise<CollectionQueryResult<z.infer<TSchema>>> {
+    const result = await this.repository.queryItems(query, snapshot)
+
+    const entry = new StateReadEntry(
+      this.tracer.nextEntryId(),
+      "queryItems",
+      await this.collectionName,
+      undefined,
+      result,
+      snapshot,
+    )
+    this.tracer.addEntry(entry)
+
+    return result
+  }
+}
+
+export const stateFixtures: Fixtures<
+  {
+    projectUnlockBackend: ProjectUnlockBackend
+    hotStateManager: HotStateManager
+    stateManager: StateManager
+  },
+  TestBaseFixtures
+> = {
+  // eslint-disable-next-line no-empty-pattern
+  projectUnlockBackend: async ({}, use) => {
+    const projectUnlockBackend = new MemoryProjectUnlockBackend()
+
+    // always "unlock" the "test" project
+    await projectUnlockBackend.unlockProject("test", "test", "test")
+
+    await use(projectUnlockBackend)
+  },
+
+  // eslint-disable-next-line no-empty-pattern
+  hotStateManager: async ({}, use) => {
+    const hotStateManager = createTestHotStateManager()
+
+    await use(hotStateManager)
+  },
+  stateManager: async ({ tracer, clock, projectUnlockBackend }, use) => {
+    const backend = new MemoryStateBackend(tracer)
+
+    const createRepository = <TSchema extends z.ZodType>(
+      collection: StateCollection,
+      schema: TSchema,
+      encryptionBackend: EncryptionBackend,
+      keyObfuscationBackend: KeyObfuscationBackend,
+    ) => {
+      const baseRepository = new StateRepository(
+        collection,
+        schema,
+        encryptionBackend,
+        keyObfuscationBackend,
+        clock,
+        tracer.logger,
+      )
+
+      return new TestStateRepository(
+        collection,
+        baseRepository,
+        tracer,
+      ) as unknown as StateRepository<TSchema>
+    }
+
+    const stateManager = new StateManager(
+      PassThroughEncryptionBackend.instance,
+      PassThroughKeyObfuscationBackend.instance,
+      backend,
+      projectUnlockBackend,
+      clock,
+      tracer.logger,
+      createRepository,
+    )
+
+    await use(stateManager)
+  },
+}

package/src/unlock/abstractions.ts

@@ -0,0 +1,49 @@
+import z from "zod"
+
+export const unlockedProjectInfo = z.object({
+  /**
+   * The base64-encoded master key for the project.
+   */
+  masterKey: z.base64(),
+
+  /**
+   * The namespace for the project, used to obfuscate keys.
+   */
+  namespace: z.string(),
+})
+
+export type UnlockedProjectInfo = z.infer<typeof unlockedProjectInfo>
+
+export interface ProjectUnlockBackend {
+  /**
+   * Checks if the project is unlocked.
+   *
+   * @param projectId The ID of the project to check.
+   * @return A promise that resolves to true if the project is unlocked, false otherwise.
+   */
+  checkProjectUnlocked(projectId: string): Promise<boolean>
+
+  /**
+   * Gets the unlocked project information.
+   *
+   * @param projectId The ID of the project to get the information for.
+   * @return A promise that resolves to the unlocked project information, or null if the project is not unlocked.
+   */
+  getUnlockedProjectInfo(projectId: string): Promise<UnlockedProjectInfo | null>
+
+  /**
+   * Unlocks the project with the given master key and namespace.
+   *
+   * @param projectId The ID of the project to unlock.
+   * @param masterKey The base64-encoded master key for the project.
+   * @param namespace The namespace for the project, used to obfuscate keys.
+   */
+  unlockProject(projectId: string, masterKey: string, namespace: string): Promise<void>
+
+  /**
+   * Locks the project, removing its unlocked state.
+   *
+   * @param projectId The ID of the project to lock.
+   */
+  lockProject(projectId: string): Promise<void>
+}

package/src/unlock/index.ts

@@ -0,0 +1,2 @@
+export * from "./abstractions"
+export * from "./memory"

package/src/unlock/memory.ts

@@ -0,0 +1,32 @@
+import type { ProjectUnlockBackend, UnlockedProjectInfo } from "./abstractions"
+
+export class MemoryProjectUnlockBackend implements ProjectUnlockBackend {
+  private readonly projects = new Map<string, UnlockedProjectInfo>()
+
+  checkProjectUnlocked(projectId: string): Promise<boolean> {
+    return Promise.resolve(this.projects.has(projectId))
+  }
+
+  getUnlockedProjectInfo(projectId: string): Promise<UnlockedProjectInfo | null> {
+    const info = this.projects.get(projectId)
+
+    return Promise.resolve(info || null)
+  }
+
+  lockProject(projectId: string): Promise<void> {
+    this.projects.delete(projectId)
+
+    return Promise.resolve()
+  }
+
+  unlockProject(projectId: string, masterKey: string, namespace: string): Promise<void> {
+    const info: UnlockedProjectInfo = {
+      masterKey,
+      namespace,
+    }
+
+    this.projects.set(projectId, info)
+
+    return Promise.resolve()
+  }
+}

package/src/worker/manager.ts

@@ -30,7 +30,15 @@ export class WorkerManager {
     )
   }
 
+  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) {
@@ -80,6 +88,9 @@ export class WorkerManager {
         logBatcher.call(log)
       })
 
+      const abortController = new AbortController()
+      this.workerAbortControllers.set(workerId, abortController)
+
       void this.workerBackend
         .run({
           projectId,
@@ -87,6 +98,7 @@ export class WorkerManager {
           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))
@@ -122,6 +134,22 @@ export class WorkerManager {
     })
   }
 
+  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()