fix: disable async_hook

Koatty · Aug 21, 2023 · ede8b32 · ede8b32
1 parent 601eefa
commit ede8b32
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 33 deletions.
diff --git a/src/index.ts b/src/index.ts
@@ -2,8 +2,11 @@
  * @Description: 
  * @Usage: 
  * @Author: richen
- * @Date: 2021-11-17 09:53:49
- * @LastEditTime: 2023-07-27 23:02:17
+ * @Date: 2020-11-20 17:37:32
+ * @LastEditors: Please set LastEditors
+ * @LastEditTime: 2023-08-21 15:45:44
+ * @License: BSD (3-Clause)
+ * @Copyright (c) - <richenlin(at)gmail.com>
  */
 
 export * from "./trace";

diff --git a/src/trace.ts b/src/trace.ts
@@ -2,14 +2,13 @@
  * @Author: richen
  * @Date: 2020-11-20 17:37:32
  * @LastEditors: Please set LastEditors
- * @LastEditTime: 2023-07-31 21:23:54
+ * @LastEditTime: 2023-08-21 15:51:19
  * @License: BSD (3-Clause)
  * @Copyright (c) - <richenlin(at)gmail.com>
  */
-import * as Koa from 'koa';
 import { FORMAT_HTTP_HEADERS, Span, Tags, Tracer } from "opentracing";
 import { v4 as uuidv4 } from "uuid";
-import { Koatty, KoattyContext } from "koatty_core";
+import { Koatty, KoattyContext, KoattyNext } from "koatty_core";
 import { asyncLocalStorage, createAsyncResource, wrapEmitter } from './wrap';
 import { httpHandler } from './handler/http';
 import { gRPCHandler } from './handler/grpc';
@@ -34,15 +33,17 @@ export function GetTraceId(options?: TraceOptions) {
  * @interface TraceOptions
  */
 export interface TraceOptions {
-  HeaderName: string;
+  RequestIdHeaderName: string;
+  RequestIdName: string;
   IdFactory: any;
 }
 
 /** 
  * defaultOptions
  */
 const defaultOptions = {
-  HeaderName: 'X-Request-Id',
+  RequestIdHeaderName: 'X-Request-Id',
+  RequestIdName: "requestId",
   IdFactory: uuidv4,
 };
 
@@ -53,13 +54,12 @@ const defaultOptions = {
  * @param {Koatty} app
  * @returns {*}  {Koa.Middleware}
  */
-export function Trace(options: TraceOptions, app: Koatty): Koa.Middleware {
+export function Trace(options: TraceOptions, app: Koatty) {
   options = { ...defaultOptions, ...options };
-  const requestIdName = options.HeaderName.toLowerCase();
   const timeout = (app.config('http_timeout') || 10) * 1000;
   const encoding = app.config('encoding') || 'utf-8';
   const openTrace = app.config("open_trace") || false;
-  const serviceName = app.name || "unknownKoattyProject";
+  const asyncHooks = app.config("async_hooks") || false;
 
   // 
   let tracer: Tracer;
@@ -70,7 +70,7 @@ export function Trace(options: TraceOptions, app: Koatty): Koa.Middleware {
     }
   }
 
-  return async (ctx: KoattyContext, next: Koa.Next) => {
+  return async (ctx: KoattyContext, next: KoattyNext) => {
     // server terminated
     let terminated = false;
     if (app.server.status === 503) {
@@ -82,59 +82,55 @@ export function Trace(options: TraceOptions, app: Koatty): Koa.Middleware {
     // 
     const respWapper = async (requestId: string, span?: Span) => {
       // metadata
-      ctx.setMetaData(options.HeaderName, requestId);
+      ctx.setMetaData(options.RequestIdName, requestId);
 
       if (ctx.protocol === "grpc") {
         // allow bypassing koa
         ctx.respond = false;
-        ctx.rpc.call.metadata.set(options.HeaderName, requestId);
+        ctx.rpc.call.metadata.set(options.RequestIdName, requestId);
         await gRPCHandler(ctx, next, { timeout, requestId, encoding, terminated, span });
       } else if (ctx.protocol === "ws" || ctx.protocol === "wss") {
         // allow bypassing koa
         ctx.respond = false;
-        ctx.set(options.HeaderName, requestId);
+        ctx.set(options.RequestIdName, requestId);
         await wsHandler(ctx, next, { timeout, requestId, encoding, terminated, span });
       } else {
         // response header
-        ctx.set(options.HeaderName, requestId);
+        ctx.set(options.RequestIdName, requestId);
         await httpHandler(ctx, next, { timeout, requestId, encoding, terminated, span });
       }
       return respond(ctx);
     }
 
     let requestId = '';
-
     if (ctx.protocol === "grpc") {
       const request: any = ctx.getMetaData("_body")[0] || {};
-      requestId = `${ctx.getMetaData(requestIdName)[0]}` || <string>request[requestIdName];
+      requestId = `${ctx.getMetaData(options.RequestIdName)[0]}` || <string>request[options.RequestIdName];
     } else {
-      requestId = <string>ctx.headers[requestIdName] || <string>ctx.query[requestIdName];
+      const requestIdHeaderName = options.RequestIdHeaderName.toLowerCase();
+      requestId = <string>ctx.headers[requestIdHeaderName] || <string>ctx.query[options.RequestIdName];
     }
     requestId = requestId || GetTraceId(options);
-
+    let span: Span;
     if (openTrace) {
-      let span: Span;
-      if (tracer) {
-        const wireCtx = tracer.extract(FORMAT_HTTP_HEADERS, ctx.req.headers);
-        if (wireCtx != null) {
-          span = tracer.startSpan(serviceName, { childOf: wireCtx });
-        } else {
-          span = tracer.startSpan(serviceName);
-        }
-        span?.addTags({ requestId });
+      const serviceName = app.name || "unknownKoattyProject";
+      const wireCtx = tracer.extract(FORMAT_HTTP_HEADERS, ctx.req.headers);
+      if (wireCtx != null) {
+        span = tracer.startSpan(serviceName, { childOf: wireCtx });
+      } else {
+        span = tracer.startSpan(serviceName);
       }
-
+      span.addTags({ requestId });
       ctx.setMetaData("tracer_span", span);
-
+    }
+    if (asyncHooks) {
       return asyncLocalStorage.run(requestId, () => {
         const asyncResource = createAsyncResource();
         wrapEmitter(ctx.req, asyncResource);
         wrapEmitter(ctx.res, asyncResource);
         return respWapper(requestId, span);
       });
     }
-
     return respWapper(requestId);
   }
-}
-
+}