Add moke test for restricted callouts.

AustinWise · AustinWise · commit 6a13d237c70d · 2022-12-30T12:11:57.000-08:00
diff --git a/src/tests/nativeaot/SmokeTests/GcRestrictedCallouts/GcRestrictedCallouts.csproj b/src/tests/nativeaot/SmokeTests/GcRestrictedCallouts/GcRestrictedCallouts.csproj
@@ -0,0 +1,16 @@
+<Project Sdk="Microsoft.NET.Sdk">
+
+  <PropertyGroup>
+    <OutputType>Exe</OutputType>
+    <CLRTestKind>BuildAndRun</CLRTestKind>
+    <CLRTestPriority>0</CLRTestPriority>
+    <AllowUnsafeBlocks>true</AllowUnsafeBlocks>
+  </PropertyGroup>
+
+  <ItemGroup>
+    <Compile Include="Program.cs" />
+    <Compile Include="RuntimeImportAttribute.cs" />
+    <Compile Include="RuntimeImports.cs" />
+  </ItemGroup>
+
+</Project>
diff --git a/src/tests/nativeaot/SmokeTests/GcRestrictedCallouts/Program.cs b/src/tests/nativeaot/SmokeTests/GcRestrictedCallouts/Program.cs
@@ -0,0 +1,38 @@
+// 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.Runtime;
+using System.Runtime.InteropServices;
+using System.Runtime.CompilerServices;
+
+static class Program
+{
+    static readonly ConditionalWeakTable<object, object> s_weakTable = new();
+    static readonly object s_inTableObject = new();
+    static readonly object s_notInTableObject = new();
+
+    static volatile bool s_testPass;
+
+    static unsafe int Main()
+    {
+        s_weakTable.Add(s_inTableObject, new object());
+
+        Console.WriteLine("RhRegisterGcCallout");
+        RuntimeImports.RhRegisterGcCallout(RuntimeImports.GcRestrictedCalloutKind.AfterMarkPhase,
+            (IntPtr)(delegate* unmanaged<uint, void>)&GcCallback);
+
+        Console.WriteLine("GC.Collect");
+        GC.Collect();
+
+        Console.WriteLine("Test passed: " + s_testPass);
+        return s_testPass ? 100 : 1;
+    }
+
+    [UnmanagedCallersOnly]
+    static void GcCallback(uint uiCondemnedGeneration)
+    {
+        s_testPass = s_weakTable.TryGetValue(s_inTableObject, out object _) &&
+                    !s_weakTable.TryGetValue(s_notInTableObject, out object _);
+    }
+}
diff --git a/src/tests/nativeaot/SmokeTests/GcRestrictedCallouts/RuntimeImportAttribute.cs b/src/tests/nativeaot/SmokeTests/GcRestrictedCallouts/RuntimeImportAttribute.cs
@@ -0,0 +1,24 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace System.Runtime
+{
+    // Exposed in Internal.CompilerServices only
+    [AttributeUsage(AttributeTargets.Method | AttributeTargets.Constructor, Inherited = false)]
+    public sealed class RuntimeImportAttribute : Attribute
+    {
+        public string DllName { get; }
+        public string EntryPoint { get; }
+
+        public RuntimeImportAttribute(string entry)
+        {
+            EntryPoint = entry;
+        }
+
+        public RuntimeImportAttribute(string dllName, string entry)
+        {
+            EntryPoint = entry;
+            DllName = dllName;
+        }
+    }
+}
diff --git a/src/tests/nativeaot/SmokeTests/GcRestrictedCallouts/RuntimeImports.cs b/src/tests/nativeaot/SmokeTests/GcRestrictedCallouts/RuntimeImports.cs
@@ -0,0 +1,31 @@
+// 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.Runtime;
+using System.Runtime.CompilerServices;
+
+namespace System.Runtime
+{
+    // Copied from src/coreclr/nativeaot/System.Private.CoreLib/src/System/Runtime/RuntimeImports.cs
+    static class RuntimeImports
+    {
+        private const string RuntimeLibrary = "*";
+
+        internal enum GcRestrictedCalloutKind
+        {
+            StartCollection = 0, // Collection is about to begin
+            EndCollection = 1, // Collection has completed
+            AfterMarkPhase = 2, // All live objects are marked (not including ready for finalization objects),
+                                // no handles have been cleared
+        }
+
+        [MethodImplAttribute(MethodImplOptions.InternalCall)]
+        [RuntimeImport(RuntimeLibrary, "RhRegisterGcCallout")]
+        internal static extern bool RhRegisterGcCallout(GcRestrictedCalloutKind eKind, IntPtr pCalloutMethod);
+
+        [MethodImplAttribute(MethodImplOptions.InternalCall)]
+        [RuntimeImport(RuntimeLibrary, "RhUnregisterGcCallout")]
+        internal static extern void RhUnregisterGcCallout(GcRestrictedCalloutKind eKind, IntPtr pCalloutMethod);
+    }
+}
