import { AsyncLocalStorage } from 'async_hooks' ; import type { Subscription } from 'rxjs' ; import type { CoreService , KibanaExecutionContext } from '../../types' ; import type { CoreContext } from '../core_context' ; import type { Logger } from '../logging' ; import type { ExecutionContextConfigType } from './execution_context_config' ; import { ExecutionContextContainer , IExecutionContextContainer , getParentContextFrom , } from './execution_context_container' ; internal export interface IExecutionContext { getParentContextFrom ( headers ) : KibanaExecutionContext | undefined ; setRequestId ( requestId ) : void ; set ( context ) : void ; /**
     * The sole purpose of this imperative internal API is to be used by the http service.
     * The event-based nature of Hapi server doesn't allow us to wrap a request handler with "withContext".
     * Since all the Hapi event lifecycle will lose the execution context.
     * Nodejs docs also recommend using AsyncLocalStorage.run() over AsyncLocalStorage.enterWith().
     * https://nodejs.org/api/async_context.html#async_context_asynclocalstorage_enterwith_store
     */ get ( ) : IExecutionContextContainer | undefined ; withContext < R > ( context , fn ) : R ; /**
     * returns serialized representation to send as a header
     **/ getAsHeader ( ) : string | undefined ; } internal export type InternalExecutionContextSetup = IExecutionContext ; internal export type InternalExecutionContextStart = IExecutionContext ; public export interface ExecutionContextSetup { /**
     * Keeps track of execution context while the passed function is executed.
     * Data are carried over all async operations spawned by the passed function.
     * The nested calls stack the registered context on top of each other.
     **/ withContext < R > ( context , fn ) : R ; } public export type ExecutionContextStart = ExecutionContextSetup ; export class ExecutionContextService implements CoreService < InternalExecutionContextSetup , InternalExecutionContextStart > { private readonly log : Logger ; private readonly contextStore : AsyncLocalStorage < IExecutionContextContainer > ; private readonly requestIdStore : AsyncLocalStorage < { requestId : string ; } > ; private enabled = false ; private configSubscription ? : Subscription ; constructor ( private readonly coreContext ) { this . log = coreContext . logger . get ( 'execution_context' ) ; this . contextStore = new AsyncLocalStorage < IExecutionContextContainer > ( ) ; this . requestIdStore = new AsyncLocalStorage < { requestId : string ; } > ( ) ; } setup ( ) { this . configSubscription = this . coreContext . configService . atPath < ExecutionContextConfigType > ( 'execution_context' ) . subscribe ( ( config ) => { this . enabled = config . enabled ; } ) ; return { getParentContextFrom , set : this . set . bind ( this ) , withContext : this . withContext . bind ( this ) , setRequestId : this . setRequestId . bind ( this ) , get : this . get . bind ( this ) , getAsHeader : this . getAsHeader . bind ( this ) , } ; } start ( ) { return { getParentContextFrom , set : this . set . bind ( this ) , setRequestId : this . setRequestId . bind ( this ) , withContext : this . withContext . bind ( this ) , get : this . get . bind ( this ) , getAsHeader : this . getAsHeader . bind ( this ) , } ; } stop ( ) { this . enabled = false ; if ( this . configSubscription ) { this . configSubscription . unsubscribe ( ) ; this . configSubscription = undefined ; } } private set ( context ) { if ( ! this . enabled ) return ; const contextContainer = new ExecutionContextContainer ( context ) ; this . contextStore . enterWith ( contextContainer ) ; this . log . debug ( JSON . stringify ( contextContainer ) ) ; } private withContext < R > ( context , fn ) { if ( ! this . enabled || ! context ) { return fn ( ) ; } const parent = this . contextStore . getStore ( ) ; const contextContainer = new ExecutionContextContainer ( context , parent ) ; this . log . debug ( JSON . stringify ( contextContainer ) ) ; return this . contextStore . run ( contextContainer , fn ) ; } private setRequestId ( requestId ) { if ( ! this . enabled ) return ; this . requestIdStore . enterWith ( { requestId } ) ; } private get ( ) { if ( ! this . enabled ) return ; return this . contextStore . getStore ( ) ; } private getAsHeader ( ) { if ( ! this . enabled ) return ; const requestId = this . requestIdStore . getStore ( ) ?. requestId ?? 'unknownId' ; const executionContext = this . contextStore . getStore ( ) ?. toString ( ) ; const executionContextStr = executionContext ? `;kibana:${ executionContext }` : '' ; return `${ requestId }${ executionContextStr }` ; } }