Unity-Technologies · scott-ferguson-unity · Jan 13, 2025
diff --git a/src/libraries/Common/src/Interop/Unix/System.Native/Interop.SocketEvent.cs b/src/libraries/Common/src/Interop/Unix/System.Native/Interop.SocketEvent.cs
@@ -47,5 +47,8 @@ internal struct SocketEvent
 
         [LibraryImport(Libraries.SystemNative, EntryPoint = "SystemNative_WaitForSocketEvents")]
         internal static unsafe partial Error WaitForSocketEvents(IntPtr port, SocketEvent* buffer, int* count);
+
+        [LibraryImport(Libraries.SystemNative, EntryPoint = "SystemNative_WakeUpSocketEventThread")]
+        internal static unsafe partial Error WakeupSocketEventThread(IntPtr port, IntPtr* handle);
     }
 }
diff --git a/src/libraries/System.Net.Sockets/src/System.Net.Sockets.csproj b/src/libraries/System.Net.Sockets/src/System.Net.Sockets.csproj
@@ -298,4 +298,7 @@
   <ItemGroup Condition="'$(TargetPlatformIdentifier)' == 'unix'">
     <Reference Include="System.Threading.Thread" />
   </ItemGroup>
+
+  <Import Project="$([MSBuild]::NormalizeDirectory('$(RepoRoot)'))\unity\unity-aot\System.Net.Sockets\System.Net.Sockets.props"  Condition="'$(UnityAot)' == 'true'" />
+
 </Project>
diff --git a/src/libraries/System.Net.Sockets/src/System/Net/Sockets/SocketAsyncEngine.Unix.cs b/src/libraries/System.Net.Sockets/src/System/Net/Sockets/SocketAsyncEngine.Unix.cs
@@ -173,12 +173,47 @@ private SocketAsyncEngine()
             }
         }
 
+#if FEATURE_RUNTIME_SHUTDOWN
+        private CancellationTokenSource? _cts;
+        private bool ContinueLoop() => !_cts!.IsCancellationRequested;
+
+        [UnsafeAccessor(UnsafeAccessorKind.StaticMethod, Name = "RegisterShutdownHandler")]
+        private static extern void ThreadRegisterShutdownHandler(Thread thread, Action action);
+
+        private void RegisterShutdownHandler()
+        {
+            _cts = new CancellationTokenSource();
+            var eventLoopThread = Thread.CurrentThread;
+            ThreadRegisterShutdownHandler(eventLoopThread, () =>
+            {
+                _cts.Cancel();
+                IntPtr handle;
+                var err = Interop.Sys.WakeupSocketEventThread(_port, &handle);
+                if (err == Interop.Error.SUCCESS)
+                {
+                    eventLoopThread.Join();
+                    _cts.Dispose();
+                }
+
+                if (handle != IntPtr.Zero)
+                    Interop.Sys.Close(handle);
+            });
+        }
+#else
+        [MethodImplAttribute(MethodImplOptions.AggressiveInlining)]
+        private bool ContinueLoop() => true;
+
+        private void RegisterShutdownHandler() {}
+#endif
+
         private void EventLoop()
         {
+            RegisterShutdownHandler();
+
             try
             {
                 SocketEventHandler handler = new SocketEventHandler(this);
-                while (true)
+                while (ContinueLoop())
                 {
                     int numEvents = EventBufferCount;
                     Interop.Error err = Interop.Sys.WaitForSocketEvents(_port, handler.Buffer, &numEvents);
@@ -200,6 +235,10 @@ private void EventLoop()
             {
                 Environment.FailFast("Exception thrown from SocketAsyncEngine event loop: " + e.ToString(), e);
             }
+            finally
+            {
+                FreeNativeResources();
+            }
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]

diff --git a/src/libraries/System.Private.CoreLib/src/System/Threading/PortableThreadPool.GateThread.cs b/src/libraries/System.Private.CoreLib/src/System/Threading/PortableThreadPool.GateThread.cs
@@ -20,13 +20,40 @@ private static class GateThread
             private static readonly AutoResetEvent RunGateThreadEvent = new AutoResetEvent(initialState: true);
             private static readonly AutoResetEvent DelayEvent = new AutoResetEvent(initialState: false);
 
+#if FEATURE_RUNTIME_SHUTDOWN
+            private static CancellationTokenSource? s_cts;
+
+            private static bool ContinueLoop() => !s_cts!.IsCancellationRequested;
+
+            private static void RegisterShutdownHandler()
+            {
+                s_cts = new CancellationTokenSource();
+                var thread = Thread.CurrentThread;
+                Thread.RegisterShutdownHandler(() =>
+                {
+                    ThreadPool.SetMinThreads(0, 0);
+                    ThreadPool.SetMaxThreads(1, 1);
+                    s_cts.Cancel();
+                    RunGateThreadEvent.Set();
+                    DelayEvent.Set();
+                    thread.Join();
+                });
+            }
+#else
+            [MethodImpl(MethodImplOptions.AggressiveInlining)]
+            private static bool ContinueLoop() => true;
+            private static void RegisterShutdownCallback() {};
+#endif
+
             private static void GateThreadStart()
             {
                 bool disableStarvationDetection =
                     AppContextConfigHelper.GetBooleanConfig("System.Threading.ThreadPool.DisableStarvationDetection", false);
                 bool debuggerBreakOnWorkStarvation =
                     AppContextConfigHelper.GetBooleanConfig("System.Threading.ThreadPool.DebugBreakOnWorkerStarvation", false);
 
+                RegisterShutdownHandler();
+
                 // The first reading is over a time range other than what we are focusing on, so we do not use the read other
                 // than to send it to any runtime-specific implementation that may also use the CPU utilization.
                 CpuUtilizationReader cpuUtilizationReader = default;
@@ -43,13 +70,13 @@ private static void GateThreadStart()
                     Gen2GcCallback.Register(threadPoolInstance.OnGen2GCCallback);
                 }
 
-                while (true)
+                while (ContinueLoop())
                 {
                     RunGateThreadEvent.WaitOne();
                     int currentTimeMs = Environment.TickCount;
                     delayHelper.SetGateActivitiesTime(currentTimeMs);
 
-                    while (true)
+                    while (ContinueLoop())
                     {
                         bool wasSignaledToWake = DelayEvent.WaitOne((int)delayHelper.GetNextDelay(currentTimeMs));
                         currentTimeMs = Environment.TickCount;

diff --git a/src/libraries/System.Private.CoreLib/src/System/Threading/TimerQueue.Portable.cs b/src/libraries/System.Private.CoreLib/src/System/Threading/TimerQueue.Portable.cs
@@ -67,6 +67,30 @@ private bool SetTimerPortable(uint actualDuration)
             return true;
         }
 
+#if FEATURE_RUNTIME_SHUTDOWN
+        private static CancellationTokenSource? s_cts;
+
+        private static bool ContinueLoop() => !s_cts!.IsCancellationRequested;
+
+        private static void RegisterShutdownHandler()
+        {
+            s_cts = new CancellationTokenSource();
+            var thread = Thread.CurrentThread;
+            Thread.RegisterShutdownHandler(
+                () =>
+                {
+                    s_cts.Cancel();
+                    s_timerEvent.Set();
+                    thread.Join();
+                });
+        }
+#else
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private static bool ContinueLoop() => true;
+
+        private static void RegisterShutdownHandler() {};
+#endif
+
         /// <summary>
         /// This method is executed on a dedicated timer thread. Its purpose is
         /// to handle timer requests and notify the TimerQueue when a timer expires.
@@ -81,8 +105,10 @@ private static void TimerThread()
                 timers = s_scheduledTimers!;
             }
 
+            RegisterShutdownHandler();
+
             int shortestWaitDurationMs = Timeout.Infinite;
-            while (true)
+            while (ContinueLoop())
             {
                 timerEvent.WaitOne(shortestWaitDurationMs);
 

diff --git a/src/native/libs/System.Native/entrypoints.c b/src/native/libs/System.Native/entrypoints.c
@@ -183,6 +183,7 @@ static const Entry s_sysNative[] =
     DllImportEntry(SystemNative_FreeSocketEventBuffer)
     DllImportEntry(SystemNative_TryChangeSocketEventRegistration)
     DllImportEntry(SystemNative_WaitForSocketEvents)
+    DllImportEntry(SystemNative_WakeUpSocketEventThread)
     DllImportEntry(SystemNative_PlatformSupportsDualModeIPv4PacketInfo)
     DllImportEntry(SystemNative_GetDomainSocketSizes)
     DllImportEntry(SystemNative_GetMaximumAddressSize)

diff --git a/src/native/libs/System.Native/pal_networking.c b/src/native/libs/System.Native/pal_networking.c
@@ -16,6 +16,7 @@
 #include <sys/time.h>
 #if HAVE_EPOLL
 #include <sys/epoll.h>
+#include <sys/eventfd.h>
 #elif HAVE_KQUEUE
 #include <sys/types.h>
 #include <sys/event.h>
@@ -2698,6 +2699,24 @@ static int32_t TryChangeSocketEventRegistrationInner(
     return err == 0 ? Error_SUCCESS : SystemNative_ConvertErrorPlatformToPal(errno);
 }
 
+static int32_t WakeUpSocketEventThreadInner(int32_t port, intptr_t* handle)
+{
+    struct epoll_event evt;
+    memset(&evt, 0, sizeof(struct epoll_event));
+
+    int fd = eventfd(0, EFD_SEMAPHORE);
+    evt.events = EPOLLIN;
+    evt.data.ptr = (void*)fd;
+    int err = epoll_ctl(port, EPOLL_CTL_ADD, fd, &evt);
+    if (err != 0)
+        return SystemNative_ConvertErrorPlatformToPal(errno);
+
+    *handle = (intptr_t)fd;
+
+    err = eventfd_write(fd, 1);
+    return err == 0 ? Error_SUCCESS : SystemNative_ConvertErrorPlatformToPal(errno);
+}
+
 static void ConvertEventEPollToSocketAsync(SocketEvent* sae, struct epoll_event* epoll)
 {
     assert(sae != NULL);
@@ -2891,6 +2910,24 @@ static int32_t TryChangeSocketEventRegistrationInner(
     return err == 0 ? Error_SUCCESS : SystemNative_ConvertErrorPlatformToPal(errno);
 }
 
+static int32_t WakeUpSocketEventThreadInner(int32_t port, intptr_t* handle)
+{
+    int pipeFds[2];
+    int err = pipe(pipeFds);
+    if (err != 0) return SystemNative_ConvertErrorPlatformToPal(errno);
+
+    struct kevent event;
+    EV_SET(&event, (uint64_t)pipeFds[0], EVFILT_READ, EV_ADD, 0, 0, GetKeventUdata(pipeFds[0]));
+    while ((err = kevent(port, &event, GetKeventNchanges(1), NULL, 0, NULL)) < 0 && errno == EINTR);
+    if (err != 0) return SystemNative_ConvertErrorPlatformToPal(errno);
+
+    *handle = (intptr_t)pipeFds[0];
+
+    err = write(pipeFds[1], &err, sizeof(err));
+    close(pipeFds[1]);
+    return err > 0 ? Error_SUCCESS : SystemNative_ConvertErrorPlatformToPal(errno);
+}
+
 static int32_t WaitForSocketEventsInner(int32_t port, SocketEvent* buffer, int32_t* count)
 {
     assert(buffer != NULL);
@@ -3017,6 +3054,14 @@ SystemNative_TryChangeSocketEventRegistration(intptr_t port, intptr_t socket, in
         portFd, socketFd, (SocketEvents)currentEvents, (SocketEvents)newEvents, data);
 }
 
+int32_t SystemNative_WakeUpSocketEventThread(intptr_t port, intptr_t* handle)
+{
+    int portFd = ToFileDescriptor(port);
+
+    *handle = 0;
+    return WakeUpSocketEventThreadInner(portFd, handle);
+}
+
 int32_t SystemNative_WaitForSocketEvents(intptr_t port, SocketEvent* buffer, int32_t* count)
 {
     if (buffer == NULL || count == NULL || *count < 0)

diff --git a/src/native/libs/System.Native/pal_networking.h b/src/native/libs/System.Native/pal_networking.h
@@ -420,3 +420,5 @@ PALEXPORT int32_t SystemNative_SendFile(intptr_t out_fd, intptr_t in_fd, int64_t
 PALEXPORT int32_t SystemNative_Disconnect(intptr_t socket);
 
 PALEXPORT uint32_t SystemNative_InterfaceNameToIndex(char* interfaceName);
+
+PALEXPORT int32_t SystemNative_WakeUpSocketEventThread(intptr_t port, intptr_t* handle);
diff --git a/unity/unity-aot/System.Net.Sockets/System.Net.Sockets.props b/unity/unity-aot/System.Net.Sockets/System.Net.Sockets.props
@@ -0,0 +1,3 @@
+<Project>
+    <Import Project="../UnityAot.Common.props" />
+</Project>
diff --git a/unity/unity-aot/System.Private.CoreLib/System.Private.CoreLib.props b/unity/unity-aot/System.Private.CoreLib/System.Private.CoreLib.props
@@ -1,8 +1,10 @@
 <Project>
+    <Import Project="../UnityAot.Common.props" />
+
     <PropertyGroup>
-        <DefineConstants>$(DefineConstants);UNITY_AOT</DefineConstants>
         <NativeAotSourcesRoot>$(CoreClrProjectRoot)\nativeaot</NativeAotSourcesRoot>
         <NativeAotBclSourcesRoot>$(NativeAotSourcesRoot)\System.Private.CoreLib\src</NativeAotBclSourcesRoot>
+        <UnityAotSourcesRoot>$(MSBuildThisFileDirectory)\src</UnityAotSourcesRoot>
     </PropertyGroup>
 
     <ItemGroup>
@@ -14,9 +16,14 @@
 
         <Compile Remove="$(CoreLibSharedDir)\System\Runtime\CompilerServices\RuntimeFeature.NonNativeAot.cs" />
         <Compile Include="$(NativeAotBclSourcesRoot)\System\Runtime\CompilerServices\RuntimeFeature.NativeAot.cs" />
+    </ItemGroup>
 
+    <ItemGroup>
+        <Compile Include="$(UnityAotSourcesRoot)\System\Threading\Thread.UnityAot.cs" />
     </ItemGroup>
+
     <ItemGroup>
         <ILLinkSubstitutionsXmls Include="$(MSBuildThisFileDirectory)\src\ILLink\ILLink.Substitutions.UnityAot.xml" />
+    	<ILLinkDescriptorsXmls Include="$(MSBuildThisFileDirectory)\src\ILLink\ILLink.Descriptors.UnityAot.xml" />
     </ItemGroup>
 </Project>
diff --git a/unity/unity-aot/System.Private.CoreLib/src/ILLink/ILLink.Descriptors.UnityAot.xml b/unity/unity-aot/System.Private.CoreLib/src/ILLink/ILLink.Descriptors.UnityAot.xml
@@ -0,0 +1,5 @@
+<linker>
+    <type fullname="System.Threading.Thread">
+        <method name="RegisterShutdownHandler" action="preserve" />
+    </type>
+</linker>
diff --git a/unity/unity-aot/System.Private.CoreLib/src/System/Threading/Thread.UnityAot.cs b/unity/unity-aot/System.Private.CoreLib/src/System/Threading/Thread.UnityAot.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.CompilerServices;
+using System.Runtime.InteropServices;
+
+namespace System.Threading;
+
+public partial class Thread
+{
+#if FEATURE_RUNTIME_SHUTDOWN
+
+    [UnmanagedCallersOnly]
+    private static void InvokeShutdown(IntPtr state)
+    {
+        GCHandle shutdownAction = GCHandle.FromIntPtr(state);
+        try
+        {
+            ((Action?)shutdownAction.Target)?.Invoke();
+        }
+        finally
+        {
+            shutdownAction.Free();
+        }
+    }
+
+    [MethodImpl(MethodImplOptions.InternalCall)]
+    private static extern unsafe void RegisterShutdownHandler(delegate* unmanaged<IntPtr, void> onShutdown, IntPtr state);
+
+    internal static unsafe void RegisterShutdownHandler(Action onShutdown)
+    {
+        var gcHandle = GCHandle.Alloc(onShutdown);
+        RegisterShutdownHandler(&InvokeShutdown, GCHandle.ToIntPtr(gcHandle));
+    }
+
+#endif
+}
diff --git a/unity/unity-aot/UnityAot.Common.props b/unity/unity-aot/UnityAot.Common.props
@@ -0,0 +1,5 @@
+<Project>
+    <PropertyGroup>
+        <DefineConstants>$(DefineConstants);UNITY_AOT;FEATURE_RUNTIME_SHUTDOWN</DefineConstants>
+    </PropertyGroup>
+</Project>
Original file line number	Diff line number	Diff line change
Expand Up		@@ -420,3 +420,5 @@ PALEXPORT int32_t SystemNative_SendFile(intptr_t out_fd, intptr_t in_fd, int64_t
		PALEXPORT int32_t SystemNative_Disconnect(intptr_t socket);

		PALEXPORT uint32_t SystemNative_InterfaceNameToIndex(char* interfaceName);

		PALEXPORT int32_t SystemNative_WakeUpSocketEventThread(intptr_t port, intptr_t* handle);