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