Merge pull request #210 from proto-kit/feature/enable-proving

Enable proving 1: Batch proving

Author: maht0rz

package-lock.json

@@ -21,7 +21,7 @@ export class ComputedBlockModel {
       blockHashes.map(
         (blockHash) => blocks.find((block) => block?.hash === blockHash)!
       ),
-      proof.proof === MOCK_PROOF ? "mock-proof" : JSON.stringify(proof)
+      proof.proof === MOCK_PROOF ? MOCK_PROOF : JSON.stringify(proof)
     );
   }
 

packages/api/src/graphql/modules/BatchStorageResolver.ts

@@ -6,7 +6,11 @@ import type {
 import { StringKeyOf, UnionToIntersection } from "../types";
 
 import { EventEmitter } from "./EventEmitter";
-import { EventEmittingComponent, EventsRecord } from "./EventEmittingComponent";
+import {
+  EventEmittingComponent,
+  EventEmittingContainer,
+  EventsRecord,
+} from "./EventEmittingComponent";
 
 export type CastToEventsRecord<Record> = Record extends EventsRecord
   ? Record
@@ -17,7 +21,13 @@ export type ModuleEvents<ModuleType extends BaseModuleType> =
     ? Events
     : InstanceType<ModuleType> extends ModuleContainer<infer NestedModules>
       ? CastToEventsRecord<ContainerEvents<NestedModules>>
-      : EventsRecord;
+      : // &
+        // (InstanceType<ModuleType> extends EventEmittingContainer<
+        //   infer ContainerEvents
+        // >
+        //   ? ContainerEvents
+        //   : {})
+        EventsRecord;
 
 export type ContainerEvents<Modules extends ModulesRecord> = {
   [Key in StringKeyOf<Modules>]: ModuleEvents<Modules[Key]>;
@@ -27,7 +37,7 @@ export type FlattenObject<Target extends Record<string, EventsRecord>> =
   UnionToIntersection<Target[keyof Target]>;
 
 export type FlattenedContainerEvents<Modules extends ModulesRecord> =
-  FlattenObject<ContainerEvents<Modules>>;
+  FlattenObject<ContainerEvents<Modules>>; // & FlattenObject<any>;
 
 export class EventEmitterProxy<
   Modules extends ModulesRecord,
@@ -45,10 +55,24 @@ export class EventEmitterProxy<
             this.emit(events, ...args);
           });
         }
+        if (this.isEventEmittingContainer(module)) {
+          module.containerEvents.onAll((events: any, args: any[]) => {
+            // eslint-disable-next-line @typescript-eslint/no-unsafe-argument
+            this.emit(events, ...args);
+          });
+        }
       }
     });
   }
 
+  private isEventEmittingContainer(
+    module: any
+  ): module is EventEmittingContainer<EventsRecord> {
+    // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+    const emitter = module.containerEvents;
+    return emitter !== undefined && emitter instanceof EventEmitter;
+  }
+
   private isEventEmitter(
     module: any
   ): module is EventEmittingComponent<EventsRecord> {

packages/common/src/events/EventEmitterProxy.ts

@@ -5,3 +5,7 @@ export type EventsRecord = Record<string, unknown[]>;
 export interface EventEmittingComponent<Events extends EventsRecord> {
   events: EventEmitter<Events>;
 }
+
+export interface EventEmittingContainer<Events extends EventsRecord> {
+  containerEvents: EventEmitter<Events>;
+}

packages/common/src/events/EventEmittingComponent.ts

@@ -25,6 +25,26 @@ function logProvable(
 }
 /* eslint-enable */
 
+const timeMap: Record<string, number> = {};
+
+function time(label = "time") {
+  timeMap[label] = Date.now();
+}
+
+function timeLog(label = "time"): string {
+  const prev = timeMap[label];
+  if (prev === undefined) {
+    return "Label not found";
+  }
+  return `${label} took ${Date.now() - prev}ms`;
+}
+
+function timeEnd(label = "time"): string {
+  const str = timeLog(label);
+  delete timeMap[label];
+  return str;
+}
+
 export const log = {
   provable: {
     info: (...args: unknown[]) => {
@@ -48,6 +68,24 @@ export const log = {
     },
   },
 
+  time,
+
+  timeLog: {
+    info: (label?: string) => loglevel.info(timeLog(label)),
+    debug: (label?: string) => loglevel.debug(timeLog(label)),
+    error: (label?: string) => loglevel.error(timeLog(label)),
+    trace: (label?: string) => loglevel.trace(timeLog(label)),
+    warn: (label?: string) => loglevel.warn(timeLog(label)),
+  },
+
+  timeEnd: {
+    info: (label?: string) => loglevel.info(timeEnd(label)),
+    debug: (label?: string) => loglevel.debug(timeEnd(label)),
+    error: (label?: string) => loglevel.error(timeEnd(label)),
+    trace: (label?: string) => loglevel.trace(timeEnd(label)),
+    warn: (label?: string) => loglevel.warn(timeEnd(label)),
+  },
+
   info: (...args: unknown[]) => {
     loglevel.info(...args);
   },

packages/common/src/log.ts

@@ -7,8 +7,8 @@ import { dummyVerificationKey } from "../dummyVerificationKey";
 import { MOCK_PROOF } from "./provableMethod";
 
 const errors = {
-  appChainNotSet: (name: string) =>
-    new Error(`Appchain was not injected for: ${name}`),
+  areProofsEnabledNotSet: (name: string) =>
+    new Error(`AreProofsEnabled was not injected for: ${name}`),
 };
 
 export interface CompileArtifact {
@@ -72,6 +72,8 @@ export function verifyToMockable<PublicInput, PublicOutput>(
       return verified;
     }
 
+    console.log("VerifyMocked");
+
     return proof.proof === MOCK_PROOF;
   };
 }
@@ -97,25 +99,29 @@ export abstract class ZkProgrammable<
   PublicInput = undefined,
   PublicOutput = void,
 > {
-  public abstract get appChain(): AreProofsEnabled | undefined;
+  public abstract get areProofsEnabled(): AreProofsEnabled | undefined;
 
   public abstract zkProgramFactory(): PlainZkProgram<
     PublicInput,
     PublicOutput
   >[];
 
+  private zkProgramSingleton?: PlainZkProgram<PublicInput, PublicOutput>[];
+
   @Memoize()
   public get zkProgram(): PlainZkProgram<PublicInput, PublicOutput>[] {
-    const zkProgram = this.zkProgramFactory();
+    if (this.zkProgramSingleton === undefined) {
+      this.zkProgramSingleton = this.zkProgramFactory();
+    }
 
-    return zkProgram.map((bucket) => {
-      if (!this.appChain) {
-        throw errors.appChainNotSet(this.constructor.name);
+    return this.zkProgramSingleton.map((bucket) => {
+      if (!this.areProofsEnabled) {
+        throw errors.areProofsEnabledNotSet(this.constructor.name);
       }
       return {
         ...bucket,
-        verify: verifyToMockable(bucket.verify, this.appChain),
-        compile: compileToMockable(bucket.compile, this.appChain),
+        verify: verifyToMockable(bucket.verify, this.areProofsEnabled),
+        compile: compileToMockable(bucket.compile, this.areProofsEnabled),
       };
     });
   }

packages/common/src/zkProgrammable/ZkProgrammable.ts

@@ -24,23 +24,23 @@ export function toProver(
   ...args: ArgumentTypes
 ) {
   return async function prover(this: ZkProgrammable<any, any>) {
-    const areProofsEnabled = this.appChain?.areProofsEnabled;
-    if (areProofsEnabled ?? false) {
-      for (const prog of this.zkProgram) {
-        if (Object.keys(prog.methods).includes(methodName)) {
-          const programProvableMethod = prog.methods[methodName];
-          // eslint-disable-next-line no-await-in-loop
-          return await Reflect.apply(programProvableMethod, this, args);
-        }
-      }
+    const { areProofsEnabled } = this.areProofsEnabled!;
+
+    const zkProgram = this.zkProgram.find((prog) =>
+      Object.keys(prog.methods).includes(methodName)
+    );
+
+    if (zkProgram === undefined) {
+      throw new Error("Correct ZkProgram not found");
+    }
+
+    if (areProofsEnabled) {
+      const programProvableMethod = zkProgram.methods[methodName];
+      return await Reflect.apply(programProvableMethod, this, args);
     }
 
-    // create a mock proof by simulating method execution in JS
+    // create a mock proof by simulating method> execution in JS
     const publicOutput = await Reflect.apply(simulatedMethod, this, args);
-    const zkProgram =
-      this.zkProgram.find((prog) => {
-        return Object.keys(prog.methods).includes(methodName);
-      }) ?? this.zkProgram[0];
 
     return new zkProgram.Proof({
       proof: MOCK_PROOF,

packages/common/src/zkProgrammable/provableMethod.ts

@@ -39,7 +39,7 @@ class TestProgrammable extends ZkProgrammable<
   TestPublicInput,
   TestPublicOutput
 > {
-  public appChain: AreProofsEnabled = appChainMock;
+  public areProofsEnabled: AreProofsEnabled = appChainMock;
 
   @provableMethod()
   public async foo(publicInput: TestPublicInput, bar: Balance) {
@@ -97,7 +97,7 @@ class TestProgrammable extends ZkProgrammable<
 }
 
 class OtherTestProgrammable extends ZkProgrammable<undefined, void> {
-  public appChain: AreProofsEnabled = appChainMock;
+  public areProofsEnabled: AreProofsEnabled = appChainMock;
 
   public constructor(public testProgrammable: TestProgrammable) {
     super();
@@ -183,7 +183,7 @@ describe("zkProgrammable", () => {
     (areProofsEnabled, { verificationKey, shouldVerifyMockProofs }) => {
       beforeAll(async () => {
         testProgrammable = new TestProgrammable();
-        testProgrammable.appChain.setProofsEnabled(areProofsEnabled);
+        testProgrammable.areProofsEnabled.setProofsEnabled(areProofsEnabled);
         zkProgramFactorySpy = jest.spyOn(testProgrammable, "zkProgramFactory");
         artifact = await testProgrammable.zkProgram[0].compile();
       }, 500_000);

packages/common/test/zkProgrammable/ZkProgrammable.test.ts

@@ -14,6 +14,7 @@ export interface BullQueueConfig {
     port: number;
     username?: string;
     password?: string;
+    db?: number;
   };
   retryAttempts?: number;
 }
@@ -25,17 +26,41 @@ export class BullQueue
   extends SequencerModule<BullQueueConfig>
   implements TaskQueue
 {
+  private activePromise?: Promise<void>;
+
   public createWorker(
     name: string,
     executor: (data: TaskPayload) => Promise<TaskPayload>,
     options?: { concurrency?: number }
   ): Closeable {
     const worker = new Worker<TaskPayload, TaskPayload>(
       name,
-      async (job) => await executor(job.data),
+      async (job) => {
+        // This weird promise logic is needed to make sure the worker is not proving in parallel
+        // This is by far not optimal - since it still picks up 1 task per queue but waits until
+        // computing them, so that leads to bad performance over multiple workers.
+        // For that we need to restructure tasks to be flowing through a single queue however
+        while (this.activePromise !== undefined) {
+          // eslint-disable-next-line no-await-in-loop
+          await this.activePromise;
+        }
+        let resOutside: () => void = () => {};
+        const promise = new Promise<void>((res) => {
+          resOutside = res;
+        });
+        this.activePromise = promise;
+
+        const result = await executor(job.data);
+        this.activePromise = undefined;
+        void resOutside();
+
+        return result;
+      },
       {
         concurrency: options?.concurrency ?? 1,
         connection: this.config.redis,
+        stalledInterval: 60000, // 1 minute
+        lockDuration: 60000, // 1 minute
 
         metrics: { maxDataPoints: MetricsTime.ONE_HOUR * 24 },
       }
@@ -68,6 +93,7 @@ export class BullQueue
       name: queueName,
 
       async addTask(payload: TaskPayload): Promise<{ taskId: string }> {
+        log.debug("Adding task: ", payload);
         const job = await queue.add(queueName, payload, {
           attempts: retryAttempts ?? 2,
         });
@@ -76,14 +102,25 @@ export class BullQueue
 
       async onCompleted(listener: (payload: TaskPayload) => Promise<void>) {
         events.on("completed", async (result) => {
-          // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
-          await listener(JSON.parse(result.returnvalue) as TaskPayload);
+          log.debug("Completed task: ", result);
+          try {
+            // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+            await listener(result.returnvalue as unknown as TaskPayload);
+          } catch (e) {
+            // Catch error explicitly since this promise is dangling,
+            // therefore any error will be voided as well
+            log.error(e);
+          }
+        });
+        events.on("error", async (error) => {
+          log.error("Error in worker", error);
         });
         await events.waitUntilReady();
       },
 
       async close(): Promise<void> {
         await events.close();
+        await queue.drain();
         await queue.close();
       },
     };

packages/deployment/src/queue/BullQueue.ts

@@ -10,16 +10,16 @@ import {
   BlockTrigger,
   Database,
   SequencerModule,
-  ProtocolStartupModule,
+  SequencerStartupModule,
 } from "@proto-kit/sequencer";
 import { TypedClass, ModulesConfig } from "@proto-kit/common";
 
 type PreconfiguredSimpleSequencerModulesRecord = {
   Mempool: typeof PrivateMempool;
   BatchProducerModule: typeof BatchProducerModule;
   BlockProducerModule: typeof BlockProducerModule;
-  ProtocolStartupModule: TypedClass<
-    ProtocolStartupModule & SequencerModule<unknown>
+  SequencerStartupModule: TypedClass<
+    SequencerStartupModule & SequencerModule<unknown>
   >;
 };
 
@@ -90,7 +90,7 @@ export class SimpleSequencerModules {
       BlockTrigger,
       TaskQueue,
       ...reducedModules,
-      ProtocolStartupModule,
+      SequencerStartupModule,
     } satisfies SimpleSequencerModulesRecord;
   }
 
@@ -102,7 +102,7 @@ export class SimpleSequencerModules {
 
       Mempool: {},
       BatchProducerModule: {},
-      ProtocolStartupModule: {},
+      SequencerStartupModule: {},
     } satisfies ModulesConfig<PreconfiguredSimpleSequencerModulesRecord>;
   }