Add memoize helper with cancellation support

Author: felixfbecker

package.json

@@ -59,9 +59,11 @@
     "@types/mocha": "2.2.32",
     "@types/mz": "0.0.30",
     "@types/node": "6.0.46",
+    "@types/sinon": "1.16.35",
     "mocha": "^3.2.0",
     "nyc": "^10.1.2",
     "rimraf": "^2.6.1",
+    "sinon": "^2.0.0",
     "source-map-support": "^0.4.11",
     "ts-node": "^1.6.1",
     "tslint": "^4.5.1",

src/cancellation.ts

@@ -1,9 +1,113 @@
-
 import * as ts from 'typescript';
-
-import { CancellationToken, CancellationTokenSource } from 'vscode-jsonrpc';
+import { CancellationToken, CancellationTokenSource, Disposable } from 'vscode-jsonrpc';
 export { CancellationToken, CancellationTokenSource };
 
+/**
+ * Provides a token that is cancelled as soon as ALL added tokens are cancelled.
+ * Useful for memoizing a function, where multiple consumers wait for the result of the same operation,
+ * which should only be cancelled if all consumers requested cancellation.
+ */
+class CancellationTokenLink {
+
+	/**
+	 * Amount of total consumers
+	 */
+	private tokens = 0;
+
+	/**
+	 * Amount of consumers that have requested cancellation
+	 */
+	private cancelled = 0;
+
+	/**
+	 * Internal token source that is cancelled when all consumers request cancellation
+	 */
+	private source = new CancellationTokenSource();
+
+	/**
+	 * A linked CancellationToken that is cancelled as soon as all consumers request cancellation
+	 */
+	get token(): CancellationToken {
+		return this.source.token;
+	}
+
+	/**
+	 * Add another CancellationToken that needs to be cancelled to trigger cancellation
+	 *
+	 * @returns Disposable that allows to remove the token again
+	 */
+	add(token: CancellationToken): Disposable {
+		this.tokens++;
+		const handlerDisposable = token.onCancellationRequested(() => {
+			this.cancelled++;
+			// If all consumers requested cancellation, cancel
+			if (this.cancelled === this.tokens) {
+				this.source.cancel();
+			}
+		});
+		return {
+			dispose: () => {
+				// Remove onCancellationRequested handler
+				handlerDisposable.dispose();
+				this.tokens--;
+				if (token.isCancellationRequested) {
+					this.cancelled--;
+				}
+			}
+		};
+	}
+}
+
+/**
+ * Memoizes the result of a promise-returning function by the first argument with support for cancellation.
+ * If the last argument is a CancellationToken, the operation is only cancelled if all calls have requested cancellation.
+ * Rejected (or cancelled) promises are automatically removed from the cache.
+ * If the operation has already finished, it will not be cancelled.
+ *
+ * @param func Function to memoize
+ * @param cache A custom Map to use for setting and getting cache items
+ *
+ * @template F The function to be memoized
+ * @template T The return type of the function, must be Promise
+ * @template K The cache key (first argument to the function)
+ */
+export function cancellableMemoize<F extends (...args: any[]) => T, T extends Promise<any>, K>(func: F, cache = new Map<K, T>()): F {
+	// Track tokens consumers provide
+	const tokenLinks = new Map<K, CancellationTokenLink>();
+	const memoized: F = <any> function (this: any, ...args: any[]) {
+		const key = args[0];
+		// Get or create CancellationTokenLink for the given first parameter
+		let tokenLink = tokenLinks.get(key);
+		if (!tokenLink) {
+			tokenLink = new CancellationTokenLink();
+			tokenLinks.set(key, tokenLink);
+		}
+		// Take last argument as CancellationToken from arguments if provided or use a token that is never cancelled
+		const token: CancellationToken = CancellationToken.is(args[args.length - 1]) ? args.pop() : CancellationToken.None;
+		// Add it to the list of tokens that need to be cancelled for final cancellation
+		tokenLink.add(token);
+		let result: T;
+		// Check if function has been called with this argument already
+		if (cache.has(key)) {
+			// Return previous result
+			result = cache.get(key);
+		} else {
+			// Call function
+			// Pass the linked cancel token
+			args.push(tokenLink.token);
+			result = <T> (<T> func.apply(this, args)).catch(err => {
+				// Don't cache rejected promises
+				cache.delete(key);
+				throw err;
+			});
+			// Save result
+			cache.set(key, result);
+		}
+		return result;
+	};
+	return memoized;
+}
+
 /**
  * Thrown when an operation was cancelled
  */

src/fs.ts

@@ -1,9 +1,10 @@
 import * as fs from 'mz/fs';
 import * as path from 'path';
-import { CancellationToken, throwIfRequested } from './cancellation';
+import { cancellableMemoize, CancellationToken, throwIfRequested } from './cancellation';
 import { LanguageClientHandler } from './lang-handler';
 import glob = require('glob');
 import Semaphore from 'semaphore-async-await';
+import * as url from 'url';
 import { InMemoryFileSystem } from './project-manager';
 import { path2uri, uri2path } from './util';
 
@@ -74,84 +75,94 @@ export class LocalFileSystem implements FileSystem {
 	}
 }
 
+/**
+ * Memoization cache that saves URIs and searches parent directories too
+ */
+class ParentUriMemoizationCache extends Map<string | undefined, Promise<void>> {
+
+	/**
+	 * Returns the value if the given URI or a parent directory of the URI is in the cache
+	 */
+	get(uri: string | undefined): Promise<void> | undefined {
+		let hit = super.get(uri);
+		if (hit) {
+			return hit;
+		}
+		// Find out if parent folder is being fetched already
+		hit = super.get(undefined);
+		if (hit) {
+			return hit;
+		}
+		if (uri) {
+			for (let parts = url.parse(uri); parts.pathname !== '/'; parts.pathname = path.dirname(parts.pathname)) {
+				hit = super.get(url.format(parts));
+				if (hit) {
+					return hit;
+				}
+			}
+		}
+		return undefined;
+	}
+
+	/**
+	 * Returns true if the given URI or a parent directory of the URI is in the cache
+	 */
+	has(key: string): boolean {
+		return this.get(key) !== undefined;
+	}
+}
+
 /**
  * Synchronizes a remote file system to an in-memory file system
  *
  * TODO: Implement Disposable with Disposer
  */
 export class FileSystemUpdater {
 
-	/**
-	 * Map from URI to Promise for a content fetch
-	 */
-	private fetches = new Map<string, Promise<void>>();
-
 	/**
 	 * Limits concurrent fetches to not fetch thousands of files in parallel
 	 */
 	private concurrencyLimit = new Semaphore(100);
 
-	private structureFetch?: Promise<void>;
-
 	constructor(private remoteFs: FileSystem, private inMemoryFs: InMemoryFileSystem) {}
 
 	/**
 	 * Fetches the file content for the given URI and adds the content to the in-memory file system
 	 *
 	 * @param uri URI of the file to fetch
 	 */
-	async fetch(uri: string, token = CancellationToken.None): Promise<void> {
+	fetch(uri: string, token = CancellationToken.None): Promise<void> {
 		// Limit concurrent fetches
-		const fetch = this.concurrencyLimit.execute(async () => {
+		return this.concurrencyLimit.execute(async () => {
 			throwIfRequested(token);
 			const content = await this.remoteFs.getTextDocumentContent(uri, token);
-			this.inMemoryFs.addFile(uri2path(uri), content);
-		});
-		// Avoid unhandled promise rejection errors (error is still propagated)
-		fetch.catch(err => {
-			// Don't cache failed or cancelled fetches
-			this.fetches.delete(uri);
+			this.inMemoryFs.add(uri, content);
 		});
-		// Track the fetch
-		this.fetches.set(uri, fetch);
-		return fetch;
 	}
 
 	/**
 	 * Returns a promise that is resolved when the given URI has been fetched (at least once) to the in-memory file system.
 	 * This function cannot be cancelled because multiple callers get the result of the same operation.
 	 *
-	 * TODO use memoize helper and support cancellation with multiple consumers
-	 *
 	 * @param uri URI of the file to ensure
 	 */
-	async ensure(uri: string): Promise<void> {
-		return this.fetches.get(uri) || this.fetch(uri);
-	}
+	ensure = cancellableMemoize(this.fetch);
 
 	/**
-	 * Fetches the file/directory structure from the remote file system and saves it in the in-memory file system
+	 * Fetches the file/directory structure for the given directory from the remote file system and saves it in the in-memory file system
+	 *
+	 * @param base The base directory which structure will be synced. Defaults to the workspace root
 	 */
-	async fetchStructure(token = CancellationToken.None): Promise<void> {
-		this.structureFetch = (async () => {
-			const uris = await this.remoteFs.getWorkspaceFiles(undefined, token);
-			for (const uri of uris) {
-				this.inMemoryFs.add(uri);
-			}
-		})();
-		this.structureFetch.catch(err => {
-			this.structureFetch = undefined;
-		});
-		return this.structureFetch;
+	async fetchStructure(base?: string, token = CancellationToken.None): Promise<void> {
+		const uris = await this.remoteFs.getWorkspaceFiles(base, token);
+		for (const uri of uris) {
+			this.inMemoryFs.add(uri);
+		}
 	}
 
 	/**
-	 * Returns a promise that is resolved as soon as the file/directory structure has been synced
+	 * Returns a promise that is resolved as soon as the file/directory structure for the given directory has been synced
 	 * from the remote file system to the in-memory file system (at least once)
-	 *
-	 * TODO use memoize helper and support cancellation with multiple consumers
 	 */
-	async ensureStructure(): Promise<void> {
-		return this.structureFetch || this.fetchStructure();
-	}
+	ensureStructure = cancellableMemoize(this.fetchStructure, new ParentUriMemoizationCache());
 }

src/project-manager.ts

@@ -292,7 +292,7 @@ export class ProjectManager implements Disposable {
 			.then(() => this.refreshConfigurations());
 
 		this.ensuredAllFiles = promise;
-		promise.catch(err => {
+		promise.catch((err: any) => {
 			console.error('Failed to fetch files for references:', err);
 			this.ensuredAllFiles = undefined;
 		});
@@ -468,7 +468,7 @@ export class ProjectManager implements Disposable {
 			await Promise.all(iterate(files).map(async path => {
 				throwIfRequested(token);
 				try {
-					await this.updater.ensure(util.path2uri('', path));
+					await this.updater.ensure(util.path2uri('', path), token);
 				} catch (err) {
 					// if cancellation was requested, break out of the loop
 					throwIfCancelledError(err);

src/test/cancellation.test.ts

@@ -0,0 +1,75 @@
+
+import * as chai from 'chai';
+import { cancellableMemoize, CancellationToken, CancellationTokenSource, CancelledError } from '../cancellation';
+import chaiAsPromised = require('chai-as-promised');
+import * as sinon from 'sinon';
+chai.use(chaiAsPromised);
+const assert = chai.assert;
+
+describe('cancellation', () => {
+	describe('cancellableMemoize()', () => {
+		it('should memoize a function by the first argument', async () => {
+			const toBeMemoized = sinon.spy((arg1: number, arg2: number, token = CancellationToken.None): Promise<number> => {
+				return Promise.resolve(Math.random());
+			});
+			const memoized = cancellableMemoize(toBeMemoized);
+			const a = await memoized(123, 456);
+			const b = await memoized(123, 456);
+			sinon.assert.calledOnce(toBeMemoized);
+			sinon.assert.calledWith(toBeMemoized, 123, 456, sinon.match.object);
+			assert.equal(a, b);
+		});
+		it('should memoize a function without parameters', async () => {
+			const toBeMemoized = sinon.spy((token = CancellationToken.None): Promise<number> => {
+				return Promise.resolve(Math.random());
+			});
+			const memoized = cancellableMemoize(toBeMemoized);
+			const a = await memoized();
+			const b = await memoized();
+			assert.equal(a, b);
+			sinon.assert.calledOnce(toBeMemoized);
+			sinon.assert.calledWith(toBeMemoized, sinon.match.object);
+		});
+		it('should not cancel the operation if there are still consumers', async () => {
+			const toBeMemoized = sinon.spy((arg: number, token = CancellationToken.None): Promise<number> => {
+				return new Promise((resolve, reject) => {
+					token.onCancellationRequested(() => {
+						reject(new CancelledError());
+					});
+					setTimeout(() => resolve(123), 500);
+				});
+			});
+			const memoized = cancellableMemoize(toBeMemoized);
+			const source1 = new CancellationTokenSource();
+			const source2 = new CancellationTokenSource();
+			const promise1 = memoized(123, source1.token);
+			const promise2 = memoized(123, source2.token);
+			source1.cancel();
+			assert.equal(await promise1, 123);
+			assert.equal(await promise2, 123);
+			sinon.assert.calledOnce(toBeMemoized);
+			sinon.assert.calledWith(toBeMemoized, 123, sinon.match.object);
+		});
+		it('should cancel the operation if all consumers requested cancellation', async () => {
+			const toBeMemoized = sinon.spy((arg: number, token: CancellationToken): Promise<number> => {
+				return new Promise((resolve, reject) => {
+					token.onCancellationRequested(() => {
+						reject(new CancelledError());
+					});
+					setTimeout(() => resolve(123), 500);
+				});
+			});
+			const memoized = cancellableMemoize(toBeMemoized);
+			const source1 = new CancellationTokenSource();
+			const source2 = new CancellationTokenSource();
+			const promise1 = memoized(123, source1.token);
+			const promise2 = memoized(123, source2.token);
+			source1.cancel();
+			source2.cancel();
+			sinon.assert.calledOnce(toBeMemoized);
+			sinon.assert.calledWith(toBeMemoized, 123, sinon.match.object);
+			await assert.isRejected(promise1, /cancel/i);
+			await assert.isRejected(promise2, /cancel/i);
+		});
+	});
+});