@highstate/backend 0.9.18 → 0.9.19

package/src/business/artifact.ts

@@ -1,87 +1,66 @@
+import type { CommonObjectMeta } from "@highstate/contract"
 import type { Logger } from "pino"
-import type { StateManager } from "../state"
-import type { ArtifactBackend } from "../artifact/abstractions"
-import type { LockManager } from "../lock"
-import type { ObjectMeta } from "@highstate/contract"
-import { v7 as uuidv7 } from "uuid"
-import { unique } from "remeda"
-import { compareArtifactUsage, type Artifact, type ArtifactUsage } from "../shared"
+import type { ArtifactBackend } from "../artifact"
+import type { Artifact, DatabaseManager, ProjectTransaction } from "../database"
 
-const artifactChunkSize = 1024 * 1024 // 1 MB
+export const artifactChunkSize = 1024 * 1024 // 1 MB
 
 /**
  * The service which handles the storage, retrieval, and management of artifacts in the system.
  */
 export class ArtifactService {
   constructor(
-    private readonly stateManager: StateManager,
+    private readonly database: DatabaseManager,
     private readonly artifactBackend: ArtifactBackend,
-    private readonly lockManager: LockManager,
     private readonly logger: Logger,
   ) {}
 
   /**
-   * Stores an artifact in the backend and indexes it in the state manager.
+   * Stores an artifact in the backend and allows caller to set up references.
    * If the artifact already exists, it does nothing.
    *
-   * Returns the artifact ID of the stored artifact.
-   *
    * @param projectId The project ID to store the artifact under.
-   * @param hash The unique hash of the artifact content.
-   * @param meta The metadata for the artifact, including name, size, and other properties.
+   * @param hash The SHA256 hash of the artifact content.
    * @param size The total size of the artifact in bytes.
-   * @param chunkSize The size of each chunk in bytes.
+   * @param meta The metadata for the artifact.
    * @param content An async iterable providing the artifact content in chunks.
-   * @param usages The list of usages for this artifact. Can be provided later.
+   * @param track A callback to set up artifact references within the transaction.
+   * @returns The model of the stored artifact.
    */
   async store(
     projectId: string,
     hash: string,
     size: number,
-    meta: ObjectMeta,
+    meta: CommonObjectMeta,
     content: AsyncIterable<Uint8Array>,
-    usages: ArtifactUsage[] = [],
-  ): Promise<string> {
-    return await this.lockManager.acquire(["artifact-hash", projectId, hash], async () => {
-      const existingArtifact = await this.stateManager
-        .getArtifactHashIndexRepository(projectId)
-        .get(hash)
-
+    track: (tx: ProjectTransaction, artifact: Artifact) => Promise<void>,
+  ): Promise<Artifact> {
+    // check if artifact already exists by hash
+    const database = await this.database.forProject(projectId)
+    const existingArtifact = await database.artifact.findUnique({ where: { hash } })
+
+    // only upload to backend if file doesn't exist there
+    if (!existingArtifact || !(await this.artifactBackend.exists(projectId, hash))) {
       if (existingArtifact) {
-        const fileExists = await this.artifactBackend.exists(projectId, hash)
-        if (fileExists) {
-          this.logger.debug(`artifact with hash "%s" already exists`, hash)
-          return existingArtifact.id
-        }
-
-        // TODO: make this check configurable
         this.logger.warn(
-          `artifact with hash "%s" exists in index but not in backend, re-uploading`,
+          `artifact with hash "%s" exists in database but not in the storage backend, re-uploading`,
           hash,
         )
       }
 
       await this.artifactBackend.store(projectId, hash, artifactChunkSize, content)
+    }
 
-      const artifact: Artifact = {
-        id: uuidv7(),
-        hash,
-        size,
-        usages,
-        meta,
-        chunkSize: artifactChunkSize,
-      }
-
-      const batch = this.stateManager.batch()
-
-      await Promise.all([
-        this.stateManager.getArtifactRepository(projectId).putItem(artifact, batch),
-        this.stateManager
-          .getArtifactHashIndexRepository(projectId)
-          .indexRepository.put(hash, artifact.id, batch),
-      ])
+    return await database.$transaction(async tx => {
+      // create or update the main artifact record
+      const artifact = await tx.artifact.upsert({
+        where: { hash },
+        create: { hash, size, meta, chunkSize: artifactChunkSize },
+        update: { meta },
+      })
 
-      await batch.write()
+      // always allow caller to set up references
+      await track(tx, artifact)
 
       this.logger.info(
         { projectId },
@@ -90,200 +69,141 @@ export class ArtifactService {
         artifact.id,
       )
 
-      return artifact.id
+      return artifact
     })
   }
 
   /**
-   * Add usages for artifacts in the project.
+   * Clears all artifact references for a specific instance and runs garbage collection.
+   * This removes all associations between the instance and its artifacts, then
+   * cleans up any unreferenced artifacts.
    *
-   * @param projectId The project ID to which the artifacts belong.
-   * @param artifactIds The IDs of the artifacts to track usages for.
-   * @param usages The list of usages to track for the artifacts.
-   * @returns
+   * @param projectId The project ID.
+   * @param instanceId The instance ID to clear artifact references for.
    */
-  async addUsages(
-    projectId: string,
-    artifactIds: string[],
-    usages: ArtifactUsage[],
-  ): Promise<void> {
-    if (artifactIds.length === 0 || usages.length === 0) {
-      return
-    }
-
-    await this.lockManager.acquire(
-      artifactIds.map(id => ["artifact", projectId, id] as const),
-      () => this._addUsages(projectId, artifactIds, usages),
-    )
-  }
-
-  private async _addUsages(
-    projectId: string,
-    artifactIds: string[],
-    usages: ArtifactUsage[],
-  ): Promise<void> {
-    this.logger.debug({
-      msg: "tracking artifact usages",
-      projectId,
-      artifactIds,
-      usages,
+  async clearInstanceArtifactReferences(projectId: string, instanceId: string): Promise<void> {
+    const database = await this.database.forProject(projectId)
+
+    this.logger.info({ projectId, instanceId }, "clearing instance artifact references")
+
+    // clear instance-artifact associations for the specific instance
+    await database.$transaction(async tx => {
+      // get all artifacts connected to this instance
+      const artifactsConnectedToInstance = await tx.artifact.findMany({
+        where: {
+          instances: {
+            some: {
+              id: instanceId,
+            },
+          },
+        },
+        select: {
+          id: true,
+        },
+      })
+
+      // disconnect the specific instance from these artifacts
+      await Promise.all(
+        artifactsConnectedToInstance.map(artifact =>
+          tx.artifact.update({
+            where: { id: artifact.id },
+            data: {
+              instances: {
+                disconnect: { id: instanceId },
+              },
+            },
+          }),
+        ),
+      )
     })
 
-    const artifacts = await this.stateManager
-      .getArtifactRepository(projectId)
-      .getManyRecord(artifactIds)
-
-    const artifactsToUpdate: Artifact[] = []
-
-    for (const artifactId of artifactIds) {
-      const artifact = artifacts[artifactId]
-      if (!artifact) {
-        this.logger.warn({
-          msg: "artifact not found during usage addition",
-          projectId,
-          artifactId,
-        })
-        continue
-      }
-
-      // add new usages to the artifact
-      for (const usage of usages) {
-        if (!artifact.usages.some(u => compareArtifactUsage(u, usage))) {
-          artifact.usages.push(usage)
-        }
-      }
-
-      artifactsToUpdate.push(artifact)
-    }
+    this.logger.info(
+      { projectId, instanceId },
+      "cleared instance artifact references, running garbage collection",
+    )
 
-    await this.stateManager.getArtifactRepository(projectId).putManyItems(artifactsToUpdate)
+    // run garbage collection to clean up now-unreferenced artifacts
+    await this.collectGarbage(projectId)
   }
 
   /**
-   * Removes usages for artifacts in the project.
-   *
-   * If an artifact has no usages left, it will be deleted immediately.
+   * Removes artifacts with no references and cleans up backend storage.
    *
-   * @param projectId The project ID to which the artifacts belong.
-   * @param artifactIds The IDs of the artifacts to remove usages for.
-   * @param usages The list of usages to remove from the artifacts.
+   * @param projectId The project ID to clean up artifacts for.
    */
-  async removeUsages(
-    projectId: string,
-    artifactIds: string[],
-    usages: ArtifactUsage[],
-  ): Promise<void> {
-    if (artifactIds.length === 0 || usages.length === 0) {
+  async collectGarbage(projectId: string): Promise<void> {
+    const database = await this.database.forProject(projectId)
+
+    // find artifacts with no references using Prisma ORM
+    const unreferencedArtifacts = await database.artifact.findMany({
+      where: {
+        AND: [
+          { serviceAccounts: { none: {} } },
+          { instances: { none: {} } },
+          { terminals: { none: {} } },
+          { pages: { none: {} } },
+        ],
+      },
+      select: {
+        id: true,
+        hash: true,
+      },
+    })
+
+    if (unreferencedArtifacts.length === 0) {
+      this.logger.debug({ projectId }, "no unreferenced artifacts found")
       return
     }
 
-    await this.lockManager.acquire(
-      artifactIds.map(id => ["artifact", projectId, id] as const),
-      () => this._removeUsages(projectId, artifactIds, usages),
+    this.logger.info(
+      { projectId, count: unreferencedArtifacts.length },
+      "collecting garbage artifacts",
     )
-  }
 
-  private async _removeUsages(
-    projectId: string,
-    artifactIds: string[],
-    usages: ArtifactUsage[],
-  ): Promise<void> {
-    this.logger.debug({
-      msg: "removing artifact usages",
-      projectId,
-      artifactIds,
-      usages,
+    // delete unreferenced artifacts in transaction
+    await database.$transaction(async tx => {
+      await tx.artifact.deleteMany({
+        where: {
+          id: { in: unreferencedArtifacts.map(a => a.id) },
+        },
+      })
     })
 
-    const artifacts = await this.stateManager
-      .getArtifactRepository(projectId)
-      .getManyRecord(artifactIds)
-
-    const artifactsToUpdate: Artifact[] = []
-    const artifactsToDelete: Artifact[] = []
-
-    for (const artifactId of artifactIds) {
-      const artifact = artifacts[artifactId]
-      if (!artifact) {
-        this.logger.warn({
-          msg: "artifact not found during usage removal",
-          projectId,
-          artifactId,
-        })
-        continue
-      }
-
-      // remove specified usages from the artifact
-      artifact.usages = artifact.usages.filter(
-        u => !usages.some(usage => compareArtifactUsage(u, usage)),
-      )
-
-      if (artifact.usages.length === 0) {
-        // if no usages left, mark for deletion
-        artifactsToDelete.push(artifact)
-      } else {
-        // otherwise, update the artifact metadata
-        artifactsToUpdate.push(artifact)
-      }
-    }
-
-    if (artifactsToUpdate.length > 0) {
-      await this.stateManager.getArtifactRepository(projectId).putManyItems(artifactsToUpdate)
-    }
+    // clean up backend storage for deleted artifacts in foreground
+    await Promise.all(
+      unreferencedArtifacts.map(async artifact => {
+        try {
+          await this.artifactBackend.delete(projectId, artifact.hash)
+        } catch (error: unknown) {
+          this.logger.warn(
+            { error, projectId, hash: artifact.hash },
+            `failed to delete artifact from backend with hash "%s"`,
+            artifact.hash,
+          )
+        }
+      }),
+    )
 
-    if (artifactsToDelete.length > 0) {
-      await this.removeArtifacts(projectId, artifactsToDelete)
-    }
+    this.logger.info(
+      { projectId, deletedCount: unreferencedArtifacts.length },
+      "garbage collection completed",
+    )
   }
 
   /**
-   * Updates the usage of artifacts which still have it and removes the usage from artifacts which no longer have it.
+   * Gets artifact entities by their IDs.
    *
-   * @param projectId The project ID to which the artifacts belong.
-   * @param usage The usage to track for the artifacts.
-   * @param oldArtifactIds The IDs of the artifacts that previously had this usage.
-   * @param newArtifactIds The IDs of the artifacts that now have this usage.
+   * @param projectId The project ID to query artifacts from.
+   * @param artifactIds The IDs of the artifacts to retrieve.
+   * @returns Array of artifact models.
    */
-  async updateUsage(
-    projectId: string,
-    usage: ArtifactUsage,
-    oldArtifactIds: string[],
-    newArtifactIds: string[],
-  ): Promise<void> {
-    const removedArtifactIds = oldArtifactIds.filter(id => !newArtifactIds.includes(id))
-    const allArtifactIds = unique([...oldArtifactIds, ...newArtifactIds])
+  async getArtifactsByIds(projectId: string, artifactIds: string[]): Promise<Artifact[]> {
+    const database = await this.database.forProject(projectId)
 
-    await this.lockManager.acquire(
-      allArtifactIds.map(id => ["artifact", projectId, id] as const),
-      async () => {
-        await this._addUsages(projectId, newArtifactIds, [usage])
-        await this._removeUsages(projectId, removedArtifactIds, [usage])
+    return await database.artifact.findMany({
+      where: {
+        id: { in: artifactIds },
       },
-    )
-  }
-
-  private async removeArtifacts(projectId: string, artifacts: Artifact[]): Promise<void> {
-    const batch = this.stateManager.batch()
-
-    await Promise.all([
-      this.stateManager.getArtifactRepository(projectId).deleteMany(
-        artifacts.map(a => a.id),
-        batch,
-      ),
-      this.stateManager.getArtifactHashIndexRepository(projectId).indexRepository.deleteMany(
-        artifacts.map(a => a.hash),
-        batch,
-      ),
-    ])
-
-    await Promise.allSettled(
-      artifacts.map(artifact => this.artifactBackend.delete(projectId, artifact.hash)),
-    )
-
-    this.logger.info({
-      msg: "deleted dangling artifacts",
-      projectId,
-      artifactHashes: artifacts.map(a => a.hash),
     })
   }
 }