Port fixes to ConditionalWeakTable (#2322)

CoreRT's ConditionalWeakTable implementation was ported to CoreCLR.  In the process, several fixes were made; this ports those fixes back.

Two issues addressed:
- Every time the table resizes, it ends up allocating a new dependency handle for every handle in the table.  That's not cheap.  Instead, we now transfer still valid handles from the old container to the new container, giving the new container the responsibility of freeing them when it's finalized... and the old container keeps it alive as long as the old container is alive via a reference.
- There's a race condition present whereby if a ConditionalWeakTable is resurrected, it's possible for a thread to try accessing a dependency handle while or after it's been freed.  This is addressed by not freeing during the first finalization, instead using that pass to clear out the table's reference to the container and then re-registering for finalization: the next time it's finalized, we can be sure it won't get resurrected and there won't be a race condition.

src/System.Private.CoreLib/src/System/Runtime/CompilerServices/ConditionalWeakTable.cs

@@ -18,7 +18,7 @@ public sealed class ConditionalWeakTable<TKey, TValue>
         #region Constructors
         public ConditionalWeakTable()
         {
-            _container = new Container();
+            _container = new Container(this);
             _lock = new Lock();
         }
         #endregion
@@ -228,7 +228,7 @@ internal void Clear()
         {
             using (LockHolder.Hold(_lock))
             {
-                _container = new Container();
+                _container = new Container(this);
             }
         }
 
@@ -310,8 +310,9 @@ private struct Entry
         //
         private class Container
         {
-            internal Container()
+            internal Container(ConditionalWeakTable<TKey, TValue> parent)
             {
+                Debug.Assert(parent != null);
                 Debug.Assert(IsPowerOfTwo(InitialCapacity));
                 int size = InitialCapacity;
                 _buckets = new int[size];
@@ -320,10 +321,18 @@ internal Container()
                     _buckets[i] = -1;
                 }
                 _entries = new Entry[size];
+
+                // Only store the parent after all of the allocations have happened successfully.
+                // Otherwise, as part of growing or clearing the container, we could end up allocating
+                // a new Container that fails (OOMs) part way through construction but that gets finalized
+                // and ends up clearing out some other container present in the associated CWT.
+                _parent = parent;
             }
-            
-            private Container(int[] buckets, Entry[] entries, int firstFreeEntry)
+
+            private Container(ConditionalWeakTable<TKey, TValue> parent, int[] buckets, Entry[] entries, int firstFreeEntry)
             {
+                Debug.Assert(parent != null);
+                _parent = parent;
                 _buckets = buckets;
                 _entries = entries;
                 _firstFreeEntry = firstFreeEntry;
@@ -484,23 +493,37 @@ internal Container Resize(int newSize)
                 {
                     int hashCode = _entries[entriesIndex].hashCode;
                     DependentHandle depHnd = _entries[entriesIndex].depHnd;
-                    if (hashCode != -1 && depHnd.IsAllocated && depHnd.GetPrimary() != null)
+                    if (hashCode != -1 && depHnd.IsAllocated)
                     {
-                        // Entry is used and has not expired. Link it into the appropriate bucket list.
-                        // Note that we have to copy the DependentHandle, since the original copy will be freed
-                        // by the Container's finalizer.  
-                        newEntries[newEntriesIndex].hashCode = hashCode;
-                        newEntries[newEntriesIndex].depHnd = depHnd.AllocateCopy();
-                        int bucket = hashCode & (newBuckets.Length - 1);
-                        newEntries[newEntriesIndex].next = newBuckets[bucket];
-                        newBuckets[bucket] = newEntriesIndex;
-                        newEntriesIndex++;
+                        if (depHnd.GetPrimary() != null)
+                        {
+                            // Entry is used and has not expired. Link it into the appropriate bucket list.
+                            newEntries[newEntriesIndex].hashCode = hashCode;
+                            newEntries[newEntriesIndex].depHnd = depHnd;
+                            int bucket = hashCode & (newBuckets.Length - 1);
+                            newEntries[newEntriesIndex].next = newBuckets[bucket];
+                            newBuckets[bucket] = newEntriesIndex;
+                            newEntriesIndex++;
+                        }
+                        else
+                        {
+                            // Pretend the item was removed, so that this container's finalizer
+                            // will clean up this dependent handle.
+                            Volatile.Write(ref _entries[entriesIndex].hashCode, -1);
+                        }
                     }
                 }
 
+                // Create the new container.  We want to transfer the responsibility of freeing the handles from
+                // the old container to the new container, and also ensure that the new container isn't finalized
+                // while the old container may still be in use.  As such, we store a reference from the old container
+                // to the new one, which will keep the new container alive as long as the old one is.
+                var newContainer = new Container(_parent, newBuckets, newEntries, newEntriesIndex);
+                _oldKeepAlive = newContainer; // once this is set, the old container's finalizer will not free transferred dependent handles
+
                 GC.KeepAlive(this); // ensure we don't get finalized while accessing DependentHandles.
 
-                return new Container(newBuckets, newEntries, newEntriesIndex);
+                return newContainer;
             }
 
             internal ICollection<TKey> Keys
@@ -610,14 +633,35 @@ private void VerifyIntegrity()
                     return;
                 }
 
-                if (_invalid)
+                //We also skip doing anything if the container is invalid, including if someone
+                // the container object was allocated but its associated table never set.
+                if (_invalid || _parent == null)
                 {
                     return;
                 }
-                Entry[] entries = _entries;
 
-                // Make sure anyone sneaking into the table post-resurrection
-                // gets booted before they can damage the native handle table.
+                // It's possible that the ConditionalWeakTable could have been resurrected, in which case code could
+                // be accessing this Container as it's being finalized.  We don't support usage after finalization,
+                // but we also don't want to potentially corrupt state by allowing dependency handles to be used as
+                // or after they've been freed.  To avoid that, if it's at all possible that another thread has a
+                // reference to this container via the CWT, we remove such a reference and then re-register for
+                // finalization: the next time around, we can be sure that no references remain to this and we can
+                // clean up the dependency handles without fear of corruption.
+                if (!_finalized)
+                {
+                    _finalized = true;
+                    using (LockHolder.Hold(_parent._lock))
+                    {
+                        if (_parent._container == this)
+                        {
+                            _parent._container = null;
+                        }
+                    }
+                    GC.ReRegisterForFinalize(this); // next time it's finalized, we'll be sure there are no remaining refs
+                    return;
+                }
+
+                Entry[] entries = _entries;
                 _invalid = true;
                 _entries = null;
                 _buckets = null;
@@ -626,15 +670,26 @@ private void VerifyIntegrity()
                 {
                     for (int entriesIndex = 0; entriesIndex < entries.Length; entriesIndex++)
                     {
-                        entries[entriesIndex].depHnd.Free();
+                        // We need to free handles in two cases:
+                        // - If this container still owns the dependency handle (meaning ownership hasn't been transferred
+                        //   to another container that replaced this one), then it should be freed.
+                        // - If this container had the entry removed, then even if in general ownership was transferred to
+                        //   another container, removed entries are not, therefore this container must free them.
+                        if (_oldKeepAlive == null || entries[entriesIndex].hashCode == -1)
+                        {
+                            entries[entriesIndex].depHnd.Free();
+                        }
                     }
                 }
             }
 
-            private int[] _buckets;             // _buckets[hashcode & (_buckets.Length - 1)] contains index of the first entry in bucket (-1 if empty)
-            private Entry[] _entries;
-            private int _firstFreeEntry;        // _firstFreeEntry < _entries.Length => table has capacity,  entries grow from the bottom of the table.
-            private bool _invalid;              // flag detects if OOM or other background exception threw us out of the lock.
+            private readonly ConditionalWeakTable<TKey, TValue> _parent;  // the ConditionalWeakTable with which this container is associated
+            private int[] _buckets;                // _buckets[hashcode & (_buckets.Length - 1)] contains index of the first entry in bucket (-1 if empty)
+            private Entry[] _entries;              // the table entries containing the stored dependency handles
+            private int _firstFreeEntry;           // _firstFreeEntry < _entries.Length => table has capacity,  entries grow from the bottom of the table.
+            private bool _invalid;                 // flag detects if OOM or other background exception threw us out of the lock.
+            private bool _finalized;               // set to true when initially finalized
+            private volatile object _oldKeepAlive; // used to ensure the next allocated container isn't finalized until this one is GC'd
         }
 
         private volatile Container _container;
@@ -717,15 +772,6 @@ public void Free()
         }
         #endregion
 
-        #region Internal Members
-        internal DependentHandle AllocateCopy()
-        {
-            object primary, secondary;
-            primary = GetPrimaryAndSecondary(out secondary);
-            return new DependentHandle(primary, secondary);
-        }
-        #endregion
-
         #region Private Data Member
         private IntPtr _handle;
         #endregion