From 4230c22a08a0438af4809d33e4cfd561f714b127 Mon Sep 17 00:00:00 2001 From: Don Jayamanne Date: Tue, 13 Sep 2022 08:56:32 +1000 Subject: [PATCH] Compress notebook output streams before rendering (#160667) * Compress notebook output streams before rendering * OOps * Combine the buffers manually * Address code review * oops * Fixes * We can have multiple stream mimes in an output * oops --- extensions/ipynb/src/serializers.ts | 27 ++++---- extensions/ipynb/src/streamCompressor.ts | 63 +++++++++++++++++++ src/vs/workbench/api/common/extHostTypes.ts | 5 +- .../view/renderers/backLayerWebView.ts | 11 ++-- .../view/renderers/stdOutErrorPreProcessor.ts | 56 +++++++++++++++++ .../contrib/notebook/common/notebookCommon.ts | 8 +++ 6 files changed, 149 insertions(+), 21 deletions(-) create mode 100644 extensions/ipynb/src/streamCompressor.ts create mode 100644 src/vs/workbench/contrib/notebook/browser/view/renderers/stdOutErrorPreProcessor.ts diff --git a/extensions/ipynb/src/serializers.ts b/extensions/ipynb/src/serializers.ts index 455fb0d2745..21e97824f18 100644 --- a/extensions/ipynb/src/serializers.ts +++ b/extensions/ipynb/src/serializers.ts @@ -7,6 +7,7 @@ import type * as nbformat from '@jupyterlab/nbformat'; import { NotebookCell, NotebookCellData, NotebookCellKind, NotebookCellOutput } from 'vscode'; import { CellMetadata, CellOutputMetadata } from './common'; import { textMimeTypes } from './deserializers'; +import { compressOutputItemStreams } from './streamCompressor'; const textDecoder = new TextDecoder(); @@ -270,21 +271,17 @@ type JupyterOutput = function convertStreamOutput(output: NotebookCellOutput): JupyterOutput { const outputs: string[] = []; - output.items - .filter((opit) => opit.mime === CellOutputMimeTypes.stderr || opit.mime === CellOutputMimeTypes.stdout) - .map((opit) => textDecoder.decode(opit.data)) - .forEach(value => { - // Ensure each line is a seprate entry in an array (ending with \n). - const lines = value.split('\n'); - // If the last item in `outputs` is not empty and the first item in `lines` is not empty, then concate them. - // As they are part of the same line. - if (outputs.length && lines.length && lines[0].length > 0) { - outputs[outputs.length - 1] = `${outputs[outputs.length - 1]}${lines.shift()!}`; - } - for (const line of lines) { - outputs.push(line); - } - }); + const compressedStream = output.items.length ? new TextDecoder().decode(compressOutputItemStreams(output.items[0].mime, output.items)) : ''; + // Ensure each line is a separate entry in an array (ending with \n). + const lines = compressedStream.split('\n'); + // If the last item in `outputs` is not empty and the first item in `lines` is not empty, then concate them. + // As they are part of the same line. + if (outputs.length && lines.length && lines[0].length > 0) { + outputs[outputs.length - 1] = `${outputs[outputs.length - 1]}${lines.shift()!}`; + } + for (const line of lines) { + outputs.push(line); + } for (let index = 0; index < (outputs.length - 1); index++) { outputs[index] = `${outputs[index]}\n`; diff --git a/extensions/ipynb/src/streamCompressor.ts b/extensions/ipynb/src/streamCompressor.ts new file mode 100644 index 00000000000..cea3184e16c --- /dev/null +++ b/extensions/ipynb/src/streamCompressor.ts @@ -0,0 +1,63 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import type { NotebookCellOutputItem } from 'vscode'; + + +/** + * Given a stream of individual stdout outputs, this function will return the compressed lines, escaping some of the common terminal escape codes. + * E.g. some terminal escape codes would result in the previous line getting cleared, such if we had 3 lines and + * last line contained such a code, then the result string would be just the first two lines. + */ +export function compressOutputItemStreams(mimeType: string, outputs: NotebookCellOutputItem[]) { + // return outputs.find(op => op.mime === mimeType)!.data.buffer; + + const buffers: Uint8Array[] = []; + let startAppending = false; + // Pick the first set of outputs with the same mime type. + for (const output of outputs) { + if (output.mime === mimeType) { + if ((buffers.length === 0 || startAppending)) { + buffers.push(output.data); + startAppending = true; + } + } else if (startAppending) { + startAppending = false; + } + } + compressStreamBuffer(buffers); + const totalBytes = buffers.reduce((p, c) => p + c.byteLength, 0); + const combinedBuffer = new Uint8Array(totalBytes); + let offset = 0; + for (const buffer of buffers) { + combinedBuffer.set(buffer, offset); + offset = offset + buffer.byteLength; + } + return combinedBuffer; +} +const MOVE_CURSOR_1_LINE_COMMAND = `${String.fromCharCode(27)}[A`; +const MOVE_CURSOR_1_LINE_COMMAND_BYTES = MOVE_CURSOR_1_LINE_COMMAND.split('').map(c => c.charCodeAt(0)); +const LINE_FEED = 10; +function compressStreamBuffer(streams: Uint8Array[]) { + streams.forEach((stream, index) => { + if (index === 0 || stream.length < MOVE_CURSOR_1_LINE_COMMAND.length) { + return; + } + + const previousStream = streams[index - 1]; + + // Remove the previous line if required. + const command = stream.subarray(0, MOVE_CURSOR_1_LINE_COMMAND.length); + if (command[0] === MOVE_CURSOR_1_LINE_COMMAND_BYTES[0] && command[1] === MOVE_CURSOR_1_LINE_COMMAND_BYTES[1] && command[2] === MOVE_CURSOR_1_LINE_COMMAND_BYTES[2]) { + const lastIndexOfLineFeed = previousStream.lastIndexOf(LINE_FEED); + if (lastIndexOfLineFeed === -1) { + return; + } + streams[index - 1] = previousStream.subarray(0, lastIndexOfLineFeed); + streams[index] = stream.subarray(MOVE_CURSOR_1_LINE_COMMAND.length); + } + }); + return streams; +} diff --git a/src/vs/workbench/api/common/extHostTypes.ts b/src/vs/workbench/api/common/extHostTypes.ts index 73c124bacf7..adca6fbc51e 100644 --- a/src/vs/workbench/api/common/extHostTypes.ts +++ b/src/vs/workbench/api/common/extHostTypes.ts @@ -17,7 +17,7 @@ import { IExtensionDescription } from 'vs/platform/extensions/common/extensions' import { FileSystemProviderErrorCode, markAsFileSystemProviderError } from 'vs/platform/files/common/files'; import { RemoteAuthorityResolverErrorCode } from 'vs/platform/remote/common/remoteAuthorityResolver'; import { IRelativePatternDto } from 'vs/workbench/api/common/extHost.protocol'; -import { CellEditType, ICellPartialMetadataEdit, IDocumentMetadataEdit } from 'vs/workbench/contrib/notebook/common/notebookCommon'; +import { CellEditType, ICellPartialMetadataEdit, IDocumentMetadataEdit, isTextStreamMime } from 'vs/workbench/contrib/notebook/common/notebookCommon'; import { checkProposedApiEnabled } from 'vs/workbench/services/extensions/common/extensions'; import type * as vscode from 'vscode'; @@ -3517,7 +3517,8 @@ export class NotebookCellOutput { for (let i = 0; i < items.length; i++) { const item = items[i]; const normalMime = normalizeMimeType(item.mime); - if (!seen.has(normalMime)) { + // We can have multiple text stream mime types in the same output. + if (!seen.has(normalMime) || isTextStreamMime(normalMime)) { seen.add(normalMime); continue; } diff --git a/src/vs/workbench/contrib/notebook/browser/view/renderers/backLayerWebView.ts b/src/vs/workbench/contrib/notebook/browser/view/renderers/backLayerWebView.ts index 69e2a7e41ff..583ba78776d 100644 --- a/src/vs/workbench/contrib/notebook/browser/view/renderers/backLayerWebView.ts +++ b/src/vs/workbench/contrib/notebook/browser/view/renderers/backLayerWebView.ts @@ -37,7 +37,7 @@ import { NOTEBOOK_WEBVIEW_BOUNDARY } from 'vs/workbench/contrib/notebook/browser import { preloadsScriptStr } from 'vs/workbench/contrib/notebook/browser/view/renderers/webviewPreloads'; import { transformWebviewThemeVars } from 'vs/workbench/contrib/notebook/browser/view/renderers/webviewThemeMapping'; import { MarkupCellViewModel } from 'vs/workbench/contrib/notebook/browser/viewModel/markupCellViewModel'; -import { CellUri, INotebookRendererInfo, NotebookSetting, RendererMessagingSpec } from 'vs/workbench/contrib/notebook/common/notebookCommon'; +import { CellUri, INotebookRendererInfo, isTextStreamMime, NotebookSetting, RendererMessagingSpec } from 'vs/workbench/contrib/notebook/common/notebookCommon'; import { INotebookKernel } from 'vs/workbench/contrib/notebook/common/notebookKernelService'; import { IScopedRendererMessaging } from 'vs/workbench/contrib/notebook/common/notebookRendererMessagingService'; import { INotebookService } from 'vs/workbench/contrib/notebook/common/notebookService'; @@ -46,6 +46,7 @@ import { WebviewWindowDragMonitor } from 'vs/workbench/contrib/webview/browser/w import { IEditorGroupsService } from 'vs/workbench/services/editor/common/editorGroupsService'; import { IWorkbenchEnvironmentService } from 'vs/workbench/services/environment/common/environmentService'; import { FromWebviewMessage, IAckOutputHeight, IClickedDataUrlMessage, ICodeBlockHighlightRequest, IContentWidgetTopRequest, IControllerPreload, ICreationContent, ICreationRequestMessage, IFindMatch, IMarkupCellInitialization, RendererMetadata, ToWebviewMessage } from './webviewMessages'; +import { compressOutputItemStreams } from 'vs/workbench/contrib/notebook/browser/view/renderers/stdOutErrorPreProcessor'; export interface ICachedInset { outputId: string; @@ -1278,12 +1279,14 @@ var requirejs = (function() { let updatedContent: ICreationContent | undefined = undefined; if (content.type === RenderOutputType.Extension) { const output = content.source.model; - const first = output.outputs.find(op => op.mime === content.mimeType)!; + const firstBuffer = isTextStreamMime(content.mimeType) ? + compressOutputItemStreams(content.mimeType, output.outputs) : + output.outputs.find(op => op.mime === content.mimeType)!.data.buffer; updatedContent = { type: RenderOutputType.Extension, outputId: outputCache.outputId, - mimeType: first.mime, - valueBytes: first.data.buffer, + mimeType: content.mimeType, + valueBytes: firstBuffer, metadata: output.metadata, }; } diff --git a/src/vs/workbench/contrib/notebook/browser/view/renderers/stdOutErrorPreProcessor.ts b/src/vs/workbench/contrib/notebook/browser/view/renderers/stdOutErrorPreProcessor.ts new file mode 100644 index 00000000000..7f20c6316fc --- /dev/null +++ b/src/vs/workbench/contrib/notebook/browser/view/renderers/stdOutErrorPreProcessor.ts @@ -0,0 +1,56 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { VSBuffer } from 'vs/base/common/buffer'; +import type { IOutputItemDto } from 'vs/workbench/contrib/notebook/common/notebookCommon'; + + +/** + * Given a stream of individual stdout outputs, this function will return the compressed lines, escaping some of the common terminal escape codes. + * E.g. some terminal escape codes would result in the previous line getting cleared, such if we had 3 lines and + * last line contained such a code, then the result string would be just the first two lines. + */ +export function compressOutputItemStreams(mimeType: string, outputs: IOutputItemDto[]) { + const buffers: Uint8Array[] = []; + let startAppending = false; + + // Pick the first set of outputs with the same mime type. + for (const output of outputs) { + if (output.mime === mimeType) { + if ((buffers.length === 0 || startAppending)) { + buffers.push(output.data.buffer); + startAppending = true; + } + } else if (startAppending) { + startAppending = false; + } + } + compressStreamBuffer(buffers); + return VSBuffer.concat(buffers.map(buffer => VSBuffer.wrap(buffer))).buffer; +} +const MOVE_CURSOR_1_LINE_COMMAND = `${String.fromCharCode(27)}[A`; +const MOVE_CURSOR_1_LINE_COMMAND_BYTES = MOVE_CURSOR_1_LINE_COMMAND.split('').map(c => c.charCodeAt(0)); +const LINE_FEED = 10; +function compressStreamBuffer(streams: Uint8Array[]) { + streams.forEach((stream, index) => { + if (index === 0 || stream.length < MOVE_CURSOR_1_LINE_COMMAND.length) { + return; + } + + const previousStream = streams[index - 1]; + + // Remove the previous line if required. + const command = stream.subarray(0, MOVE_CURSOR_1_LINE_COMMAND.length); + if (command[0] === MOVE_CURSOR_1_LINE_COMMAND_BYTES[0] && command[1] === MOVE_CURSOR_1_LINE_COMMAND_BYTES[1] && command[2] === MOVE_CURSOR_1_LINE_COMMAND_BYTES[2]) { + const lastIndexOfLineFeed = previousStream.lastIndexOf(LINE_FEED); + if (lastIndexOfLineFeed === -1) { + return; + } + streams[index - 1] = previousStream.subarray(0, lastIndexOfLineFeed); + streams[index] = stream.subarray(MOVE_CURSOR_1_LINE_COMMAND.length); + } + }); + return streams; +} diff --git a/src/vs/workbench/contrib/notebook/common/notebookCommon.ts b/src/vs/workbench/contrib/notebook/common/notebookCommon.ts index 361c0baa2b1..af7f455c8af 100644 --- a/src/vs/workbench/contrib/notebook/common/notebookCommon.ts +++ b/src/vs/workbench/contrib/notebook/common/notebookCommon.ts @@ -947,3 +947,11 @@ export interface NotebookExtensionDescription { readonly id: ExtensionIdentifier; readonly location: UriComponents | undefined; } + +/** + * Whether the provided mime type is a text streamn like `stdout`, `stderr`. + */ +export function isTextStreamMime(mimeType: string) { + return ['application/vnd.code.notebook.stdout', 'application/x.notebook.stdout', 'application/x.notebook.stream', 'application/vnd.code.notebook.stderr', 'application/x.notebook.stderr'].includes(mimeType); +} +