Fix filebrowser to handle cancellation and multiple open/expand requests (#499)

* add cancel filebrowser request

* update expand request contract

* cleanup

* add queue

* add expand queue and update cancellation

* dispose cancelsource

* remove unnecessary field

* cleanup using directives

* address pr comment

* more changes

* change contract

src/Microsoft.SqlTools.ServiceLayer/FileBrowser/Contracts/FileBrowserOpenRequest.cs

@@ -26,6 +26,11 @@ namespace Microsoft.SqlTools.ServiceLayer.FileBrowser.Contracts
         /// File extension filter (e.g. *.bak)
         /// </summary>
         public string[] FileFilters;
+
+        /// <summary>
+        /// True if this is a request to change file filter
+        /// </summary>
+        public bool ChangeFilter;
     }
 
     /// <summary>

src/Microsoft.SqlTools.ServiceLayer/FileBrowser/FileBrowserOperation.cs

@@ -8,6 +8,7 @@ using System.Collections.Generic;
 using System.Data.SqlClient;
 using System.Globalization;
 using System.Text.RegularExpressions;
+using System.Threading;
 using Microsoft.SqlTools.ServiceLayer.FileBrowser.Contracts;
 
 namespace Microsoft.SqlTools.ServiceLayer.FileBrowser
@@ -15,11 +16,14 @@ namespace Microsoft.SqlTools.ServiceLayer.FileBrowser
     /// <summary>
     /// Implementation for file browser operation
     /// </summary>
-    internal class FileBrowserOperation : FileBrowserBase
+    internal class FileBrowserOperation : FileBrowserBase, IDisposable
     {
         private FileTree fileTree;
         private string expandPath;
         private string[] fileFilters;
+        private bool fileTreeCreated;
+        private CancellationTokenSource cancelSource;
+        private CancellationToken cancelToken;
 
         #region Constructors
 
@@ -29,6 +33,8 @@ namespace Microsoft.SqlTools.ServiceLayer.FileBrowser
         public FileBrowserOperation()
         {
             this.fileTree = new FileTree();
+            this.cancelSource = new CancellationTokenSource();
+            this.cancelToken = cancelSource.Token;
         }
 
         /// <summary>
@@ -39,15 +45,7 @@ namespace Microsoft.SqlTools.ServiceLayer.FileBrowser
         public FileBrowserOperation(SqlConnection connectionInfo, string expandPath, string[] fileFilters = null): this()
         {
             this.sqlConnection = connectionInfo;
-            this.expandPath = expandPath;
+            this.Initialize(expandPath, fileFilters);
-            if (fileFilters == null)
-            {
-                this.fileFilters = new string[1] { "*" };
-            }
-            else
-            {
-                this.fileFilters = fileFilters;
-            }
         }
 
         #endregion
@@ -70,13 +68,65 @@ namespace Microsoft.SqlTools.ServiceLayer.FileBrowser
             }
         }
 
+        public bool FileTreeCreated
+        {
+            get
+            {
+                return this.fileTreeCreated;
+            }
+        }
+
+        public SqlConnection SqlConnection
+        {
+            get
+            {
+                return this.sqlConnection;
+            }
+        }
+
+        public bool IsCancellationRequested
+        {
+            get
+            {
+                return this.cancelToken.IsCancellationRequested;
+            }
+        }
+
+        public void Cancel()
+        {
+            this.cancelSource.Cancel();
+        }
         #endregion
 
+        public void Initialize(string expandPath, string[] fileFilters)
+        {
+            this.expandPath = expandPath;
+            if (fileFilters == null)
+            {
+                this.fileFilters = new string[1] { "*" };
+            }
+            else
+            {
+                this.fileFilters = fileFilters;
+            }
+        }
+
+        public void Dispose()
+        {
+            if (this.sqlConnection != null)
+            {
+                this.sqlConnection.Close();
+            }
+            this.cancelSource.Dispose();
+        }
+
         public void PopulateFileTree()
         {
+            this.fileTreeCreated = false;
             this.PathSeparator = GetPathSeparator(this.Enumerator, this.sqlConnection);
             PopulateDrives();
             ExpandSelectedNode(this.expandPath);
+            this.fileTreeCreated = true;
         }
 
         /// <summary>
@@ -94,6 +144,11 @@ namespace Microsoft.SqlTools.ServiceLayer.FileBrowser
 
                 foreach (string dir in dirs)
                 {
+                    if (cancelToken.IsCancellationRequested)
+                    {
+                        break;
+                    }
+
                     FileTreeNode currentNode = null;
                     foreach (FileTreeNode node in currentChildren)
                     {
@@ -131,10 +186,42 @@ namespace Microsoft.SqlTools.ServiceLayer.FileBrowser
             }
         }
 
-        public List<FileTreeNode> GetChildren(string filepath)
+        public void PopulateDrives()
+        {
+            bool first = true;
+            if (!cancelToken.IsCancellationRequested)
+            {
+                foreach (var fileInfo in EnumerateDrives(Enumerator, sqlConnection))
+                {
+                    if (cancelToken.IsCancellationRequested)
+                    {
+                        break;
+                    }
+
+                    // Windows drive letter paths have a '\' at the end. Linux drive paths won't have a '\'.
+                    var node = new FileTreeNode
+                    {
+                        Name = Convert.ToString(fileInfo.path, CultureInfo.InvariantCulture).TrimEnd('\\'),
+                        FullPath = fileInfo.path
+                    };
+
+                    this.fileTree.RootNode.AddChildNode(node);
+
+                    if (first)
+                    {
+                        this.fileTree.SelectedNode = node;
+                        first = false;
+                    }
+
+                    node.Children = this.GetChildren(node.FullPath);
+                }
+            }
+        }
+
+        public List<FileTreeNode> GetChildren(string filePath)
         { 
             List<FileTreeNode> children = new List<FileTreeNode>();
-            foreach (var file in EnumerateFilesInFolder(Enumerator, sqlConnection, filepath))
+            foreach (var file in EnumerateFilesInFolder(Enumerator, sqlConnection, filePath))
             {
                 bool isFile = !string.IsNullOrEmpty(file.fileName);
                 FileTreeNode treeNode = new FileTreeNode();
@@ -157,34 +244,9 @@ namespace Microsoft.SqlTools.ServiceLayer.FileBrowser
                     children.Add(treeNode);
                 }
             }
-
             return children;
         }
 
-        public void PopulateDrives()
-        {
-            bool first = true;
-            foreach (var fileInfo in EnumerateDrives(Enumerator, sqlConnection))
-            {
-                // Windows drive letter paths have a '\' at the end. Linux drive paths won't have a '\'.
-                var node = new FileTreeNode
-                {
-                    Name = Convert.ToString(fileInfo.path, CultureInfo.InvariantCulture).TrimEnd('\\'),
-                    FullPath = fileInfo.path
-                };
-
-                this.fileTree.RootNode.AddChildNode(node);
-
-                if (first)
-                {
-                    this.fileTree.SelectedNode = node;
-                    first = false;
-                }
-
-                node.Children = this.GetChildren(node.FullPath);
-            }
-        }
-
         /// <summary>
         /// Filter a filename based on the full mask provide.  The full mask may be a collection a masks seperated by semi-colons.
         /// For example: *; *.txt

src/Microsoft.SqlTools.ServiceLayer/FileBrowser/FileBrowserService.cs

@@ -5,14 +5,13 @@
 
 using System;
 using System.Collections.Concurrent;
-using System.Data.Common;
 using System.Data.SqlClient;
 using System.Threading.Tasks;
 using Microsoft.SqlTools.Hosting.Protocol;
 using Microsoft.SqlTools.ServiceLayer.Connection;
-using Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection;
 using Microsoft.SqlTools.ServiceLayer.FileBrowser.Contracts;
 using Microsoft.SqlTools.ServiceLayer.Hosting;
+using Microsoft.SqlTools.ServiceLayer.LanguageServices;
 using Microsoft.SqlTools.ServiceLayer.Utility;
 
 namespace Microsoft.SqlTools.ServiceLayer.FileBrowser
@@ -20,7 +19,7 @@ namespace Microsoft.SqlTools.ServiceLayer.FileBrowser
     /// <summary>
     /// Main class for file browser service
     /// </summary>
-    public sealed class FileBrowserService
+    public sealed class FileBrowserService: IDisposable
     {
         private static readonly Lazy<FileBrowserService> LazyInstance = new Lazy<FileBrowserService>(() => new FileBrowserService());
         public static FileBrowserService Instance => LazyInstance.Value;
@@ -29,6 +28,9 @@ namespace Microsoft.SqlTools.ServiceLayer.FileBrowser
         private readonly ConcurrentDictionary<string, FileBrowserOperation> ownerToFileBrowserMap = new ConcurrentDictionary<string, FileBrowserOperation>();
         private readonly ConcurrentDictionary<string, ValidatePathsCallback> validatePathsCallbackMap = new ConcurrentDictionary<string, ValidatePathsCallback>();
         private ConnectionService connectionService;
+        private ConnectedBindingQueue expandNodeQueue = new ConnectedBindingQueue(needsMetadata: false);
+        private static int DefaultExpandTimeout = 120000;
+        private string serviceName = "FileBrowser";
 
         /// <summary>
         /// Signature for callback method that validates the selected file paths
@@ -43,6 +45,7 @@ namespace Microsoft.SqlTools.ServiceLayer.FileBrowser
                 if (connectionService == null)
                 {
                     connectionService = ConnectionService.Instance;
+                    connectionService.RegisterConnectedQueue(this.serviceName, this.expandNodeQueue);
                 }
                 return connectionService;
             }
@@ -133,54 +136,98 @@ namespace Microsoft.SqlTools.ServiceLayer.FileBrowser
             FileBrowserOperation removedOperation;
             response.Succeeded = ownerToFileBrowserMap.TryRemove(fileBrowserParams.OwnerUri, out removedOperation);
 
+            if (removedOperation != null && this.expandNodeQueue != null)
+            {
+                bool hasPendingQueueItems = this.expandNodeQueue.HasPendingQueueItems;
+                if (removedOperation.FileTreeCreated && !hasPendingQueueItems)
+                {
+                    removedOperation.Dispose();
+                    this.expandNodeQueue.CloseConnections(removedOperation.SqlConnection.DataSource, removedOperation.SqlConnection.Database, DefaultExpandTimeout);
+                }
+                else if (!removedOperation.FileTreeCreated)
+                {
+                    removedOperation.Cancel();
+                }
+                else if (hasPendingQueueItems)
+                { 
+                    this.expandNodeQueue.StopQueueProcessor(DefaultExpandTimeout);
+                }
+            }
+
             await requestContext.SendResult(response);
         }
 
         #endregion
 
+        public void Dispose()
+        {
+            this.expandNodeQueue.Dispose();
+        }
+
         internal async Task RunFileBrowserOpenTask(FileBrowserOpenParams fileBrowserParams, RequestContext<bool> requestContext)
         {
             FileBrowserOpenedParams result = new FileBrowserOpenedParams();
+            SqlConnection conn = null;
+            FileBrowserOperation browser = null;
+            bool isCancelRequested = false;
+
+            if (this.expandNodeQueue.IsCancelRequested)
+            {
+                this.expandNodeQueue.StartQueueProcessor();
+            }
 
             try
             {
-                ConnectionInfo connInfo;
+                if (!fileBrowserParams.ChangeFilter)
-                this.ConnectionServiceInstance.TryFindConnection(fileBrowserParams.OwnerUri, out connInfo);
-                SqlConnection conn = null;
-
-                if (connInfo != null)
                 {
-                    DbConnection dbConn;
+                    ConnectionInfo connInfo;
-                    connInfo.TryGetConnection(ConnectionType.Default, out dbConn);
+                    this.ConnectionServiceInstance.TryFindConnection(fileBrowserParams.OwnerUri, out connInfo);
-                    if (dbConn != null)
+                    if (connInfo != null)
                     {
-                        conn = ReliableConnectionHelper.GetAsSqlConnection(dbConn);
+                        // Open new connection for each Open request
+                        conn = ConnectionService.OpenSqlConnection(connInfo, this.serviceName);
+                        browser = new FileBrowserOperation(conn, fileBrowserParams.ExpandPath, fileBrowserParams.FileFilters);
                     }
                 }
-
-                if (conn != null)
-                {
-                    FileBrowserOperation browser = new FileBrowserOperation(conn, fileBrowserParams.ExpandPath, fileBrowserParams.FileFilters);
-                    browser.PopulateFileTree();
-
-                    ownerToFileBrowserMap.AddOrUpdate(fileBrowserParams.OwnerUri, browser, (key, value) => browser);
-
-                    result.OwnerUri = fileBrowserParams.OwnerUri;
-                    result.FileTree = browser.FileTree;
-                    result.Succeeded = true;
-                }
                 else
                 {
-                    result.Succeeded = false;
+                    ownerToFileBrowserMap.TryGetValue(fileBrowserParams.OwnerUri, out browser);
+                    if (browser != null)
+                    {
+                        browser.Initialize(fileBrowserParams.ExpandPath, fileBrowserParams.FileFilters);
+                    }
+                }
+
+                if (browser != null)
+                {
+                    ownerToFileBrowserMap.AddOrUpdate(fileBrowserParams.OwnerUri, browser, (key, value) => browser);
+
+                    // Create file browser tree
+                    browser.PopulateFileTree();
+
+                    // Check if cancel was requested
+                    if (browser.IsCancellationRequested)
+                    {
+                        browser.Dispose();
+                        isCancelRequested = true;
+                    }
+                    else
+                    {
+                        result.OwnerUri = fileBrowserParams.OwnerUri;
+                        result.FileTree = browser.FileTree;
+                        result.Succeeded = true;
+                    }
                 }
             }
             catch (Exception ex)
             {
-                result.Succeeded = false;
                 result.Message = ex.Message;
             }
 
-            await requestContext.SendEvent(FileBrowserOpenedNotification.Type, result);
+            if (!isCancelRequested)
+            {
+                await requestContext.SendEvent(FileBrowserOpenedNotification.Type, result);
+            }
         }
 
         internal async Task RunFileBrowserExpandTask(FileBrowserExpandParams fileBrowserParams, RequestContext<bool> requestContext)
@@ -188,13 +235,35 @@ namespace Microsoft.SqlTools.ServiceLayer.FileBrowser
             FileBrowserExpandedParams result = new FileBrowserExpandedParams();
             try
             {
-                FileBrowserOperation browser;
+                FileBrowserOperation operation;
-                result.Succeeded = ownerToFileBrowserMap.TryGetValue(fileBrowserParams.OwnerUri, out browser);
+                ConnectionInfo connInfo;
-                if (result.Succeeded && browser != null)
+                result.Succeeded = ownerToFileBrowserMap.TryGetValue(fileBrowserParams.OwnerUri, out operation);
+                this.ConnectionServiceInstance.TryFindConnection(fileBrowserParams.OwnerUri, out connInfo);
+
+                if (result.Succeeded && operation != null && connInfo != null)
                 {
-                    result.Children = browser.GetChildren(fileBrowserParams.ExpandPath).ToArray();
+                    QueueItem queueItem = expandNodeQueue.QueueBindingOperation(
-                    result.ExpandPath = fileBrowserParams.ExpandPath;
+                        key: expandNodeQueue.AddConnectionContext(connInfo, this.serviceName),
-                    result.OwnerUri = fileBrowserParams.OwnerUri;
+                        bindingTimeout: DefaultExpandTimeout,
+                        waitForLockTimeout: DefaultExpandTimeout,
+                        bindOperation: (bindingContext, cancelToken) =>
+                        {
+                            result.ExpandPath = fileBrowserParams.ExpandPath;
+                            result.Children = operation.GetChildren(fileBrowserParams.ExpandPath).ToArray();
+                            result.OwnerUri = fileBrowserParams.OwnerUri;
+                            return result;
+                        });
+
+                    queueItem.ItemProcessed.WaitOne();
+
+                    if (this.expandNodeQueue.IsCancelRequested)
+                    {
+                        this.expandNodeQueue.CloseConnections(operation.SqlConnection.DataSource, operation.SqlConnection.Database, DefaultExpandTimeout);
+                    }
+                    else if (queueItem.GetResultAsT<FileBrowserExpandedParams>() != null)
+                    {
+                        result = queueItem.GetResultAsT<FileBrowserExpandedParams>();
+                    }
                 }
             }
             catch (Exception ex)
@@ -238,7 +307,6 @@ namespace Microsoft.SqlTools.ServiceLayer.FileBrowser
             }
             catch (Exception ex)
             {
-                result.Succeeded = false;
                 result.Message = ex.Message;
             }
 

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

@@ -20,7 +20,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
     {
         internal const int QueueThreadStackSize = 5 * 1024 * 1024;
 
-        private CancellationTokenSource processQueueCancelToken = new CancellationTokenSource();
+        private CancellationTokenSource processQueueCancelToken = null;
 
         private ManualResetEvent itemQueuedEvent = new ManualResetEvent(initialState: false);
 
@@ -44,8 +44,12 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
         public BindingQueue()
         {
             this.BindingContextMap = new Dictionary<string, IBindingContext>();
+            this.StartQueueProcessor();
+        }
 
-            this.queueProcessorTask = StartQueueProcessor();
+        public void StartQueueProcessor()
+        {
+            this.queueProcessorTask = StartQueueProcessorAsync();
         }
 
         /// <summary>
@@ -58,6 +62,18 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             return this.queueProcessorTask.Wait(timeout);
         }
 
+        /// <summary>
+        /// Returns true if cancellation is requested
+        /// </summary>
+        /// <returns></returns>
+        public bool IsCancelRequested
+        {
+            get
+            {
+                return this.processQueueCancelToken.IsCancellationRequested;
+            }
+        }
+
         /// <summary>
         /// Queue a binding request item
         /// </summary>
@@ -182,7 +198,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             }
         }
 
-        private bool HasPendingQueueItems
+        public bool HasPendingQueueItems
         {
             get
             {
@@ -214,8 +230,14 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
         /// <summary>
         /// Starts the queue processing thread
         /// </summary>        
-        private Task StartQueueProcessor()
+        private Task StartQueueProcessorAsync()
         {
+            if (this.processQueueCancelToken != null)
+            {
+                this.processQueueCancelToken.Dispose();
+            }
+            this.processQueueCancelToken = new CancellationTokenSource();
+
             return Task.Factory.StartNew(
                 ProcessQueue,
                 this.processQueueCancelToken.Token,
@@ -368,6 +390,11 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
 
         public void Dispose()
         {
+            if (this.processQueueCancelToken != null)
+            {
+                this.processQueueCancelToken.Dispose();
+            }
+
             if (itemQueuedEvent != null)
             {
                 itemQueuedEvent.Dispose();