add lock on qcc to fix sync issue

snowflakedb · Nov 17, 2023 · 8137f56 · 8137f56
1 parent 62f2206
commit 8137f56
Show file tree

Hide file tree

Showing 2 changed files with 84 additions and 28 deletions.
diff --git a/Snowflake.Data.Tests/IntegrationTests/SFDbCommandIT.cs b/Snowflake.Data.Tests/IntegrationTests/SFDbCommandIT.cs
@@ -60,6 +60,51 @@ public void TestExecAsyncAPI()
             }
         }
 
+        [Test]
+        public void TestExecAsyncAPIParallel()
+        {
+            SnowflakeDbConnectionPool.ClearAllPools();
+            using (DbConnection conn = new SnowflakeDbConnection())
+            {
+                conn.ConnectionString = ConnectionString;
+
+                Task connectTask = conn.OpenAsync(CancellationToken.None);
+                connectTask.Wait();
+                Assert.AreEqual(ConnectionState.Open, conn.State);
+
+                Task[] taskArray = new Task[100];
+                for (int i = 0; i < taskArray.Length; i++)
+                {
+                    taskArray[i] = Task.Factory.StartNew(() =>
+                    {
+                        using (DbCommand cmd = conn.CreateCommand())
+                        {
+                            long queryResult = 0;
+                            cmd.CommandText = "select count(seq4()) from table(generator(timelimit => 3)) v";
+                            Task<DbDataReader> execution = cmd.ExecuteReaderAsync();
+                            Task readCallback = execution.ContinueWith((t) =>
+                            {
+                                using (DbDataReader reader = t.Result)
+                                {
+                                    Assert.IsTrue(reader.Read());
+                                    queryResult = reader.GetInt64(0);
+                                    Assert.IsFalse(reader.Read());
+                                }
+                            });
+                            // query is not finished yet, result is still 0;
+                            Assert.AreEqual(0, queryResult);
+                            // block till query finished
+                            readCallback.Wait();
+                            // queryResult should be updated by callback
+                            Assert.AreNotEqual(0, queryResult);
+                        }
+                    });
+                }
+                Task.WaitAll(taskArray);
+                conn.Close();
+            }
+        }
+
         [Test]
         public void TestCancelExecuteAsync()
         {

diff --git a/Snowflake.Data/Core/QueryContextCache.cs b/Snowflake.Data/Core/QueryContextCache.cs
@@ -69,6 +69,7 @@ public int Compare(QueryContextElement x, QueryContextElement y)
 
     internal class QueryContextCache
     {
+        private readonly object _qccLock;
         private int _capacity; // Capacity of the cache
         private Dictionary<long, QueryContextElement> _idMap; // Map for id and QCC
         private Dictionary<long, QueryContextElement> _priorityMap; // Map for priority and QCC
@@ -78,6 +79,7 @@ internal class QueryContextCache
 
         public QueryContextCache(int capacity)
         {
+            _qccLock = new object();
             _capacity = capacity;
             _idMap = new Dictionary<long, QueryContextElement>();
             _priorityMap = new Dictionary<long, QueryContextElement>();
@@ -189,14 +191,17 @@ public void ClearCache()
 
         public void SetCapacity(int cap)
         {
-            // check without locking first for performance reason
-            if (_capacity == cap)
-                return;
-
-            _logger.Debug($"set capacity from {_capacity} to {cap}");
-            _capacity = cap;
-            CheckCacheCapacity();
-            LogCacheEntries();
+            lock(_qccLock)
+            {
+                // check without locking first for performance reason
+                if (_capacity == cap)
+                    return;
+
+                _logger.Debug($"set capacity from {_capacity} to {cap}");
+                _capacity = cap;
+                CheckCacheCapacity();
+                LogCacheEntries();
+            }
         }
 
         /**
@@ -221,26 +226,29 @@ public int GetSize()
          */
         public void Update(ResponseQueryContext queryContext)
         {
-            // Log existing cache entries
-            LogCacheEntries();
-
-            if (queryContext == null || queryContext.Entries == null)
-            {
-                // Clear the cache
-                ClearCache();
-                return;
-            }
-            foreach (ResponseQueryContextElement entry in queryContext.Entries)
+            lock(_qccLock)
             {
-                Merge(entry.Id, entry.ReadTimestamp, entry.Priority, entry.Context);
-            }
+                // Log existing cache entries
+                LogCacheEntries();
 
-            SyncPriorityMap();
+                if (queryContext == null || queryContext.Entries == null)
+                {
+                    // Clear the cache
+                    ClearCache();
+                    return;
+                }
+                foreach (ResponseQueryContextElement entry in queryContext.Entries)
+                {
+                    Merge(entry.Id, entry.ReadTimestamp, entry.Priority, entry.Context);
+                }
 
-            // After merging all entries, truncate to capacity
-            CheckCacheCapacity();
-            // Log existing cache entries
-            LogCacheEntries();
+                SyncPriorityMap();
+
+                // After merging all entries, truncate to capacity
+                CheckCacheCapacity();
+                // Log existing cache entries
+                LogCacheEntries();
+            }
         }
 
         /**
@@ -251,10 +259,13 @@ public RequestQueryContext GetQueryContextRequest()
         {
             RequestQueryContext reqQCC = new RequestQueryContext();
             reqQCC.Entries = new List<RequestQueryContextElement>();
-            foreach (QueryContextElement elem in _cacheSet)
+            lock(_qccLock)
             {
-                RequestQueryContextElement reqElem = new RequestQueryContextElement(elem);
-                reqQCC.Entries.Add(reqElem);
+                foreach (QueryContextElement elem in _cacheSet)
+                {
+                    RequestQueryContextElement reqElem = new RequestQueryContextElement(elem);
+                    reqQCC.Entries.Add(reqElem);
+                }
             }
 
             return reqQCC;