Partition the ComWrappers RCW cache into per-processor buckets

src/libraries/System.Private.CoreLib/src/System/Runtime/InteropServices/ComWrappers.cs

@@ -5,6 +5,7 @@
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
+using System.Numerics;
 using System.Runtime.CompilerServices;
 using System.Runtime.Versioning;
 using System.Threading;
@@ -1296,10 +1297,56 @@ internal void RemoveWrappersFromCache(IEnumerable<NativeObjectWrapper> wrappers)
             _rcwCache.RemoveAll(wrappers);
         }
 
-        private sealed class RcwCache
+        /// <summary>
+        /// The cache mapping COM instances to the <see cref="NativeObjectWrapper"/> objects tracking their RCWs.
+        /// </summary>
+        /// <remarks>
+        /// The cache is partitioned into several independent buckets, each with its own lock, so that operations on
+        /// COM instances that map to different buckets don't contend with one another. Reducing that contention is
+        /// important because the cache is consulted on essentially every transition from native to managed code, and
+        /// because the finalizer thread concurrently takes write locks to remove entries for collected RCWs.
+        /// </remarks>
+        private readonly struct RcwCache
         {
-            private readonly ReaderWriterLockSlim _lock = new ReaderWriterLockSlim();
-            private readonly Dictionary<IntPtr, GCHandle> _cache = [];
+            private readonly Bucket[] _buckets;
+
+            public RcwCache()
+            {
+                // Use as many buckets as there are processors, matching the default concurrency level of
+                // 'ConcurrentDictionary'. The count is rounded up to a power of two so that the bucket for a
+                // given COM instance can be selected with a mask rather than a division.
+                int bucketCount = (int)BitOperations.RoundUpToPowerOf2((uint)Environment.ProcessorCount);
+                Bucket[] buckets = new Bucket[bucketCount];
+
+                for (int i = 0; i < buckets.Length; i++)
+                {
+                    buckets[i] = new Bucket();
+                }
+
+                _buckets = buckets;
+            }
+
+            /// <summary>
+            /// Gets the bucket owning the entries for a given COM instance.
+            /// </summary>
+            /// <param name="comPointer">The com instance to get the bucket for.</param>
+            /// <returns>The bucket owning the entries for <paramref name="comPointer"/>.</returns>
+            [MethodImpl(MethodImplOptions.AggressiveInlining)]
+            private ref readonly Bucket GetBucket(IntPtr comPointer)
+            {
+                Bucket[] buckets = _buckets;
+
+                // COM instances are heap allocated, so they're always at least pointer aligned (and 16-byte aligned
+                // in practice). That means their low bits are constant and can't be used to select a bucket directly.
+                // Multiplying by a large odd constant (2^64 divided by the golden ratio) mixes every input bit into
+                // the high half of the product, which is then masked to produce the index. The whole sequence lowers
+                // to a multiply, a shift and a mask, which is negligible next to the lookup that follows.
+                ulong hash = (ulong)(nuint)comPointer * 0x9E3779B97F4A7C15;
+                uint index = (uint)(hash >> 32) & (uint)(buckets.Length - 1);
+
+                // Return the bucket by reference, so that it's addressed in place in the array rather than copied
+                return ref buckets[index];
+            }
 
             /// <summary>
             /// Gets the current RCW proxy object for <paramref name="comPointer"/> if it exists in the cache or inserts a new entry with <paramref name="comProxy"/>.
@@ -1310,139 +1357,183 @@ private sealed class RcwCache
             /// <returns>The proxy object currently in the cache for <paramref name="comPointer"/> or the proxy object owned by <paramref name="wrapper"/> if no entry exists and the corresponding native wrapper.</returns>
             public (NativeObjectWrapper actualWrapper, object actualProxy) GetOrAddProxyForComInstance(IntPtr comPointer, NativeObjectWrapper wrapper, object comProxy)
             {
-                _lock.EnterWriteLock();
-                try
-                {
-                    Debug.Assert(wrapper.ProxyHandle.Target == comProxy);
-                    ref GCHandle rcwEntry = ref CollectionsMarshal.GetValueRefOrAddDefault(_cache, comPointer, out bool exists);
-                    if (!exists)
-                    {
-                        // Someone else didn't beat us to adding the entry to the cache.
-                        // Add our entry here.
-                        rcwEntry = GCHandle.Alloc(wrapper, GCHandleType.Weak);
-                    }
-                    else if (rcwEntry.Target is not (NativeObjectWrapper cachedWrapper))
-                    {
-                        Debug.Assert(rcwEntry.IsAllocated);
-                        // The target was collected, so we need to update the cache entry.
-                        rcwEntry.Target = wrapper;
-                    }
-                    else
-                    {
-                        object? existingProxy = cachedWrapper.ProxyHandle.Target;
-                        // The target NativeObjectWrapper was not collected, but we need to make sure
-                        // that the proxy object is still alive.
-                        if (existingProxy is not null)
-                        {
-                            // The existing proxy object is still alive, we will use that.
-                            return (cachedWrapper, existingProxy);
-                        }
+                return GetBucket(comPointer).GetOrAddProxyForComInstance(comPointer, wrapper, comProxy);
+            }
 
-                        // The proxy object was collected, so we need to update the cache entry.
-                        rcwEntry.Target = wrapper;
-                    }
+            /// <summary>
+            /// Gets the current RCW proxy object for <paramref name="comPointer"/>, if it exists in the cache and is still alive.
+            /// </summary>
+            /// <param name="comPointer">The com instance we want to get the RCW for.</param>
+            /// <returns>The proxy object currently in the cache for <paramref name="comPointer"/>, if any.</returns>
+            public object? FindProxyForComInstance(IntPtr comPointer)
+            {
+                return GetBucket(comPointer).FindProxyForComInstance(comPointer);
+            }
 
-                    // We either added an entry to the cache or updated an existing entry that was dead.
-                    // Return our target object.
-                    return (wrapper, comProxy);
-                }
-                finally
+            /// <summary>
+            /// Removes the entry associating <paramref name="comPointer"/> with <paramref name="wrapper"/>, if present.
+            /// </summary>
+            /// <param name="comPointer">The com instance to remove the entry for.</param>
+            /// <param name="wrapper">The <see cref="NativeObjectWrapper"/> that is expected to be in the cache.</param>
+            public void Remove(IntPtr comPointer, NativeObjectWrapper wrapper)
+            {
+                GetBucket(comPointer).Remove(comPointer, wrapper);
+            }
+
+            /// <summary>
+            /// Removes the entries for all input <see cref="NativeObjectWrapper"/> objects, if present.
+            /// </summary>
+            /// <param name="wrappers">The <see cref="NativeObjectWrapper"/> objects to remove the entries for.</param>
+            public void RemoveAll(IEnumerable<NativeObjectWrapper> wrappers)
+            {
+                // The wrappers can span multiple buckets, so they're removed one at a time. This is only used when
+                // tearing down an apartment, so the extra lock acquisitions don't matter. Note that entries are not
+                // removed atomically as a batch anymore, but that was never something callers could rely on: the
+                // cache is only ever observed one entry at a time.
+                foreach (NativeObjectWrapper wrapper in wrappers)
                 {
-                    _lock.ExitWriteLock();
+                    IntPtr comPointer = wrapper.ExternalComObject;
+
+                    GetBucket(comPointer).Remove(comPointer, wrapper);
                 }
             }
 
-            public object? FindProxyForComInstance(IntPtr comPointer)
+            /// <summary>
+            /// A single partition of the RCW cache, holding the entries for all COM instances that map to it.
+            /// </summary>
+            /// <remarks>
+            /// This is a struct so that buckets are stored inline in the containing array. That saves a dereference
+            /// on each lookup, and lets several buckets share a cache line. There is no false sharing to worry about,
+            /// as the fields are only ever read: all mutable state lives in the referenced lock and dictionary.
+            /// </remarks>
+            private readonly struct Bucket
             {
-                _lock.EnterReadLock();
-                try
+                private readonly ReaderWriterLockSlim _lock;
+                private readonly Dictionary<IntPtr, WeakGCHandle<NativeObjectWrapper>> _cache;
+
+                public Bucket()
+                {
+                    _lock = new ReaderWriterLockSlim();
+                    _cache = [];
+                }
+
+                /// <inheritdoc cref="RcwCache.GetOrAddProxyForComInstance"/>
+                public (NativeObjectWrapper actualWrapper, object actualProxy) GetOrAddProxyForComInstance(IntPtr comPointer, NativeObjectWrapper wrapper, object comProxy)
                 {
-                    if (!_cache.TryGetValue(comPointer, out GCHandle existingHandle))
+                    _lock.EnterWriteLock();
+                    try
                     {
-                        // No entry in the cache.
-                        return null;
+                        Debug.Assert(wrapper.ProxyHandle.Target == comProxy);
+                        ref WeakGCHandle<NativeObjectWrapper> rcwEntry = ref CollectionsMarshal.GetValueRefOrAddDefault(_cache, comPointer, out bool exists);
+                        if (!exists)
+                        {
+                            // Someone else didn't beat us to adding the entry to the cache.
+                            // Add our entry here.
+                            rcwEntry = new WeakGCHandle<NativeObjectWrapper>(wrapper);
+                        }
+                        else if (!rcwEntry.TryGetTarget(out NativeObjectWrapper? cachedWrapper))
+                        {
+                            Debug.Assert(rcwEntry.IsAllocated);
+                            // The target was collected, so we need to update the cache entry.
+                            rcwEntry.SetTarget(wrapper);
+                        }
+                        else
+                        {
+                            object? existingProxy = cachedWrapper.ProxyHandle.Target;
+                            // The target NativeObjectWrapper was not collected, but we need to make sure
+                            // that the proxy object is still alive.
+                            if (existingProxy is not null)
+                            {
+                                // The existing proxy object is still alive, we will use that.
+                                return (cachedWrapper, existingProxy);
+                            }
+
+                            // The proxy object was collected, so we need to update the cache entry.
+                            rcwEntry.SetTarget(wrapper);
+                        }
+
+                        // We either added an entry to the cache or updated an existing entry that was dead.
+                        // Return our target object.
+                        return (wrapper, comProxy);
                     }
-                    if (existingHandle.Target is NativeObjectWrapper { ProxyHandle.Target: object cachedProxy })
+                    finally
                     {
-                        // The target exists and is still alive. Return it.
-                        return cachedProxy;
+                        _lock.ExitWriteLock();
                     }
-                    // The target was collected, so we need to remove the entry from the cache.
-                    // We'll do this in a write lock after we exit the read lock.
-                    // We don't use an upgradeable lock here as only one thread can hold an upgradeable lock at a time,
-                    // effectively eliminating the benefit of using a reader-writer lock.
-                }
-                finally
-                {
-                    _lock.ExitReadLock();
                 }
 
-                _lock.EnterWriteLock();
-                try
+                /// <inheritdoc cref="RcwCache.FindProxyForComInstance"/>
+                public object? FindProxyForComInstance(IntPtr comPointer)
                 {
-                    // Someone else could have removed the entry or added a new one in the time
-                    // between us releasing the read lock and acquiring the write lock.
-                    if (_cache.TryGetValue(comPointer, out GCHandle existingHandle)
-                        && existingHandle.Target is null)
+                    _lock.EnterReadLock();
+                    try
                     {
-                        // There's still a dead entry in the cache,
-                        // remove it.
-                        _cache.Remove(comPointer);
-                        existingHandle.Free();
+                        if (!_cache.TryGetValue(comPointer, out WeakGCHandle<NativeObjectWrapper> existingHandle))
+                        {
+                            // No entry in the cache.
+                            return null;
+                        }
+                        if (existingHandle.TryGetTarget(out NativeObjectWrapper? cachedWrapper)
+                            && cachedWrapper.ProxyHandle.Target is object cachedProxy)
+                        {
+                            // The target exists and is still alive. Return it.
+                            return cachedProxy;
+                        }
+                        // The target was collected, so we need to remove the entry from the cache.
+                        // We'll do this in a write lock after we exit the read lock.
+                        // We don't use an upgradeable lock here as only one thread can hold an upgradeable lock at a time,
+                        // effectively eliminating the benefit of using a reader-writer lock.
+                    }
+                    finally
+                    {
+                        _lock.ExitReadLock();
                     }
-                }
-                finally
-                {
-                    _lock.ExitWriteLock();
-                }
 
-                return null;
-            }
+                    _lock.EnterWriteLock();
+                    try
+                    {
+                        // Someone else could have removed the entry or added a new one in the time
+                        // between us releasing the read lock and acquiring the write lock.
+                        if (_cache.TryGetValue(comPointer, out WeakGCHandle<NativeObjectWrapper> existingHandle)
+                            && !existingHandle.TryGetTarget(out _))
+                        {
+                            // There's still a dead entry in the cache,
+                            // remove it.
+                            _cache.Remove(comPointer);
+                            existingHandle.Dispose();
+                        }
+                    }
+                    finally
+                    {
+                        _lock.ExitWriteLock();
+                    }
 
-            public void Remove(IntPtr comPointer, NativeObjectWrapper wrapper)
-            {
-                _lock.EnterWriteLock();
-                try
-                {
-                    Remove_Locked(comPointer, wrapper);
-                }
-                finally
-                {
-                    _lock.ExitWriteLock();
+                    return null;
                 }
-            }
 
-            public void RemoveAll(IEnumerable<NativeObjectWrapper> wrappers)
-            {
-                _lock.EnterWriteLock();
-                try
+                /// <inheritdoc cref="RcwCache.Remove"/>
+                public void Remove(IntPtr comPointer, NativeObjectWrapper wrapper)
                 {
-                    foreach (NativeObjectWrapper wrapper in wrappers)
+                    _lock.EnterWriteLock();
+                    try
                     {
-                        Remove_Locked(wrapper.ExternalComObject, wrapper);
+                        // TryGetOrCreateObjectForComInstanceInternal may have put a new entry into the cache
+                        // in the time between the GC cleared the contents of the GC handle but before the
+                        // NativeObjectWrapper finalizer ran.
+                        // Only remove the entry if the target of the GC handle is the NativeObjectWrapper
+                        // or is null (indicating that the corresponding NativeObjectWrapper has been scheduled for finalization).
+                        if (_cache.TryGetValue(comPointer, out WeakGCHandle<NativeObjectWrapper> cachedRef)
+                            && (!cachedRef.TryGetTarget(out NativeObjectWrapper? cachedWrapper)
+                                || cachedWrapper == wrapper))
+                        {
+                            _cache.Remove(comPointer);
+                            cachedRef.Dispose();
+                        }
+                    }
+                    finally
+                    {
+                        _lock.ExitWriteLock();
                     }
-                }
-                finally
-                {
-                    _lock.ExitWriteLock();
-                }
-            }
-
-            private void Remove_Locked(IntPtr comPointer, NativeObjectWrapper wrapper)
-            {
-                Debug.Assert(_lock.IsWriteLockHeld);
-                // This method is used in a scenario where we already have a lock on the cache, so we can skip acquiring the lock again.
-                // TryGetOrCreateObjectForComInstanceInternal may have put a new entry into the cache
-                // in the time between the GC cleared the contents of the GC handle but before the
-                // NativeObjectWrapper finalizer ran.
-                // Only remove the entry if the target of the GC handle is the NativeObjectWrapper
-                // or is null (indicating that the corresponding NativeObjectWrapper has been scheduled for finalization).
-                if (_cache.TryGetValue(comPointer, out GCHandle cachedRef)
-                    && (wrapper == cachedRef.Target
-                        || cachedRef.Target is null))
-                {
-                    _cache.Remove(comPointer);
-                    cachedRef.Free();
                 }
             }
         }