Send search results back from EH in batches - #47058

src/vs/platform/search/common/search.ts

@@ -33,7 +33,7 @@ export interface ISearchResultProvider {
 	search(query: ISearchQuery): PPromise<ISearchComplete, ISearchProgressItem>;
 }
 
-export interface IFolderQuery<U=uri> {
+export interface IFolderQuery<U extends UriComponents=uri> {
 	folder: U;
 	excludePattern?: glob.IExpression;
 	includePattern?: glob.IExpression;
@@ -65,7 +65,7 @@ export interface IQueryOptions extends ICommonQueryOptions<uri> {
 	includePattern?: string;
 }
 
-export interface ISearchQueryProps<U> extends ICommonQueryOptions<U> {
+export interface ISearchQueryProps<U extends UriComponents> extends ICommonQueryOptions<U> {
 	type: QueryType;
 
 	excludePattern?: glob.IExpression;
@@ -103,11 +103,13 @@ export interface IPatternInfo {
 	isSmartCase?: boolean;
 }
 
-export interface IFileMatch {
-	resource?: uri;
+export interface IFileMatch<U extends UriComponents = uri> {
+	resource?: U;
 	lineMatches?: ILineMatch[];
 }
 
+export type IRawFileMatch = IFileMatch<UriComponents>;
+
 export interface ILineMatch {
 	preview: string;
 	lineNumber: number;

src/vs/workbench/api/electron-browser/mainThreadSearch.ts

@@ -9,7 +9,7 @@ import { IDisposable, dispose } from 'vs/base/common/lifecycle';
 import { values } from 'vs/base/common/map';
 import URI, { UriComponents } from 'vs/base/common/uri';
 import { PPromise, TPromise } from 'vs/base/common/winjs.base';
-import { IFileMatch, ILineMatch, ISearchComplete, ISearchProgressItem, ISearchQuery, ISearchResultProvider, ISearchService, QueryType } from 'vs/platform/search/common/search';
+import { IFileMatch, ISearchComplete, ISearchProgressItem, ISearchQuery, ISearchResultProvider, ISearchService, QueryType, IRawFileMatch } from 'vs/platform/search/common/search';
 import { extHostNamedCustomer } from 'vs/workbench/api/electron-browser/extHostCustomers';
 import { ExtHostContext, ExtHostSearchShape, IExtHostContext, MainContext, MainThreadSearchShape } from '../node/extHost.protocol';
 
@@ -40,7 +40,7 @@ export class MainThreadSearch implements MainThreadSearchShape {
 		this._searchProvider.delete(handle);
 	}
 
-	$handleFindMatch(handle: number, session, data: UriComponents | [UriComponents, ILineMatch]): void {
+	$handleFindMatch(handle: number, session, data: UriComponents | IRawFileMatch[]): void {
 		this._searchProvider.get(handle).handleFindMatch(session, data);
 	}
 }
@@ -57,14 +57,15 @@ class SearchOperation {
 		//
 	}
 
-	addMatch(resource: URI, match: ILineMatch): void {
-		if (!this.matches.has(resource.toString())) {
-			this.matches.set(resource.toString(), { resource, lineMatches: [] });
-		}
-		if (match) {
-			this.matches.get(resource.toString()).lineMatches.push(match);
+	addMatch(match: IFileMatch): void {
+		if (this.matches.has(match.resource.toString())) {
+			// Merge with previous IFileMatches
+			this.matches.get(match.resource.toString()).lineMatches.push(...match.lineMatches);
+		} else {
+			this.matches.set(match.resource.toString(), match);
 		}
-		this.progress(this.matches.get(resource.toString()));
+
+		this.progress(this.matches.get(match.resource.toString()));
 	}
 }
 
@@ -124,22 +125,23 @@ class RemoteSearchProvider implements ISearchResultProvider {
 		});
 	}
 
-	handleFindMatch(session: number, dataOrUri: UriComponents | [UriComponents, ILineMatch]): void {
+	handleFindMatch(session: number, dataOrUri: UriComponents | IRawFileMatch[]): void {
 		if (!this._searches.has(session)) {
 			// ignore...
 			return;
 		}
-		let resource: URI;
-		let match: ILineMatch;
 
+		const searchOp = this._searches.get(session);
 		if (Array.isArray(dataOrUri)) {
-			resource = URI.revive(dataOrUri[0]);
-			match = dataOrUri[1];
+			dataOrUri.forEach(m => {
+				searchOp.addMatch({
+					resource: URI.revive(m.resource),
+					lineMatches: m.lineMatches
+				});
+			});
 		} else {
-			resource = URI.revive(dataOrUri);
+			searchOp.addMatch({ resource: URI.revive(dataOrUri) });
 		}
-
-		this._searches.get(session).addMatch(resource, match);
 	}
 }
 

src/vs/workbench/api/node/extHost.protocol.ts

@@ -46,7 +46,7 @@ import { IStat, FileChangeType, FileOpenFlags, IWatchOptions, FileSystemProvider
 import { ConfigurationScope } from 'vs/platform/configuration/common/configurationRegistry';
 import { CommentRule, CharacterPair, EnterAction } from 'vs/editor/common/modes/languageConfiguration';
 import { ISingleEditOperation } from 'vs/editor/common/model';
-import { ILineMatch, IPatternInfo, IRawSearchQuery } from 'vs/platform/search/common/search';
+import { IPatternInfo, IRawSearchQuery, IRawFileMatch } from 'vs/platform/search/common/search';
 import { LogLevel } from 'vs/platform/log/common/log';
 import { TaskExecutionDTO, TaskDTO, TaskHandleDTO } from 'vs/workbench/api/shared/tasks';
 
@@ -400,7 +400,7 @@ export interface MainThreadFileSystemShape extends IDisposable {
 export interface MainThreadSearchShape extends IDisposable {
 	$registerSearchProvider(handle: number, scheme: string): void;
 	$unregisterProvider(handle: number): void;
-	$handleFindMatch(handle: number, session, data: UriComponents | [UriComponents, ILineMatch]): void;
+	$handleFindMatch(handle: number, session, data: UriComponents | IRawFileMatch[]): void;
 }
 
 export interface MainThreadTaskShape extends IDisposable {

src/vs/workbench/api/node/extHostSearch.ts

@@ -6,7 +6,7 @@
 
 import { asWinJsPromise } from 'vs/base/common/async';
 import { TPromise } from 'vs/base/common/winjs.base';
-import { IPatternInfo, IFolderQuery, IRawSearchQuery } from 'vs/platform/search/common/search';
+import { IPatternInfo, IFolderQuery, IRawSearchQuery, IRawFileMatch, IFileMatch } from 'vs/platform/search/common/search';
 import * as vscode from 'vscode';
 import { ExtHostSearchShape, IMainContext, MainContext, MainThreadSearchShape } from './extHost.protocol';
 import URI, { UriComponents } from 'vs/base/common/uri';
@@ -81,15 +81,148 @@ export class ExtHostSearch implements ExtHostSearchShape {
 			encoding: query.fileEncoding
 		};
 
+		const collector = new TextSearchResultsCollector(handle, session, this._proxy);
 		const progress = {
 			report: (data: vscode.TextSearchResult) => {
-				this._proxy.$handleFindMatch(handle, session, [data.uri, {
-					lineNumber: data.range.start.line,
-					preview: data.preview.leading + data.preview.matching + data.preview.trailing,
-					offsetAndLengths: [[data.preview.leading.length, data.preview.matching.length]]
-				}]);
+				collector.add(data);
 			}
 		};
-		return asWinJsPromise(token => provider.provideTextSearchResults(pattern, searchOptions, progress, token));
+		return asWinJsPromise(token => provider.provideTextSearchResults(pattern, searchOptions, progress, token))
+			.then(() => collector.flush());
+	}
+}
+
+class TextSearchResultsCollector {
+	private _batchedCollector: BatchedCollector<IRawFileMatch>;
+
+	private _currentFileMatch: IFileMatch;
+
+	constructor(private _handle: number, private _session: number, private _proxy: MainThreadSearchShape) {
+		this._batchedCollector = new BatchedCollector<IRawFileMatch>(512, items => this.sendItems(items));
+	}
+
+	add(data: vscode.TextSearchResult): void {
+		// Collects TextSearchResults into IRawFileMatches and collates using BatchedCollector.
+		// This is efficient for ripgrep which sends results back one file at a time. It wouldn't be efficient for other search
+		// providers that send results in random order. We could do this step afterwards instead.
+		if (this._currentFileMatch && this._currentFileMatch.resource.toString() !== data.uri.toString()) {
+			this.pushToCollector();
+			this._currentFileMatch = null;
+		}
+
+		if (!this._currentFileMatch) {
+			this._currentFileMatch = {
+				resource: data.uri,
+				lineMatches: []
+			};
+		}
+
+		// TODO@roblou - line text is sent for every match
+		this._currentFileMatch.lineMatches.push({
+			lineNumber: data.range.start.line,
+			preview: data.preview.leading + data.preview.matching + data.preview.trailing,
+			offsetAndLengths: [[data.preview.leading.length, data.preview.matching.length]]
+		});
+	}
+
+	private pushToCollector(): void {
+		const size = this._currentFileMatch.lineMatches.reduce((acc, match) => acc + match.offsetAndLengths.length, 0);
+		this._batchedCollector.addItem(this._currentFileMatch, size);
+	}
+
+	flush(): void {
+		this.pushToCollector();
+		this._batchedCollector.flush();
+	}
+
+	private sendItems(items: IRawFileMatch | IRawFileMatch[]): void {
+		items = Array.isArray(items) ? items : [items];
+		this._proxy.$handleFindMatch(this._handle, this._session, items);
+	}
+}
+
+/**
+ * Collects items that have a size - before the cumulative size of collected items reaches START_BATCH_AFTER_COUNT, the callback is called for every
+ * set of items collected.
+ * But after that point, the callback is called with batches of maxBatchSize.
+ * If the batch isn't filled within some time, the callback is also called.
+ */
+class BatchedCollector<T> {
+	private static readonly TIMEOUT = 4000;
+
+	// After START_BATCH_AFTER_COUNT items have been collected, stop flushing on timeout
+	private static readonly START_BATCH_AFTER_COUNT = 50;
+
+	private totalNumberCompleted = 0;
+	private batch: T[] = [];
+	private batchSize = 0;
+	private timeoutHandle: number;
+
+	constructor(private maxBatchSize: number, private cb: (items: T | T[]) => void) {
+	}
+
+	addItem(item: T, size: number): void {
+		if (!item) {
+			return;
+		}
+
+		if (this.maxBatchSize > 0) {
+			this.addItemToBatch(item, size);
+		} else {
+			this.cb(item);
+		}
+	}
+
+	addItems(items: T[], size: number): void {
+		if (!items) {
+			return;
+		}
+
+		if (this.maxBatchSize > 0) {
+			this.addItemsToBatch(items, size);
+		} else {
+			this.cb(items);
+		}
+	}
+
+	private addItemToBatch(item: T, size: number): void {
+		this.batch.push(item);
+		this.batchSize += size;
+		this.onUpdate();
+	}
+
+	private addItemsToBatch(item: T[], size: number): void {
+		this.batch = this.batch.concat(item);
+		this.batchSize += size;
+		this.onUpdate();
+	}
+
+	private onUpdate(): void {
+		if (this.totalNumberCompleted < BatchedCollector.START_BATCH_AFTER_COUNT) {
+			// Flush because we aren't batching yet
+			this.flush();
+		} else if (this.batchSize >= this.maxBatchSize) {
+			// Flush because the batch is full
+			this.flush();
+		} else if (!this.timeoutHandle) {
+			// No timeout running, start a timeout to flush
+			this.timeoutHandle = setTimeout(() => {
+				this.flush();
+			}, BatchedCollector.TIMEOUT);
+		}
+	}
+
+	flush(): void {
+		if (this.batchSize) {
+			this.totalNumberCompleted += this.batchSize;
+			this.cb(this.batch);
+			this.batch = [];
+			this.batchSize = 0;
+
+			if (this.timeoutHandle) {
+				clearTimeout(this.timeoutHandle);
+				this.timeoutHandle = 0;
+			}
+		}
 	}
 }