import { callback , delay } from "awaiting" ; import { EventEmitter } from "events" ; import { JupyterKernel , VERSION } from "./jupyter" ; import { MessageType } from "@nteract/messaging" ; import { bind_methods , copy_with , deep_copy , uuid , trunc , } from "@cocalc/util/misc" ; import { CodeExecutionEmitterInterface , ExecOpts , StdinFunction , } from "@cocalc/frontend/jupyter/project-interface" ; type State = "init" | "closed" | "running" ; export class CodeExecutionEmitter extends EventEmitter implements CodeExecutionEmitterInterface { readonly kernel : JupyterKernel ; readonly code : string ; readonly id ? : string ; readonly stdin ? : StdinFunction ; readonly halt_on_error : boolean ; private iopub_done : boolean = false ; private shell_done : boolean = false ; private state : State = "init" ; private all_output : object [ ] = [ ] ; private _message : any ; private _go_cb : Function | undefined = undefined ; private timeout_ms ? : number ; private timer ? : any ; private killing : string = "" ; constructor ( kernel , opts ) { super ( ) ; this . kernel = kernel ; this . code = opts . code ; this . id = opts . id ; this . stdin = opts . stdin ; this . halt_on_error = ! ! opts . halt_on_error ; this . timeout_ms = opts . timeout_ms ; this . _message = { parent_header : { } , metadata : { } , channel : "shell" , header : { msg_id : `execute_${ uuid ( ) }` , username : "" , session : "" , msg_type : "execute_request" as MessageType , version : VERSION , date : new Date ( ) . toISOString ( ) , } , content : { code : this . code , silent : false , store_history : true , user_expressions : { } , allow_stdin : this . stdin != null , } , } ; bind_methods ( this ) ; } emit_output ( output ) { this . all_output . push ( output ) ; this . emit ( "output" , output ) ; } cancel ( ) { this . emit ( "canceled" ) ; } close ( ) { if ( this . state == "closed" ) return ; if ( this . timer != null ) { clearTimeout ( this . timer ) ; delete this . timer ; } this . state = "closed" ; this . emit ( "closed" ) ; this . removeAllListeners ( ) ; } throw_error ( err ) { this . emit ( "error" , err ) ; this . close ( ) ; } async _handle_stdin ( mesg ) { const dbg = this . kernel . dbg ( "_handle_stdin" ) ; if ( ! this . stdin ) { throw Error ( "BUG -- stdin handling not supported" ) ; } dbg ( `STDIN kernel --> server: ${ JSON . stringify ( mesg ) }` ) ; if ( mesg . parent_header . msg_id !== this . _message . header . msg_id ) { dbg ( `STDIN msg_id mismatch: ${ mesg . parent_header . msg_id }!=${ this . _message . header . msg_id }` ) ; return ; } let response ; try { response = await this . stdin ( mesg . content . prompt ? mesg . content . prompt : "" , ! ! mesg . content . password ) ; } catch ( err ) { response = `ERROR -- ${ err }` ; } dbg ( `STDIN client --> server ${ JSON . stringify ( response ) }` ) ; const m = { channel : "stdin" , parent_header : this . _message . header , metadata : { } , header : { msg_id : uuid ( ) , username : "" , session : "" , msg_type : "input_reply" as MessageType , version : VERSION , date : new Date ( ) . toISOString ( ) , } , content : { value : response , } , } ; dbg ( `STDIN server --> kernel: ${ JSON . stringify ( m ) }` ) ; this . kernel . channel ?. next ( m ) ; } _handle_shell ( mesg ) { if ( mesg . parent_header . msg_id !== this . _message . header . msg_id ) { return ; } const dbg = this . kernel . dbg ( "_handle_shell" ) ; dbg ( `got SHELL message -- ${ JSON . stringify ( mesg ) }` ) ; if ( mesg . content ?. status == "error" || mesg . content ?. status == "abort" ) { if ( this . halt_on_error ) { this . kernel . clear_execute_code_queue ( ) ; } this . set_shell_done ( true ) ; } else if ( mesg . content ?. status == "ok" ) { this . _push_mesg ( mesg ) ; this . set_shell_done ( true ) ; } } private set_shell_done ( value ) { this . shell_done = value ; if ( this . iopub_done && this . shell_done ) { this . _finish ( ) ; } } private set_iopub_done ( value ) { this . iopub_done = value ; if ( this . iopub_done && this . shell_done ) { this . _finish ( ) ; } } _handle_iopub ( mesg ) { if ( mesg . parent_header . msg_id !== this . _message . header . msg_id ) { return ; } const dbg = this . kernel . dbg ( "_handle_iopub" ) ; dbg ( `got IOPUB message -- ${ JSON . stringify ( mesg ) }` ) ; if ( mesg . content ?. comm_id != null ) { } else { this . _push_mesg ( mesg ) ; } this . set_iopub_done ( ! ! this . killing || mesg . content ?. execution_state == "idle" ) ; } private handle_closed ( ) { const dbg = this . kernel . dbg ( "CodeExecutionEmitter.handle_closed" ) ; dbg ( "kernel closed" ) ; this . killing = "kernel crashed" ; this . _finish ( ) ; } _finish ( ) { if ( this . state == "closed" ) { return ; } this . kernel . removeListener ( "iopub" , this . _handle_iopub ) ; if ( this . stdin != null ) { this . kernel . removeListener ( "stdin" , this . _handle_stdin ) ; } this . kernel . removeListener ( "shell" , this . _handle_shell ) ; if ( this . kernel . _execute_code_queue != null ) { this . kernel . _execute_code_queue . shift ( ) ; this . kernel . _process_execute_code_queue ( ) ; } this . kernel . removeListener ( "close" , this . handle_closed ) ; this . _push_mesg ( { done : true } ) ; this . close ( ) ; this . _go_cb ?. ( this . killing ) ; this . _go_cb = undefined ; } _push_mesg ( mesg ) { const header = mesg . header ; mesg = copy_with ( mesg , [ "metadata" , "content" , "buffers" , "done" ] ) ; mesg = deep_copy ( mesg ) ; if ( header !== undefined ) { mesg . msg_type = header . msg_type ; } this . emit_output ( mesg ) ; } async go ( ) { await callback ( this . _go ) ; return this . all_output ; } _go ( cb ) { if ( this . state != "init" ) { cb ( "may only run once" ) ; return ; } this . state = "running" ; const dbg = this . kernel . dbg ( `_execute_code('${ trunc ( this . code , 15 ) }')` ) ; dbg ( `code='${ this . code }'` ) ; if ( this . kernel . get_state ( ) === "closed" ) { this . close ( ) ; cb ( "closed" ) ; return ; } this . _go_cb = cb ; if ( this . stdin != null ) { this . kernel . on ( "stdin" , this . _handle_stdin ) ; } this . kernel . on ( "shell" , this . _handle_shell ) ; this . kernel . on ( "iopub" , this . _handle_iopub ) ; dbg ( "send the message to get things rolling" ) ; this . kernel . channel ?. next ( this . _message ) ; this . kernel . on ( "closed" , this . handle_closed ) ; if ( this . timeout_ms ) { this . timer = setTimeout ( this . timeout , this . timeout_ms ) ; } } private async timeout ( ) { const dbg = this . kernel . dbg ( "CodeExecutionEmitter.timeout" ) ; if ( this . state == "closed" ) { dbg ( "already finished, so nothing to worry about" ) ; return ; } this . killing = "Timeout Error: execution time limit = " + Math . round ( ( this . timeout_ms ?? 0 ) / 1000 ) + " seconds" ; let tries = 3 ; let d = 1000 ; while ( this . state != ( "closed" as State ) && tries > 0 ) { dbg ( "code still running, so try to interrupt it" ) ; this . kernel . signal ( "SIGINT" ) ; await delay ( d ) ; d *= 1.3 ; tries -= 1 ; } if ( this . state != ( "closed" as State ) ) { dbg ( "now try SIGKILL, which should kill things for sure." ) ; this . kernel . signal ( "SIGKILL" ) ; this . _finish ( ) ; } } }