Block Initializing Edit Sessions When In Progress (#265)

* Implementation of a wait handle for initialize

* WIP

* Adding more initialize unit tests

src/Microsoft.SqlTools.ServiceLayer/EditData/EditDataService.cs

@@ -57,6 +57,10 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
         private readonly Lazy<ConcurrentDictionary<string, EditSession>> editSessions = new Lazy<ConcurrentDictionary<string, EditSession>>(
             () => new ConcurrentDictionary<string, EditSession>());
 
+        private readonly Lazy<ConcurrentDictionary<string, TaskCompletionSource<bool>>> initializeWaitHandles =
+            new Lazy<ConcurrentDictionary<string, TaskCompletionSource<bool>>>(
+                () => new ConcurrentDictionary<string, TaskCompletionSource<bool>>());
+
         #endregion
 
         #region Properties
@@ -66,6 +70,12 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
         /// </summary>
         internal ConcurrentDictionary<string, EditSession> ActiveSessions => editSessions.Value;
 
+        /// <summary>
+        /// Dictionary mapping OwnerURIs to wait handlers for initialize tasks. Pretty much only
+        /// provided for unit test scenarios.
+        /// </summary>
+        internal ConcurrentDictionary<string, TaskCompletionSource<bool>> InitializeWaitHandles => initializeWaitHandles.Value;
+
         #endregion
 
         /// <summary>
@@ -160,7 +170,14 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
                 // Make sure we have info to process this request
                 Validate.IsNotNullOrWhitespaceString(nameof(initParams.OwnerUri), initParams.OwnerUri);
                 Validate.IsNotNullOrWhitespaceString(nameof(initParams.ObjectName), initParams.ObjectName);
+                Validate.IsNotNullOrWhitespaceString(nameof(initParams.ObjectType), initParams.ObjectType);
 
+                // Try to add a new wait handler to the 
+                if (!InitializeWaitHandles.TryAdd(initParams.OwnerUri, new TaskCompletionSource<bool>()))
+                {
+                    throw new InvalidOperationException(SR.EditDataInitializeInProgress);
+                }
+    
                 // Setup a callback for when the query has successfully created
                 Func<Query, Task<bool>> queryCreateSuccessCallback = async query =>
                 {
@@ -169,21 +186,26 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
                 };
 
                 // Setup a callback for when the query failed to be created
-                Func<string, Task> queryCreateFailureCallback = requestContext.SendError;
+                Func<string, Task> queryCreateFailureCallback = async message =>
+                {
+                    await requestContext.SendError(message);
+                    CompleteInitializeWaitHandler(initParams.OwnerUri, false);
+                };
 
                 // Setup a callback for when the query completes execution successfully
                 Query.QueryAsyncEventHandler queryCompleteSuccessCallback =
                     q => QueryCompleteCallback(q, initParams, requestContext);
 
                 // Setup a callback for when the query completes execution with failure
-                Query.QueryAsyncEventHandler queryCompleteFailureCallback = query =>
+                Query.QueryAsyncEventHandler queryCompleteFailureCallback = async query =>
                 {
                     EditSessionReadyParams readyParams = new EditSessionReadyParams
                     {
                         OwnerUri = initParams.OwnerUri,
                         Success = false
                     };
-                    return requestContext.SendEvent(EditSessionReadyEvent.Type, readyParams);
+                    await requestContext.SendEvent(EditSessionReadyEvent.Type, readyParams);
+                    CompleteInitializeWaitHandler(initParams.OwnerUri, false);
                 };
 
                 // Put together a query for the results and execute it
@@ -199,6 +221,7 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
             catch (Exception e)
             {
                 await requestContext.SendError(e.Message);
+                CompleteInitializeWaitHandler(initParams.OwnerUri, false);
             }
         }
 
@@ -304,6 +327,17 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
 
             // Send the edit session ready notification
             await requestContext.SendEvent(EditSessionReadyEvent.Type, readyParams);
+            CompleteInitializeWaitHandler(initParams.OwnerUri, true);
+        }
+
+        private void CompleteInitializeWaitHandler(string ownerUri, bool result)
+        {
+            // If there isn't a wait handler, just ignore it
+            TaskCompletionSource<bool> initializeWaiter;
+            if (ownerUri != null && InitializeWaitHandles.TryRemove(ownerUri, out initializeWaiter))
+            {
+                initializeWaiter.SetResult(result);
+            }
         }
 
         #endregion

src/Microsoft.SqlTools.ServiceLayer/Localization/sr.cs

@@ -501,6 +501,14 @@ namespace Microsoft.SqlTools.ServiceLayer
             }
         }
 
+        public static string EditDataInitializeInProgress
+        {
+            get
+            {
+                return Keys.GetString(Keys.EditDataInitializeInProgress);
+            }
+        }
+
         public static string EE_BatchSqlMessageNoProcedureInfo
         {
             get
@@ -1046,6 +1054,9 @@ namespace Microsoft.SqlTools.ServiceLayer
             public const string EditDataCommitInProgress = "EditDataCommitInProgress";
 
 
+            public const string EditDataInitializeInProgress = "EditDataInitializeInProgress";
+
+
             public const string EE_BatchSqlMessageNoProcedureInfo = "EE_BatchSqlMessageNoProcedureInfo";
 
 

src/Microsoft.SqlTools.ServiceLayer/Localization/sr.resx

@@ -423,6 +423,10 @@
     <value>A commit task is in progress. Please wait for completion.</value>
     <comment></comment>
   </data>
+  <data name="EditDataInitializeInProgress" xml:space="preserve">
+    <value>Another edit data initialize is in progress for this owner URI. Please wait for completion.</value>
+    <comment></comment>
+  </data>
   <data name="EE_BatchSqlMessageNoProcedureInfo" xml:space="preserve">
     <value>Msg {0}, Level {1}, State {2}, Line {3}</value>
     <comment></comment>

src/Microsoft.SqlTools.ServiceLayer/Localization/sr.strings

@@ -200,6 +200,8 @@ EditDataScriptFilePathNull = An output filename must be provided
 
 EditDataCommitInProgress = A commit task is in progress. Please wait for completion.
 
+EditDataInitializeInProgress = Another edit data initialize is in progress for this owner URI. Please wait for completion.
+
 ############################################################################
 # DacFx Resources
 

test/Microsoft.SqlTools.ServiceLayer.UnitTests/EditData/ServiceIntegrationTests.cs

@@ -215,6 +215,142 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             edit.Verify(e => e.SetCell(It.IsAny<int>(), It.IsAny<string>()), Times.Once);
         }
 
+        [Theory]
+        [InlineData(null, "table", "table")]            // Null owner URI
+        [InlineData(Common.OwnerUri, null, "table")]    // Null object name
+        [InlineData(Common.OwnerUri, "table", null)]    // Null object type
+        public async Task InitializeNullParams(string ownerUri, string objName, string objType)
+        {
+            // Setup: Create an edit data service without a session
+            var eds = new EditDataService(null, null, null);
+
+            // If:
+            // ... I have init params with a null parameter
+            var initParams = new EditInitializeParams
+            {
+                ObjectName = objName,
+                OwnerUri = ownerUri,
+                ObjectType = objType
+            };
+
+            // ... And I initialize an edit session with that
+            var efv = new EventFlowValidator<EditInitializeResult>()
+                .AddErrorValidation<string>(Assert.NotNull)
+                .Complete();
+            await eds.HandleInitializeRequest(initParams, efv.Object);
+
+            // Then:
+            // ... An error event should have been raised
+            efv.Validate();
+
+            // ... There should not be a session
+            Assert.Empty(eds.ActiveSessions);
+
+            // ... There should not be a wait handler
+            Assert.Empty(eds.InitializeWaitHandles);
+        }
+
+        [Fact]
+        public async Task InitializeInProgress()
+        {
+            // Setup: Create an edit data service with an "in-progress initialize"
+            var eds = new EditDataService(null, null, null);
+            eds.InitializeWaitHandles[Common.OwnerUri] = new TaskCompletionSource<bool>();
+
+            // If:
+            // ... I ask to initialize a session when an initialize task is already in progress
+            var initParams = new EditInitializeParams
+            {
+                ObjectName = "table",
+                OwnerUri = Common.OwnerUri,
+                ObjectType = "table"
+            };
+            var efv = new EventFlowValidator<EditInitializeResult>()
+                .AddErrorValidation<string>(Assert.NotNull)
+                .Complete();
+            await eds.HandleInitializeRequest(initParams, efv.Object);
+
+            // Then:
+            // ... An error event should have been raised
+            efv.Validate();
+
+            // ... There should not be a session
+            Assert.Empty(eds.ActiveSessions);
+        }
+
+        [Fact]
+        public async Task InitializeQueryCreateFailed()
+        {
+            // Setup:
+            // ... Create a query execution service that will throw on creation of the query
+            var qes = QueryExecution.Common.GetPrimedExecutionService(null, false, false, null);
+
+            // ... Create an edit data service that uses the mocked up query service
+            var eds = new EditDataService(qes, null, null);
+
+            // If:
+            // ... I initialize a session
+            var initParams = new EditInitializeParams
+            {
+                ObjectName = "table",
+                OwnerUri = Common.OwnerUri,
+                ObjectType = "table"
+            };
+            var efv = new EventFlowValidator<EditInitializeResult>()
+                .AddErrorValidation<string>(Assert.NotEmpty)
+                .Complete();
+            await eds.HandleInitializeRequest(initParams, efv.Object);
+
+            // Then:
+            // ... We should have gotten an error back
+            efv.Validate();
+
+            // ... There should not be any sessions created
+            Assert.Empty(eds.ActiveSessions);
+
+            // ... There should not be a wait handle
+            Assert.Empty(eds.InitializeWaitHandles);
+        }
+
+        [Fact]
+        public async Task InitializeQueryExecutionFails()
+        {
+            // Setup:
+            // ... Create a query execution service that will throw on execution of the query
+            var qes = QueryExecution.Common.GetPrimedExecutionService(null, true, true, null);
+
+            // ... Create an edit data service that uses the mocked up query service
+            var eds = new EditDataService(qes, null, null);
+
+            // If:
+            // ... I initialize a session
+            var initParams = new EditInitializeParams
+            {
+                ObjectName = "table",
+                OwnerUri = Common.OwnerUri,
+                ObjectType = "table"
+            };
+            var efv = new EventFlowValidator<EditInitializeResult>()
+                .AddResultValidation(Assert.NotNull)
+                .AddEventValidation(EditSessionReadyEvent.Type, esrp =>
+                {
+                    Assert.NotNull(esrp);
+                    Assert.False(esrp.Success);
+                }).Complete();
+            await eds.HandleInitializeRequest(initParams, efv.Object);
+            await eds.InitializeWaitHandles[Common.OwnerUri].Task;
+
+            // Then:
+            // ... We should have started execution, but failed
+            efv.Validate();
+
+            // ... There should not be any sessions created
+            Assert.Empty(eds.ActiveSessions);
+
+            // ... There should not be a wait handle. It should have been cleaned up by now
+            Assert.Empty(eds.InitializeWaitHandles);
+        }
+
         private static EditSession GetDefaultSession()
         {
             // ... Create a session with a proper query and metadata