Fix OE and Binding Queue reliability bugs (#702)

src/Microsoft.SqlTools.ServiceLayer/LanguageServices/BindingQueue.cs

@@ -39,6 +39,8 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
         /// </summary>
         internal Dictionary<string, IBindingContext> BindingContextMap { get; set; }
 
+        internal Dictionary<IBindingContext, Task> BindingContextTasks { get; set; } = new Dictionary<IBindingContext, Task>();
+
         /// <summary>
         /// Constructor for a binding queue instance
         /// </summary>
@@ -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;
+                        var bindingContextTask = this.BindingContextTasks[bindingContext];
-                        try
-                        {                                                    
-                            // prefer the queue item binding item, otherwise use the context default timeout
-                            int bindTimeout = queueItem.BindingTimeout ?? bindingContext.BindingTimeout;
 
-                            // handle the case a previous binding operation is still running                                                 
+                        // Run in the binding context task in case this task has to wait for a previous binding operation
-                            if (!bindingContext.BindingLock.WaitOne(queueItem.WaitForLockTimeout ?? 0))
+                        this.BindingContextTasks[bindingContext] = bindingContextTask.ContinueWith((task) =>
-                            {
+                        {
-                                queueItem.Result = queueItem.TimeoutOperation != null
+                            bool lockTaken = false;
-                                    ? queueItem.TimeoutOperation(bindingContext)
+                            try
-                                    : null;
+                            {                                                    
+                                // prefer the queue item binding item, otherwise use the context default timeout
+                                int bindTimeout = queueItem.BindingTimeout ?? bindingContext.BindingTimeout;
 
-                                continue;
+                                // handle the case a previous binding operation is still running
-                            }
+                                if (!bindingContext.BindingLock.WaitOne(queueItem.WaitForLockTimeout ?? 0))
-
-                            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
                                 {
-                                    result = queueItem.BindOperation(
+                                    try
-                                        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);
+                                        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
+                                    return;
-                                if (queueItem.TimeoutOperation != null)
-                                {                                    
-                                    queueItem.Result = queueItem.TimeoutOperation(bindingContext);                              
                                 }
 
-                                lockTaken = false;
+                                bindingContext.BindingLock.Reset();  
 
-                                bindTask
+                                lockTaken = true;
-                                    .ContinueWith((a) => bindingContext.BindingLock.Set())
+
-                                    .ContinueWithOnFaulted(t => Logger.Write(TraceEventType.Error, "Binding queue threw exception " + t.Exception.ToString()));
+                                // 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 (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)
                             {
-                                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();                          
                             }
-
+                        }, TaskContinuationOptions.None);
-                            queueItem.ItemProcessed.Set();
-                        }
 
                         // if a queue processing cancellation was requested then exit the loop
                         if (token.IsCancellationRequested)

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
         /// <param name="refresh">force to refresh</param>
         /// <param name="refresh">name of the sql object to filter</param>
         /// <returns></returns>
-        public abstract IEnumerable<TreeNode> Expand(TreeNode parent, bool refresh, string name, bool includeSystemObjects);
+        public abstract IEnumerable<TreeNode> Expand(TreeNode parent, bool refresh, string name, bool includeSystemObjects, CancellationToken cancellationToken);
 
         /// <summary>
         /// The list of filters that should be applied on the smo object list

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
         /// </summary>
         /// <returns>Children as an IList. This is the raw children collection, not a copy</returns>
-        public IList<TreeNode> Expand(string name= null)
+        public IList<TreeNode> 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
         /// </summary>
         /// <returns>Children as an IList. This is the raw children collection, not a copy</returns>
-        public IList<TreeNode> Expand()
+        public IList<TreeNode> Expand(CancellationToken cancellationToken)
         {
-            return Expand(null);
+            return Expand(null, cancellationToken);
         }
 
         /// <summary>
         /// Refresh this node and returns its children
         /// </summary>
         /// <returns>Children as an IList. This is the raw children collection, not a copy</returns>
-        public virtual IList<TreeNode> Refresh()
+        public virtual IList<TreeNode> 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<TreeNode> items = factory.Expand(this, refresh, name, includeSystemObjects);
+                            IEnumerable<TreeNode> items = factory.Expand(this, refresh, name, includeSystemObjects, cancellationToken);
                             if (items != null)
                             {
                                 foreach (TreeNode item in items)

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;

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))

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<SmoQueryContext>();
             if (IsAccessible(context))
             {
-                base.PopulateChildren(refresh, name);
+                base.PopulateChildren(refresh, name, cancellationToken);
             }
             else
             {

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<TreeNode> Expand(TreeNode parent, bool refresh, string name, bool includeSystemObjects)
+        public override IEnumerable<TreeNode> Expand(TreeNode parent, bool refresh, string name, bool includeSystemObjects, CancellationToken cancellationToken)
         {
             List<TreeNode> allChildren = new List<TreeNode>();
 
@@ -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
         /// </summary>
         /// <param name="allChildren">List to which nodes should be added</param>
         /// <param name="parent">Parent the nodes are being added to</param>
-        protected virtual void OnExpandPopulateNonFolders(IList<TreeNode> allChildren, TreeNode parent, bool refresh, string name)
+        protected virtual void OnExpandPopulateNonFolders(IList<TreeNode> 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");

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)

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);
         }
+
+        /// <summary>
+        /// Queue an task with a long operation causing a timeout and make sure subsequent tasks still execute
+        /// </summary>
+        [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);
+        }
     }
 }

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<TreeNode> children = node.Expand();
+            IList<TreeNode> 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<FolderNode>(children[2], "Folder", SR.SchemaHierarchy_ServerObjects);
             // And the database is contained under it
             TreeNode databases = children[0];
-            IList<TreeNode> dbChildren = databases.Expand();
+            IList<TreeNode> dbChildren = databases.Expand(new CancellationToken());
             Assert.Equal(2, dbChildren.Count);
             Assert.Equal(SR.SchemaHierarchy_SystemDatabases, dbChildren[0].NodeValue);
 

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<SqlConnectionOpener>();
             connectedBindingQueue.SetConnectionOpener(mockConnectionOpener.Object);
             service.ConnectedBindingQueue = connectedBindingQueue;
@@ -271,7 +273,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer
         public void FindNodeCanExpandParentNodes()
         {
             var mockTreeNode = new Mock<TreeNode>();
-            object[] populateChildrenArguments = { ItExpr.Is<bool>(x => x == false), ItExpr.IsNull<string>() };
+            object[] populateChildrenArguments = { ItExpr.Is<bool>(x => x == false), ItExpr.IsNull<string>(), new CancellationToken() };
             mockTreeNode.Protected().Setup("PopulateChildren", populateChildrenArguments);
             mockTreeNode.Object.IsAlwaysLeaf = false;