Make Throttler async (microsoft#185688)

Author: roblourens

src/vs/base/common/async.ts

@@ -166,26 +166,36 @@ export interface ITask<T> {
  * 			throttler.queue(deliver);
  * 		}
  */
-export class Throttler {
+export class Throttler implements IDisposable {
 
 	private activePromise: Promise<any> | null;
 	private queuedPromise: Promise<any> | null;
 	private queuedPromiseFactory: ITask<Promise<any>> | null;
 
+	private isDisposed = false;
+
 	constructor() {
 		this.activePromise = null;
 		this.queuedPromise = null;
 		this.queuedPromiseFactory = null;
 	}
 
 	queue<T>(promiseFactory: ITask<Promise<T>>): Promise<T> {
+		if (this.isDisposed) {
+			throw new Error('Throttler is disposed');
+		}
+
 		if (this.activePromise) {
 			this.queuedPromiseFactory = promiseFactory;
 
 			if (!this.queuedPromise) {
 				const onComplete = () => {
 					this.queuedPromise = null;
 
+					if (this.isDisposed) {
+						return;
+					}
+
 					const result = this.queue(this.queuedPromiseFactory!);
 					this.queuedPromiseFactory = null;
 
@@ -214,6 +224,10 @@ export class Throttler {
 			});
 		});
 	}
+
+	dispose(): void {
+		this.isDisposed = true;
+	}
 }
 
 export class Sequencer {

src/vs/base/test/common/async.test.ts

@@ -155,6 +155,40 @@ suite('Async', () => {
 
 			return Promise.all(promises);
 		});
+
+		test('disposal after queueing', async () => {
+			let factoryCalls = 0;
+			const factory = async () => {
+				factoryCalls++;
+				return async.timeout(0);
+			};
+
+			const throttler = new async.Throttler();
+			const promises: Promise<any>[] = [];
+
+			promises.push(throttler.queue(factory));
+			promises.push(throttler.queue(factory));
+			throttler.dispose();
+
+			await Promise.all(promises);
+			assert.strictEqual(factoryCalls, 1);
+		});
+
+		test('disposal before queueing', async () => {
+			let factoryCalls = 0;
+			const factory = async () => {
+				factoryCalls++;
+				return async.timeout(0);
+			};
+
+			const throttler = new async.Throttler();
+			const promises: Promise<any>[] = [];
+
+			throttler.dispose();
+			assert.throws(() => promises.push(throttler.queue(factory)));
+			assert.strictEqual(factoryCalls, 0);
+			await Promise.all(promises);
+		});
 	});
 
 	suite('Delayer', function () {

src/vs/workbench/contrib/notebook/browser/contrib/cellStatusBar/contributedStatusBarItemController.ts

@@ -76,9 +76,8 @@ class CellStatusBarHelper extends Disposable {
 	private _currentItemLists: INotebookCellStatusBarItemList[] = [];
 
 	private _activeToken: CancellationTokenSource | undefined;
-	private _isDisposed: boolean = false;
 
-	private readonly _updateThrottler = new Throttler();
+	private readonly _updateThrottler = this._register(new Throttler());
 
 	constructor(
 		private readonly _notebookViewModel: INotebookViewModel,
@@ -102,18 +101,7 @@ class CellStatusBarHelper extends Disposable {
 	private _updateSoon(): void {
 		// Wait a tick to make sure that the event is fired to the EH before triggering status bar providers
 		this._register(disposableTimeout(() => {
-			this._updateThrottler.queue(async () => {
-				if (this._isDisposed) {
-					// This order of events can happen
-					// - Start one update
-					// - Start a second update, its queued
-					// - This class is disposed, cancelling the first update
-					// - The second update runs, and we're disposed. So bail at this point.
-					return;
-				}
-
-				return this._update();
-			});
+			this._updateThrottler.queue(() => this._update());
 		}, 0));
 	}
 
@@ -140,7 +128,6 @@ class CellStatusBarHelper extends Disposable {
 
 	override dispose() {
 		super.dispose();
-		this._isDisposed = true;
 		this._activeToken?.dispose(true);
 
 		this._notebookViewModel.deltaCellStatusBarItems(this._currentItemIds, [{ handle: this._cell.handle, items: [] }]);