Make ExceptWith atomic

apache · May 15, 2024 · 2682952 · 2682952
1 parent 3635caa
commit 2682952
Showing 1 changed file with 61 additions and 13 deletions.
diff --git a/src/Lucene.Net/Support/ConcurrentHashSet.cs b/src/Lucene.Net/Support/ConcurrentHashSet.cs
@@ -74,9 +74,9 @@ public int Count
                 {
                     AcquireAllLocks(ref acquiredLocks);
 
-                    for (var i = 0; i < _tables.CountPerLock.Length; i++)
+                    foreach (var t in _tables.CountPerLock)
                     {
-                        count += _tables.CountPerLock[i];
+                        count += t;
                     }
                 }
                 finally
@@ -88,6 +88,21 @@ public int Count
             }
         }
 
+        private int CountInternal
+        {
+            get
+            {
+                int count = 0;
+
+                foreach (var t in _tables.CountPerLock)
+                {
+                    count += t;
+                }
+
+                return count;
+            }
+        }
+
         /// <summary>
         /// Gets a value that indicates whether the <see cref="ConcurrentHashSet{T}"/> is empty.
         /// </summary>
@@ -298,16 +313,21 @@ public void Clear()
             {
                 AcquireAllLocks(ref locksAcquired);
 
-                var newTables = new Tables(new Node[DefaultCapacity], _tables.Locks, new int[_tables.CountPerLock.Length]);
-                _tables = newTables;
-                _budget = Math.Max(1, newTables.Buckets.Length / newTables.Locks.Length);
+                ClearInternal();
             }
             finally
             {
                 ReleaseLocks(0, locksAcquired);
             }
         }
 
+        private void ClearInternal()
+        {
+            var newTables = new Tables(new Node[DefaultCapacity], _tables.Locks, new int[_tables.CountPerLock.Length]);
+            _tables = newTables;
+            _budget = Math.Max(1, newTables.Buckets.Length / newTables.Locks.Length);
+        }
+
         /// <summary>
         /// Determines whether the <see cref="ConcurrentHashSet{T}"/> contains the specified
         /// item.
@@ -347,16 +367,25 @@ public bool Contains(T item)
         public bool TryRemove(T item)
         {
             var hashcode = _comparer.GetHashCode(item);
+            return TryRemoveInternal(item, hashcode, acquireLock: true);
+        }
+
+        private bool TryRemoveInternal(T item, int hashcode, bool acquireLock)
+        {
             while (true)
             {
                 var tables = _tables;
 
                 GetBucketAndLockNo(hashcode, out int bucketNo, out int lockNo, tables.Buckets.Length, tables.Locks.Length);
 
                 object syncRoot = tables.Locks[lockNo];
-                UninterruptableMonitor.Enter(syncRoot);
+                var lockTaken = false;
+
                 try
                 {
+                    if (acquireLock)
+                        UninterruptableMonitor.Enter(syncRoot, ref lockTaken);
+
                     // If the table just got resized, we may not be holding the right lock, and must retry.
                     // This should be a rare occurrence.
                     if (tables != _tables)
@@ -388,7 +417,8 @@ public bool TryRemove(T item)
                 }
                 finally
                 {
-                    UninterruptableMonitor.Exit(syncRoot);
+                    if (lockTaken)
+                        UninterruptableMonitor.Exit(syncRoot);
                 }
 
                 return false;
@@ -753,15 +783,33 @@ private void CopyToItems(T[] array, int index)
 
         public void ExceptWith(IEnumerable<T> other)
         {
-            if (ReferenceEquals(this, other))
+            if (other is null)
+                throw new ArgumentNullException(nameof(other));
+
+            var locksAcquired = 0;
+            try
             {
-                Clear();
-                return;
-            }
+                AcquireAllLocks(ref locksAcquired);
+
+                if (CountInternal == 0)
+                {
+                    return;
+                }
+
+                if (ReferenceEquals(this, other))
+                {
+                    ClearInternal();
+                    return;
+                }
 
-            foreach (var item in other)
+                foreach (var item in other)
+                {
+                    TryRemoveInternal(item, _comparer.GetHashCode(item), acquireLock: false);
+                }
+            }
+            finally
             {
-                TryRemove(item);
+                ReleaseLocks(0, locksAcquired);
             }
         }