fix: add termined

Koatty · Nov 16, 2022 · fdec6bb · fdec6bb
1 parent c1b35f9
commit fdec6bb
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 25 deletions.
diff --git a/src/handler/grpc.ts b/src/handler/grpc.ts
@@ -3,7 +3,7 @@
  * @Usage: 
  * @Author: richen
  * @Date: 2021-11-19 00:23:06
- * @LastEditTime: 2022-11-01 15:42:27
+ * @LastEditTime: 2022-11-16 16:17:21
  */
 import * as Helper from "koatty_lib";
 import { KoattyContext } from "koatty_core";
@@ -43,12 +43,14 @@ export async function grpcHandler(ctx: KoattyContext, next: Function, ext?: any)
     const response: any = {};
 
     try {
-        response.timeout = null;
-        // promise.race
-        await Promise.race([new Promise((resolve, reject) => {
-            response.timeout = setTimeout(reject, timeout, new Exception('Deadline exceeded', 1, 4));
-            return;
-        }), next()]);
+        if (!ext.termined) {
+            response.timeout = null;
+            // promise.race
+            await Promise.race([new Promise((resolve, reject) => {
+                response.timeout = setTimeout(reject, timeout, new Exception('Deadline exceeded', 1, 4));
+                return;
+            }), next()]);
+        }
 
         if (ctx.body !== undefined && ctx.status === 404) {
             ctx.status = 200;

diff --git a/src/handler/http.ts b/src/handler/http.ts
@@ -3,7 +3,7 @@
  * @Usage: 
  * @Author: richen
  * @Date: 2021-11-19 00:14:59
- * @LastEditTime: 2022-11-01 15:42:41
+ * @LastEditTime: 2022-11-16 16:17:37
  */
 import { Helper } from "koatty_lib";
 import { catcher } from "../catcher";
@@ -43,12 +43,14 @@ export async function httpHandler(ctx: KoattyContext, next: Function, ext?: any)
     // try /catch
     const response: any = ctx.res;
     try {
-        response.timeout = null;
-        // promise.race
-        await Promise.race([new Promise((resolve, reject) => {
-            response.timeout = setTimeout(reject, timeout, new Exception('Request Timeout', 1, 408));
-            return;
-        }), next()]);
+        if (!ext.termined) {
+            response.timeout = null;
+            // promise.race
+            await Promise.race([new Promise((resolve, reject) => {
+                response.timeout = setTimeout(reject, timeout, new Exception('Request Timeout', 1, 408));
+                return;
+            }), next()]);
+        }
 
         if (ctx.body !== undefined && ctx.status === 404) {
             ctx.status = 200;

diff --git a/src/handler/ws.ts b/src/handler/ws.ts
@@ -3,7 +3,7 @@
  * @Usage: 
  * @Author: richen
  * @Date: 2021-11-19 00:24:43
- * @LastEditTime: 2022-11-01 15:42:24
+ * @LastEditTime: 2022-11-16 16:17:49
 */
 import { inspect } from "util";
 import * as Helper from "koatty_lib";
@@ -54,12 +54,14 @@ export async function wsHandler(ctx: KoattyContext, next: Function, ext?: any):
     // try /catch
     const response: any = ctx.res;
     try {
-        response.timeout = null;
-        // promise.race
-        await Promise.race([new Promise((resolve, reject) => {
-            response.timeout = setTimeout(reject, timeout, new Exception('Request Timeout', 1, 408));
-            return;
-        }), next()]);
+        if (!ext.termined) {
+            response.timeout = null;
+            // promise.race
+            await Promise.race([new Promise((resolve, reject) => {
+                response.timeout = setTimeout(reject, timeout, new Exception('Request Timeout', 1, 408));
+                return;
+            }), next()]);
+        }
 
         if (ctx.body !== undefined && ctx.status === 404) {
             ctx.status = 200;

diff --git a/src/index.ts b/src/index.ts
@@ -2,7 +2,7 @@
  * @Author: richen
  * @Date: 2020-11-20 17:37:32
  * @LastEditors: Please set LastEditors
- * @LastEditTime: 2022-03-01 15:56:09
+ * @LastEditTime: 2022-11-16 16:16:28
  * @License: BSD (3-Clause)
  * @Copyright (c) - <richenlin(at)gmail.com>
  */
@@ -58,6 +58,14 @@ export function Trace(options: TraceOptions, app: Koatty): Koa.Middleware {
     const encoding = app.config('encoding') || 'utf-8';
     const openTrace = app.config("open_trace") || false;
     return async (ctx: KoattyContext, next: Koa.Next) => {
+        // server termined
+        let termined = false;
+        if (app.server.status === 503) {
+            ctx.status = 503;
+            ctx.set('Connection', 'close');
+            ctx.body = 'Server is in the process of shutting down';
+            termined = true;
+        }
         // 
         const respWapper = async (currTraceId: string) => {
             // metadata
@@ -66,16 +74,16 @@ export function Trace(options: TraceOptions, app: Koatty): Koa.Middleware {
                 // allow bypassing koa
                 ctx.respond = false;
                 ctx.rpc.call.metadata.set(options.HeaderName, currTraceId);
-                await grpcHandler(ctx, next, { timeout, currTraceId, encoding });
+                await grpcHandler(ctx, next, { timeout, currTraceId, encoding, termined });
             } else if (ctx.protocol === "ws" || ctx.protocol === "wss") {
                 // allow bypassing koa
                 ctx.respond = false;
                 ctx.set(options.HeaderName, currTraceId);
-                await wsHandler(ctx, next, { timeout, currTraceId, encoding });
+                await wsHandler(ctx, next, { timeout, currTraceId, encoding, termined });
             } else {
                 // response header
                 ctx.set(options.HeaderName, currTraceId);
-                await httpHandler(ctx, next, { timeout, currTraceId, encoding });
+                await httpHandler(ctx, next, { timeout, currTraceId, encoding, termined });
             }
             return respond(ctx);
         }