import config from '../../config/AgentConfig' ; import Context from '../../trace/context/Context' ; import Span from '../../trace/span/Span' ; import SpanContext from '../../trace/context/SpanContext' ; import DummyContext from '../../trace/context/DummyContext' ; import async_hooks from 'async_hooks' ; type AsyncState = { spans : Span [ ] ; } ; let store ; if ( async_hooks . AsyncLocalStorage ) { store = new async_hooks . AsyncLocalStorage ( ) ; } else { const executionAsyncId = async_hooks . executionAsyncId ; const asyncLocalStore = { } ; store = { getStore ( ) { return asyncLocalStore [ executionAsyncId ( ) ] as unknown as AsyncState ; } , enterWith ( s ) { asyncLocalStore [ executionAsyncId ( ) ] = s ; } , } ; async_hooks . createHook ( { init ( asyncId , type , triggerId ) { asyncLocalStore [ asyncId ] = asyncLocalStore [ triggerId ] ; } , destroy ( asyncId ) { delete asyncLocalStore [ asyncId ] ; } , } ) . enable ( ) ; } class ContextManager { isCold = true ; cosntructor ( ) { setTimeout ( ( ) => this . isCold = false , 1000 ) . unref ( ) ; } checkCold ( ) { const isCold = this . isCold ; this . isCold = false ; return isCold ; } get asyncState ( ) : AsyncState { let asyncState = store . getStore ( ) ; if ( ! asyncState ) { asyncState = { spans : [ ] } ; store . enterWith ( asyncState ) ; } return asyncState ; } get currentSpan ( ) : Span { const spans = store . getStore ( ) ?. spans ; return spans ?. [ spans . length - 1 ] as Span ; } ; get hasContext ( ) : boolean | undefined { return Boolean ( store . getStore ( ) ?. spans . length ) ; } get current ( ) : Context { const asyncState = this . asyncState ; if ( asyncState . spans . length ) return asyncState . spans [ asyncState . spans . length - 1 ] . context ; if ( SpanContext . nActiveSegments < config . maxBufferSize ) return new SpanContext ( ) ; return new DummyContext ( ) ; } get spans ( ) : Span [ ] { return this . asyncState . spans ; } spansDup ( ) { let asyncState = store . getStore ( ) ; if ( ! asyncState ) { asyncState = { spans : [ ] } ; } else { asyncState = { spans : [ ... asyncState . spans ] } ; } store . enterWith ( asyncState ) ; return asyncState . spans ; } clear ( span ) { const spans = this . spansDup ( ) ; const idx = spans . indexOf ( span ) ; if ( idx !== - 1 ) spans . splice ( idx , 1 ) ; } restore ( span ) { const spans = this . spansDup ( ) ; if ( spans . indexOf ( span ) === - 1 ) spans . push ( span ) ; } withSpan ( span , callback , ... args ) { if ( ! span . startTime ) span . start ( ) ; try { return callback ( ... args ) ; } catch ( e ) { span . error ( e ) ; throw e ; } finally { span . stop ( ) ; } } async withSpanAwait ( span , callback , ... args ) { if ( ! span . startTime ) span . start ( ) ; try { return await callback ( ... args ) ; } catch ( e ) { span . error ( e ) ; throw e ; } finally { span . stop ( ) ; } } } export default new ContextManager ( ) ;