Even more robust handling of extension tree loading (#9007)

* Even more robust handling of extension tree loading

* Fix initial loading and add new test for CMS

* Fix compile errors

* Fix logic

* Remove debug log statements

extensions/big-data-cluster/src/bigDataCluster/tree/controllerTreeDataProvider.ts

@@ -32,9 +32,11 @@ export class ControllerTreeDataProvider implements vscode.TreeDataProvider<TreeN
 	public readonly onDidChangeTreeData: vscode.Event<TreeNode> = this._onDidChangeTreeData.event;
 	private root: ControllerRootNode;
 	private credentialProvider: azdata.CredentialProvider;
+	private initialized: boolean = false;
 
 	constructor(private memento: vscode.Memento) {
 		this.root = new ControllerRootNode(this);
+		this.root.addChild(new LoadingControllerNode());
 	}
 
 	public async getChildren(element?: TreeNode): Promise<TreeNode[]> {
@@ -42,14 +44,11 @@ export class ControllerTreeDataProvider implements vscode.TreeDataProvider<TreeN
 			return element.getChildren();
 		}
 
-		if (this.root.hasChildren) {
-			return this.root.getChildren();
+		if (!this.initialized) {
+			this.loadSavedControllers().catch(err => { vscode.window.showErrorMessage(localize('bdc.controllerTreeDataProvider.error', "Unexpected error loading saved controllers: {0}", err)); });
 		}
 
-		// Kick off loading the saved controllers but then immediately return the loading node so
-		// the user isn't left with an empty tree while we load the nodes
-		this.loadSavedControllers().catch(err => { vscode.window.showErrorMessage(localize('bdc.controllerTreeDataProvider.error', "Unexpected error loading saved controllers: {0}", err)); });
-		return [new LoadingControllerNode()];
+		return this.root.getChildren();
 	}
 
 	public getTreeItem(element: TreeNode): vscode.TreeItem | Thenable<vscode.TreeItem> {
@@ -89,12 +88,11 @@ export class ControllerTreeDataProvider implements vscode.TreeDataProvider<TreeN
 	}
 
 	private removeNonControllerNodes(): void {
-		this.removePlaceholderNodes();
-		this.removeDefectiveControllerNodes();
+		this.removePlaceholderNodes(this.root.children);
+		this.removeDefectiveControllerNodes(this.root.children);
 	}
 
-	private removePlaceholderNodes(): void {
-		let nodes = this.root.children;
+	private removePlaceholderNodes(nodes: TreeNode[]): void {
 		if (nodes.length > 0) {
 			for (let i = 0; i < nodes.length; ++i) {
 				if (nodes[i] instanceof AddControllerNode ||
@@ -106,8 +104,7 @@ export class ControllerTreeDataProvider implements vscode.TreeDataProvider<TreeN
 		}
 	}
 
-	private removeDefectiveControllerNodes(): void {
-		let nodes = this.root.children;
+	private removeDefectiveControllerNodes(nodes: TreeNode[]): void {
 		if (nodes.length > 0) {
 			for (let i = 0; i < nodes.length; ++i) {
 				if (nodes[i] instanceof ControllerNode) {
@@ -121,31 +118,42 @@ export class ControllerTreeDataProvider implements vscode.TreeDataProvider<TreeN
 	}
 
 	private async loadSavedControllers(): Promise<void> {
-		this.root.clearChildren();
-		let controllers: IControllerInfoSlim[] = this.memento.get('controllers');
-		if (controllers) {
-			for (const c of controllers) {
-				let password = undefined;
-				if (c.rememberPassword) {
-					password = await this.getPassword(c.url, c.username);
+		// Optimistically set to true so we don't double-load the tree
+		this.initialized = true;
+		try {
+			let controllers: IControllerInfoSlim[] = this.memento.get('controllers');
+			let treeNodes: TreeNode[] = [];
+			if (controllers) {
+				for (const c of controllers) {
+					let password = undefined;
+					if (c.rememberPassword) {
+						password = await this.getPassword(c.url, c.username);
+					}
+					if (!c.auth) {
+						// Added before we had added authentication
+						c.auth = 'basic';
+					}
+					treeNodes.push(new ControllerNode(
+						c.url, c.auth, c.username, password, c.rememberPassword,
+						undefined, this.root, this, undefined
+					));
 				}
-				if (!c.auth) {
-					// Added before we had added authentication
-					c.auth = 'basic';
-				}
-				this.root.addChild(new ControllerNode(
-					c.url, c.auth, c.username, password, c.rememberPassword,
-					undefined, this.root, this, undefined
-				));
+				this.removeDefectiveControllerNodes(treeNodes);
 			}
-			this.removeDefectiveControllerNodes();
+
+			this.root.clearChildren();
+			if (treeNodes.length === 0) {
+				this.root.addChild(new AddControllerNode());
+			} else {
+				treeNodes.forEach(node => this.root.addChild(node));
+			}
+			this.notifyNodeChanged();
+		} catch (err) {
+			// Reset so we can try again if the tree refreshes
+			this.initialized = false;
+			throw err;
 		}
 
-		if (!this.root.hasChildren) {
-			this.root.addChild(new AddControllerNode());
-		}
-
-		this.notifyNodeChanged();
 	}
 
 	public async saveControllers(): Promise<void> {

extensions/cms/src/cmsResource/tree/treeProvider.ts

@@ -17,6 +17,7 @@ import { CmsResourceTreeNode } from './cmsResourceTreeNode';
 export class CmsResourceTreeProvider implements TreeDataProvider<TreeNode>, ICmsResourceTreeChangeHandler {
 
 	private _appContext: AppContext;
+	private _children: TreeNode[] = [CmsResourceMessageTreeNode.create(CmsResourceTreeProvider.loadingLabel, undefined)];
 
 	public constructor(
 		public readonly appContext: AppContext
@@ -31,16 +32,14 @@ export class CmsResourceTreeProvider implements TreeDataProvider<TreeNode>, ICms
 		}
 
 		if (!this.isSystemInitialized) {
-			// Kick off loading the saved servers but then immediately return the loading node so
-			// the user isn't left with an empty tree while we load the nodes
-			this.loadSavedServers().catch(err => this._appContext.apiWrapper.showErrorMessage(localize('cms.resource.tree.treeProvider.loadError', "Unexpected error occured while loading saved servers {0}", err)));
-			return [CmsResourceMessageTreeNode.create(CmsResourceTreeProvider.loadingLabel, undefined)];
+			this.loadSavedServers().catch(err => this._appContext.apiWrapper.showErrorMessage(localize('cms.resource.tree.treeProvider.loadError', "Unexpected error occurred while loading saved servers {0}", err)));
+			return this._children;
 		}
 		try {
 			let registeredCmsServers = this.appContext.cmsUtils.registeredCmsServers;
 			if (registeredCmsServers && registeredCmsServers.length > 0) {
 				this.isSystemInitialized = true;
-				return registeredCmsServers.map((server) => {
+				this._children = registeredCmsServers.map((server) => {
 					return new CmsResourceTreeNode(
 						server.name,
 						server.description,
@@ -49,11 +48,12 @@ export class CmsResourceTreeProvider implements TreeDataProvider<TreeNode>, ICms
 						this._appContext, this, null);
 				}).sort((a, b) => a.name.localeCompare(b.name));
 			} else {
-				return [new CmsResourceEmptyTreeNode()];
+				this._children = [new CmsResourceEmptyTreeNode()];
 			}
 		} catch (error) {
-			return [new CmsResourceEmptyTreeNode()];
+			this._children = [new CmsResourceEmptyTreeNode()];
 		}
+		return this._children;
 	}
 
 	public get onDidChangeTreeData(): Event<TreeNode> {
@@ -93,6 +93,10 @@ export class CmsResourceTreeProvider implements TreeDataProvider<TreeNode>, ICms
 					await this.appContext.cmsUtils.cacheRegisteredCmsServer(server.name, server.description,
 						server.ownerUri, server.connection);
 				}
+				this._children = servers;
+			} else {
+				// No saved servers so just show the Add Server node since we're done loading
+				this._children = [new CmsResourceEmptyTreeNode()];
 			}
 			this._onDidChangeTreeData.fire(undefined);
 		} catch (error) {

extensions/cms/src/test/cmsResource/tree/treeProvider.test.ts

@@ -13,6 +13,7 @@ import { CmsResourceTreeProvider } from '../../../cmsResource/tree/treeProvider'
 import { CmsResourceMessageTreeNode } from '../../../cmsResource/messageTreeNode';
 import { CmsResourceEmptyTreeNode } from '../../../cmsResource/tree/cmsResourceEmptyTreeNode';
 import { CmsUtils } from '../../../cmsUtils';
+import { sleep } from '../../utils';
 
 // Mock services
 let mockAppContext: AppContext;
@@ -29,22 +30,42 @@ describe('CmsResourceTreeProvider.getChildren', function (): void {
 		mockAppContext = new AppContext(mockExtensionContext.object, mockApiWrapper.object, mockCmsUtils.object);
 	});
 
-	it('Should not be initialized.', async function (): Promise<void> {
+	it('Should be loading while waiting for saved servers to load', async function (): Promise<void> {
 		const treeProvider = new CmsResourceTreeProvider(mockAppContext);
-		should.notEqual(treeProvider.isSystemInitialized, true);
+		// We need to return at least one node so the async loading part is hit
+		mockCmsUtils.setup(x => x.getSavedServers).returns(() => {
+			return () => [{name: 'name',
+				description: 'desc',
+				ownerUri: 'uri',
+				connection: undefined}];
+		});
+		// Set up so loading the servers doesn't return immediately - thus we'll still have the Loading node
+		mockCmsUtils.setup(x => x.cacheRegisteredCmsServer).returns(() => {
+			return async () => { await sleep(600000); return undefined; };
+		});
+		should.notEqual(treeProvider.isSystemInitialized, true, 'Expected isSystemInitialized not to be true');
 		const children = await treeProvider.getChildren(undefined);
-		should.equal(children.length, 1);
-		should.equal(children[0].parent, undefined);
-		should.equal(children[0] instanceof CmsResourceMessageTreeNode, true);
+		should.equal(children.length, 1, 'Expected exactly one child node');
+		should.equal(children[0].parent, undefined, 'Expected node to not have a parent');
+		should.equal(children[0] instanceof CmsResourceMessageTreeNode, true, 'Expected node to be a CmsResourceMessageTreeNode');
+	});
+
+	it('Should be empty resource node when no servers to load', async function (): Promise<void> {
+		const treeProvider = new CmsResourceTreeProvider(mockAppContext);
+		should.notEqual(treeProvider.isSystemInitialized, true, 'Expected isSystemInitialized not to be true');
+		const children = await treeProvider.getChildren(undefined);
+		should.equal(children.length, 1, 'Expected exactly one child node');
+		should.equal(children[0].parent, undefined, 'Expected node to not have a parent');
+		should.equal(children[0] instanceof CmsResourceEmptyTreeNode, true, 'Expected node to be a CmsResourceEmptyTreeNode');
 	});
 
 	it('Should not be loading after initialized.', async function (): Promise<void> {
 		const treeProvider = new CmsResourceTreeProvider(mockAppContext);
 		treeProvider.isSystemInitialized = true;
-		should.equal(true, treeProvider.isSystemInitialized);
+		should.equal(true, treeProvider.isSystemInitialized, 'Expected isSystemInitialized to be true');
 		mockCmsUtils.setup(x => x.registeredCmsServers).returns(() => []);
 		const children = await treeProvider.getChildren(undefined);
-		should.equal(children[0] instanceof CmsResourceEmptyTreeNode, true);
+		should.equal(children[0] instanceof CmsResourceEmptyTreeNode, true, 'Expected child node to be CmsResourceEmptyTreeNode');
 	});
 
 	it('Should show CMS nodes if there are cached servers', async function (): Promise<void> {
@@ -59,6 +80,6 @@ describe('CmsResourceTreeProvider.getChildren', function (): void {
 			}];
 		});
 		const children = await treeProvider.getChildren(undefined);
-		should.equal(children[0] !== null, true);
+		should.exist(children[0], 'Child node did not exist');
 	});
 });

extensions/cms/src/test/utils.ts

@@ -0,0 +1,8 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+export async function sleep(ms: number): Promise<{}> {
+	return new Promise(resolve => setTimeout(resolve, ms));
+}