diff --git a/src/Microsoft.SqlTools.ServiceLayer/LanguageServices/BindingQueue.cs b/src/Microsoft.SqlTools.ServiceLayer/LanguageServices/BindingQueue.cs index 93b1a9f6..7cb0d592 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/LanguageServices/BindingQueue.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/LanguageServices/BindingQueue.cs @@ -39,6 +39,8 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices /// internal Dictionary BindingContextMap { get; set; } + internal Dictionary BindingContextTasks { get; set; } = new Dictionary(); + /// /// Constructor for a binding queue instance /// @@ -145,7 +147,9 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices { if (!this.BindingContextMap.ContainsKey(key)) { - this.BindingContextMap.Add(key, new T()); + var bindingContext = new T(); + this.BindingContextMap.Add(key, bindingContext); + this.BindingContextTasks.Add(bindingContext, Task.Run(() => null)); } return this.BindingContextMap[key]; @@ -190,11 +194,17 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices var bindingContext = this.BindingContextMap[key]; if (bindingContext.ServerConnection != null && bindingContext.ServerConnection.IsOpen) { - bindingContext.ServerConnection.Disconnect(); + // Disconnecting can take some time so run it in a separate task so that it doesn't block removal + Task.Run(() => + { + bindingContext.ServerConnection.Cancel(); + bindingContext.ServerConnection.Disconnect(); + }); } // remove key from the map this.BindingContextMap.Remove(key); + this.BindingContextTasks.Remove(bindingContext); } } } @@ -286,86 +296,119 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices continue; } - bool lockTaken = false; - try - { - // prefer the queue item binding item, otherwise use the context default timeout - int bindTimeout = queueItem.BindingTimeout ?? bindingContext.BindingTimeout; + var bindingContextTask = this.BindingContextTasks[bindingContext]; - // handle the case a previous binding operation is still running - if (!bindingContext.BindingLock.WaitOne(queueItem.WaitForLockTimeout ?? 0)) - { - queueItem.Result = queueItem.TimeoutOperation != null - ? queueItem.TimeoutOperation(bindingContext) - : null; + // Run in the binding context task in case this task has to wait for a previous binding operation + this.BindingContextTasks[bindingContext] = bindingContextTask.ContinueWith((task) => + { + bool lockTaken = false; + try + { + // prefer the queue item binding item, otherwise use the context default timeout + int bindTimeout = queueItem.BindingTimeout ?? bindingContext.BindingTimeout; - continue; - } - - bindingContext.BindingLock.Reset(); - - lockTaken = true; - - // execute the binding operation - object result = null; - CancellationTokenSource cancelToken = new CancellationTokenSource(); - - // run the operation in a separate thread - var bindTask = Task.Run(() => - { - try + // handle the case a previous binding operation is still running + if (!bindingContext.BindingLock.WaitOne(queueItem.WaitForLockTimeout ?? 0)) { - result = queueItem.BindOperation( - bindingContext, - cancelToken.Token); - } - catch (Exception ex) - { - Logger.Write(TraceEventType.Error, "Unexpected exception on the binding queue: " + ex.ToString()); - if (queueItem.ErrorHandler != null) + try { - result = queueItem.ErrorHandler(ex); + Logger.Write(TraceEventType.Warning, "Binding queue operation timed out waiting for previous operation to finish"); + queueItem.Result = queueItem.TimeoutOperation != null + ? queueItem.TimeoutOperation(bindingContext) + : null; + } + catch (Exception ex) + { + Logger.Write(TraceEventType.Error, "Exception running binding queue lock timeout handler: " + ex.ToString()); + } + finally + { + queueItem.ItemProcessed.Set(); } - } - }); - - // check if the binding tasks completed within the binding timeout - if (bindTask.Wait(bindTimeout)) - { - queueItem.Result = result; - } - else - { - cancelToken.Cancel(); - // if the task didn't complete then call the timeout callback - if (queueItem.TimeoutOperation != null) - { - queueItem.Result = queueItem.TimeoutOperation(bindingContext); + return; } - lockTaken = false; + bindingContext.BindingLock.Reset(); - bindTask - .ContinueWith((a) => bindingContext.BindingLock.Set()) - .ContinueWithOnFaulted(t => Logger.Write(TraceEventType.Error, "Binding queue threw exception " + t.Exception.ToString())); + lockTaken = true; + + // execute the binding operation + object result = null; + CancellationTokenSource cancelToken = new CancellationTokenSource(); + + // run the operation in a separate thread + var bindTask = Task.Run(() => + { + try + { + result = queueItem.BindOperation( + bindingContext, + cancelToken.Token); + } + catch (Exception ex) + { + Logger.Write(TraceEventType.Error, "Unexpected exception on the binding queue: " + ex.ToString()); + if (queueItem.ErrorHandler != null) + { + result = queueItem.ErrorHandler(ex); + } + } + }); + + Task.Run(() => + { + try + { + // check if the binding tasks completed within the binding timeout + if (bindTask.Wait(bindTimeout)) + { + queueItem.Result = result; + } + else + { + cancelToken.Cancel(); + + // if the task didn't complete then call the timeout callback + if (queueItem.TimeoutOperation != null) + { + queueItem.Result = queueItem.TimeoutOperation(bindingContext); + } + + bindTask.ContinueWithOnFaulted(t => Logger.Write(TraceEventType.Error, "Binding queue threw exception " + t.Exception.ToString())); + + // Give the task a chance to cancel before moving on to the next operation + Task.WaitAny(bindTask, Task.Delay(bindingContext.BindingTimeout)); + } + } + catch (Exception ex) + { + Logger.Write(TraceEventType.Error, "Binding queue task completion threw exception " + ex.ToString()); + } + finally + { + // set item processed to avoid deadlocks + if (lockTaken) + { + bindingContext.BindingLock.Set(); + } + queueItem.ItemProcessed.Set(); + } + }); } - } - catch (Exception ex) - { - // catch and log any exceptions raised in the binding calls - // set item processed to avoid deadlocks - Logger.Write(TraceEventType.Error, "Binding queue threw exception " + ex.ToString()); - } - finally - { - if (lockTaken) + catch (Exception ex) { - bindingContext.BindingLock.Set(); + // catch and log any exceptions raised in the binding calls + // set item processed to avoid deadlocks + Logger.Write(TraceEventType.Error, "Binding queue threw exception " + ex.ToString()); + // set item processed to avoid deadlocks + if (lockTaken) + { + bindingContext.BindingLock.Set(); + } + queueItem.ItemProcessed.Set(); } - - queueItem.ItemProcessed.Set(); - } + }, TaskContinuationOptions.None); // if a queue processing cancellation was requested then exit the loop if (token.IsCancellationRequested) diff --git a/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Nodes/ChildFactory.cs b/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Nodes/ChildFactory.cs index 8e536065..43151021 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Nodes/ChildFactory.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Nodes/ChildFactory.cs @@ -4,6 +4,7 @@ // using System.Collections.Generic; +using System.Threading; using Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel; namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes @@ -30,7 +31,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes /// force to refresh /// name of the sql object to filter /// - public abstract IEnumerable Expand(TreeNode parent, bool refresh, string name, bool includeSystemObjects); + public abstract IEnumerable Expand(TreeNode parent, bool refresh, string name, bool includeSystemObjects, CancellationToken cancellationToken); /// /// The list of filters that should be applied on the smo object list diff --git a/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Nodes/TreeNode.cs b/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Nodes/TreeNode.cs index 823791f4..d5f0dc42 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Nodes/TreeNode.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Nodes/TreeNode.cs @@ -8,6 +8,7 @@ using System.Collections.Generic; using System.Collections.ObjectModel; using System.Diagnostics; using System.Globalization; +using System.Threading; using Microsoft.SqlTools.ServiceLayer.Metadata.Contracts; using Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Contracts; using Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel; @@ -236,14 +237,14 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes /// Expands this node and returns its children /// /// Children as an IList. This is the raw children collection, not a copy - public IList Expand(string name= null) + public IList Expand(string name, CancellationToken cancellationToken) { // TODO consider why solution explorer has separate Children and Items options if (children.IsInitialized) { return children; } - PopulateChildren(false, name); + PopulateChildren(false, name, cancellationToken); return children; } @@ -251,19 +252,19 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes /// Expands this node and returns its children /// /// Children as an IList. This is the raw children collection, not a copy - public IList Expand() + public IList Expand(CancellationToken cancellationToken) { - return Expand(null); + return Expand(null, cancellationToken); } /// /// Refresh this node and returns its children /// /// Children as an IList. This is the raw children collection, not a copy - public virtual IList Refresh() + public virtual IList Refresh(CancellationToken cancellationToken) { // TODO consider why solution explorer has separate Children and Items options - PopulateChildren(true, null); + PopulateChildren(true, null, cancellationToken); return children; } @@ -315,7 +316,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes return Parent as T; } - protected virtual void PopulateChildren(bool refresh, string name = null) + protected virtual void PopulateChildren(bool refresh, string name, CancellationToken cancellationToken) { Logger.Write(TraceEventType.Verbose, string.Format(CultureInfo.InvariantCulture, "Populating oe node :{0}", this.GetNodePath())); Debug.Assert(IsAlwaysLeaf == false); @@ -337,9 +338,10 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes { foreach (var factory in childFactories) { + cancellationToken.ThrowIfCancellationRequested(); try { - IEnumerable items = factory.Expand(this, refresh, name, includeSystemObjects); + IEnumerable items = factory.Expand(this, refresh, name, includeSystemObjects, cancellationToken); if (items != null) { foreach (TreeNode item in items) diff --git a/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerService.cs b/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerService.cs index 4f59976b..c8730ea0 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerService.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerService.cs @@ -418,11 +418,11 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer { if (forceRefresh) { - nodes = node.Refresh().Select(x => x.ToNodeInfo()).ToArray(); + nodes = node.Refresh(cancelToken).Select(x => x.ToNodeInfo()).ToArray(); } else { - nodes = node.Expand().Select(x => x.ToNodeInfo()).ToArray(); + nodes = node.Expand(cancelToken).Select(x => x.ToNodeInfo()).ToArray(); } response.Nodes = nodes; response.ErrorMessage = node.ErrorMessage; diff --git a/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerUtils.cs b/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerUtils.cs index cf3c658e..30c66503 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerUtils.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerUtils.cs @@ -4,6 +4,7 @@ // using System; +using System.Threading; using Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes; using Microsoft.SqlTools.ServiceLayer.Utility; @@ -60,7 +61,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer { return node; } - var children = expandIfNeeded && !node.IsAlwaysLeaf ? node.Expand() : node.GetChildren(); + var children = expandIfNeeded && !node.IsAlwaysLeaf ? node.Expand(new CancellationToken()) : node.GetChildren(); foreach (var child in children) { if (filter != null && filter(child)) diff --git a/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/SmoModel/DatabaseTreeNode.cs b/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/SmoModel/DatabaseTreeNode.cs index ff5a782d..78d7b7ea 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/SmoModel/DatabaseTreeNode.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/SmoModel/DatabaseTreeNode.cs @@ -7,6 +7,7 @@ using System; using System.Diagnostics; using System.Globalization; +using System.Threading; using Microsoft.SqlServer.Management.Smo; using Microsoft.SqlTools.Utility; @@ -40,12 +41,12 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel } } - protected override void PopulateChildren(bool refresh, string name = null) + protected override void PopulateChildren(bool refresh, string name, CancellationToken cancellationToken) { SmoQueryContext context = this.GetContextAs(); if (IsAccessible(context)) { - base.PopulateChildren(refresh, name); + base.PopulateChildren(refresh, name, cancellationToken); } else { diff --git a/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/SmoModel/SmoChildFactoryBase.cs b/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/SmoModel/SmoChildFactoryBase.cs index 78f33922..8da3b18b 100644 --- a/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/SmoModel/SmoChildFactoryBase.cs +++ b/src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/SmoModel/SmoChildFactoryBase.cs @@ -9,6 +9,7 @@ using System.Diagnostics; using System.Globalization; using System.Linq; using System.Reflection; +using System.Threading; using Microsoft.SqlServer.Management.Smo; using Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes; using Microsoft.SqlTools.Utility; @@ -23,7 +24,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel return null; } - public override IEnumerable Expand(TreeNode parent, bool refresh, string name, bool includeSystemObjects) + public override IEnumerable Expand(TreeNode parent, bool refresh, string name, bool includeSystemObjects, CancellationToken cancellationToken) { List allChildren = new List(); @@ -31,7 +32,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel { OnExpandPopulateFoldersAndFilter(allChildren, parent, includeSystemObjects); RemoveFoldersFromInvalidSqlServerVersions(allChildren, parent); - OnExpandPopulateNonFolders(allChildren, parent, refresh, name); + OnExpandPopulateNonFolders(allChildren, parent, refresh, name, cancellationToken); OnBeginAsyncOperations(parent); } catch(Exception ex) @@ -83,7 +84,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel /// /// List to which nodes should be added /// Parent the nodes are being added to - protected virtual void OnExpandPopulateNonFolders(IList allChildren, TreeNode parent, bool refresh, string name) + protected virtual void OnExpandPopulateNonFolders(IList allChildren, TreeNode parent, bool refresh, string name, CancellationToken cancellationToken) { Logger.Write(TraceEventType.Verbose, string.Format(CultureInfo.InvariantCulture, "child factory parent :{0}", parent.GetNodePath())); @@ -115,6 +116,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel } foreach (var querier in queriers) { + cancellationToken.ThrowIfCancellationRequested(); if (!querier.IsValidFor(serverValidFor)) { continue; @@ -125,6 +127,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel var smoObjectList = querier.Query(context, propertyFilter, refresh, smoProperties).ToList(); foreach (var smoObject in smoObjectList) { + cancellationToken.ThrowIfCancellationRequested(); if (smoObject == null) { Logger.Write(TraceEventType.Error, "smoObject should not be null"); diff --git a/test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/ObjectExplorer/ObjectExplorerServiceTests.cs b/test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/ObjectExplorer/ObjectExplorerServiceTests.cs index 791c5fd6..df7cef15 100644 --- a/test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/ObjectExplorer/ObjectExplorerServiceTests.cs +++ b/test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/ObjectExplorer/ObjectExplorerServiceTests.cs @@ -8,6 +8,7 @@ using System.Globalization; using System.IO; using System.Linq; using System.Text; +using System.Threading; using System.Threading.Tasks; using Microsoft.SqlTools.ServiceLayer.Connection.Contracts; using Microsoft.SqlTools.ServiceLayer.ObjectExplorer; @@ -303,7 +304,7 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.ObjectExplorer if (serverNode) { Assert.Equal(nodeInfo.NodeType, NodeTypes.Server.ToString()); - var children = session.Root.Expand(); + var children = session.Root.Expand(new CancellationToken()); //All server children should be folder nodes foreach (var item in children) diff --git a/test/Microsoft.SqlTools.ServiceLayer.UnitTests/LanguageServer/BindingQueueTests.cs b/test/Microsoft.SqlTools.ServiceLayer.UnitTests/LanguageServer/BindingQueueTests.cs index 38f17df9..14aeeb2d 100644 --- a/test/Microsoft.SqlTools.ServiceLayer.UnitTests/LanguageServer/BindingQueueTests.cs +++ b/test/Microsoft.SqlTools.ServiceLayer.UnitTests/LanguageServer/BindingQueueTests.cs @@ -141,7 +141,6 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.LanguageServer public void QueueWithUnhandledExceptionTest() { InitializeTestSettings(); - ManualResetEvent mre = new ManualResetEvent(false); bool isExceptionHandled = false; object defaultReturnObject = new object(); var queueItem = this.bindingQueue.QueueBindingOperation( @@ -150,11 +149,10 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.LanguageServer timeoutOperation: TestTimeoutOperation, errorHandler: (exception) => { isExceptionHandled = true; - mre.Set(); return defaultReturnObject; }); - mre.WaitOne(10000); + queueItem.ItemProcessed.WaitOne(10000); this.bindingQueue.StopQueueProcessor(15000); @@ -213,5 +211,58 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.LanguageServer Assert.Equal(1, this.timeoutCallCount); Assert.True(this.isCancelationRequested); } + + /// + /// Queue an task with a long operation causing a timeout and make sure subsequent tasks still execute + /// + [Fact] + public void QueueWithTimeoutRunsNextTask() + { + string operationKey = "testkey"; + ManualResetEvent firstEventExecuted = new ManualResetEvent(false); + ManualResetEvent secondEventExecuted = new ManualResetEvent(false); + bool firstOperationCanceled = false; + bool secondOperationExecuted = false; + InitializeTestSettings(); + + this.bindCallbackDelay = 1000; + var totalTimeout = (this.bindCallbackDelay + this.bindingContext.BindingTimeout) * 2; + + this.bindingQueue.QueueBindingOperation( + key: operationKey, + bindingTimeout: bindCallbackDelay / 2, + bindOperation: (bindingContext, cancellationToken) => + { + secondEventExecuted.WaitOne(); + if (cancellationToken.IsCancellationRequested) + { + firstOperationCanceled = true; + } + firstEventExecuted.Set(); + return null; + }, + timeoutOperation: TestTimeoutOperation); + + this.bindingQueue.QueueBindingOperation( + key: operationKey, + bindingTimeout: bindCallbackDelay, + bindOperation: (bindingContext, cancellationToken) => + { + secondOperationExecuted = true; + secondEventExecuted.Set(); + return null; + }, + waitForLockTimeout: totalTimeout + ); + + var result = firstEventExecuted.WaitOne(totalTimeout); + Assert.True(result); + + this.bindingQueue.StopQueueProcessor(15000); + + Assert.Equal(1, this.timeoutCallCount); + Assert.True(firstOperationCanceled); + Assert.True(secondOperationExecuted); + } } } diff --git a/test/Microsoft.SqlTools.ServiceLayer.UnitTests/ObjectExplorer/NodeTests.cs b/test/Microsoft.SqlTools.ServiceLayer.UnitTests/ObjectExplorer/NodeTests.cs index aef8bad8..8f139844 100644 --- a/test/Microsoft.SqlTools.ServiceLayer.UnitTests/ObjectExplorer/NodeTests.cs +++ b/test/Microsoft.SqlTools.ServiceLayer.UnitTests/ObjectExplorer/NodeTests.cs @@ -7,6 +7,7 @@ using System; using System.Collections.Generic; using System.Data.SqlClient; using System.Globalization; +using System.Threading; using Microsoft.SqlServer.Management.Common; using Microsoft.SqlServer.Management.Smo; using Microsoft.SqlTools.Extensibility; @@ -400,7 +401,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer ServerNode node = SetupServerNodeWithServer(smoServer); // When I populate its children - IList children = node.Expand(); + IList children = node.Expand(new CancellationToken()); // Then I expect it to contain server-level folders Assert.Equal(3, children.Count); @@ -409,7 +410,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer VerifyTreeNode(children[2], "Folder", SR.SchemaHierarchy_ServerObjects); // And the database is contained under it TreeNode databases = children[0]; - IList dbChildren = databases.Expand(); + IList dbChildren = databases.Expand(new CancellationToken()); Assert.Equal(2, dbChildren.Count); Assert.Equal(SR.SchemaHierarchy_SystemDatabases, dbChildren[0].NodeValue); diff --git a/test/Microsoft.SqlTools.ServiceLayer.UnitTests/ObjectExplorer/ObjectExplorerServiceTests.cs b/test/Microsoft.SqlTools.ServiceLayer.UnitTests/ObjectExplorer/ObjectExplorerServiceTests.cs index 1974df80..49f6ed4c 100644 --- a/test/Microsoft.SqlTools.ServiceLayer.UnitTests/ObjectExplorer/ObjectExplorerServiceTests.cs +++ b/test/Microsoft.SqlTools.ServiceLayer.UnitTests/ObjectExplorer/ObjectExplorerServiceTests.cs @@ -6,6 +6,7 @@ using System; using System.Data.SqlClient; using System.Linq; +using System.Threading; using System.Threading.Tasks; using Microsoft.SqlTools.Hosting.Protocol; using Microsoft.SqlTools.ServiceLayer.Connection; @@ -52,6 +53,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer connectedBindingContext.ServerConnection = new ServerConnection(new SqlConnection(fakeConnectionString)); connectedBindingQueue = new ConnectedBindingQueue(false); connectedBindingQueue.BindingContextMap.Add($"{details.ServerName}_{details.DatabaseName}_{details.UserName}_NULL", connectedBindingContext); + connectedBindingQueue.BindingContextTasks.Add(connectedBindingContext, Task.Run(() => null)); mockConnectionOpener = new Mock(); connectedBindingQueue.SetConnectionOpener(mockConnectionOpener.Object); service.ConnectedBindingQueue = connectedBindingQueue; @@ -271,7 +273,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer public void FindNodeCanExpandParentNodes() { var mockTreeNode = new Mock(); - object[] populateChildrenArguments = { ItExpr.Is(x => x == false), ItExpr.IsNull() }; + object[] populateChildrenArguments = { ItExpr.Is(x => x == false), ItExpr.IsNull(), new CancellationToken() }; mockTreeNode.Protected().Setup("PopulateChildren", populateChildrenArguments); mockTreeNode.Object.IsAlwaysLeaf = false;