Skip to content

Commit

Permalink
Merge pull request #234381 from microsoft/rebornix/weak-tick
Browse files Browse the repository at this point in the history 
Notebook snapshot (cells and buffer)
  • Loading branch information
@rebornix
rebornix authored Nov 21, 2024
2 parents ec1e84b + 20ecaf9 commit 143f300
Show file tree
Hide file tree
Showing 6 changed files with 136 additions and 62 deletions.
53 changes: 51 additions & 2 deletions src/vs/workbench/contrib/notebook/browser/services/notebookServiceImpl.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ import { INotebookEditorOptions } from '../notebookBrowser.js';
import { NotebookDiffEditorInput } from '../../common/notebookDiffEditorInput.js';
import { NotebookCellTextModel } from '../../common/model/notebookCellTextModel.js';
import { NotebookTextModel } from '../../common/model/notebookTextModel.js';
import { ACCESSIBLE_NOTEBOOK_DISPLAY_ORDER, CellUri, NotebookSetting, INotebookContributionData, INotebookExclusiveDocumentFilter, INotebookRendererInfo, INotebookTextModel, IOrderedMimeType, IOutputDto, MimeTypeDisplayOrder, NotebookEditorPriority, NotebookRendererMatch, NOTEBOOK_DISPLAY_ORDER, RENDERER_EQUIVALENT_EXTENSIONS, RENDERER_NOT_AVAILABLE, NotebookExtensionDescription, INotebookStaticPreloadInfo } from '../../common/notebookCommon.js';
import { ACCESSIBLE_NOTEBOOK_DISPLAY_ORDER, CellUri, NotebookSetting, INotebookContributionData, INotebookExclusiveDocumentFilter, INotebookRendererInfo, INotebookTextModel, IOrderedMimeType, IOutputDto, MimeTypeDisplayOrder, NotebookEditorPriority, NotebookRendererMatch, NOTEBOOK_DISPLAY_ORDER, RENDERER_EQUIVALENT_EXTENSIONS, RENDERER_NOT_AVAILABLE, NotebookExtensionDescription, INotebookStaticPreloadInfo, NotebookData } from '../../common/notebookCommon.js';
import { NotebookEditorInput } from '../../common/notebookEditorInput.js';
import { INotebookEditorModelResolverService } from '../../common/notebookEditorModelResolverService.js';
import { NotebookOutputRendererInfo, NotebookStaticPreloadInfo as NotebookStaticPreloadInfo } from '../../common/notebookOutputRenderer.js';
Expand All @@ -42,9 +42,12 @@ import { IUriIdentityService } from '../../../../../platform/uriIdentity/common/
import { INotebookDocument, INotebookDocumentService } from '../../../../services/notebook/common/notebookDocumentService.js';
import { MergeEditorInput } from '../../../mergeEditor/browser/mergeEditorInput.js';
import type { EditorInputWithOptions, IResourceDiffEditorInput, IResourceMergeEditorInput } from '../../../../common/editor.js';
import { streamToBuffer, VSBuffer, VSBufferReadableStream } from '../../../../../base/common/buffer.js';
import { bufferToStream, streamToBuffer, VSBuffer, VSBufferReadableStream } from '../../../../../base/common/buffer.js';
import type { IEditorGroup } from '../../../../services/editor/common/editorGroupsService.js';
import { NotebookMultiDiffEditorInput } from '../diff/notebookMultiDiffEditorInput.js';
import { SnapshotContext } from '../../../../services/workingCopy/common/fileWorkingCopy.js';
import { CancellationToken } from '../../../../../base/common/cancellation.js';
import { CancellationError } from '../../../../../base/common/errors.js';

export class NotebookProviderInfoStore extends Disposable {

Expand Down Expand Up @@ -776,6 +779,52 @@ export class NotebookService extends Disposable implements INotebookService {
return notebookModel;
}

async createNotebookTextDocumentSnapshot(uri: URI, context: SnapshotContext, token: CancellationToken): Promise<VSBufferReadableStream> {
const model = this.getNotebookTextModel(uri);

if (!model) {
throw new Error(`notebook for ${uri} doesn't exist`);
}

const info = await this.withNotebookDataProvider(model.viewType);

if (!(info instanceof SimpleNotebookProviderInfo)) {
throw new Error('CANNOT open file notebook with this provider');
}

const serializer = info.serializer;
const outputSizeLimit = this._configurationService.getValue<number>(NotebookSetting.outputBackupSizeLimit) * 1024;
const data: NotebookData = model.createSnapshot({ context: context, outputSizeLimit: outputSizeLimit, transientOptions: serializer.options });
const bytes = await serializer.notebookToData(data);

if (token.isCancellationRequested) {
throw new CancellationError();
}
return bufferToStream(bytes);
}

async restoreNotebookTextModelFromSnapshot(uri: URI, viewType: string, snapshot: VSBufferReadableStream): Promise<NotebookTextModel> {
const model = this.getNotebookTextModel(uri);

if (!model) {
throw new Error(`notebook for ${uri} doesn't exist`);
}

const info = await this.withNotebookDataProvider(model.viewType);

if (!(info instanceof SimpleNotebookProviderInfo)) {
throw new Error('CANNOT open file notebook with this provider');
}

const serializer = info.serializer;

const bytes = await streamToBuffer(snapshot);
const data = await info.serializer.dataToNotebook(bytes);
model.restoreSnapshot(data, serializer.options);

return model;
}

getNotebookTextModel(uri: URI): NotebookTextModel | undefined {
return this._models.get(uri)?.model;
}
Expand Down
70 changes: 57 additions & 13 deletions src/vs/workbench/contrib/notebook/common/model/notebookTextModel.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,29 +3,31 @@
* Licensed under the MIT License. See License.txt in the project root for license information.
*--------------------------------------------------------------------------------------------*/

import { Emitter, Event, PauseableEmitter } from '../../../../../base/common/event.js';
import { Disposable, dispose, IDisposable } from '../../../../../base/common/lifecycle.js';
import { URI } from '../../../../../base/common/uri.js';
import { NotebookCellTextModel } from './notebookCellTextModel.js';
import { INotebookTextModel, NotebookCellOutputsSplice, NotebookDocumentMetadata, NotebookCellMetadata, ICellEditOperation, CellEditType, CellUri, diff, NotebookCellsChangeType, ICellDto2, TransientOptions, NotebookTextModelChangedEvent, IOutputDto, ICellOutput, IOutputItemDto, ISelectionState, NullablePartialNotebookCellMetadata, NotebookCellInternalMetadata, NullablePartialNotebookCellInternalMetadata, NotebookTextModelWillAddRemoveEvent, NotebookCellTextModelSplice, ICell, NotebookCellCollapseState, NotebookCellDefaultCollapseConfig, CellKind, NotebookCellExecutionState } from '../notebookCommon.js';
import { IUndoRedoService, UndoRedoElementType, IUndoRedoElement, IResourceUndoRedoElement, UndoRedoGroup, IWorkspaceUndoRedoElement } from '../../../../../platform/undoRedo/common/undoRedo.js';
import { MoveCellEdit, SpliceCellsEdit, CellMetadataEdit } from './cellEdit.js';
import { ISequence, LcsDiff } from '../../../../../base/common/diff/diff.js';
import { Emitter, Event, PauseableEmitter } from '../../../../../base/common/event.js';
import { hash } from '../../../../../base/common/hash.js';
import { NotebookCellOutputTextModel } from './notebookCellOutputTextModel.js';
import { IModelService } from '../../../../../editor/common/services/model.js';
import { Disposable, dispose, IDisposable } from '../../../../../base/common/lifecycle.js';
import { Schemas } from '../../../../../base/common/network.js';
import { filter } from '../../../../../base/common/objects.js';
import { isEqual } from '../../../../../base/common/resources.js';
import { ILanguageService } from '../../../../../editor/common/languages/language.js';
import { FindMatch, ITextModel } from '../../../../../editor/common/model.js';
import { TextModel } from '../../../../../editor/common/model/textModel.js';
import { isDefined } from '../../../../../base/common/types.js';
import { ILanguageDetectionService } from '../../../../services/languageDetection/common/languageDetectionWorkerService.js';
import { URI } from '../../../../../base/common/uri.js';
import { Position } from '../../../../../editor/common/core/position.js';
import { Range } from '../../../../../editor/common/core/range.js';
import { ILanguageService } from '../../../../../editor/common/languages/language.js';
import { FindMatch, ITextModel } from '../../../../../editor/common/model.js';
import { TextModel } from '../../../../../editor/common/model/textModel.js';
import { SearchParams } from '../../../../../editor/common/model/textModelSearch.js';
import { IModelService } from '../../../../../editor/common/services/model.js';
import { IModelContentChangedEvent } from '../../../../../editor/common/textModelEvents.js';
import { IResourceUndoRedoElement, IUndoRedoElement, IUndoRedoService, IWorkspaceUndoRedoElement, UndoRedoElementType, UndoRedoGroup } from '../../../../../platform/undoRedo/common/undoRedo.js';
import { ILanguageDetectionService } from '../../../../services/languageDetection/common/languageDetectionWorkerService.js';
import { SnapshotContext } from '../../../../services/workingCopy/common/fileWorkingCopy.js';
import { CellEditType, CellKind, CellUri, diff, ICell, ICellDto2, ICellEditOperation, ICellOutput, INotebookSnapshotOptions, INotebookTextModel, IOutputDto, IOutputItemDto, ISelectionState, NotebookCellCollapseState, NotebookCellDefaultCollapseConfig, NotebookCellExecutionState, NotebookCellInternalMetadata, NotebookCellMetadata, NotebookCellOutputsSplice, NotebookCellsChangeType, NotebookCellTextModelSplice, NotebookData, NotebookDocumentMetadata, NotebookTextModelChangedEvent, NotebookTextModelWillAddRemoveEvent, NullablePartialNotebookCellInternalMetadata, NullablePartialNotebookCellMetadata, TransientOptions } from '../notebookCommon.js';
import { INotebookExecutionStateService } from '../notebookExecutionStateService.js';
import { CellMetadataEdit, MoveCellEdit, SpliceCellsEdit } from './cellEdit.js';
import { NotebookCellOutputTextModel } from './notebookCellOutputTextModel.js';
import { NotebookCellTextModel } from './notebookCellTextModel.js';

class StackOperation implements IWorkspaceUndoRedoElement {
type: UndoRedoElementType.Workspace;
Expand Down Expand Up @@ -441,6 +443,48 @@ export class NotebookTextModel extends Disposable implements INotebookTextModel
);
}

createSnapshot(options: INotebookSnapshotOptions): NotebookData {
const transientOptions = options.transientOptions ?? this.transientOptions;
const data: NotebookData = {
metadata: filter(this.metadata, key => !transientOptions.transientDocumentMetadata[key]),
cells: [],
};

let outputSize = 0;
for (const cell of this.cells) {
const cellData: ICellDto2 = {
cellKind: cell.cellKind,
language: cell.language,
mime: cell.mime,
source: cell.getValue(),
outputs: [],
internalMetadata: cell.internalMetadata
};

if (options.context === SnapshotContext.Backup && options.outputSizeLimit > 0) {
cell.outputs.forEach(output => {
output.outputs.forEach(item => {
outputSize += item.data.byteLength;
});
});
if (outputSize > options.outputSizeLimit) {
throw new Error('Notebook too large to backup');
}
}

cellData.outputs = !transientOptions.transientOutputs ? cell.outputs : [];
cellData.metadata = filter(cell.metadata, key => !transientOptions.transientCellMetadata[key]);

data.cells.push(cellData);
}

return data;
}

restoreSnapshot(snapshot: NotebookData, transientOptions?: TransientOptions): void {
this.reset(snapshot.cells, snapshot.metadata, transientOptions ?? this.transientOptions);
}

static computeEdits(model: NotebookTextModel, cells: ICellDto2[], executingHandles: number[] = []): ICellEditOperation[] {
const edits: ICellEditOperation[] = [];
const isExecuting = (cell: NotebookCellTextModel) => executingHandles.includes(cell.handle);
Expand Down
9 changes: 9 additions & 0 deletions src/vs/workbench/contrib/notebook/common/notebookCommon.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ import { ICellRange } from './notebookRange.js';
import { RegisteredEditorPriority } from '../../../services/editor/common/editorResolverService.js';
import { generateMetadataUri, generate as generateUri, parseMetadataUri, parse as parseUri } from '../../../services/notebook/common/notebookDocumentService.js';
import { IWorkingCopyBackupMeta, IWorkingCopySaveEvent } from '../../../services/workingCopy/common/workingCopy.js';
import { SnapshotContext } from '../../../services/workingCopy/common/fileWorkingCopy.js';

export const NOTEBOOK_EDITOR_ID = 'workbench.editor.notebook';
export const NOTEBOOK_DIFF_EDITOR_ID = 'workbench.editor.notebookTextDiffEditor';
Expand Down Expand Up @@ -276,6 +277,12 @@ export interface ICell {
onDidChangeInternalMetadata: Event<CellInternalMetadataChangedEvent>;
}

export interface INotebookSnapshotOptions {
context: SnapshotContext;
outputSizeLimit: number;
transientOptions?: TransientOptions;
}

export interface INotebookTextModel extends INotebookTextModelLike {
readonly notebookType: string;
readonly viewType: string;
Expand All @@ -286,6 +293,8 @@ export interface INotebookTextModel extends INotebookTextModelLike {
readonly length: number;
readonly cells: readonly ICell[];
reset(cells: ICellDto2[], metadata: NotebookDocumentMetadata, transientOptions: TransientOptions): void;
createSnapshot(options: INotebookSnapshotOptions): NotebookData;
restoreSnapshot(snapshot: NotebookData, transientOptions?: TransientOptions): void;
applyEdits(rawEdits: ICellEditOperation[], synchronous: boolean, beginSelectionState: ISelectionState | undefined, endSelectionsComputer: () => ISelectionState | undefined, undoRedoGroup: UndoRedoGroup | undefined, computeUndoRedo?: boolean): boolean;
onDidChangeContent: Event<NotebookTextModelChangedEvent>;
onWillDispose: Event<void>;
Expand Down
47 changes: 3 additions & 44 deletions src/vs/workbench/contrib/notebook/common/notebookEditorModel.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,14 +3,13 @@
* Licensed under the MIT License. See License.txt in the project root for license information.
*--------------------------------------------------------------------------------------------*/

import { VSBufferReadableStream, bufferToStream, streamToBuffer } from '../../../../base/common/buffer.js';
import { VSBufferReadableStream, streamToBuffer } from '../../../../base/common/buffer.js';
import { CancellationToken } from '../../../../base/common/cancellation.js';
import { CancellationError } from '../../../../base/common/errors.js';
import { Emitter, Event } from '../../../../base/common/event.js';
import { IMarkdownString } from '../../../../base/common/htmlContent.js';
import { Disposable, DisposableStore } from '../../../../base/common/lifecycle.js';
import { Schemas } from '../../../../base/common/network.js';
import { filter } from '../../../../base/common/objects.js';
import { assertType } from '../../../../base/common/types.js';
import { URI } from '../../../../base/common/uri.js';
import { IConfigurationService } from '../../../../platform/configuration/common/configuration.js';
Expand All @@ -19,7 +18,7 @@ import { ITelemetryService } from '../../../../platform/telemetry/common/telemet
import { IRevertOptions, ISaveOptions, IUntypedEditorInput } from '../../../common/editor.js';
import { EditorModel } from '../../../common/editor/editorModel.js';
import { NotebookTextModel } from './model/notebookTextModel.js';
import { ICellDto2, INotebookEditorModel, INotebookLoadOptions, IResolvedNotebookEditorModel, NotebookCellsChangeType, NotebookData, NotebookSetting } from './notebookCommon.js';
import { INotebookEditorModel, INotebookLoadOptions, IResolvedNotebookEditorModel, NotebookCellsChangeType, NotebookSetting } from './notebookCommon.js';
import { INotebookLoggingService } from './notebookLoggingService.js';
import { INotebookSerializer, INotebookService, SimpleNotebookProviderInfo } from './notebookService.js';
import { IFilesConfigurationService } from '../../../services/filesConfiguration/common/filesConfigurationService.js';
Expand Down Expand Up @@ -294,47 +293,7 @@ export class NotebookFileWorkingCopyModel extends Disposable implements IStoredF
}

async snapshot(context: SnapshotContext, token: CancellationToken): Promise<VSBufferReadableStream> {
const serializer = await this.getNotebookSerializer();

const data: NotebookData = {
metadata: filter(this._notebookModel.metadata, key => !serializer.options.transientDocumentMetadata[key]),
cells: [],
};

let outputSize = 0;
for (const cell of this._notebookModel.cells) {
const cellData: ICellDto2 = {
cellKind: cell.cellKind,
language: cell.language,
mime: cell.mime,
source: cell.getValue(),
outputs: [],
internalMetadata: cell.internalMetadata
};

const outputSizeLimit = this._configurationService.getValue<number>(NotebookSetting.outputBackupSizeLimit) * 1024;
if (context === SnapshotContext.Backup && outputSizeLimit > 0) {
cell.outputs.forEach(output => {
output.outputs.forEach(item => {
outputSize += item.data.byteLength;
});
});
if (outputSize > outputSizeLimit) {
throw new Error('Notebook too large to backup');
}
}

cellData.outputs = !serializer.options.transientOutputs ? cell.outputs : [];
cellData.metadata = filter(cell.metadata, key => !serializer.options.transientCellMetadata[key]);

data.cells.push(cellData);
}

const bytes = await serializer.notebookToData(data);
if (token.isCancellationRequested) {
throw new CancellationError();
}
return bufferToStream(bytes);
return this._notebookService.createNotebookTextDocumentSnapshot(this._notebookModel.uri, context, token);
}

async update(stream: VSBufferReadableStream, token: CancellationToken): Promise<void> {
Expand Down
3 changes: 3 additions & 0 deletions src/vs/workbench/contrib/notebook/common/notebookService.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ import { IFileStatWithMetadata, IWriteFileOptions } from '../../../../platform/f
import { ITextQuery } from '../../../services/search/common/search.js';
import { NotebookPriorityInfo } from '../../search/common/search.js';
import { INotebookFileMatchNoModel } from '../../search/common/searchNotebookHelpers.js';
import { SnapshotContext } from '../../../services/workingCopy/common/fileWorkingCopy.js';


export const INotebookService = createDecorator<INotebookService>('notebookService');
Expand Down Expand Up @@ -80,6 +81,8 @@ export interface INotebookService {
saveMimeDisplayOrder(target: ConfigurationTarget): void;

createNotebookTextModel(viewType: string, uri: URI, stream?: VSBufferReadableStream): Promise<NotebookTextModel>;
createNotebookTextDocumentSnapshot(uri: URI, context: SnapshotContext, token: CancellationToken): Promise<VSBufferReadableStream>;
restoreNotebookTextModelFromSnapshot(uri: URI, viewType: string, snapshot: VSBufferReadableStream): Promise<NotebookTextModel>;
getNotebookTextModel(uri: URI): NotebookTextModel | undefined;
getNotebookTextModels(): Iterable<NotebookTextModel>;
listNotebookDocuments(): readonly NotebookTextModel[];
Expand Down
16 changes: 13 additions & 3 deletions src/vs/workbench/contrib/notebook/test/browser/notebookEditorModel.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
*--------------------------------------------------------------------------------------------*/

import assert from 'assert';
import { VSBuffer } from '../../../../../base/common/buffer.js';
import { bufferToStream, VSBuffer, VSBufferReadableStream } from '../../../../../base/common/buffer.js';
import { CancellationToken } from '../../../../../base/common/cancellation.js';
import { DisposableStore } from '../../../../../base/common/lifecycle.js';
import { Mimes } from '../../../../../base/common/mime.js';
Expand Down Expand Up @@ -246,7 +246,8 @@ suite('NotebookFileWorkingCopyModel', function () {
assert.strictEqual(notebook.cells[0].metadata!.bar, undefined);
return VSBuffer.fromString('');
}
}
},
configurationService
),
configurationService,
telemetryService,
Expand Down Expand Up @@ -309,7 +310,7 @@ suite('NotebookFileWorkingCopyModel', function () {
});
});

function mockNotebookService(notebook: NotebookTextModel, notebookSerializer: Promise<INotebookSerializer> | INotebookSerializer) {
function mockNotebookService(notebook: NotebookTextModel, notebookSerializer: Promise<INotebookSerializer> | INotebookSerializer, configurationService: TestConfigurationService = new TestConfigurationService()): INotebookService {
return new class extends mock<INotebookService>() {
private serializer: INotebookSerializer | undefined = undefined;
override async withNotebookDataProvider(viewType: string): Promise<SimpleNotebookProviderInfo> {
Expand All @@ -336,5 +337,14 @@ function mockNotebookService(notebook: NotebookTextModel, notebookSerializer: Pr
}
);
}
override async createNotebookTextDocumentSnapshot(uri: URI, context: SnapshotContext, token: CancellationToken): Promise<VSBufferReadableStream> {
const info = await this.withNotebookDataProvider(notebook.viewType);
const serializer = info.serializer;
const outputSizeLimit = configurationService.getValue(NotebookSetting.outputBackupSizeLimit) ?? 1024;
const data: NotebookData = notebook.createSnapshot({ context: context, outputSizeLimit: outputSizeLimit, transientOptions: serializer.options });
const bytes = await serializer.notebookToData(data);

return bufferToStream(bytes);
}
};
}

0 comments on commit 143f300

Please sign in to comment.