import { UpdateResponse } from '@elastic/elasticsearch/lib/api/typesWithBodyKey' ; import moment from 'moment' ; import * as Rx from 'rxjs' ; import { timeout } from 'rxjs/operators' ; import { finished , Writable } from 'stream' ; import { promisify } from 'util' ; import { getContentStream , LevelLogger } from '../' ; import { ReportingCore } from '../../' ; import { RunContext , TaskManagerStartContract , TaskRunCreatorFunction , } from '../../../../task_manager/server' ; import { CancellationToken } from '../../../common/cancellation_token' ; import { durationToNumber , numberToDuration } from '../../../common/schema_utils' ; import type { ReportOutput } from '../../../common/types' ; import type { ReportingConfigType } from '../../config' ; import type { BasePayload , ExportTypeDefinition , RunTaskFn } from '../../types' ; import type { ReportDocument , ReportingStore } from '../store' ; import { Report , SavedReport } from '../store' ; import type { ReportFailedFields , ReportProcessingFields } from '../store/store' ; import { ReportingTask , ReportingTaskStatus , REPORTING_EXECUTE_TYPE , ReportTaskParams , TaskRunResult , } from './' ; import { errorLogger } from './error_logger' ; type CompletedReportOutput = Omit < ReportOutput , 'content' > ; interface ReportingExecuteTaskInstance { state : object ; taskType : string ; params : ReportTaskParams ; runAt ? : Date ; } interface TaskExecutor extends Pick < ExportTypeDefinition , 'jobContentEncoding' > { jobExecutor : RunTaskFn < BasePayload > ; } function isOutput ( output ) { return ( output as CompletedReportOutput ) . size != null ; } function reportFromTask ( task ) { return new Report ( { ... task , _id : task . id , _index : task . index } ) ; } export class ExecuteReportTask implements ReportingTask { public TYPE = REPORTING_EXECUTE_TYPE ; private logger : LevelLogger ; private taskManagerStart ? : TaskManagerStartContract ; private taskExecutors ? : Map < string , TaskExecutor > ; private kibanaId ? : string ; private kibanaName ? : string ; private store ? : ReportingStore ; constructor ( private reporting , private config , logger ) { this . logger = logger . clone ( [ 'runTask' ] ) ; } public async init ( taskManager ) { this . taskManagerStart = taskManager ; const { reporting } = this ; const exportTypesRegistry = reporting . getExportTypesRegistry ( ) ; const executors = new Map < string , TaskExecutor > ( ) ; for ( const exportType of exportTypesRegistry . getAll ( ) ) { const exportTypeLogger = this . logger . clone ( [ exportType . id ] ) ; const jobExecutor = exportType . runTaskFnFactory ( reporting , exportTypeLogger ) ; executors . set ( exportType . jobType , { jobExecutor , jobContentEncoding : exportType . jobContentEncoding , } ) ; } this . taskExecutors = executors ; const config = reporting . getConfig ( ) ; this . kibanaId = config . kbnConfig . get ( 'server' , 'uuid' ) ; this . kibanaName = config . kbnConfig . get ( 'server' , 'name' ) ; } private async getStore ( ) { if ( this . store ) { return this . store ; } const { store } = await this . reporting . getPluginStartDeps ( ) ; this . store = store ; return store ; } private getTaskManagerStart ( ) { if ( ! this . taskManagerStart ) { throw new Error ( 'Reporting task runner has not been initialized!' ) ; } return this . taskManagerStart ; } private getJobContentEncoding ( jobType ) { return this . taskExecutors ?. get ( jobType ) ?. jobContentEncoding ; } public async _claimJob ( task ) { if ( this . kibanaId == null ) { throw new Error ( `Kibana instance ID is undefined!` ) ; } if ( this . kibanaName == null ) { throw new Error ( `Kibana instance name is undefined!` ) ; } const store = await this . getStore ( ) ; const report = await store . findReportFromTask ( task ) ; if ( report . status === 'completed' ) { throw new Error ( `Can not claim the report job: it is already completed!` ) ; } const m = moment ( ) ; const maxAttempts = this . config . capture . maxAttempts ; if ( report . attempts >= maxAttempts ) { const err = new Error ( `Max attempts reached (${ maxAttempts }). Queue timeout reached.` ) ; await this . _failJob ( report , err ) ; throw err ; } const queueTimeout = durationToNumber ( this . config . queue . timeout ) ; const startTime = m . toISOString ( ) ; const expirationTime = m . add ( queueTimeout ) . toISOString ( ) ; const doc = { kibana_id : this . kibanaId , kibana_name : this . kibanaName , attempts : report . attempts + 1 , max_attempts : maxAttempts , started_at : startTime , timeout : queueTimeout , process_expiration : expirationTime , } ; const claimedReport = new SavedReport ( { ... report , ... doc , } ) ; this . logger . debug ( `Claiming ${ claimedReport . jobtype } ${ report . _id } ` + `[_index: ${ report . _index }]  ` + `[_seq_no: ${ report . _seq_no }]  ` + `[_primary_term: ${ report . _primary_term }]  ` + `[attempts: ${ report . attempts }]  ` + `[process_expiration: ${ expirationTime }]` ) ; const resp = await store . setReportClaimed ( claimedReport , doc ) ; claimedReport . _seq_no = resp . _seq_no ; claimedReport . _primary_term = resp . _primary_term ; return claimedReport ; } private async _failJob ( report , error ? ) { const message = `Failing ${ report . jobtype } job ${ report . _id }` ; let docOutput ; if ( error ) { errorLogger ( this . logger , message , error ) ; docOutput = this . _formatOutput ( error ) ; } else { errorLogger ( this . logger , message ) ; } const store = await this . getStore ( ) ; const completedTime = moment ( ) . toISOString ( ) ; const doc = { completed_at : completedTime , output : docOutput ?? null , } ; return await store . setReportFailed ( report , doc ) ; } private _formatOutput ( output ) { const docOutput = { } as ReportOutput ; const unknownMime = null ; if ( isOutput ( output ) ) { docOutput . content_type = output . content_type || unknownMime ; docOutput . max_size_reached = output . max_size_reached ; docOutput . csv_contains_formulas = output . csv_contains_formulas ; docOutput . size = output . size ; docOutput . warnings = output . warnings && output . warnings . length > 0 ? output . warnings : undefined ; } else { const defaultOutput = null ; docOutput . content = output . toString ( ) || defaultOutput ; docOutput . content_type = unknownMime ; docOutput . warnings = [ output . toString ( ) ] ; } return docOutput ; } public async _performJob ( task , cancellationToken , stream ) { if ( ! this . taskExecutors ) { throw new Error ( `Task run function factories have not been called yet!` ) ; } const runner = this . taskExecutors . get ( task . jobtype ) ; if ( ! runner ) { throw new Error ( `No defined task runner function for ${ task . jobtype }!` ) ; } const queueTimeout = durationToNumber ( this . config . queue . timeout ) ; return Rx . from ( runner . jobExecutor ( task . id , task . payload , cancellationToken , stream ) ) . pipe ( timeout ( queueTimeout ) ) . toPromise ( ) ; } public async _completeJob ( report , output ) { let docId = `/${ report . _index }/_doc/${ report . _id }` ; this . logger . debug ( `Saving ${ report . jobtype } to ${ docId }.` ) ; const completedTime = moment ( ) . toISOString ( ) ; const docOutput = this . _formatOutput ( output ) ; const store = await this . getStore ( ) ; const doc = { completed_at : completedTime , output : docOutput , } ; docId = `/${ report . _index }/_doc/${ report . _id }` ; const resp = await store . setReportCompleted ( report , doc ) ; this . logger . info ( `Saved ${ report . jobtype } job ${ docId }` ) ; report . _seq_no = resp . _seq_no ; report . _primary_term = resp . _primary_term ; return report ; } private getTaskRunner ( ) { return ( context ) => { let jobId ; const cancellationToken = new CancellationToken ( ) ; return { run : async ( ) => { let report ; const task = context . taskInstance . params as ReportTaskParams ; jobId = task ?. id ; try { if ( ! jobId ) { throw new Error ( 'Invalid report data provided in scheduled task!' ) ; } this . reporting . trackReport ( jobId ) ; report = await this . _claimJob ( task ) ; } catch ( failedToClaim ) { errorLogger ( this . logger , `Error in claiming ${ jobId }` , failedToClaim ) ; } if ( ! report ) { this . reporting . untrackReport ( jobId ) ; errorLogger ( this . logger , `Job ${ jobId } could not be claimed. Exiting...` ) ; return ; } const { jobtype : jobType , attempts } = report ; const maxAttempts = this . config . capture . maxAttempts ; this . logger . debug ( `Starting ${ jobType } report ${ jobId }: attempt ${ attempts } of ${ maxAttempts }.` ) ; this . logger . debug ( `Reports running: ${ this . reporting . countConcurrentReports ( ) }.` ) ; const eventLog = this . reporting . getEventLogger ( new Report ( { ... task , _id : task . id , _index : task . index } ) ) ; try { const jobContentEncoding = this . getJobContentEncoding ( jobType ) ; const stream = await getContentStream ( this . reporting , { id : report . _id , index : report . _index , if_primary_term : report . _primary_term , if_seq_no : report . _seq_no , } , { encoding : jobContentEncoding === 'base64' ? 'base64' : 'raw' , } ) ; eventLog . logExecutionStart ( ) ; const output = await this . _performJob ( task , cancellationToken , stream ) ; stream . end ( ) ; eventLog . logExecutionComplete ( { byteSize : stream . bytesWritten } ) ; await promisify ( finished ) ( stream , { readable : false } ) ; report . _seq_no = stream . getSeqNo ( ) ! ; report . _primary_term = stream . getPrimaryTerm ( ) ! ; if ( output ) { this . logger . debug ( `Job output size: ${ stream . bytesWritten } bytes.` ) ; report = await this . _completeJob ( report , { ... output , size : stream . bytesWritten , } ) ; } this . logger . debug ( `Stopping ${ jobId }.` ) ; } catch ( failedToExecuteErr ) { eventLog . logError ( failedToExecuteErr ) ; cancellationToken . cancel ( ) ; if ( attempts < maxAttempts ) { try { if ( report == null ) { throw new Error ( `Report ${ jobId } is null!` ) ; } const remainingAttempts = maxAttempts - report . attempts ; errorLogger ( this . logger , `Scheduling retry for job ${ jobId }. Retries remaining: ${ remainingAttempts }.` , failedToExecuteErr ) ; await this . rescheduleTask ( reportFromTask ( task ) . toReportTaskJSON ( ) , this . logger ) ; } catch ( rescheduleErr ) { errorLogger ( this . logger , `Could not reschedule the errored job ${ jobId }!` , rescheduleErr ) ; } } else { try { const maxAttemptsMsg = `Max attempts (${ attempts }) reached for job ${ jobId }. Failed with: ${ failedToExecuteErr }` ; if ( report == null ) { throw new Error ( `Report ${ jobId } is null!` ) ; } const resp = await this . _failJob ( report , new Error ( maxAttemptsMsg ) ) ; report . _seq_no = resp . _seq_no ; report . _primary_term = resp . _primary_term ; } catch ( failedToFailError ) { errorLogger ( this . logger , `Could not fail ${ jobId }!` , failedToFailError ) ; } } } finally { this . reporting . untrackReport ( jobId ) ; this . logger . debug ( `Reports running: ${ this . reporting . countConcurrentReports ( ) }.` ) ; } } , cancel : async ( ) => { if ( jobId ) { this . logger . warn ( `Cancelling job ${ jobId }...` ) ; } cancellationToken . cancel ( ) ; } , } ; } ; } public getTaskDefinition ( ) { const queueTimeout = Math . ceil ( numberToDuration ( this . config . queue . timeout ) . asSeconds ( ) ) + 's' ; const maxConcurrency = this . config . queue . pollEnabled ? 1 : 0 ; return { type : REPORTING_EXECUTE_TYPE , title : 'Reporting: execute job' , createTaskRunner : this . getTaskRunner ( ) , maxAttempts : 1 , timeout : queueTimeout , maxConcurrency , } ; } public async scheduleTask ( params ) { const taskInstance = { taskType : REPORTING_EXECUTE_TYPE , state : { } , params , } ; return await this . getTaskManagerStart ( ) . schedule ( taskInstance ) ; } private async rescheduleTask ( task , logger ) { logger . info ( `Rescheduling task:${ task . id } to retry after error.` ) ; const oldTaskInstance = { taskType : REPORTING_EXECUTE_TYPE , state : { } , params : task , } ; const newTask = await this . getTaskManagerStart ( ) . schedule ( oldTaskInstance ) ; logger . debug ( `Rescheduled task:${ task . id }. New task: task:${ newTask . id }` ) ; return newTask ; } public getStatus ( ) { if ( this . taskManagerStart ) { return ReportingTaskStatus . INITIALIZED ; } return ReportingTaskStatus . UNINITIALIZED ; } }