add deserialize timing and per-session caching

microsoft · May 24, 2023 · cc6cae4 · cc6cae4
1 parent 8c96a2a
commit cc6cae4
Show file tree

Hide file tree

Showing 2 changed files with 65 additions and 19 deletions.
diff --git a/src/vs/workbench/contrib/search/browser/searchActionsNotebook.ts b/src/vs/workbench/contrib/search/browser/searchActionsNotebook.ts
@@ -57,6 +57,7 @@ registerAction2(class NotebookDeserializeTest extends Action2 {
 		let processedBytes = 0;
 		let processedCells = 0;
 		let matchCount = 0;
+		let deserializeTime = 0;
 		const start = Date.now();
 		const pattern = 'start_index';
 		let i = 0;
@@ -65,6 +66,7 @@ registerAction2(class NotebookDeserializeTest extends Action2 {
 				break;
 			}
 			i++;
+			const deserializeStart = Date.now();
 			const uri = fileMatch.resource;
 			const content = await fileService.readFileStream(uri);
 			try {
@@ -83,7 +85,8 @@ registerAction2(class NotebookDeserializeTest extends Action2 {
 					_data = await info.serializer.dataToNotebook(bytes);
 				}
 
-
+				const deserializeEnd = Date.now();
+				deserializeTime += deserializeEnd - deserializeStart;
 				_data.cells.forEach((cell, index) => {
 					const input = cell.source;
 					const matches = this.getMatches(input, uri, index, pattern);
@@ -101,7 +104,7 @@ registerAction2(class NotebookDeserializeTest extends Action2 {
 		const end = Date.now();
 		logService.info(`${matchCount} matches found`);
 		logService.info(`notebook deserialize END | ${end - start}ms | ${((processedBytes / 1024) / 1024).toFixed(2)}MB | Number of Files: ${processedFiles} | Number of Cells: ${processedCells}`);
-
+		logService.info(`notebook deserialize time | ${deserializeTime}ms | ${(deserializeTime / (end - start)) * 100}% of total time`);
 	}
 
 	getMatches(source: string, uri: URI, cellIndex: number, target: string) {

diff --git a/src/vs/workbench/contrib/search/browser/searchModel.ts b/src/vs/workbench/contrib/search/browser/searchModel.ts
@@ -1927,6 +1927,7 @@ export class SearchModel extends Disposable {
 
 	private currentCancelTokenSource: CancellationTokenSource | null = null;
 	private searchCancelledForNewSearch: boolean = false;
+	private _notebookDataCache: NotebookDataCache;
 
 	constructor(
 		@ISearchService private readonly searchService: ISearchService,
@@ -1935,12 +1936,11 @@ export class SearchModel extends Disposable {
 		@IInstantiationService private readonly instantiationService: IInstantiationService,
 		@IUriIdentityService private readonly uriIdentityService: IUriIdentityService,
 		@INotebookEditorService private readonly notebookEditorService: INotebookEditorService,
-		@INotebookService private readonly notebookService: INotebookService,
-		@IFileService private readonly fileService: IFileService,
 		@ILogService private readonly logService: ILogService,
 	) {
 		super();
 		this._searchResult = this.instantiationService.createInstance(SearchResult, this);
+		this._notebookDataCache = this.instantiationService.createInstance(NotebookDataCache);
 	}
 
 	isReplaceActive(): boolean {
@@ -1986,6 +1986,8 @@ export class SearchModel extends Disposable {
 			folderQueries: textQuery.folderQueries
 		};
 
+		let deserializeTime = 0;
+		const start = Date.now();
 
 		const searchComplete = await this.searchService.fileSearch(
 			query,
@@ -1996,22 +1998,12 @@ export class SearchModel extends Disposable {
 		for (const fileMatch of filesToSearch) {
 			const cellMatches: ICellMatch[] = [];
 			const uri = fileMatch.resource;
-			const content = await this.fileService.readFileStream(uri);
-			try {
-				const info = await this.notebookService.withNotebookDataProvider('jupyter-notebook');
-				if (!(info instanceof SimpleNotebookProviderInfo)) {
-					throw new Error('CANNOT open file notebook with this provider');
-				}
-
-				let _data: NotebookData = {
-					metadata: {},
-					cells: []
-				};
-				if (uri.scheme !== Schemas.vscodeInteractive) {
-					const bytes = await streamToBuffer(content.value);
-					_data = await info.serializer.dataToNotebook(bytes);
-				}
 
+			try {
+				const deserializeStart = Date.now();
+				const _data = await this._notebookDataCache.getNotebookData(uri);
+				const deserializeEnd = Date.now();
+				deserializeTime += deserializeEnd - deserializeStart;
 
 				_data.cells.forEach((cell, index) => {
 					const source = cell.source;
@@ -2043,6 +2035,11 @@ export class SearchModel extends Disposable {
 			}
 		}
 
+		const end = Date.now();
+		this.logService.info(`query: ${textQuery.contentPattern.pattern}`);
+		this.logService.info(`notebook deserialize END | ${end - start}ms`);
+		this.logService.info(`notebook deserialize time | ${deserializeTime}ms | ${(deserializeTime / (end - start)) * 100}% of total time`);
+
 		return {
 			results: results,
 			limitHit: false
@@ -2388,6 +2385,52 @@ export class RangeHighlightDecorations implements IDisposable {
 	});
 }
 
+interface INotebookDataEditInfo {
+	notebookData: NotebookData;
+	mTime: number;
+}
+class NotebookDataCache {
+	private _entries: ResourceMap<INotebookDataEditInfo>;
+
+	constructor(
+		@IUriIdentityService private readonly uriIdentityService: IUriIdentityService,
+		@IFileService private readonly fileService: IFileService,
+		@INotebookService private readonly notebookService: INotebookService,
+	) {
+		this._entries = new ResourceMap<INotebookDataEditInfo>(uri => this.uriIdentityService.extUri.getComparisonKey(uri));
+	}
+
+	async getNotebookData(notebookUri: URI): Promise<NotebookData> {
+		if (notebookUri.scheme === Schemas.vscodeInteractive) {
+			//unsupported
+			throw new Error(`CANNOT parse notebook uri ${notebookUri.toString()}`);
+		}
+		const mTime = (await this.fileService.stat(notebookUri)).mtime;
+
+		const entry = this._entries.get(notebookUri);
+
+		if (entry && entry.mTime === mTime) {
+			return entry.notebookData;
+		} else {
+			const info = await this.notebookService.withNotebookDataProvider('jupyter-notebook');
+			if (!(info instanceof SimpleNotebookProviderInfo)) {
+				throw new Error('CANNOT open file notebook with this provider');
+			}
+
+			let _data: NotebookData = {
+				metadata: {},
+				cells: []
+			};
+
+			const content = await this.fileService.readFileStream(notebookUri);
+			const bytes = await streamToBuffer(content.value);
+			_data = await info.serializer.dataToNotebook(bytes);
+			this._entries.set(notebookUri, { notebookData: _data, mTime });
+			return _data;
+		}
+	}
+
+}
 function textSearchResultToMatches(rawMatch: ITextSearchMatch, fileMatch: FileMatch): Match[] {
 	const previewLines = rawMatch.preview.text.split('\n');
 	if (Array.isArray(rawMatch.ranges)) {