Native Masking for AsyncStateServices

nx.json

@@ -17,6 +17,7 @@
     "prisma:generate": {
       "inputs": [
         "{projectRoot}/package.json",
+        "{projectRoot}/prisma/*",
         "{projectRoot}/prisma/**/*"
       ]
     },

packages/indexer/package.json

@@ -22,7 +22,7 @@
   "dependencies": {
     "@envelop/extended-validation": "^4.1.0",
     "@inkjs/ui": "^1.0.0",
-    "@prisma/client": "^5.18.0",
+    "@prisma/client": "^5.22.0",
     "@types/yargs": "^17.0.29",
     "figlet": "^1.7.0",
     "ink": "^4.4.1",

packages/persistance/package.json

@@ -20,9 +20,9 @@
     "access": "public"
   },
   "dependencies": {
-    "@prisma/client": "^5.18.0",
+    "@prisma/client": "^5.22.0",
     "lodash": "^4.17.21",
-    "prisma": "^5.18.0",
+    "prisma": "^5.22.0",
     "redis": "^4.6.12",
     "reflect-metadata": "^0.1.13"
   },

packages/persistance/prisma/migrations/20250129151707_state_masking/migration.sql

@@ -0,0 +1,29 @@
+/*
+  Warnings:
+
+  - The primary key for the `State` table will be changed. If it partially fails, the table could be left without primary key constraint.
+  - You are about to drop the column `mask` on the `State` table. All the data in the column will be lost.
+  - Added the required column `maskId` to the `State` table without a default value. This is not possible if the table is not empty.
+
+*/
+-- AlterTable
+TRUNCATE TABLE "State";
+ALTER TABLE "State" DROP CONSTRAINT "State_pkey",
+DROP COLUMN "mask",
+ADD COLUMN     "maskId" INTEGER NOT NULL,
+ADD CONSTRAINT "State_pkey" PRIMARY KEY ("path", "maskId");
+
+-- CreateTable
+CREATE TABLE "Mask" (
+    "id" SERIAL NOT NULL,
+    "name" VARCHAR(256) NOT NULL,
+    "parent" INTEGER,
+
+    CONSTRAINT "Mask_pkey" PRIMARY KEY ("id")
+);
+
+-- AddForeignKey
+ALTER TABLE "State" ADD CONSTRAINT "State_maskId_fkey" FOREIGN KEY ("maskId") REFERENCES "Mask"("id") ON DELETE RESTRICT ON UPDATE CASCADE;
+
+-- AddForeignKey
+ALTER TABLE "Mask" ADD CONSTRAINT "Mask_parent_fkey" FOREIGN KEY ("parent") REFERENCES "Mask"("id") ON DELETE SET NULL ON UPDATE CASCADE;

packages/persistance/prisma/schema.prisma

@@ -3,9 +3,6 @@
 
 generator client {
   provider = "prisma-client-js"
-  // output   = "./../node_modules/.prisma/client"
-  // Enable after upgrade to 5.9.0
-  // previewFeatures = ["relationJoins"]
 }
 
 datasource db {
@@ -16,9 +13,21 @@ datasource db {
 model State {
   path   Decimal   @db.Decimal(78, 0)
   values Decimal[] @db.Decimal(78, 0)
-  mask   String    @db.VarChar(256)
 
-  @@id([path, mask])
+  maskId Int
+  mask   Mask @relation(fields: [maskId], references: [id])
+
+  @@id([path, maskId])
+}
+
+model Mask {
+  id   Int    @id @default(autoincrement())
+  name String @db.VarChar(256)
+
+  parent     Int?
+  parentMask Mask?   @relation("ParentMasks", fields: [parent], references: [id])
+  children   Mask[]  @relation("ParentMasks")
+  State      State[]
 }
 
 model Transaction {

packages/persistance/src/PrismaDatabaseConnection.ts

@@ -69,9 +69,6 @@ export class PrismaDatabaseConnection
       blockStorage: {
         useClass: PrismaBlockStorage,
       },
-      unprovenStateService: {
-        useFactory: () => new PrismaStateService(this, this.tracer, "block"),
-      },
       settlementStorage: {
         useClass: PrismaSettlementStorage,
       },
@@ -81,6 +78,9 @@ export class PrismaDatabaseConnection
       transactionStorage: {
         useClass: PrismaTransactionStorage,
       },
+      unprovenStateService: {
+        useFactory: () => new PrismaStateService(this, this.tracer, "block"),
+      },
     };
   }
 
@@ -95,6 +95,7 @@ export class PrismaDatabaseConnection
       "Settlement",
       "IncomingMessageBatch",
       "IncomingMessageBatchTransaction",
+      "Mask",
     ];
 
     await this.prismaClient.$transaction(

packages/persistance/src/services/prisma/PrismaStateService.ts

@@ -11,6 +11,10 @@ import { injectable } from "tsyringe";
 
 import type { PrismaConnection } from "../../PrismaDatabaseConnection";
 
+import { readState } from "./sql/readState";
+import { deleteCollisionsFromParentMask } from "./sql/deleteCollisionsFromParentMask";
+import { mergeIntoParent } from "./sql/mergeIntoParent";
+
 // We need to create a correctly configured Decimal constructor
 // with our parameters
 const Decimal = Prisma.Decimal.clone({
@@ -21,35 +25,71 @@ const Decimal = Prisma.Decimal.clone({
 export class PrismaStateService implements AsyncStateService {
   private cache: StateEntry[] = [];
 
+  private maskId?: number;
+
   /**
    * @param connection
-   * @param mask A indicator to which masking level the values belong
    * @param tracer
+   * @param mask A indicator to which masking level the values belong.
+   * This name has to be unique
+   * @param parent
    */
   public constructor(
     private readonly connection: PrismaConnection,
     public readonly tracer: Tracer,
-    private readonly mask: string
+    private readonly mask: string,
+    private readonly parent?: number
   ) {}
 
+  private async getMaskId(): Promise<number> {
+    if (this.maskId === undefined) {
+      this.maskId = await this.initializeMask(this.mask, this.parent);
+    }
+    return this.maskId;
+  }
+
+  private async initializeMask(mask: string, parent?: number): Promise<number> {
+    const { prismaClient } = this.connection;
+
+    const found = await prismaClient.mask.findFirst({
+      where: {
+        name: mask,
+        parent,
+      },
+    });
+
+    if (found === null) {
+      const createdMask = await prismaClient.mask.create({
+        data: {
+          parent,
+          name: mask,
+        },
+      });
+      return createdMask.id;
+    }
+    return found.id;
+  }
+
   @trace("db.state.commit")
   public async commit(): Promise<void> {
     const { prismaClient } = this.connection;
 
+    const maskId = await this.getMaskId();
+
     const data = this.cache
       .filter((entry) => entry.value !== undefined)
       .map((entry) => ({
         path: new Decimal(entry.key.toString()),
         values: entry.value!.map((field) => new Decimal(field.toString())),
-        mask: this.mask,
+        maskId,
       }));
 
     await prismaClient.state.deleteMany({
       where: {
         path: {
           in: this.cache.map((x) => new Decimal(x.key.toString())),
         },
-        mask: this.mask,
+        maskId,
       },
     });
     await prismaClient.state.createMany({
@@ -60,23 +100,20 @@ export class PrismaStateService implements AsyncStateService {
   }
 
   public async getMany(keys: Field[]): Promise<StateEntry[]> {
-    const records = await this.connection.prismaClient.state.findMany({
-      where: {
-        AND: [
-          {
-            path: {
-              in: keys.map((key) => new Decimal(key.toString())),
-            },
-          },
-          {
-            mask: this.mask,
-          },
-        ],
-      },
-    });
+    const maskId = await this.getMaskId();
+    const paths = keys.map((key) => new Decimal(key.toString()));
+
+    const records: {
+      path: Prisma.Decimal;
+      values: Prisma.Decimal[] | null;
+    }[] = await this.connection.prismaClient.$queryRaw(
+      readState(maskId, paths)
+    );
+
     return records.map((record) => ({
       key: Field(record.path.toFixed()),
-      value: record.values.map((x) => Field(x.toFixed())),
+      // TODO Figure out why that is nullable
+      value: record.values?.map((x) => Field(x.toFixed())) ?? [],
     }));
   }
 
@@ -92,4 +129,62 @@ export class PrismaStateService implements AsyncStateService {
   public writeStates(entries: StateEntry[]): void {
     this.cache.push(...entries);
   }
+
+  public async createMask(name: string): Promise<AsyncStateService> {
+    const maskId = await this.getMaskId();
+    return new PrismaStateService(this.connection, this.tracer, name, maskId);
+  }
+
+  public async mergeIntoParent(): Promise<void> {
+    const maskId = await this.getMaskId();
+
+    const client = this.connection.prismaClient;
+
+    if (this.parent !== undefined) {
+      // Rough strategy here:
+      // 1. Delete all entries that are bound to be overwritten from the parent mask
+      // 2. Update this mask's entries to parent mask id
+      // 3. Re-link all children of this mask to this mask's parent
+      // 4. Delete mask
+
+      await client.$transaction([
+        client.$queryRaw(deleteCollisionsFromParentMask(maskId)),
+        // MergeIntoParent could be a prisma query, but it isn't, because eventually,
+        // this Service should be stateless, therefore the parentId wouldn't be on hand
+        // anymore, so we need to inline it's retrieval into the query
+        client.$queryRaw(mergeIntoParent(maskId)),
+        client.mask.updateMany({
+          where: {
+            parent: maskId,
+          },
+          data: {
+            parent: this.parent,
+          },
+        }),
+        client.mask.delete({
+          where: {
+            id: maskId,
+          },
+        }),
+      ]);
+    } else {
+      throw new Error("Can't merge into parent without a parent");
+    }
+  }
+
+  public async drop(): Promise<void> {
+    const maskId = await this.getMaskId();
+
+    await this.connection.prismaClient.state.deleteMany({
+      where: {
+        maskId,
+      },
+    });
+
+    await this.connection.prismaClient.mask.delete({
+      where: {
+        id: maskId,
+      },
+    });
+  }
 }

packages/persistance/src/services/prisma/sql/deleteCollisionsFromParentMask.ts

@@ -0,0 +1,24 @@
+import { Prisma } from "@prisma/client";
+
+/*
+ DELETE PATHS FROM PARENT MASK
+
+ Strategy on this query:
+ 1. Find the mask we are interested in, as a record with the id and it's parent id
+ 2. Select all paths that exist on that mask
+ 3. Delete all records whose mask is the parent of (1), and whose path is contained in (2)
+ */
+
+export function deleteCollisionsFromParentMask(maskId: number) {
+  return Prisma.sql`
+    WITH mask AS (
+        SELECT id, parent FROM "Mask" m WHERE m.id = ${maskId}),
+    paths AS (
+        SELECT path FROM "State" s
+        JOIN mask m on s."maskId" = m.id
+     )
+    DELETE FROM "State" s
+        WHERE s."maskId" IN (SELECT parent FROM mask)
+            AND s.path IN (SELECT path FROM paths)
+  `;
+}

packages/persistance/src/services/prisma/sql/mergeIntoParent.ts

@@ -0,0 +1,16 @@
+import { Prisma } from "@prisma/client";
+
+/*
+Strategy:
+ 1. Find the mask we are interested in, as a record with the id and it's parent id
+ 2. Update the maskId to the masks parent for all the entries of that particular mask
+*/
+
+export function mergeIntoParent(maskId: number) {
+  return Prisma.sql`
+    WITH mask AS (
+        SELECT id, parent FROM "Mask" m WHERE m.id = ${maskId})
+    UPDATE "State" SET "maskId" = (SELECT parent FROM mask)
+    WHERE "maskId" = (SELECT id FROM mask)
+  `;
+}

packages/persistance/src/services/prisma/sql/readState.ts

@@ -0,0 +1,16 @@
+import { Prisma } from "@prisma/client";
+
+export function readState(maskId: number, paths: Prisma.Decimal[]) {
+  return Prisma.sql`
+    WITH RECURSIVE find_masks(id, parent) AS (
+        SELECT m.id, m.parent FROM "Mask" m WHERE m.id = ${maskId}
+        UNION ALL
+        SELECT s.id, s.parent FROM "Mask" s, find_masks f
+        WHERE f.parent = s.id
+    )
+    SELECT distinct on (state.path) state.path, state.values FROM find_masks
+        JOIN "State" state ON state."maskId" = find_masks.id
+    WHERE state.path = ANY(${paths})
+    ORDER BY state.path, state."maskId" DESC;
+  `;
+}

packages/processor/package.json

@@ -24,7 +24,7 @@
   "main": "dist/index.js",
   "dependencies": {
     "@inkjs/ui": "^1.0.0",
-    "@prisma/client": "^5.19.1",
+    "@prisma/client": "^5.22.0",
     "@types/yargs": "^17.0.29",
     "figlet": "^1.7.0",
     "ink": "^4.4.1",

packages/protocol/src/index.ts

@@ -39,7 +39,7 @@ export * from "./state/context/RuntimeMethodExecutionContext";
 export * from "./state/protocol/ProtocolState";
 export * from "./state/State";
 export * from "./state/StateMap";
-export * from "./state/StateService";
+export * from "./state/SimpleAsyncStateService";
 export * from "./state/StateServiceProvider";
 export * from "./state/assert/assert";
 export * from "./settlement/contracts/authorizations/ContractAuthorization";

packages/protocol/src/protocol/Protocol.ts

@@ -13,7 +13,7 @@ import { DependencyContainer, Lifecycle } from "tsyringe";
 import { BlockProvable } from "../prover/block/BlockProvable";
 import { StateTransitionProvable } from "../prover/statetransition/StateTransitionProvable";
 import { StateServiceProvider } from "../state/StateServiceProvider";
-import { SimpleAsyncStateService } from "../state/StateService";
+import { SimpleAsyncStateService } from "../state/SimpleAsyncStateService";
 import { NoopBlockHook } from "../hooks/NoopBlockHook";
 import { BlockHeightHook } from "../hooks/BlockHeightHook";
 import { LastStateRootBlockHook } from "../hooks/LastStateRootBlockHook";