Merge pull request microsoft#187766 from microsoft/aamunger/notebookOutput

perf optimizations for large stream notebook outputs

Author: amunger

extensions/notebook-renderers/src/index.ts

@@ -4,8 +4,8 @@
  *--------------------------------------------------------------------------------------------*/
 
 import type { ActivationFunction, OutputItem, RendererContext } from 'vscode-notebook-renderer';
-import { createOutputContent, scrollableClass } from './textHelper';
-import { HtmlRenderingHook, IDisposable, IRichRenderContext, JavaScriptRenderingHook, RenderOptions } from './rendererTypes';
+import { createOutputContent, appendOutput, scrollableClass } from './textHelper';
+import { HtmlRenderingHook, IDisposable, IRichRenderContext, JavaScriptRenderingHook, OutputWithAppend, RenderOptions } from './rendererTypes';
 import { ttPolicy } from './htmlHelper';
 
 function clearContainer(container: HTMLElement) {
@@ -152,7 +152,7 @@ function renderError(
 	outputInfo: OutputItem,
 	outputElement: HTMLElement,
 	ctx: IRichRenderContext,
-	trustHTML: boolean
+	trustHtml: boolean
 ): IDisposable {
 	const disposableStore = createDisposableStore();
 
@@ -172,7 +172,7 @@ function renderError(
 		outputElement.classList.add('traceback');
 
 		const outputScrolling = scrollingEnabled(outputInfo, ctx.settings);
-		const content = createOutputContent(outputInfo.id, [err.stack ?? ''], ctx.settings.lineLimit, outputScrolling, trustHTML);
+		const content = createOutputContent(outputInfo.id, err.stack ?? '', { linesLimit: ctx.settings.lineLimit, scrollable: outputScrolling, trustHtml });
 		const contentParent = document.createElement('div');
 		contentParent.classList.toggle('word-wrap', ctx.settings.outputWordWrap);
 		disposableStore.push(ctx.onDidChangeSettings(e => {
@@ -270,29 +270,23 @@ function scrollingEnabled(output: OutputItem, options: RenderOptions) {
 //      div.output.output-stream		<-- outputElement parameter
 //        div.scrollable? tabindex="0" 	<-- contentParent
 //          div output-item-id="{guid}"	<-- content from outputItem parameter
-function renderStream(outputInfo: OutputItem, outputElement: HTMLElement, error: boolean, ctx: IRichRenderContext): IDisposable {
+function renderStream(outputInfo: OutputWithAppend, outputElement: HTMLElement, error: boolean, ctx: IRichRenderContext): IDisposable {
 	const disposableStore = createDisposableStore();
 	const outputScrolling = scrollingEnabled(outputInfo, ctx.settings);
+	const outputOptions = { linesLimit: ctx.settings.lineLimit, scrollable: outputScrolling, trustHtml: false, error };
 
 	outputElement.classList.add('output-stream');
 
-	const text = outputInfo.text();
-	const newContent = createOutputContent(outputInfo.id, [text], ctx.settings.lineLimit, outputScrolling, false);
-	newContent.setAttribute('output-item-id', outputInfo.id);
-	if (error) {
-		newContent.classList.add('error');
-	}
-
 	const scrollTop = outputScrolling ? findScrolledHeight(outputElement) : undefined;
 
 	const previousOutputParent = getPreviousMatchingContentGroup(outputElement);
 	// If the previous output item for the same cell was also a stream, append this output to the previous
 	if (previousOutputParent) {
 		const existingContent = previousOutputParent.querySelector(`[output-item-id="${outputInfo.id}"]`) as HTMLElement | null;
 		if (existingContent) {
-			existingContent.replaceWith(newContent);
-
+			appendOutput(outputInfo, existingContent, outputOptions);
 		} else {
+			const newContent = createOutputContent(outputInfo.id, outputInfo.text(), outputOptions);
 			previousOutputParent.appendChild(newContent);
 		}
 		previousOutputParent.classList.toggle('scrollbar-visible', previousOutputParent.scrollHeight > previousOutputParent.clientHeight);
@@ -301,12 +295,9 @@ function renderStream(outputInfo: OutputItem, outputElement: HTMLElement, error:
 		const existingContent = outputElement.querySelector(`[output-item-id="${outputInfo.id}"]`) as HTMLElement | null;
 		let contentParent = existingContent?.parentElement;
 		if (existingContent && contentParent) {
-			existingContent.replaceWith(newContent);
-			while (newContent.nextSibling) {
-				// clear out any stale content if we had previously combined streaming outputs into this one
-				newContent.nextSibling.remove();
-			}
+			appendOutput(outputInfo, existingContent, outputOptions);
 		} else {
+			const newContent = createOutputContent(outputInfo.id, outputInfo.text(), outputOptions);
 			contentParent = document.createElement('div');
 			contentParent.appendChild(newContent);
 			while (outputElement.firstChild) {
@@ -333,7 +324,7 @@ function renderText(outputInfo: OutputItem, outputElement: HTMLElement, ctx: IRi
 
 	const text = outputInfo.text();
 	const outputScrolling = scrollingEnabled(outputInfo, ctx.settings);
-	const content = createOutputContent(outputInfo.id, [text], ctx.settings.lineLimit, outputScrolling, false);
+	const content = createOutputContent(outputInfo.id, text, { linesLimit: ctx.settings.lineLimit, scrollable: outputScrolling, trustHtml: false });
 	content.classList.add('output-plaintext');
 	if (ctx.settings.outputWordWrap) {
 		content.classList.add('word-wrap');

extensions/notebook-renderers/src/rendererTypes.ts

@@ -35,3 +35,14 @@ export interface RenderOptions {
 }
 
 export type IRichRenderContext = RendererContext<void> & { readonly settings: RenderOptions; readonly onDidChangeSettings: Event<RenderOptions> };
+
+export type OutputElementOptions = {
+	linesLimit: number;
+	scrollable?: boolean;
+	error?: boolean;
+	trustHtml?: boolean;
+};
+
+export interface OutputWithAppend extends OutputItem {
+	appendedText?(): string | undefined;
+}

extensions/notebook-renderers/src/test/notebookRenderer.test.ts

@@ -5,8 +5,8 @@
 
 import * as assert from 'assert';
 import { activate } from '..';
-import { OutputItem, RendererApi } from 'vscode-notebook-renderer';
-import { IDisposable, IRichRenderContext, RenderOptions } from '../rendererTypes';
+import { RendererApi } from 'vscode-notebook-renderer';
+import { IDisposable, IRichRenderContext, OutputWithAppend, RenderOptions } from '../rendererTypes';
 import { JSDOM } from "jsdom";
 
 const dom = new JSDOM();
@@ -116,10 +116,13 @@ suite('Notebook builtin output renderer', () => {
 		}
 	}
 
-	function createOutputItem(text: string, mime: string, id: string = '123'): OutputItem {
+	function createOutputItem(text: string, mime: string, id: string = '123', appendedText?: string): OutputWithAppend {
 		return {
 			id: id,
 			mime: mime,
+			appendedText() {
+				return appendedText;
+			},
 			text() {
 				return text;
 			},
@@ -177,9 +180,9 @@ suite('Notebook builtin output renderer', () => {
 			assert.ok(renderer, 'Renderer not created');
 
 			const outputElement = new OutputHtml().getFirstOuputElement();
-			const outputItem = createOutputItem('content', 'text/plain');
+			const outputItem = createOutputItem('content', mimeType);
 			await renderer!.renderOutputItem(outputItem, outputElement);
-			const outputItem2 = createOutputItem('replaced content', 'text/plain');
+			const outputItem2 = createOutputItem('replaced content', mimeType);
 			await renderer!.renderOutputItem(outputItem2, outputElement);
 
 			const inserted = outputElement.firstChild as HTMLElement;
@@ -189,6 +192,87 @@ suite('Notebook builtin output renderer', () => {
 
 	});
 
+	test('Append streaming output', async () => {
+		const context = createContext({ outputWordWrap: false, outputScrolling: true });
+		const renderer = await activate(context);
+		assert.ok(renderer, 'Renderer not created');
+
+		const outputElement = new OutputHtml().getFirstOuputElement();
+		const outputItem = createOutputItem('content', stdoutMimeType, '123', 'ignoredAppend');
+		await renderer!.renderOutputItem(outputItem, outputElement);
+		const outputItem2 = createOutputItem('content\nappended', stdoutMimeType, '123', '\nappended');
+		await renderer!.renderOutputItem(outputItem2, outputElement);
+
+		const inserted = outputElement.firstChild as HTMLElement;
+		assert.ok(inserted.innerHTML.indexOf('>content</') !== -1, `Previous content should still exist: ${outputElement.innerHTML}`);
+		assert.ok(inserted.innerHTML.indexOf('ignoredAppend') === -1, `Append value should not be used on first render: ${outputElement.innerHTML}`);
+		assert.ok(inserted.innerHTML.indexOf('>appended</') !== -1, `Content was not appended to output element: ${outputElement.innerHTML}`);
+		assert.ok(inserted.innerHTML.indexOf('>content</') === inserted.innerHTML.lastIndexOf('>content</'), `Original content should not be duplicated: ${outputElement.innerHTML}`);
+	});
+
+	test(`Appending multiple streaming outputs`, async () => {
+		const context = createContext({ outputScrolling: true });
+		const renderer = await activate(context);
+		assert.ok(renderer, 'Renderer not created');
+
+		const outputHtml = new OutputHtml();
+		const firstOutputElement = outputHtml.getFirstOuputElement();
+		const outputItem1 = createOutputItem('first stream content', stdoutMimeType, '1');
+		const outputItem2 = createOutputItem(JSON.stringify(error), errorMimeType, '2');
+		const outputItem3 = createOutputItem('second stream content', stdoutMimeType, '3');
+		await renderer!.renderOutputItem(outputItem1, firstOutputElement);
+		const secondOutputElement = outputHtml.appendOutputElement();
+		await renderer!.renderOutputItem(outputItem2, secondOutputElement);
+		const thirdOutputElement = outputHtml.appendOutputElement();
+		await renderer!.renderOutputItem(outputItem3, thirdOutputElement);
+
+		const appendedItem1 = createOutputItem('', stdoutMimeType, '1', ' appended1');
+		await renderer!.renderOutputItem(appendedItem1, firstOutputElement);
+		const appendedItem3 = createOutputItem('', stdoutMimeType, '3', ' appended3');
+		await renderer!.renderOutputItem(appendedItem3, thirdOutputElement);
+
+		assert.ok(firstOutputElement.innerHTML.indexOf('>first stream content') > -1, `Content was not added to output element: ${outputHtml.cellElement.innerHTML}`);
+		assert.ok(firstOutputElement.innerHTML.indexOf('appended1') > -1, `Content was not appended to output element: ${outputHtml.cellElement.innerHTML}`);
+		assert.ok(secondOutputElement.innerHTML.indexOf('>NameError</') > -1, `Content was not added to output element: ${outputHtml.cellElement.innerHTML}`);
+		assert.ok(thirdOutputElement.innerHTML.indexOf('>second stream content') > -1, `Content was not added to output element: ${outputHtml.cellElement.innerHTML}`);
+		assert.ok(thirdOutputElement.innerHTML.indexOf('appended3') > -1, `Content was not appended to output element: ${outputHtml.cellElement.innerHTML}`);
+	});
+
+	test('Append large streaming outputs', async () => {
+		const context = createContext({ outputWordWrap: false, outputScrolling: true });
+		const renderer = await activate(context);
+		assert.ok(renderer, 'Renderer not created');
+
+		const outputElement = new OutputHtml().getFirstOuputElement();
+		const lotsOfLines = new Array(4998).fill('line').join('\n');
+		const firstOuput = lotsOfLines + 'expected1';
+		const outputItem = createOutputItem(firstOuput, stdoutMimeType, '123');
+		await renderer!.renderOutputItem(outputItem, outputElement);
+		const appended = '\n' + lotsOfLines + 'expectedAppend';
+		const outputItem2 = createOutputItem(firstOuput + appended, stdoutMimeType, '123', appended);
+		await renderer!.renderOutputItem(outputItem2, outputElement);
+
+		const inserted = outputElement.firstChild as HTMLElement;
+		assert.ok(inserted.innerHTML.indexOf('expected1') !== -1, `Last bit of previous content should still exist`);
+		assert.ok(inserted.innerHTML.indexOf('expectedAppend') !== -1, `Content was not appended to output element`);
+	});
+
+	test('Streaming outputs larger than the line limit are truncated', async () => {
+		const context = createContext({ outputWordWrap: false, outputScrolling: true });
+		const renderer = await activate(context);
+		assert.ok(renderer, 'Renderer not created');
+
+		const outputElement = new OutputHtml().getFirstOuputElement();
+		const lotsOfLines = new Array(11000).fill('line').join('\n');
+		const firstOuput = 'shouldBeTruncated' + lotsOfLines + 'expected1';
+		const outputItem = createOutputItem(firstOuput, stdoutMimeType, '123');
+		await renderer!.renderOutputItem(outputItem, outputElement);
+
+		const inserted = outputElement.firstChild as HTMLElement;
+		assert.ok(inserted.innerHTML.indexOf('expected1') !== -1, `Last bit of content should exist`);
+		assert.ok(inserted.innerHTML.indexOf('shouldBeTruncated') === -1, `Beginning content should be truncated`);
+	});
+
 	test(`Render with wordwrap and scrolling for error output`, async () => {
 		const context = createContext({ outputWordWrap: true, outputScrolling: true });
 		const renderer = await activate(context);
@@ -268,6 +352,29 @@ suite('Notebook builtin output renderer', () => {
 		assert.ok(inserted.innerHTML.indexOf('>second stream content</') === -1, `Content was not replaced in output element: ${outputHtml.cellElement.innerHTML}`);
 	});
 
+	test(`Consolidated streaming outputs should append matching outputs correctly`, async () => {
+		const context = createContext({ outputScrolling: true });
+		const renderer = await activate(context);
+		assert.ok(renderer, 'Renderer not created');
+
+		const outputHtml = new OutputHtml();
+		const outputElement = outputHtml.getFirstOuputElement();
+		const outputItem1 = createOutputItem('first stream content', stdoutMimeType, '1');
+		const outputItem2 = createOutputItem('second stream content', stdoutMimeType, '2');
+		await renderer!.renderOutputItem(outputItem1, outputElement);
+		const secondOutput = outputHtml.appendOutputElement();
+		await renderer!.renderOutputItem(outputItem2, secondOutput);
+		const appendingOutput = createOutputItem('', stdoutMimeType, '2', ' appended');
+		await renderer!.renderOutputItem(appendingOutput, secondOutput);
+
+
+		const inserted = outputElement.firstChild as HTMLElement;
+		assert.ok(inserted, `nothing appended to output element: ${outputElement.innerHTML}`);
+		assert.ok(inserted.innerHTML.indexOf('>first stream content</') > -1, `Content was not added to output element: ${outputHtml.cellElement.innerHTML}`);
+		assert.ok(inserted.innerHTML.indexOf('>second stream content') > -1, `Second content was not added to ouptut element: ${outputHtml.cellElement.innerHTML}`);
+		assert.ok(inserted.innerHTML.indexOf('appended') > -1, `Content was not appended to ouptut element: ${outputHtml.cellElement.innerHTML}`);
+	});
+
 	test(`Streaming outputs interleaved with other mime types will produce separate outputs`, async () => {
 		const context = createContext({ outputScrolling: false });
 		const renderer = await activate(context);

extensions/notebook-renderers/src/textHelper.ts

@@ -4,9 +4,12 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { handleANSIOutput } from './ansi';
-
+import { OutputElementOptions, OutputWithAppend } from './rendererTypes';
 export const scrollableClass = 'scrollable';
 
+const softScrollableLineLimit = 5000;
+const hardScrollableLineLimit = 8000;
+
 /**
  * Output is Truncated. View as a [scrollable element] or open in a [text editor]. Adjust cell output [settings...]
  */
@@ -91,22 +94,70 @@ function truncatedArrayOfString(id: string, buffer: string[], linesLimit: number
 
 function scrollableArrayOfString(id: string, buffer: string[], trustHtml: boolean) {
 	const element = document.createElement('div');
-	if (buffer.length > 5000) {
+	if (buffer.length > softScrollableLineLimit) {
 		element.appendChild(generateNestedViewAllElement(id));
 	}
 
-	element.appendChild(handleANSIOutput(buffer.slice(-5000).join('\n'), trustHtml));
+	element.appendChild(handleANSIOutput(buffer.slice(-1 * softScrollableLineLimit).join('\n'), trustHtml));
 
 	return element;
 }
 
-export function createOutputContent(id: string, outputs: string[], linesLimit: number, scrollable: boolean, trustHtml: boolean): HTMLElement {
+const outputLengths: Record<string, number> = {};
 
-	const buffer = outputs.join('\n').split(/\r\n|\r|\n/g);
+function appendScrollableOutput(element: HTMLElement, id: string, appended: string, trustHtml: boolean) {
+	if (!outputLengths[id]) {
+		outputLengths[id] = 0;
+	}
 
+	const buffer = appended.split(/\r\n|\r|\n/g);
+	const appendedLength = buffer.length + outputLengths[id];
+	// Only append outputs up to the hard limit of lines, then replace it with the last softLimit number of lines
+	if (appendedLength > hardScrollableLineLimit) {
+		return false;
+	}
+	else {
+		element.appendChild(handleANSIOutput(buffer.join('\n'), trustHtml));
+		outputLengths[id] = appendedLength;
+	}
+	return true;
+}
+
+export function createOutputContent(id: string, outputText: string, options: OutputElementOptions): HTMLElement {
+	const { linesLimit, error, scrollable, trustHtml } = options;
+	const buffer = outputText.split(/\r\n|\r|\n/g);
+	outputLengths[id] = outputLengths[id] = Math.min(buffer.length, softScrollableLineLimit);
+
+	let outputElement: HTMLElement;
 	if (scrollable) {
-		return scrollableArrayOfString(id, buffer, trustHtml);
+		outputElement = scrollableArrayOfString(id, buffer, !!trustHtml);
 	} else {
-		return truncatedArrayOfString(id, buffer, linesLimit, trustHtml);
+		outputElement = truncatedArrayOfString(id, buffer, linesLimit, !!trustHtml);
+	}
+
+	outputElement.setAttribute('output-item-id', id);
+	if (error) {
+		outputElement.classList.add('error');
 	}
+
+	return outputElement;
 }
+
+export function appendOutput(outputInfo: OutputWithAppend, existingContent: HTMLElement, options: OutputElementOptions) {
+	const appendedText = outputInfo.appendedText?.();
+	// appending output only supported for scrollable ouputs currently
+	if (appendedText && options.scrollable) {
+		if (appendScrollableOutput(existingContent, outputInfo.id, appendedText, false)) {
+			return;
+		}
+	}
+
+	const newContent = createOutputContent(outputInfo.id, outputInfo.text(), options);
+	existingContent.replaceWith(newContent);
+	while (newContent.nextSibling) {
+		// clear out any stale content if we had previously combined streaming outputs into this one
+		newContent.nextSibling.remove();
+	}
+
+}
+

src/vs/workbench/api/common/extHostNotebookDocument.ts

@@ -132,7 +132,7 @@ export class ExtHostCell {
 					const compressed = notebookCommon.compressOutputItemStreams(mimeOutputs.get(mime)!);
 					output.items.push({
 						mime,
-						data: compressed.buffer
+						data: compressed.data.buffer
 					});
 				});
 			}