Added closeable decorators to all services needing closing

rpanic · rpanic · commit 8baeee965551 · 2024-12-13T18:22:24.000+01:00
diff --git a/packages/api/src/graphql/GraphqlSequencerModule.ts b/packages/api/src/graphql/GraphqlSequencerModule.ts
@@ -1,6 +1,6 @@
 import assert from "node:assert";
 
-import { SequencerModule } from "@proto-kit/sequencer";
+import { Closeable, closeable, SequencerModule } from "@proto-kit/sequencer";
 import {
   ChildContainerProvider,
   Configurable,
@@ -29,9 +29,10 @@ export interface GraphqlModulesDefintion<
   config?: ModulesConfig<GraphQLModules>;
 }
 
+@closeable()
 export class GraphqlSequencerModule<GraphQLModules extends GraphqlModulesRecord>
   extends ModuleContainer<GraphQLModules>
-  implements Configurable<unknown>, SequencerModule<unknown>
+  implements Configurable<unknown>, SequencerModule<unknown>, Closeable
 {
   public static from<GraphQLModules extends GraphqlModulesRecord>(
     definition: GraphqlModulesDefintion<GraphQLModules>
@@ -93,4 +94,10 @@ export class GraphqlSequencerModule<GraphQLModules extends GraphqlModulesRecord>
     }
     await this.graphqlServer.startServer();
   }
+
+  public async close() {
+    if (this.graphqlServer !== undefined) {
+      await this.graphqlServer.close();
+    }
+  }
 }
diff --git a/packages/api/src/graphql/GraphqlServer.ts b/packages/api/src/graphql/GraphqlServer.ts
@@ -145,7 +145,18 @@ export class GraphqlServer extends SequencerModule<GraphqlServerOptions> {
     });
   }
 
-  public close() {
-    this.server?.close();
+  public async close() {
+    if (this.server !== undefined) {
+      const { server } = this;
+
+      await new Promise<void>((res) => {
+        server.close((error) => {
+          if (error !== undefined) {
+            log.error(error);
+          }
+          res();
+        });
+      });
+    }
   }
 }
diff --git a/packages/deployment/src/queue/BullQueue.ts b/packages/deployment/src/queue/BullQueue.ts
@@ -6,6 +6,7 @@ import {
   InstantiatedQueue,
   TaskQueue,
   AbstractTaskQueue,
+  closeable,
 } from "@proto-kit/sequencer";
 
 import { InstantiatedBullQueue } from "./InstantiatedBullQueue";
@@ -28,13 +29,15 @@ interface BullWorker extends Closeable {
 /**
  * TaskQueue implementation for BullMQ
  */
+@closeable()
 export class BullQueue
   extends AbstractTaskQueue<BullQueueConfig>
-  implements TaskQueue
+  implements TaskQueue, Closeable
 {
   private activePromise?: Promise<void>;
 
   private activeWorkers: Record<string, BullWorker> = {};
+
   private activeJobs = 0;
 
   public createWorker(
@@ -131,4 +134,10 @@ export class BullQueue
   public async start() {
     noop();
   }
+
+  public async close() {
+    await this.closeQueues();
+
+    // Closing of active workers is handled by the LocalTaskWorkerModule
+  }
 }
diff --git a/packages/indexer/test/GeneratedResolverFactoryGraphqlModule.test.ts b/packages/indexer/test/GeneratedResolverFactoryGraphqlModule.test.ts
@@ -139,7 +139,7 @@ describe("GeneratedResolverFactoryGraphqlModule", () => {
     }
   });
 
-  afterAll(() => {
-    indexer.resolve("GraphqlServer").close();
+  afterAll(async () => {
+    await indexer.resolve("GraphqlServer").close();
   });
 });
diff --git a/packages/persistance/src/PrismaRedisDatabase.ts b/packages/persistance/src/PrismaRedisDatabase.ts
@@ -3,6 +3,7 @@ import {
   SequencerModule,
   StorageDependencyMinimumDependencies,
   Database,
+  closeable,
 } from "@proto-kit/sequencer";
 import { ChildContainerProvider } from "@proto-kit/common";
 import { PrismaClient } from "@prisma/client";
@@ -25,6 +26,7 @@ export interface PrismaRedisCombinedConfig {
 }
 
 @sequencerModule()
+@closeable()
 export class PrismaRedisDatabase
   extends SequencerModule<PrismaRedisCombinedConfig>
   implements PrismaConnection, RedisConnection, Database
diff --git a/packages/persistance/test-integration/PrismaBlockProduction.test.ts b/packages/persistance/test-integration/PrismaBlockProduction.test.ts
@@ -46,7 +46,7 @@ describe("prisma integration", () => {
   };
 
   const teardown = async () => {
-    await appChain.sequencer.resolve("Database").close();
+    await appChain.sequencer.close();
   };
 
   describe("produce fuzzed block", () => {
diff --git a/packages/sdk/src/appChain/AppChain.ts b/packages/sdk/src/appChain/AppChain.ts
@@ -348,5 +348,9 @@ export class AppChain<
       .resolve(WorkerReadyModule)
       .waitForReady();
   }
+
+  public async close() {
+    await this.sequencer.close();
+  }
 }
 /* eslint-enable @typescript-eslint/consistent-type-assertions */
diff --git a/packages/sequencer/src/protocol/production/trigger/TimedBlockTrigger.ts b/packages/sequencer/src/protocol/production/trigger/TimedBlockTrigger.ts
@@ -2,7 +2,7 @@ import { inject, injectable } from "tsyringe";
 import { injectOptional, log } from "@proto-kit/common";
 import gcd from "compute-gcd";
 
-import { Closeable } from "../../../sequencer/builder/Closeable";
+import { closeable, Closeable } from "../../../sequencer/builder/Closeable";
 import { BatchProducerModule } from "../BatchProducerModule";
 import { Mempool } from "../../../mempool/Mempool";
 import { BlockQueue } from "../../../storage/repositories/BlockStorage";
@@ -32,6 +32,7 @@ export interface TimedBlockTriggerEvent extends BlockEvents {
 }
 
 @injectable()
+@closeable()
 export class TimedBlockTrigger
   extends BlockTriggerBase<TimedBlockTriggerConfig, TimedBlockTriggerEvent>
   implements BlockTrigger, Closeable
diff --git a/packages/sequencer/src/sequencer/SequencerStartupModule.ts b/packages/sequencer/src/sequencer/SequencerStartupModule.ts
@@ -21,9 +21,14 @@ import {
 import { VerificationKeyService } from "../protocol/runtime/RuntimeVerificationKeyService";
 
 import { SequencerModule, sequencerModule } from "./builder/SequencerModule";
+import { Closeable, closeable } from "./builder/Closeable";
 
 @sequencerModule()
-export class SequencerStartupModule extends SequencerModule {
+@closeable()
+export class SequencerStartupModule
+  extends SequencerModule
+  implements Closeable
+{
   public constructor(
     private readonly flowCreator: FlowCreator,
     @inject("Protocol")
@@ -144,4 +149,8 @@ export class SequencerStartupModule extends SequencerModule {
 
     log.info("Protocol circuits compiled successfully, commencing startup");
   }
+
+  public async close() {
+    await this.registrationFlow.close();
+  }
 }
diff --git a/packages/sequencer/src/sequencer/executor/Sequencer.ts b/packages/sequencer/src/sequencer/executor/Sequencer.ts
@@ -92,7 +92,9 @@ export class Sequencer<Modules extends SequencerModulesRecord>
     log.info("Closing sequencer...");
     const closeables = this.container.resolveAll<Closeable>("Closeable");
     await Promise.all(
-      closeables.map(async (closeable) => await closeable.close())
+      closeables.map(async (closeable) => {
+        await closeable.close();
+      })
     );
     log.info("Sequencer closed");
   }
diff --git a/packages/sequencer/src/storage/Database.ts b/packages/sequencer/src/storage/Database.ts
@@ -1,6 +1,8 @@
+import { Closeable } from "../sequencer/builder/Closeable";
+
 import { StorageDependencyFactory } from "./StorageDependencyFactory";
 
-export interface Database extends StorageDependencyFactory {
+export interface Database extends StorageDependencyFactory, Closeable {
   /**
    * Prunes all data from the database connection.
    * Note: This function should only be called immediately at startup,
diff --git a/packages/sequencer/src/storage/inmemory/InMemoryDatabase.ts b/packages/sequencer/src/storage/inmemory/InMemoryDatabase.ts
@@ -7,6 +7,7 @@ import {
 } from "../../sequencer/builder/SequencerModule";
 import { StorageDependencyMinimumDependencies } from "../StorageDependencyFactory";
 import { Database } from "../Database";
+import { closeable } from "../../sequencer/builder/Closeable";
 
 import { InMemoryBlockStorage } from "./InMemoryBlockStorage";
 import { InMemoryAsyncMerkleTreeStore } from "./InMemoryAsyncMerkleTreeStore";
@@ -16,6 +17,7 @@ import { InMemorySettlementStorage } from "./InMemorySettlementStorage";
 import { InMemoryTransactionStorage } from "./InMemoryTransactionStorage";
 
 @sequencerModule()
+@closeable()
 export class InMemoryDatabase extends SequencerModule implements Database {
   public dependencies(): StorageDependencyMinimumDependencies {
     return {
@@ -59,6 +61,10 @@ export class InMemoryDatabase extends SequencerModule implements Database {
     noop();
   }
 
+  public async close() {
+    noop();
+  }
+
   public async pruneDatabase(): Promise<void> {
     // Figure out how to implement this nicely.
     // However, this would only be a op when pruneDatabase will be called
diff --git a/packages/sequencer/src/worker/queue/AbstractTaskQueue.ts b/packages/sequencer/src/worker/queue/AbstractTaskQueue.ts
@@ -16,4 +16,10 @@ export abstract class AbstractTaskQueue<
     }
     return this.queues[name];
   }
+
+  protected async closeQueues() {
+    await Promise.all(
+      Object.values(this.queues).map(async (queue) => await queue.close())
+    );
+  }
 }
diff --git a/packages/sequencer/src/worker/worker/LocalTaskWorkerModule.ts b/packages/sequencer/src/worker/worker/LocalTaskWorkerModule.ts
@@ -7,6 +7,8 @@ import {
   ModulesRecord,
   NoConfig,
   Presets,
+  ResolvableModules,
+  StringKeyOf,
   TypedClass,
 } from "@proto-kit/common";
 import { ReturnType } from "@proto-kit/protocol";
@@ -29,6 +31,7 @@ import {
   StateTransitionTask,
 } from "../../protocol/production/tasks/StateTransitionTask";
 import { CircuitCompilerTask } from "../../protocol/production/tasks/CircuitCompilerTask";
+import { closeable } from "../../sequencer/builder/Closeable";
 
 import { FlowTaskWorker } from "./FlowTaskWorker";
 import { TaskWorkerModule } from "./TaskWorkerModule";
@@ -53,6 +56,7 @@ type LocalTaskWorkerModuleEvents = { ready: [boolean] };
  * cloud workers.
  */
 @sequencerModule()
+@closeable()
 export class LocalTaskWorkerModule<Tasks extends TaskWorkerModulesRecord>
   extends ModuleContainer<Tasks>
   implements
@@ -63,6 +67,10 @@ export class LocalTaskWorkerModule<Tasks extends TaskWorkerModulesRecord>
 
   public containerEvents = new EventEmitter<LocalTaskWorkerModuleEvents>();
 
+  private worker?: FlowTaskWorker<
+    InstanceType<ResolvableModules<Tasks>[StringKeyOf<Tasks>]>[]
+  > = undefined;
+
   public static from<Tasks extends TaskWorkerModulesRecord>(
     modules: Tasks
   ): TypedClass<LocalTaskWorkerModule<Tasks>> {
@@ -103,6 +111,8 @@ export class LocalTaskWorkerModule<Tasks extends TaskWorkerModulesRecord>
     });
 
     const worker = new FlowTaskWorker(this.taskQueue(), [...tasks]);
+    this.worker = worker;
+
     await worker.start();
 
     void worker
@@ -114,6 +124,12 @@ export class LocalTaskWorkerModule<Tasks extends TaskWorkerModulesRecord>
         log.error("Error occurring waiting for the ready event", e);
       });
   }
+
+  public async close() {
+    if (this.worker !== undefined) {
+      await this.worker.close();
+    }
+  }
 }
 
 export class VanillaTaskWorkerModules {
diff --git a/packages/sequencer/test-integration/workers/ChildProcessWorker.ts b/packages/sequencer/test-integration/workers/ChildProcessWorker.ts
@@ -26,6 +26,7 @@ export class ChildProcessWorker {
   }
 
   kill() {
-    this?.process?.kill();
+    this.process!.kill("SIGKILL");
+    console.log("Killed", this.process!.killed);
   }
 }
diff --git a/packages/sequencer/test-integration/workers/workers-proven.test.ts b/packages/sequencer/test-integration/workers/workers-proven.test.ts
@@ -1,16 +1,10 @@
 import "reflect-metadata";
-import { expectDefined, log, sleep } from "@proto-kit/common";
-import { AppChain } from "@proto-kit/sdk";
 import { container } from "tsyringe";
 import { PrivateKey, UInt64 } from "o1js";
-import { BlockTestService } from "../../test/integration/services/BlockTestService";
+import { expectDefined, log } from "@proto-kit/common";
+import { AppChain } from "@proto-kit/sdk";
 import { BullQueue } from "@proto-kit/deployment";
-import {
-  BullConfig,
-  protocolClass,
-  runtimeClass,
-  runtimeProtocolConfig,
-} from "./modules";
+
 import {
   BatchProducerModule,
   BlockProducerModule,
@@ -21,7 +15,15 @@ import {
   Sequencer,
   SequencerStartupModule,
 } from "../../src";
+import { BlockTestService } from "../../test/integration/services/BlockTestService";
 import { ConstantFeeStrategy } from "../../src/protocol/baselayer/fees/ConstantFeeStrategy";
+
+import {
+  BullConfig,
+  protocolClass,
+  runtimeClass,
+  runtimeProtocolConfig,
+} from "./modules";
 import { ChildProcessWorker } from "./ChildProcessWorker";
 
 const timeout = 300000;
@@ -41,8 +43,10 @@ describe("worker-proven", () => {
       worker.start(true);
     });
 
-    afterAll(() => {
+    afterAll(async () => {
       worker.kill();
+
+      await appChain.close();
     });
 
     it(
diff --git a/packages/stack/test/graphql/graphql.test.ts b/packages/stack/test/graphql/graphql.test.ts
@@ -114,7 +114,7 @@ describe("graphql client test", () => {
   }, 20_000);
 
   afterAll(async () => {
-    server.sequencer.resolveOrFail("GraphqlServer", GraphqlServer).close();
+    await server.sequencer.close();
   }, 20_000);
 
   it("should retrieve state", async () => {

Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,6 @@`
`1`	`1`	`import assert from "node:assert";`
`2`	`2`
`3`		`-import { SequencerModule } from "@proto-kit/sequencer";`
	`3`	`+import { Closeable, closeable, SequencerModule } from "@proto-kit/sequencer";`
`4`	`4`	`import {`
`5`	`5`	`ChildContainerProvider,`
`6`	`6`	`Configurable,`
`@@ -29,9 +29,10 @@ export interface GraphqlModulesDefintion<`
`29`	`29`	`config?: ModulesConfig<GraphQLModules>;`
`30`	`30`	`}`
`31`	`31`
	`32`	`+@closeable()`
`32`	`33`	`export class GraphqlSequencerModule<GraphQLModules extends GraphqlModulesRecord>`
`33`	`34`	`extends ModuleContainer<GraphQLModules>`
`34`		`- implements Configurable<unknown>, SequencerModule<unknown>`
	`35`	`+ implements Configurable<unknown>, SequencerModule<unknown>, Closeable`
`35`	`36`	`{`
`36`	`37`	`public static from<GraphQLModules extends GraphqlModulesRecord>(`
`37`	`38`	`definition: GraphqlModulesDefintion<GraphQLModules>`
`@@ -93,4 +94,10 @@ export class GraphqlSequencerModule<GraphQLModules extends GraphqlModulesRecord>`
`93`	`94`	`}`
`94`	`95`	`await this.graphqlServer.startServer();`
`95`	`96`	`}`
	`97`	`+`
	`98`	`+ public async close() {`
	`99`	`+ if (this.graphqlServer !== undefined) {`
	`100`	`+ await this.graphqlServer.close();`
	`101`	`+ }`
	`102`	`+ }`
`96`	`103`	`}`
Original file line number	Diff line number	Diff line change
`@@ -139,7 +139,7 @@ describe("GeneratedResolverFactoryGraphqlModule", () => {`
`139`	`139`	`}`
`140`	`140`	`});`
`141`	`141`
`142`		`- afterAll(() => {`
`143`		`- indexer.resolve("GraphqlServer").close();`
	`142`	`+ afterAll(async () => {`
	`143`	`+ await indexer.resolve("GraphqlServer").close();`
`144`	`144`	`});`
`145`	`145`	`});`
Original file line number	Diff line number	Diff line change
`@@ -348,5 +348,9 @@ export class AppChain<`
`348`	`348`	`.resolve(WorkerReadyModule)`
`349`	`349`	`.waitForReady();`
`350`	`350`	`}`
	`351`	`+`
	`352`	`+ public async close() {`
	`353`	`+ await this.sequencer.close();`
	`354`	`+ }`
`351`	`355`	`}`
`352`	`356`	`/* eslint-enable @typescript-eslint/consistent-type-assertions */`