From 3e9ba19fa2f3e1b4903650f253fb4cda129b66dd Mon Sep 17 00:00:00 2001
From: Gao Sun <gao@silverhand.io>
Date: Mon, 13 Mar 2023 15:44:01 +0800
Subject: [PATCH] refactor(core): optimize tenant disposal

---
 packages/core/src/app/init.ts       | 32 ++++++++++++++------
 packages/core/src/tenants/Tenant.ts | 46 +++++++++++++++++++++++++++++
 packages/core/src/tenants/index.ts  |  4 +--
 3 files changed, 71 insertions(+), 11 deletions(-)

diff --git a/packages/core/src/app/init.ts b/packages/core/src/app/init.ts
index 723fb0c7a..52c9878a7 100644
--- a/packages/core/src/app/init.ts
+++ b/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;
     }
   });
diff --git a/packages/core/src/tenants/Tenant.ts b/packages/core/src/tenants/Tenant.ts
index 5a27c7cda..b3f3bded0 100644
--- a/packages/core/src/tenants/Tenant.ts
+++ b/packages/core/src/tenants/Tenant.ts
@@ -34,6 +34,9 @@ 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;
@@ -130,4 +133,47 @@ export default class Tenant implements TenantContext {
     this.app = app;
     this.provider = provider;
   }
+
+  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);
+      };
+    });
+  }
 }
diff --git a/packages/core/src/tenants/index.ts b/packages/core/src/tenants/index.ts
index f265e2ce0..1f6f0b39d 100644
--- a/packages/core/src/tenants/index.ts
+++ b/packages/core/src/tenants/index.ts
@@ -5,8 +5,8 @@ import Tenant from './Tenant.js';
 export class TenantPool {
   protected cache = new LRUCache<string, Tenant>({
     max: 100,
-    dispose: async (tenant) => {
-      await tenant.envSet.end();
+    dispose: (tenant) => {
+      void tenant.dispose();
     },
   });