Merge pull request #246 from proto-kit/feature/closeable

Added closeable collection and shutdown mechanism

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();
+    }
+  }
 }

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();
+        });
+      });
+    }
   }
 }

packages/common/src/config/ModuleContainer.ts

@@ -28,6 +28,7 @@ import {
 } from "./ConfigurableModule";
 import { ChildContainerProvider } from "./ChildContainerProvider";
 import { ChildContainerCreatable } from "./ChildContainerCreatable";
+import { getInjectAliases } from "./injectAlias";
 
 const errors = {
   configNotSetInContainer: (moduleName: string) =>
@@ -228,6 +229,16 @@ export class ModuleContainer<
     }
   }
 
+  protected registerAliases(originalToken: string, clas: TypedClass<any>) {
+    const aliases = getInjectAliases(clas);
+
+    aliases.forEach((alias) =>
+      this.container.register(alias, {
+        useToken: originalToken,
+      })
+    );
+  }
+
   /**
    * Register modules into the current container, and registers
    * a respective resolution hook in order to decorate the module
@@ -250,6 +261,8 @@ export class ModuleContainer<
           { lifecycle: Lifecycle.ContainerScoped }
         );
         this.onAfterModuleResolution(moduleName);
+
+        this.registerAliases(moduleName, useClass);
       }
     });
   }
@@ -412,7 +425,11 @@ export class ModuleContainer<
           this.container.register(key, declaration, {
             lifecycle: Lifecycle.Singleton,
           });
-          // eslint-disable-next-line sonarjs/no-duplicated-branches
+          this.registerAliases(
+            key,
+            // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+            declaration.useClass as TypedClass<unknown>
+          );
         } else if (isTokenProvider(declaration)) {
           this.container.register(key, declaration, {
             lifecycle: Lifecycle.Singleton,

packages/common/src/config/injectAlias.ts

@@ -0,0 +1,70 @@
+import { TypedClass } from "../types";
+
+export const injectAliasMetadataKey = "protokit-inject-alias";
+
+/**
+ * Attaches metadata to the class that the ModuleContainer can pick up
+ * and inject this class in the DI container under the specified aliases.
+ * This method supports inheritance, therefore also gets aliases defined
+ * on superclasses
+ */
+export function injectAlias(aliases: string[]) {
+  return (target: TypedClass<unknown>) => {
+    // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+    const superAliases = Reflect.getMetadata(
+      injectAliasMetadataKey,
+      // eslint-disable-next-line @typescript-eslint/no-unsafe-argument
+      Object.getPrototypeOf(target)
+    ) as string[] | undefined;
+
+    // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+    const existingAliases = Reflect.getMetadata(
+      injectAliasMetadataKey,
+      target
+    ) as string[] | undefined;
+
+    let allAliases = aliases;
+
+    if (superAliases !== undefined) {
+      allAliases = allAliases.concat(superAliases);
+    }
+    if (existingAliases !== undefined) {
+      allAliases = allAliases.concat(existingAliases);
+    }
+
+    Reflect.defineMetadata(
+      injectAliasMetadataKey,
+      allAliases.filter(
+        (value, index, array) => array.indexOf(value) === index
+      ),
+      target
+    );
+  };
+}
+
+/**
+ * Marks the class to implement a certain interface T, while also attaching
+ * a DI-injection alias as metadata, that will be picked up by the ModuleContainer
+ * to allow resolving by that interface name
+ * @param name The name of the injection alias, convention is to use the same as the name of T
+ */
+export function implement<T>(name: string) {
+  return (
+    /**
+     * Check if the target class extends RuntimeModule, while
+     * also providing static config presets
+     */
+    target: TypedClass<T>
+  ) => {
+    injectAlias([name])(target);
+  };
+}
+
+export function getInjectAliases(target: TypedClass<unknown>): string[] {
+  // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+  const aliases = Reflect.getMetadata(
+    injectAliasMetadataKey,
+    target
+  ) as string[];
+  return aliases ?? [];
+}

packages/common/src/index.ts

@@ -22,3 +22,4 @@ export * from "./compiling/AtomicCompileHelper";
 export * from "./compiling/CompileRegistry";
 export * from "./compiling/CompilableModule";
 export * from "./compiling/services/ChildVerificationKeyService";
+export * from "./config/injectAlias";

packages/common/test/config/ModuleContainer.test.ts

@@ -10,7 +10,8 @@ import {
   ModulesRecord,
 } from "../../src/config/ModuleContainer";
 import { TypedClass } from "../../src/types";
-import { DependencyFactory } from "../../src";
+import { DependencyFactory, expectDefined } from "../../src";
+import { injectAlias } from "../../src/config/injectAlias";
 
 // module container will accept modules that extend this type
 class BaseTestModule<Config> extends ConfigurableModule<Config> {}
@@ -24,6 +25,7 @@ interface TestModuleConfig {
 }
 
 @injectable()
+@injectAlias(["child-alias", "multi-alias"])
 class ChildModule extends BaseTestModule<NoConfig> {
   public constructor(@inject("TestModule") public readonly testModule: any) {
     super();
@@ -34,6 +36,7 @@ class ChildModule extends BaseTestModule<NoConfig> {
   }
 }
 
+@injectAlias(["base-alias", "multi-alias"])
 class TestModule
   extends BaseTestModule<TestModuleConfig>
   implements DependencyFactory
@@ -66,7 +69,11 @@ class WrongTestModule {}
 
 class TestModuleContainer<
   Modules extends TestModulesRecord,
-> extends ModuleContainer<Modules> {}
+> extends ModuleContainer<Modules> {
+  public get dependencyContainer() {
+    return this.container;
+  }
+}
 
 describe("moduleContainer", () => {
   let container: TestModuleContainer<{
@@ -169,4 +176,40 @@ describe("moduleContainer", () => {
     expect(config.testConfigProperty2).toBe(2);
     expect(config.testConfigProperty3).toBe(undefined);
   });
+
+  it("should resolve dependencies correctly via alias", () => {
+    container.configure({
+      TestModule: {
+        testConfigProperty,
+      },
+
+      OtherTestModule: {
+        otherTestConfigProperty: testConfigProperty,
+      },
+    });
+
+    container.create(() => tsyringeContainer.createChildContainer());
+
+    // Unfortunately we still need this so that the dependencies are registered
+    // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+    const m1 = container.resolve("base-alias" as any);
+    const m2 = container.resolve("TestModule");
+
+    expectDefined(m1);
+    // Check if its the same reference
+    expect(m1).toBe(m2);
+
+    const dm1 = container.resolve("child-alias" as any) as ChildModule;
+    const dm2 = container.resolve("DependencyModule1");
+
+    expect(dm1.x()).toBe("dependency factory works");
+    expect(dm1.testModule).toBeDefined();
+    expect(dm1).toBe(dm2);
+
+    const multi =
+      container.dependencyContainer.resolveAll<BaseTestModule<unknown>>(
+        "multi-alias"
+      );
+    expect(multi).toHaveLength(2);
+  });
 });

packages/common/test/config/injectAlias.test.ts

@@ -0,0 +1,28 @@
+import "reflect-metadata";
+import { getInjectAliases, injectAlias } from "../../src/config/injectAlias";
+
+@injectAlias(["foo", "bar"])
+class TestClass {}
+
+@injectAlias(["ayy"])
+class TestClass2 extends TestClass {}
+
+describe("injectAlias metadata", () => {
+  it("set and retrieve", () => {
+    expect.assertions(2);
+
+    const aliases = getInjectAliases(TestClass);
+
+    expect(aliases).toHaveLength(2);
+    expect(aliases).toStrictEqual(["foo", "bar"]);
+  });
+
+  it("recursive", () => {
+    expect.assertions(2);
+
+    const aliases = getInjectAliases(TestClass2);
+
+    expect(aliases).toHaveLength(3);
+    expect(aliases).toStrictEqual(["ayy", "foo", "bar"]);
+  });
+});

packages/deployment/src/queue/BullQueue.ts

@@ -6,6 +6,7 @@ import {
   InstantiatedQueue,
   TaskQueue,
   AbstractTaskQueue,
+  closeable,
 } from "@proto-kit/sequencer";
 
 import { InstantiatedBullQueue } from "./InstantiatedBullQueue";
@@ -24,9 +25,10 @@ export interface BullQueueConfig {
 /**
  * TaskQueue implementation for BullMQ
  */
+@closeable()
 export class BullQueue
   extends AbstractTaskQueue<BullQueueConfig>
-  implements TaskQueue
+  implements TaskQueue, Closeable
 {
   private activePromise?: Promise<void>;
 
@@ -101,4 +103,10 @@ export class BullQueue
   public async start() {
     noop();
   }
+
+  public async close() {
+    await this.closeQueues();
+
+    // Closing of active workers is handled by the LocalTaskWorkerModule
+  }
 }

packages/indexer/test/GeneratedResolverFactoryGraphqlModule.test.ts

@@ -139,7 +139,7 @@ describe("GeneratedResolverFactoryGraphqlModule", () => {
     }
   });
 
-  afterAll(() => {
-    indexer.resolve("GraphqlServer").close();
+  afterAll(async () => {
+    await indexer.resolve("GraphqlServer").close();
   });
 });

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";
@@ -26,6 +27,7 @@ export interface PrismaRedisCombinedConfig {
 }
 
 @sequencerModule()
+@closeable()
 export class PrismaRedisDatabase
   extends SequencerModule<PrismaRedisCombinedConfig>
   implements PrismaConnection, RedisConnection, Database

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", () => {

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 */

packages/sequencer/src/index.ts

@@ -6,6 +6,7 @@ export * from "./mempool/private/PrivateMempool";
 export * from "./sequencer/executor/Sequencer";
 export * from "./sequencer/executor/Sequenceable";
 export * from "./sequencer/builder/SequencerModule";
+export * from "./sequencer/builder/Closeable";
 export * from "./worker/flow/Flow";
 export * from "./worker/flow/Task";
 export * from "./worker/flow/JSONTaskSerializer";

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 "../../../worker/queue/TaskQueue";
+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