[feat] Add support for cleaning up disposable instances, part deux (#183)

* [feat] Add support for cleaning up disposable instances

* [feat] Add support for asynchronous disposables

Co-authored-by: Steven Hobson-Campbell <[email protected]>
Co-authored-by: Lee C <[email protected]>

Author: 3 people

README.md

@@ -33,6 +33,7 @@ constructor injection.
 - [Circular dependencies](#circular-dependencies)
   - [The `delay` helper function](#the-delay-helper-function)
   - [Interfaces and circular dependencies](#interfaces-and-circular-dependencies)
+- [Disposable instances](#disposable-instances)
 - [Full examples](#full-examples)
   - [Example without interfaces](#example-without-interfaces)
   - [Example with interfaces](#example-with-interfaces)
@@ -681,6 +682,20 @@ export class Bar implements IBar {
 }
 ```
 
+# Disposable instances
+All instances created by the container that implement the [`Disposable`](./src/types/disposable.ts)
+interface will automatically be disposed of when the container is disposed.
+
+```typescript
+container.dispose();
+```
+
+or to await all asynchronous disposals:
+
+```typescript
+await container.dispose();
+```
+
 # Full examples
 
 ## Example without interfaces

src/__tests__/disposable.test.ts

@@ -0,0 +1,27 @@
+import Disposable, {isDisposable} from "../types/disposable";
+
+describe("Disposable", () => {
+  describe("isDisposable", () => {
+    it("returns false for non-disposable object", () => {
+      const nonDisposable = {};
+
+      expect(isDisposable(nonDisposable)).toBeFalsy();
+    });
+
+    it("returns false when dispose method takes too many args", () => {
+      const specialDisposable = {
+        dispose(_: any) {}
+      };
+
+      expect(isDisposable(specialDisposable)).toBeFalsy();
+    });
+
+    it("returns true for disposable object", () => {
+      const disposable: Disposable = {
+        dispose() {}
+      };
+
+      expect(isDisposable(disposable)).toBeTruthy();
+    });
+  });
+});

src/__tests__/global-container.test.ts

@@ -11,6 +11,7 @@ import {instance as globalContainer} from "../dependency-container";
 import injectAll from "../decorators/inject-all";
 import Lifecycle from "../types/lifecycle";
 import {ValueProvider} from "../providers";
+import Disposable from "../types/disposable";
 
 interface IBar {
   value: string;
@@ -820,3 +821,87 @@ test("predicateAwareClassFactory returns new instances each call with caching of
 
   expect(factory(globalContainer)).not.toBe(factory(globalContainer));
 });
+
+describe("dispose", () => {
+  class Foo implements Disposable {
+    disposed = false;
+    dispose(): void {
+      this.disposed = true;
+    }
+  }
+  class Bar implements Disposable {
+    disposed = false;
+    dispose(): void {
+      this.disposed = true;
+    }
+  }
+  class Baz implements Disposable {
+    disposed = false;
+    async dispose(): Promise<void> {
+      return new Promise(resolve => {
+        process.nextTick(() => {
+          this.disposed = true;
+          resolve();
+        });
+      });
+    }
+  }
+
+  it("renders the container useless", () => {
+    const container = globalContainer.createChildContainer();
+    container.dispose();
+
+    expect(() => container.register("Bar", {useClass: Bar})).toThrow(
+      /disposed/
+    );
+    expect(() => container.reset()).toThrow(/disposed/);
+    expect(() => container.resolve("indisposed")).toThrow(/disposed/);
+  });
+
+  it("disposes all child disposables", () => {
+    const container = globalContainer.createChildContainer();
+
+    const foo = container.resolve(Foo);
+    const bar = container.resolve(Bar);
+
+    container.dispose();
+
+    expect(foo.disposed).toBeTruthy();
+    expect(bar.disposed).toBeTruthy();
+  });
+
+  it("disposes asynchronous disposables", async () => {
+    const container = globalContainer.createChildContainer();
+
+    const foo = container.resolve(Foo);
+    const baz = container.resolve(Baz);
+
+    await container.dispose();
+
+    expect(foo.disposed).toBeTruthy();
+    expect(baz.disposed).toBeTruthy();
+  });
+
+  it("disposes all instances of the same type", () => {
+    const container = globalContainer.createChildContainer();
+
+    const foo1 = container.resolve(Foo);
+    const foo2 = container.resolve(Foo);
+
+    container.dispose();
+
+    expect(foo1.disposed).toBeTruthy();
+    expect(foo2.disposed).toBeTruthy();
+  });
+
+  it("doesn't dispose of instances created external to the container", () => {
+    const foo = new Foo();
+    const container = globalContainer.createChildContainer();
+
+    container.registerInstance(Foo, foo);
+    container.resolve(Foo);
+    container.dispose();
+
+    expect(foo.disposed).toBeFalsy();
+  });
+});

src/dependency-container.ts

@@ -28,6 +28,7 @@ import Lifecycle from "./types/lifecycle";
 import ResolutionContext from "./resolution-context";
 import {formatErrorCtor} from "./error-helpers";
 import {DelayedConstructor} from "./lazy-helpers";
+import Disposable, {isDisposable} from "./types/disposable";
 import InterceptorOptions from "./types/interceptor-options";
 import Interceptors from "./interceptors";
 
@@ -45,6 +46,8 @@ export const typeInfo = new Map<constructor<any>, ParamInfo[]>();
 class InternalDependencyContainer implements DependencyContainer {
   private _registry = new Registry();
   private interceptors = new Interceptors();
+  private disposed = false;
+  private disposables = new Set<Disposable>();
 
   public constructor(private parent?: InternalDependencyContainer) {}
 
@@ -81,6 +84,8 @@ class InternalDependencyContainer implements DependencyContainer {
     providerOrConstructor: Provider<T> | constructor<T>,
     options: RegistrationOptions = {lifecycle: Lifecycle.Transient}
   ): InternalDependencyContainer {
+    this.ensureNotDisposed();
+
     let provider: Provider<T>;
 
     if (!isProvider(providerOrConstructor)) {
@@ -139,6 +144,8 @@ class InternalDependencyContainer implements DependencyContainer {
     from: InjectionToken<T>,
     to: InjectionToken<T>
   ): InternalDependencyContainer {
+    this.ensureNotDisposed();
+
     if (isNormalToken(to)) {
       return this.register(from, {
         useToken: to
@@ -154,6 +161,8 @@ class InternalDependencyContainer implements DependencyContainer {
     token: InjectionToken<T>,
     instance: T
   ): InternalDependencyContainer {
+    this.ensureNotDisposed();
+
     return this.register(token, {
       useValue: instance
     });
@@ -171,6 +180,8 @@ class InternalDependencyContainer implements DependencyContainer {
     from: InjectionToken<T>,
     to?: InjectionToken<T>
   ): InternalDependencyContainer {
+    this.ensureNotDisposed();
+
     if (isNormalToken(from)) {
       if (isNormalToken(to)) {
         return this.register(
@@ -213,6 +224,8 @@ class InternalDependencyContainer implements DependencyContainer {
     token: InjectionToken<T>,
     context: ResolutionContext = new ResolutionContext()
   ): T {
+    this.ensureNotDisposed();
+
     const registration = this.getRegistration(token);
 
     if (!registration && isNormalToken(token)) {
@@ -282,6 +295,8 @@ class InternalDependencyContainer implements DependencyContainer {
     registration: Registration,
     context: ResolutionContext
   ): T {
+    this.ensureNotDisposed();
+
     // If we have already resolved this scoped dependency, return it
     if (
       registration.options.lifecycle === Lifecycle.ResolutionScoped &&
@@ -334,6 +349,8 @@ class InternalDependencyContainer implements DependencyContainer {
     token: InjectionToken<T>,
     context: ResolutionContext = new ResolutionContext()
   ): T[] {
+    this.ensureNotDisposed();
+
     const registrations = this.getAllRegistrations(token);
 
     if (!registrations && isNormalToken(token)) {
@@ -360,6 +377,8 @@ class InternalDependencyContainer implements DependencyContainer {
   }
 
   public isRegistered<T>(token: InjectionToken<T>, recursive = false): boolean {
+    this.ensureNotDisposed();
+
     return (
       this._registry.has(token) ||
       (recursive &&
@@ -369,12 +388,15 @@ class InternalDependencyContainer implements DependencyContainer {
   }
 
   public reset(): void {
+    this.ensureNotDisposed();
     this._registry.clear();
     this.interceptors.preResolution.clear();
     this.interceptors.postResolution.clear();
   }
 
   public clearInstances(): void {
+    this.ensureNotDisposed();
+
     for (const [token, registrations] of this._registry.entries()) {
       this._registry.setAll(
         token,
@@ -391,6 +413,8 @@ class InternalDependencyContainer implements DependencyContainer {
   }
 
   public createChildContainer(): DependencyContainer {
+    this.ensureNotDisposed();
+
     const childContainer = new InternalDependencyContainer(this);
 
     for (const [token, registrations] of this._registry.entries()) {
@@ -443,6 +467,21 @@ class InternalDependencyContainer implements DependencyContainer {
     });
   }
 
+  public async dispose(): Promise<void> {
+    this.disposed = true;
+
+    const promises: Promise<unknown>[] = [];
+    this.disposables.forEach(disposable => {
+      const maybePromise = disposable.dispose();
+
+      if (maybePromise) {
+        promises.push(maybePromise);
+      }
+    });
+
+    await Promise.all(promises);
+  }
+
   private getRegistration<T>(token: InjectionToken<T>): Registration | null {
     if (this.isRegistered(token)) {
       return this._registry.get(token)!;
@@ -478,18 +517,27 @@ class InternalDependencyContainer implements DependencyContainer {
         this.resolve(target, context)
       );
     }
-    const paramInfo = typeInfo.get(ctor);
-    if (!paramInfo || paramInfo.length === 0) {
-      if (ctor.length === 0) {
-        return new ctor();
-      } else {
-        throw new Error(`TypeInfo not known for "${ctor.name}"`);
+
+    const instance: T = (() => {
+      const paramInfo = typeInfo.get(ctor);
+      if (!paramInfo || paramInfo.length === 0) {
+        if (ctor.length === 0) {
+          return new ctor();
+        } else {
+          throw new Error(`TypeInfo not known for "${ctor.name}"`);
+        }
       }
-    }
 
-    const params = paramInfo.map(this.resolveParams(context, ctor));
+      const params = paramInfo.map(this.resolveParams(context, ctor));
+
+      return new ctor(...params);
+    })();
 
-    return new ctor(...params);
+    if (isDisposable(instance)) {
+      this.disposables.add(instance);
+    }
+
+    return instance;
   }
 
   private resolveParams<T>(context: ResolutionContext, ctor: constructor<T>) {
@@ -523,6 +571,14 @@ class InternalDependencyContainer implements DependencyContainer {
       }
     };
   }
+
+  private ensureNotDisposed(): void {
+    if (this.disposed) {
+      throw new Error(
+        "This container has been disposed, you cannot interact with a disposed container"
+      );
+    }
+  }
 }
 
 export const instance: DependencyContainer = new InternalDependencyContainer();

src/index.ts

@@ -6,6 +6,7 @@ if (typeof Reflect === "undefined" || !Reflect.getMetadata) {
 
 export {
   DependencyContainer,
+  Disposable,
   Lifecycle,
   RegistrationOptions,
   Frequency

src/types/dependency-container.ts

@@ -5,6 +5,7 @@ import ValueProvider from "../providers/value-provider";
 import ClassProvider from "../providers/class-provider";
 import constructor from "./constructor";
 import RegistrationOptions from "./registration-options";
+import Disposable from "./disposable";
 import InterceptorOptions from "./interceptor-options";
 
 export type ResolutionType = "Single" | "All";
@@ -30,7 +31,7 @@ export interface PostResolutionInterceptorCallback<T = any> {
   ): void;
 }
 
-export default interface DependencyContainer {
+export default interface DependencyContainer extends Disposable {
   register<T>(
     token: InjectionToken<T>,
     provider: ValueProvider<T>
@@ -120,4 +121,10 @@ export default interface DependencyContainer {
     callback: PostResolutionInterceptorCallback<T>,
     options?: InterceptorOptions
   ): void;
+
+  /**
+   * Calls `.dispose()` on all disposable instances created by the container.
+   * After calling this, the container may no longer be used.
+   */
+  dispose(): Promise<void> | void;
 }

src/types/disposable.ts

@@ -0,0 +1,16 @@
+export default interface Disposable {
+  dispose(): Promise<void> | void;
+}
+
+export function isDisposable(value: any): value is Disposable {
+  if (typeof value.dispose !== "function") return false;
+
+  const disposeFun: Function = value.dispose;
+
+  // `.dispose()` takes in no arguments
+  if (disposeFun.length > 0) {
+    return false;
+  }
+
+  return true;
+}