Fix generic virtual dispatch hang after multiple AssemblyLoadContext unloads

src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/CastCache.cs

@@ -22,6 +22,9 @@ internal unsafe struct CastCache
         private const uint VERSION_NUM_MASK = (1 << VERSION_NUM_SIZE) - 1;
         private const int BUCKET_SIZE = 8;
 
+        // The number of elements in the sentinel table (see s_sentinelTable).
+        private const int SENTINEL_TABLE_SIZE = 2;
+
         // nothing is ever stored into this, so we can use a static instance.
         private static int[]? s_sentinelTable;
 
@@ -36,7 +39,7 @@ internal unsafe struct CastCache
 
         public CastCache(int initialCacheSize, int maxCacheSize)
         {
-            Debug.Assert(BitOperations.PopCount((uint)initialCacheSize) == 1 && initialCacheSize > 1);
+            Debug.Assert(BitOperations.PopCount((uint)initialCacheSize) == 1 && initialCacheSize > SENTINEL_TABLE_SIZE);
             Debug.Assert(BitOperations.PopCount((uint)maxCacheSize) == 1 && maxCacheSize >= initialCacheSize);
 
             _initialCacheSize = initialCacheSize;
@@ -45,7 +48,7 @@ public CastCache(int initialCacheSize, int maxCacheSize)
             // A trivial 2-elements table used for "flushing" the cache.
             // Nothing is ever stored in such a small table and identity of the sentinel is not important.
             // It is required that we are able to allocate this, we may need this in OOM cases.
-            s_sentinelTable ??= CreateCastCache(2, throwOnFail: true);
+            s_sentinelTable ??= CreateCastCache(SENTINEL_TABLE_SIZE, throwOnFail: true);
 
             _table =
 #if !DEBUG
@@ -123,6 +126,13 @@ private static ref int TableMask(ref int tableData)
             return ref Unsafe.Add(ref tableData, 1);
         }
 
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private static bool IsSentinel(ref int tableData)
+        {
+            // The sentinel is the only table with SENTINEL_TABLE_SIZE elements.
+            return TableMask(ref tableData) == SENTINEL_TABLE_SIZE - 1;
+        }
+
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         private static ref uint VictimCounter(ref int tableData)
         {
@@ -256,9 +266,9 @@ internal void TrySet(nuint source, nuint target, bool result)
             do
             {
                 tableData = ref TableData(_table);
-                if (TableMask(ref tableData) == 1)
+                if (IsSentinel(ref tableData))
                 {
-                    // 2-element table is used as a sentinel.
+                    // the sentinel table is used to indicate that
                     // we did not allocate a real table yet or have flushed it.
                     // try replacing the table, but do not insert anything.
                     MaybeReplaceCacheWithLarger(_lastFlushSize);
@@ -331,7 +341,7 @@ internal void TrySet(nuint source, nuint target, bool result)
             // reread tableData after TryGrow.
             tableData = ref TableData(_table);
 
-            if (TableMask(ref tableData) == 1)
+            if (IsSentinel(ref tableData))
             {
                 // do not insert into a sentinel.
                 return;

src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/GenericCache.cs

@@ -57,6 +57,9 @@ private struct Entry
         private const uint VERSION_NUM_MASK = (1 << VERSION_NUM_SIZE) - 1;
         private const int BUCKET_SIZE = 8;
 
+        // The number of elements in the sentinel table (see _sentinelTable).
+        private const int SENTINEL_TABLE_SIZE = 2;
+
         // The fields of this structure are known to coreclr, so if they are updated, you must also update object.h
 
         // The actual storage.
@@ -74,7 +77,7 @@ private struct Entry
         // creates a new cache instance
         public GenericCache(int initialCacheSize, int maxCacheSize)
         {
-            Debug.Assert(BitOperations.PopCount((uint)initialCacheSize) == 1 && initialCacheSize > 1);
+            Debug.Assert(BitOperations.PopCount((uint)initialCacheSize) == 1 && initialCacheSize > SENTINEL_TABLE_SIZE);
             Debug.Assert(BitOperations.PopCount((uint)maxCacheSize) == 1 && maxCacheSize >= initialCacheSize);
 
             _initialCacheSize = initialCacheSize;
@@ -83,7 +86,7 @@ public GenericCache(int initialCacheSize, int maxCacheSize)
             // A trivial 2-elements table used for "flushing" the cache.
             // Nothing is ever stored in such a small table and identity of the sentinel is not important.
             // It is required that we are able to allocate this, we may need this in OOM cases.
-            _sentinelTable = CreateCacheTable(2, throwOnFail: true)!;
+            _sentinelTable = CreateCacheTable(SENTINEL_TABLE_SIZE, throwOnFail: true)!;
 
             _table =
 #if !DEBUG
@@ -138,6 +141,14 @@ private static ref Entry Element(Entry[] table, int index)
             return ref Unsafe.Add(ref Unsafe.As<byte, Entry>(ref Unsafe.As<RawArrayData>(table).Data), index + 1);
         }
 
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private static bool IsSentinel(Entry[] table)
+        {
+            // The sentinel is the only table with SENTINEL_TABLE_SIZE elements.
+            // NOTE: the actual array is one element longer, since element 0 is used for aux data.
+            return table.Length == SENTINEL_TABLE_SIZE + 1;
+        }
+
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         internal bool TryGet(TKey key, out TValue? value)
         {
@@ -246,9 +257,9 @@ internal void TrySet(TKey key, TValue value)
             do
             {
                 table = _table;
-                if (table.Length == 2)
+                if (IsSentinel(table))
                 {
-                    // 2-element table is used as a sentinel.
+                    // sentinel table is used to indicate that
                     // we did not allocate a real table yet or have flushed it.
                     // try replacing the table, but do not insert anything.
                     MaybeReplaceCacheWithLarger(_lastFlushSize);
@@ -322,7 +333,7 @@ internal void TrySet(TKey key, TValue value)
             // reread tableData after TryGrow.
             table = _table;
 
-            if (table.Length == 2)
+            if (IsSentinel(table))
             {
                 // do not insert into a sentinel.
                 return;

src/tests/Loader/CollectibleAssemblies/GenericVirtualMethod/GenericVirtualMethod.cs

@@ -0,0 +1,96 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.IO;
+using System.Reflection;
+using System.Runtime.CompilerServices;
+using System.Runtime.Loader;
+using Xunit;
+using TestLibrary;
+
+// Repeatedly loads an assembly into a collectible AssemblyLoadContext, invokes a generic
+// virtual method on a type from that assembly and unloads the context.
+// The results of the generic virtual dispatch are stored in a process-wide cache, which is
+// flushed when a collectible context is unloaded. If the flush does not take effect, a stale
+// entry could be returned for a new type that happens to reuse the addresses of an unloaded
+// one, which results in a hang or a crash.
+public class GenericVirtualMethodUnloading
+{
+    private class TestALC : AssemblyLoadContext
+    {
+        public TestALC(int id) : base($"GenericVirtualMethod{id}", isCollectible: true)
+        {
+        }
+    }
+
+    [ActiveIssue("https://github.com/dotnet/runtimelab/issues/155: Collectible assemblies", typeof(Utilities), nameof(Utilities.IsNativeAot))]
+    [ActiveIssue("https://github.com/dotnet/runtime/issues/34072", TestRuntimes.Mono)]
+    [Fact]
+    public static void CallGenericVirtualMethodAcrossUnloads()
+    {
+        string payloadPath = Path.Combine(Path.GetDirectoryName(Assembly.GetExecutingAssembly().Location), "GenericVirtualMethodUnloaded.dll");
+
+        for (int iteration = 0; iteration < 10; iteration++)
+        {
+            WeakReference context = InvokeAndUnload(payloadPath, iteration);
+            for (int collection = 0; context.IsAlive && collection < 10; collection++)
+            {
+                GC.Collect();
+                GC.WaitForPendingFinalizers();
+            }
+
+            Assert.False(context.IsAlive, $"Context {iteration} did not unload.");
+            VerifyVirtualDispatchCacheIsEmpty();
+        }
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static WeakReference InvokeAndUnload(string payloadPath, int iteration)
+    {
+        TestALC alc = new TestALC(iteration);
+        WeakReference weakAlc = new WeakReference(alc, trackResurrection: false);
+
+        Assembly payload = alc.LoadFromAssemblyPath(payloadPath);
+        Type machineType = payload.GetType("DerivedMachine", throwOnError: true);
+        Type stateType = payload.GetType("MarkerState", throwOnError: true);
+        Type baseType = payload.GetType("Machine", throwOnError: true);
+
+        object machine = Activator.CreateInstance(machineType)!;
+        MethodInfo change = baseType.GetMethod("Change")!.MakeGenericMethod(stateType);
+
+        TargetInvocationException ex = Assert.Throws<TargetInvocationException>(() => change.Invoke(machine, null));
+        Assert.NotNull(ex.InnerException);
+        Assert.Equal("ExpectedException", ex.InnerException!.GetType().FullName);
+
+        alc.Unload();
+        return weakAlc;
+    }
+
+    // Unloading a collectible context flushes the virtual function pointer cache, so that
+    // targets belonging to the unloaded context cannot be returned by a later lookup.
+    // Check that the cache is indeed empty after the unload.
+    private static void VerifyVirtualDispatchCacheIsEmpty()
+    {
+        Type helpersType = typeof(object).Assembly.GetType("System.Runtime.CompilerServices.VirtualDispatchHelpers");
+        if (helpersType == null)
+        {
+            // The cache is specific to CoreCLR.
+            return;
+        }
+
+        object cache = helpersType.GetField("s_virtualFunctionPointerCache", BindingFlags.NonPublic | BindingFlags.Static).GetValue(null);
+        Array table = (Array)cache.GetType().GetField("_table", BindingFlags.NonPublic | BindingFlags.Instance).GetValue(cache);
+
+        FieldInfo infoField = table.GetType().GetElementType().GetField("_info", BindingFlags.NonPublic | BindingFlags.Instance);
+        FieldInfo versionField = infoField.FieldType.GetField("_version", BindingFlags.NonPublic | BindingFlags.Instance);
+
+        // element 0 of the table holds the auxiliary data of the table, the entries start at index 1.
+        // A nonzero version means the entry is in use.
+        for (int i = 1; i < table.Length; i++)
+        {
+            object entryInfo = infoField.GetValue(table.GetValue(i));
+            Assert.Equal(0u, (uint)versionField.GetValue(entryInfo));
+        }
+    }
+}

src/tests/Loader/CollectibleAssemblies/GenericVirtualMethod/GenericVirtualMethod.csproj

@@ -0,0 +1,13 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <PropertyGroup>
+    <!-- Needed for GC.WaitForPendingFinalizers and to exercise the process-wide cache in isolation -->
+    <RequiresProcessIsolation>true</RequiresProcessIsolation>
+  </PropertyGroup>
+  <ItemGroup>
+    <ProjectReference Include="GenericVirtualMethodUnloaded.csproj" />
+    <Compile Include="GenericVirtualMethod.cs" />
+  </ItemGroup>
+  <ItemGroup>
+    <ProjectReference Include="$(TestLibraryProjectPath)" />
+  </ItemGroup>
+</Project>

src/tests/Loader/CollectibleAssemblies/GenericVirtualMethod/GenericVirtualMethodUnloaded.cs

@@ -0,0 +1,43 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Collections.Generic;
+
+public class State;
+public sealed class MarkerState : State;
+public sealed class ExpectedException : Exception;
+
+public abstract class Machine
+{
+    private readonly Dictionary<Type, State> _states = new();
+
+    public void Change<T>() where T : State
+    {
+        _ = GetOrCreate<T>();
+    }
+
+    private T GetOrCreate<T>() where T : State
+    {
+        if (!_states.TryGetValue(typeof(T), out State state))
+        {
+            state = Construct<T>();
+            _states[typeof(T)] = state;
+        }
+
+        return (T)state;
+    }
+
+    protected virtual T Construct<T>() where T : State
+    {
+        return Activator.CreateInstance<T>();
+    }
+}
+
+public sealed class DerivedMachine : Machine
+{
+    protected override T Construct<T>()
+    {
+        throw new ExpectedException();
+    }
+}

src/tests/Loader/CollectibleAssemblies/GenericVirtualMethod/GenericVirtualMethodUnloaded.csproj

@@ -0,0 +1,11 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <PropertyGroup>
+    <OutputType>Library</OutputType>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="GenericVirtualMethodUnloaded.cs" />
+  </ItemGroup>
+  <ItemGroup>
+    <ProjectReference Include="$(TestLibraryProjectPath)" />
+  </ItemGroup>
+</Project>