Merge pull request #3383 from logto-io/gao-optimize-tenant-disposal

refactor(core): optimize tenant disposal

packages/core/src/app/init.ts

@@ -17,6 +17,18 @@ const logListening = (type: 'core' | 'admin' = 'core') => {
   }
 };
 
+const getTenant = async (tenantId: string) => {
+  try {
+    return await tenantPool.get(tenantId);
+  } catch (error: unknown) {
+    if (error instanceof TenantNotFoundError) {
+      return error;
+    }
+
+    throw error;
+  }
+};
+
 export default async function initApp(app: Koa): Promise<void> {
   app.use(async (ctx, next) => {
     if (EnvSet.values.isDomainBasedMultiTenancy && ctx.URL.pathname === '/status') {
@@ -33,18 +45,20 @@ export default async function initApp(app: Koa): Promise<void> {
       return next();
     }
 
+    const tenant = await getTenant(tenantId);
+
+    if (tenant instanceof TenantNotFoundError) {
+      ctx.status = 404;
+
+      return next();
+    }
+
     try {
-      const tenant = await tenantPool.get(tenantId);
+      tenant.requestStart();
       await tenant.run(ctx, next);
-
-      return;
+      tenant.requestEnd();
     } catch (error: unknown) {
-      if (error instanceof TenantNotFoundError) {
-        ctx.status = 404;
-
-        return next();
-      }
-
+      tenant.requestEnd();
       throw error;
     }
   });

packages/core/src/env-set/index.ts

@@ -52,10 +52,6 @@ export class EnvSet {
     return this.#pool;
   }
 
-  get poolSafe() {
-    return this.#pool;
-  }
-
   get oidc() {
     if (!this.#oidc) {
       return throwNotLoadedError();

packages/core/src/tenants/Tenant.ts

@@ -34,23 +34,15 @@ export default class Tenant implements TenantContext {
     return new Tenant(envSet, id);
   }
 
+  #requestCount = 0;
+  #onRequestEmpty?: () => Promise<void>;
+
   public readonly provider: Provider;
   public readonly queries: Queries;
   public readonly libraries: Libraries;
+  public readonly run: MiddlewareType;
 
-  public readonly app: Koa;
-
-  get run(): MiddlewareType {
-    if (
-      EnvSet.values.isPathBasedMultiTenancy &&
-      // If admin URL Set is specified, consider that URL first
-      !(EnvSet.values.adminUrlSet.deduplicated().length > 0 && this.id === adminTenantId)
-    ) {
-      return mount('/' + this.id, this.app);
-    }
-
-    return mount(this.app);
-  }
+  private readonly app: Koa;
 
   private constructor(public readonly envSet: EnvSet, public readonly id: string) {
     const queries = new Queries(envSet.pool);
@@ -129,5 +121,56 @@ export default class Tenant implements TenantContext {
 
     this.app = app;
     this.provider = provider;
+
+    const { isPathBasedMultiTenancy, adminUrlSet } = EnvSet.values;
+    this.run =
+      isPathBasedMultiTenancy &&
+      // If admin URL Set is specified, consider that URL first
+      !(adminUrlSet.deduplicated().length > 0 && this.id === adminTenantId)
+        ? mount('/' + this.id, this.app)
+        : mount(this.app);
+  }
+
+  public requestStart() {
+    this.#requestCount += 1;
+  }
+
+  public requestEnd() {
+    if (this.#requestCount > 0) {
+      this.#requestCount -= 1;
+
+      if (this.#requestCount === 0) {
+        void this.#onRequestEmpty?.();
+      }
+    }
+  }
+
+  /**
+   * Try to dispose the tenant resources. If there are any pending requests, this function will wait for them to end with 5s timeout.
+   *
+   * Currently this function only ends the database pool.
+   *
+   * @returns Resolves `true` for a normal disposal and `'timeout'` for a timeout.
+   */
+  public async dispose() {
+    if (this.#requestCount <= 0) {
+      await this.envSet.end();
+
+      return true;
+    }
+
+    return new Promise<true | 'timeout'>((resolve) => {
+      const timeout = setTimeout(async () => {
+        this.#onRequestEmpty = undefined;
+        await this.envSet.end();
+        resolve('timeout');
+      }, 5000);
+
+      this.#onRequestEmpty = async () => {
+        clearTimeout(timeout);
+        await this.envSet.end();
+        resolve(true);
+      };
+    });
   }
 }

packages/core/src/tenants/index.ts

@@ -3,10 +3,11 @@ import LRUCache from 'lru-cache';
 import Tenant from './Tenant.js';
 
 export class TenantPool {
-  protected cache = new LRUCache<string, Tenant>({
+  protected cache = new LRUCache<string, Promise<Tenant>>({
     max: 100,
-    dispose: async (tenant) => {
-      await tenant.envSet.end();
+    dispose: async (entry) => {
+      const tenant = await entry;
+      void tenant.dispose();
     },
   });
 
@@ -18,7 +19,7 @@ export class TenantPool {
     }
 
     console.log('Init tenant:', tenantId);
-    const newTenant = await Tenant.create(tenantId);
+    const newTenant = Tenant.create(tenantId);
     this.cache.set(tenantId, newTenant);
 
     return newTenant;
@@ -26,10 +27,10 @@ export class TenantPool {
 
   async endAll(): Promise<void> {
     await Promise.all(
-      this.cache.dump().map(([, tenant]) => {
-        const { poolSafe } = tenant.value.envSet;
+      this.cache.dump().map(async ([, entry]) => {
+        const tenant = await entry.value;
 
-        return poolSafe?.end();
+        return tenant.envSet.end();
       })
     );
   }