Skip to content

Commit 23003ec

Browse files
koesie10koesie10
committed
Extract process output handling in CLI server
1 parent c32065e commit 23003ec

1 file changed

Lines changed: 169 additions & 102 deletions

File tree

  • extensions/ql-vscode/src/codeql-cli

‎extensions/ql-vscode/src/codeql-cli/cli.ts‎

Lines changed: 169 additions & 102 deletions
Original file line numberDiff line numberDiff line change
@@ -369,9 +369,6 @@ export class CodeQLCliServer implements Disposable {
369369
onLine?: OnLineCallback,
370370
silent?: boolean,
371371
): Promise<string> {
372-
const stderrBuffers: Buffer[] = [];
373-
// The current buffer of stderr of a single line. To be used for logging.
374-
let currentLineStderrBuffer: Buffer = Buffer.alloc(0);
375372
if (this.commandInProcess) {
376373
throw new Error("runCodeQlCliInternal called while cli was running");
377374
}
@@ -383,8 +380,6 @@ export class CodeQLCliServer implements Disposable {
383380
}
384381
// Grab the process so that typescript know that it is always defined.
385382
const process = this.process;
386-
// The array of fragments of stdout
387-
const stdoutBuffers: Buffer[] = [];
388383

389384
// Compute the full args array
390385
const args = command.concat(LOGGING_FLAGS).concat(commandArgs);
@@ -396,26 +391,83 @@ export class CodeQLCliServer implements Disposable {
396391
);
397392
}
398393
try {
399-
await new Promise<void>((resolve, reject) => {
400-
// Start listening to stdout
401-
process.stdout.addListener("data", (newData: Buffer) => {
402-
if (onLine) {
403-
void (async () => {
404-
const response = await onLine(newData.toString("utf-8"));
405-
406-
if (!response) {
407-
return;
408-
}
409-
410-
process.stdin.write(`${response}${EOL}`);
411-
412-
// Remove newData from stdoutBuffers because the data has been consumed
413-
// by the onLine callback.
414-
stdoutBuffers.splice(stdoutBuffers.indexOf(newData), 1);
415-
})();
416-
}
394+
return await this.handleProcessOutput(process, {
395+
handleNullTerminator: true,
396+
onListenStart: (process) => {
397+
// Write the command followed by a null terminator.
398+
process.stdin.write(JSON.stringify(args), "utf8");
399+
process.stdin.write(this.nullBuffer);
400+
},
401+
description,
402+
args,
403+
silent,
404+
onLine,
405+
});
406+
} catch (err) {
407+
// Kill the process if it isn't already dead.
408+
this.killProcessIfRunning();
409+
410+
throw err;
411+
}
412+
} finally {
413+
this.commandInProcess = false;
414+
// start running the next command immediately
415+
this.runNext();
416+
}
417+
}
418+
419+
private async handleProcessOutput(
420+
process: ChildProcessWithoutNullStreams,
421+
{
422+
handleNullTerminator,
423+
args,
424+
description,
425+
onLine,
426+
onListenStart,
427+
silent,
428+
}: {
429+
handleNullTerminator: boolean;
430+
args: string[];
431+
description: string;
432+
onLine?: OnLineCallback;
433+
onListenStart?: (process: ChildProcessWithoutNullStreams) => void;
434+
silent?: boolean;
435+
},
436+
): Promise<string> {
437+
const stderrBuffers: Buffer[] = [];
438+
// The current buffer of stderr of a single line. To be used for logging.
439+
let currentLineStderrBuffer: Buffer = Buffer.alloc(0);
440+
441+
// The listeners of the process. Declared here so they can be removed in the finally block.
442+
let stdoutListener: ((newData: Buffer) => void) | undefined = undefined;
443+
let stderrListener: ((newData: Buffer) => void) | undefined = undefined;
444+
let closeListener: ((code: number | null) => void) | undefined = undefined;
445+
446+
try {
447+
// The array of fragments of stdout
448+
const stdoutBuffers: Buffer[] = [];
417449

418-
stdoutBuffers.push(newData);
450+
await new Promise<void>((resolve, reject) => {
451+
stdoutListener = (newData: Buffer) => {
452+
if (onLine) {
453+
void (async () => {
454+
const response = await onLine(newData.toString("utf-8"));
455+
456+
if (!response) {
457+
return;
458+
}
459+
460+
process.stdin.write(`${response}${EOL}`);
461+
462+
// Remove newData from stdoutBuffers because the data has been consumed
463+
// by the onLine callback.
464+
stdoutBuffers.splice(stdoutBuffers.indexOf(newData), 1);
465+
})();
466+
}
467+
468+
stdoutBuffers.push(newData);
469+
470+
if (handleNullTerminator) {
419471
// If the buffer ends in '0' then exit.
420472
// We don't have to check the middle as no output will be written after the null until
421473
// the next command starts
@@ -425,89 +477,104 @@ export class CodeQLCliServer implements Disposable {
425477
) {
426478
resolve();
427479
}
428-
});
429-
// Listen to stderr
430-
process.stderr.addListener("data", (newData: Buffer) => {
431-
stderrBuffers.push(newData);
432-
433-
if (!silent) {
434-
currentLineStderrBuffer = Buffer.concat([
435-
currentLineStderrBuffer,
436-
newData,
437-
]);
438-
439-
// Print the stderr to the logger as it comes in. We need to ensure that
440-
// we don't split messages on the same line, so we buffer the stderr and
441-
// split it on EOLs.
442-
const eolBuffer = Buffer.from(EOL);
443-
444-
let hasCreatedSubarray = false;
445-
446-
let eolIndex;
447-
while (
448-
(eolIndex = currentLineStderrBuffer.indexOf(eolBuffer)) !== -1
449-
) {
450-
const line = currentLineStderrBuffer.subarray(0, eolIndex);
451-
void this.logger.log(line.toString("utf-8"));
452-
currentLineStderrBuffer = currentLineStderrBuffer.subarray(
453-
eolIndex + eolBuffer.length,
454-
);
455-
hasCreatedSubarray = true;
456-
}
480+
}
481+
};
482+
stderrListener = (newData: Buffer) => {
483+
stderrBuffers.push(newData);
484+
485+
if (!silent) {
486+
currentLineStderrBuffer = Buffer.concat([
487+
currentLineStderrBuffer,
488+
newData,
489+
]);
490+
491+
// Print the stderr to the logger as it comes in. We need to ensure that
492+
// we don't split messages on the same line, so we buffer the stderr and
493+
// split it on EOLs.
494+
const eolBuffer = Buffer.from(EOL);
495+
496+
let hasCreatedSubarray = false;
497+
498+
let eolIndex;
499+
while (
500+
(eolIndex = currentLineStderrBuffer.indexOf(eolBuffer)) !== -1
501+
) {
502+
const line = currentLineStderrBuffer.subarray(0, eolIndex);
503+
void this.logger.log(line.toString("utf-8"));
504+
currentLineStderrBuffer = currentLineStderrBuffer.subarray(
505+
eolIndex + eolBuffer.length,
506+
);
507+
hasCreatedSubarray = true;
508+
}
457509

458-
// We have created a subarray, which means that the complete original buffer is now referenced
459-
// by the subarray. We need to create a new buffer to avoid memory leaks.
460-
if (hasCreatedSubarray) {
461-
currentLineStderrBuffer = Buffer.from(currentLineStderrBuffer);
462-
}
510+
// We have created a subarray, which means that the complete original buffer is now referenced
511+
// by the subarray. We need to create a new buffer to avoid memory leaks.
512+
if (hasCreatedSubarray) {
513+
currentLineStderrBuffer = Buffer.from(currentLineStderrBuffer);
463514
}
464-
});
465-
// Listen for process exit.
466-
process.addListener("close", (code) =>
467-
reject(new ExitCodeError(code)),
468-
);
469-
// Write the command followed by a null terminator.
470-
process.stdin.write(JSON.stringify(args), "utf8");
471-
process.stdin.write(this.nullBuffer);
472-
});
473-
// Join all the data together
474-
const fullBuffer = Buffer.concat(stdoutBuffers);
475-
// Make sure we remove the terminator;
476-
const data = fullBuffer.toString("utf8", 0, fullBuffer.length - 1);
477-
if (!silent) {
478-
void this.logger.log(currentLineStderrBuffer.toString("utf8"));
479-
currentLineStderrBuffer = Buffer.alloc(0);
480-
void this.logger.log("CLI command succeeded.");
481-
}
482-
return data;
483-
} catch (err) {
484-
// Kill the process if it isn't already dead.
485-
this.killProcessIfRunning();
515+
}
516+
};
517+
closeListener = (code) => {
518+
if (handleNullTerminator) {
519+
reject(new ExitCodeError(code));
520+
} else {
521+
if (code === 0) {
522+
resolve();
523+
} else {
524+
reject(new ExitCodeError(code));
525+
}
526+
}
527+
};
486528

487-
// Report the error (if there is a stderr then use that otherwise just report the error code or nodejs error)
488-
const cliError = getCliError(
489-
err,
490-
stderrBuffers.length > 0
491-
? Buffer.concat(stderrBuffers).toString("utf8")
492-
: undefined,
493-
description,
494-
args,
495-
);
496-
cliError.stack += getErrorStack(err);
497-
throw cliError;
498-
} finally {
499-
if (!silent && currentLineStderrBuffer.length > 0) {
500-
void this.logger.log(currentLineStderrBuffer.toString("utf8"));
501-
}
502-
// Remove the listeners we set up.
503-
process.stdout.removeAllListeners("data");
504-
process.stderr.removeAllListeners("data");
505-
process.removeAllListeners("close");
529+
// Start listening to stdout
530+
process.stdout.addListener("data", stdoutListener);
531+
// Listen to stderr
532+
process.stderr.addListener("data", stderrListener);
533+
// Listen for process exit.
534+
process.addListener("close", closeListener);
535+
536+
onListenStart?.(process);
537+
});
538+
// Join all the data together
539+
const fullBuffer = Buffer.concat(stdoutBuffers);
540+
// Make sure we remove the terminator
541+
const data = fullBuffer.toString(
542+
"utf8",
543+
0,
544+
handleNullTerminator ? fullBuffer.length - 1 : fullBuffer.length,
545+
);
546+
if (!silent) {
547+
void this.logger.log(currentLineStderrBuffer.toString("utf8"));
548+
currentLineStderrBuffer = Buffer.alloc(0);
549+
void this.logger.log("CLI command succeeded.");
506550
}
551+
return data;
552+
} catch (err) {
553+
// Report the error (if there is a stderr then use that otherwise just report the error code or nodejs error)
554+
const cliError = getCliError(
555+
err,
556+
stderrBuffers.length > 0
557+
? Buffer.concat(stderrBuffers).toString("utf8")
558+
: undefined,
559+
description,
560+
args,
561+
);
562+
cliError.stack += getErrorStack(err);
563+
throw cliError;
507564
} finally {
508-
this.commandInProcess = false;
509-
// start running the next command immediately
510-
this.runNext();
565+
if (!silent && currentLineStderrBuffer.length > 0) {
566+
void this.logger.log(currentLineStderrBuffer.toString("utf8"));
567+
}
568+
// Remove the listeners we set up.
569+
if (stdoutListener) {
570+
process.stdout.removeListener("data", stdoutListener);
571+
}
572+
if (stderrListener) {
573+
process.stderr.removeListener("data", stderrListener);
574+
}
575+
if (closeListener) {
576+
process.removeListener("close", closeListener);
577+
}
511578
}
512579
}
513580

0 commit comments

Comments
 (0)