Refactor TryCatchFinallyHooksBuilder class and callStack hook

Author: pmunin

src/TryCatchFinallyHooks.ts

@@ -45,11 +45,10 @@ export class TryCatchFinallyHooksBuilder<FullHookContext extends FunctionContext
     const _this = this
     const beforeHooksTry = this.beforeHooksTry.bind(this)
     const afterHooksTry = this.afterHooksTry?.bind(this)
-    type TContext = FullHookContext & FunctionContext
     return (func:any)=>{
       return createTryCatchFinally(func, {
         onTry(funcCtx) {
-          const ctx:TContext = funcCtx as any
+          const ctx:FullHookContext = funcCtx as any
           if(args)
             (ctx as any).args = args
           const bht = beforeHooksTry(ctx)
@@ -122,6 +121,12 @@ export class TryCatchFinallyHooksBuilder<FullHookContext extends FunctionContext
   }
 
   private _currentContext:(FullHookContext & FunctionContext)|undefined = undefined
+
+
+  /**
+   * Only safe to use in sync functions or in async functions before any awaited code.
+   * Otherwise, the context may be changed by another async function - in that case use callStack hook instead
+   */
   get current(){
     return this._currentContext
   }

src/callStack.test.ts

@@ -7,13 +7,13 @@ function createTrack(log:any){
   .add(callStack)
   .add({
     onTry(ctx) {
-      log("onTry", [...ctx.callstack.map(c=>c.args.name), ctx.args.name].join('/'))
+      log("onTry", [...ctx.getCallStack().map(c=>c.name), ctx.name].join('/'))
       return {
         onFinally() {
-          log("onFinally",[...ctx.callstack.map(c=>c.args.name), ctx.args.name].join('/'))
+          log("onFinally",[...ctx.getCallStack().map(c=>c.name), ctx.name].join('/'))
         },
         onCatch() {
-          log("onCatch",[...ctx.callstack.map(c=>c.args.name), ctx.args.name].join('/'))
+          log("onCatch",[...ctx.getCallStack().map(c=>c.name), ctx.name].join('/'))
         }
       }
     }

src/callStack.ts

@@ -11,36 +11,45 @@ import { FunctionContext } from "./tryCatchFinally";
  * During onTry Context.func (that is about to be invoked) is replaced with a special function wrapper that has assigned unique name `callstack_callId:<context.callId>`.
  * Now those function names in Error.captureCallstack() can be associated context.callstack.callId via list of context objects that are currently active.
  * 
- * If you don't see actions in your callstack, increase Error.stackTraceLimit
+ * If you don't see actions in your callstack, increase Error.stackTraceLimit, but keep in mind that it impacts performance
  */
 
 const callPrefix = 'callstack_callId:'
 
-export type CallstackContext = { args: { name?: string;}, name:string, id: string, callstack: CallstackContext[] };
+export type CallStackContext = { 
+  args: {
+    name?: string;
+    isVisibleInCallStack?:boolean
+  },
+  name:string,
+  id: string,
+  
+  /**
+   * Returns current tracked callstack. Filters Error.captureStackTrace by tracked functions.
+   * If you don't see actions in your callstack, increase Error.stackTraceLimit
+   */
+  getCallStack(): CallStackContext[] 
+};
 
 
-export const callStack = TryCatchFinallyHooksBuilder.createHook<CallstackContext>(
+export const callStack = TryCatchFinallyHooksBuilder.createHook<CallStackContext>(
   (ctx)=>{
     if(!ctx.name)
       ctx.name = ctx.args.name || ctx.func.name || '<anonymous>'
 
     if(!ctx.id)
       ctx.id = Math.random().toString(16).slice(2)
     getOrUpdateAllActiveContexts(actions=>[...actions, ctx])
-    
 
-    const origFunc = ctx.func
-    const callstackCallWrapper = function (this:any,...args:any) { return origFunc.apply(this, args) }
-    Object.defineProperty(callstackCallWrapper, 'name', {value:callPrefix+ctx.id})
-    ctx.func = callstackCallWrapper
-
-    Object.defineProperty(ctx, 'callstack', { 
-      get() {
-        const contextsMap = new Map(getOrUpdateAllActiveContexts().map(ctx=>[ctx.id, ctx]))
-        const actionCallIdStack = getActionCallIdsStack()
-        return actionCallIdStack.map(ctxId=>contextsMap.get(ctxId)!)
-      }
-    });
+    const isVisible = 'isVisibleInCallStack' in ctx.args ? ctx.args.isVisibleInCallStack : true
+    
+    if(isVisible){
+      const origFunc = ctx.func
+      const callstackCallWrapper = function (this:any,...args:any) { return origFunc.apply(this, args) }
+      Object.defineProperty(callstackCallWrapper, 'name', {value:callPrefix+ctx.id})
+      ctx.func = callstackCallWrapper
+    }
+    ctx.getCallStack = getCurrentCallstack
 
     return {
       onFinally() {
@@ -51,6 +60,15 @@ export const callStack = TryCatchFinallyHooksBuilder.createHook<CallstackContext
   }
 )
 
+/**
+ * Returns current tracked callstack. Filters Error.captureStackTrace by tracked functions.
+ * If you don't see actions in your callstack, increase Error.stackTraceLimit
+ */
+export function getCurrentCallstack():CallStackContext[]{
+  const contextsMap = new Map(getOrUpdateAllActiveContexts().map(ctx=>[ctx.id, ctx]))
+  return getActionCallIdsStack().map(ctxId=>contextsMap.get(ctxId)).filter(c=>c) as CallStackContext[]
+}
+
 function prepareActionStackTrace(error: Error, stack: NodeJS.CallSite[]): string[] {
   return stack
     .map(s=>s.getFunctionName()!)
@@ -76,16 +94,16 @@ function getActionCallIdsStack():string[]
 
 const allActiveContexts// mimicking AsyncLocalStorage
  = {
-  _items:undefined as CallstackContext[]|undefined,
-  getStore():CallstackContext[]|undefined{
+  _items:undefined as CallStackContext[]|undefined,
+  getStore():CallStackContext[]|undefined{
     return [...this._items||[]]
   },
-  enterWith(items:CallstackContext[]){
+  enterWith(items:CallStackContext[]){
     this._items = items||[]
   }
  }
 
-function getOrUpdateAllActiveContexts(update?:(old:CallstackContext[])=>undefined|CallstackContext[]):CallstackContext[]{
+function getOrUpdateAllActiveContexts(update?:(old:CallStackContext[])=>undefined|CallStackContext[]):CallStackContext[]{
   const store = [...allActiveContexts.getStore()||[]]
   if(update)
   {

src/example/tracker.ts

@@ -11,7 +11,7 @@ const track = new TryCatchFinallyHooksBuilder()
     return {
       onFinally() {
         const datadog: any = {};
-        datadog.histogram(`action-${ctx.args.name}-duration`, ctx.duration, { tags: ['action:' + ctx.name, 'callstack:' + ctx.callstack.map(c => c.name).join('/')] });
+        datadog.histogram(`action-${ctx.args.name}-duration`, ctx.duration, { tags: ['action:' + ctx.name, 'callstack:' + ctx.getCallStack().map(c => c.name).join('/')] });
       },
     };
   })