Implement SafeComponentHandle and earlier managed cleanup

src/SampSharp.OpenMp.Core/IStartupContext.cs

@@ -41,4 +41,9 @@ public interface IStartupContext
     /// Occurs when the application has been initialized.
     /// </summary>
     event EventHandler? Initialized;
+
+    /// <summary>
+    /// Occurs when an open.mp component is being freed. The component which is being freed is passed as an argument.
+    /// </summary>
+    event EventHandler<IComponent>? ComponentFreed;
 }

src/SampSharp.OpenMp.Core/SampSharpInitParams.cs

@@ -31,4 +31,14 @@ public readonly ref struct SampSharpInitParams
     /// Gets information about the SampSharp open.mp component.
     /// </summary>
     public SampSharpInfo Info => _info.Value;
-}
+
+    /// <summary>
+    /// Gets a pointer to an unmanaged function that configures a cleanup callback to be invoked.
+    /// </summary>
+    public readonly unsafe delegate* unmanaged[Cdecl]<nint, void> SetOnCleanup;
+
+    /// <summary>
+    /// Gets a pointer to an unmanaged function that configures a callback to be invoked when a component is being freed.
+    /// </summary>
+    public readonly unsafe delegate* unmanaged[Cdecl]<nint, void> SetOnFreeComponent;
+}

src/SampSharp.OpenMp.Core/StartupContext.cs

@@ -1,4 +1,5 @@
 using SampSharp.OpenMp.Core.Api;
+using System.Runtime.InteropServices;
 
 namespace SampSharp.OpenMp.Core;
 
@@ -7,6 +8,12 @@ namespace SampSharp.OpenMp.Core;
 /// </summary>
 public sealed class StartupContext : IStartupContext
 {
+    // Delegates must be kept in memory to prevent GC
+    // ReSharper disable PrivateFieldCanBeConvertedToLocalVariable
+    private readonly Action _cleanup;
+    private readonly OnFreeComponentDelegate _freeComponent;
+    // ReSharper restore PrivateFieldCanBeConvertedToLocalVariable
+
     /// <summary>
     /// The API version of the native open.mp component which supports this version of the managed API. This is used to check for version mismatches between the managed and native
     /// components.
@@ -32,6 +39,19 @@ public StartupContext(SampSharpInitParams init)
             Core.LogLine(LogLevel.Error, ex.ToString());
         };
         SampSharpExceptionHandler.SetExceptionHandler(_unhandledExceptionHandler);
+
+        // Keep delegates in memory to prevent GC
+        _cleanup = OnCleanup;
+        _freeComponent = OnFreeComponent;
+
+        var cleanupPtr = Marshal.GetFunctionPointerForDelegate(_cleanup);
+        var freeComponentPtr = Marshal.GetFunctionPointerForDelegate(_freeComponent);
+
+        unsafe
+        {
+            init.SetOnCleanup(cleanupPtr);
+            init.SetOnFreeComponent(freeComponentPtr);
+        }
     }
 
     /// <inheritdoc />
@@ -63,6 +83,9 @@ public ExceptionHandler UnhandledExceptionHandler
     /// <inheritdoc />
     public event EventHandler? Initialized;
 
+    /// <inheritdoc />
+    public event EventHandler<IComponent>? ComponentFreed;
+
     /// <summary>
     /// Internal method. Do not invoke manually.
     /// </summary>
@@ -74,14 +97,6 @@ public void InitializeUsing(IStartup configurator)
         Initialized?.Invoke(this, EventArgs.Empty);
     }
 
-    /// <summary>
-    /// Internal method. Do not invoke manually.
-    /// </summary>
-    public void InvokeCleanup()
-    {
-        Cleanup?.Invoke(this, EventArgs.Empty);
-    }
-
     /// <summary>
     /// Internal method. Do not invoke manually.
     /// </summary>
@@ -116,4 +131,17 @@ private static void VersionCheck(SampSharpInitParams init)
             Environment.FailFast(message);
         }
     }
+
+    private void OnCleanup()
+    {
+        Cleanup?.Invoke(this, EventArgs.Empty);
+    }
+
+    private void OnFreeComponent(IComponent component)
+    {
+        ComponentFreed?.Invoke(this, component);
+    }
+
+    [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
+    private delegate void OnFreeComponentDelegate(IComponent component);
 }

src/SampSharp.OpenMp.Entities/Hosting/EcsHost.cs

@@ -7,32 +7,40 @@ namespace SampSharp.Entities;
 [Extension(0x57e43771d28c5e7e)]
 internal sealed partial class EcsHost(IServiceProvider serviceProvider, UnhandledExceptionHandler? exceptionHandler) : Extension
 {
+    private IStartupContext? _context;
     private IServiceProvider? _serviceProvider = serviceProvider;
 
     public IServiceProvider ServiceProvider => _serviceProvider ?? throw new InvalidOperationException();
 
     public void Start(IStartupContext context)
     {
+        _context = context;
+
         context.UseSynchronizationContext();
 
         context.UnhandledExceptionHandler = UnhandledExceptionHandler;
 
+        context.Cleanup += ContextOnCleanup;
+
         LoadSystems();
 
         // Fire initial event
         OnGameModeInit();
     }
 
-    protected override void Cleanup()
+    private void ContextOnCleanup(object? sender, EventArgs e)
     {
+        _context?.Cleanup -= ContextOnCleanup;
+        _context = null;
+
         OnGameModeExit();
 
         if (_serviceProvider is IDisposable disposable)
         {
-            //  TODO: This cleanup is called so late - we can't unsubscribe event handlers anymore, but the disposables in registered systems will try to unsubscribe them. This may cause a System.ExecutionEngineException on shutdown.
             disposable.Dispose();
-            _serviceProvider = null;
         }
+
+        _serviceProvider = null;
     }
 
     private void UnhandledExceptionHandler(string context, Exception exception)

src/SampSharp.OpenMp.Entities/Hosting/EcsHostBuilder.cs

@@ -74,7 +74,7 @@ internal EcsHost Build(IStartupContext context)
 
     private IServiceProvider BuildServiceProvider(IStartupContext context)
     {
-        var environment = new SampSharpEnvironment(context.Configurator.GetType().Assembly, context.Core, context.ComponentList);
+        var environment = new SampSharpEnvironment(context.Configurator.GetType().Assembly, context.Core, context.ComponentList, new SafeComponentHandleProvider(context));
 
         var services = new ServiceCollection();
 

src/SampSharp.OpenMp.Entities/Hosting/ISafeComponentHandle.cs

@@ -0,0 +1,7 @@
+namespace SampSharp.Entities;
+
+internal interface ISafeComponentHandle
+{
+    nint Handle { get; }
+    void Free();
+}

src/SampSharp.OpenMp.Entities/Hosting/ISafeComponentHandleProvider.cs

@@ -0,0 +1,16 @@
+using SampSharp.OpenMp.Core.Api;
+
+namespace SampSharp.Entities;
+
+/// <summary>
+/// Provides handles to open.mp components which are automatically cleared when the component is freed. This allows for safe access to components without risking access to freed components, which can lead to crashes or undefined behavior.
+/// </summary>
+public interface ISafeComponentHandleProvider
+{
+    /// <summary>
+    /// Retrieves a handle to a component of the specified type.
+    /// </summary>
+    /// <typeparam name="T">The type of the component to retrieve.</typeparam>
+    /// <returns>A <see cref="SafeComponentHandle{T}"/> representing a handle to the requested component. The handle will be cleared when the open.mp component is freed.</returns>
+    SafeComponentHandle<T> Get<T>() where T : unmanaged, IComponent.IManagedInterface;
+}

src/SampSharp.OpenMp.Entities/Hosting/SafeComponentHandle.cs

@@ -0,0 +1,65 @@
+using System.Diagnostics.CodeAnalysis;
+using SampSharp.OpenMp.Core.Api;
+
+namespace SampSharp.Entities;
+
+/// <summary>
+/// Provides a safe pointer to an open.mp component of type <typeparamref name="T"/>.
+/// </summary>
+/// <typeparam name="T">The unmanaged open.mp component type.</typeparam>
+public sealed class SafeComponentHandle<T> : ISafeComponentHandle where T : unmanaged, IComponent.IManagedInterface
+{
+    private nint _componentHandle;
+    private T _value;
+
+    internal SafeComponentHandle(T value, nint componentHandle)
+    {
+        _componentHandle = componentHandle;
+        Value = value;
+    }
+
+    /// <summary>
+    /// Gets the current value stored in the container.
+    /// </summary>
+    public T Value
+    {
+        get
+        {
+            if (_value.HasValue)
+            {
+                return _value;
+            }
+
+            ThrowDisposed();
+            return default;
+        }
+        private set => _value = value;
+    }
+
+    /// <summary>
+    /// Gets a value indicating whether the current instance contains a valid value.
+    /// </summary>
+    public bool HasValue => _value.HasValue;
+
+    nint ISafeComponentHandle.Handle => _componentHandle;
+
+    void ISafeComponentHandle.Free()
+    {
+        Value = default;
+        _componentHandle = 0;
+    }
+
+    [DoesNotReturn]
+    private static void ThrowDisposed()
+    {
+        throw new ObjectDisposedException(nameof(T));
+    }
+
+    /// <summary>
+    /// Defines an implicit conversion from <see cref="SafeComponentHandle{T}"/> to <typeparamref name="T"/>, allowing for seamless access to the underlying component value while ensuring safety against freed components.
+    /// </summary>
+    public static implicit operator T(SafeComponentHandle<T> safeHandle)
+    {
+        return safeHandle?.Value ?? default;
+    }
+}

src/SampSharp.OpenMp.Entities/Hosting/SafeComponentHandleProvider.cs

@@ -0,0 +1,52 @@
+using SampSharp.OpenMp.Core;
+using SampSharp.OpenMp.Core.Api;
+
+namespace SampSharp.Entities;
+
+internal class SafeComponentHandleProvider : ISafeComponentHandleProvider
+{
+    private readonly IStartupContext _startupContext;
+    private readonly Dictionary<UID, ISafeComponentHandle> _safeHandles = [];
+
+    public SafeComponentHandleProvider(IStartupContext startupContext)
+    {
+        _startupContext = startupContext;
+
+        startupContext.ComponentFreed += OnComponentFreed;
+        startupContext.Cleanup += OnCleanup;
+    }
+
+    public SafeComponentHandle<T> Get<T>() where T : unmanaged, IComponent.IManagedInterface
+    {
+        var uid = T.ComponentId;
+
+        if (_safeHandles.TryGetValue(uid, out var existing))
+        {
+            return (SafeComponentHandle<T>)existing;
+        }
+
+        unsafe
+        {
+            var component = _startupContext.ComponentList.QueryComponent(uid);
+
+            var typedHandle = T.FromComponentHandle(component.Handle);
+            var typedComponent = *(T*)&typedHandle;
+
+            var safeHandle = new SafeComponentHandle<T>(typedComponent, component.Handle);
+            _safeHandles[uid] = safeHandle;
+            return safeHandle;
+        }
+    }
+
+    private void OnCleanup(object? sender, EventArgs e)
+    {
+        _startupContext.ComponentFreed -= OnComponentFreed;
+        _startupContext.Cleanup -= OnCleanup;
+    }
+
+    private void OnComponentFreed(object? sender, IComponent e)
+    {
+        var handle = e.Handle;
+        _safeHandles.Values.FirstOrDefault(x => x.Handle == handle)?.Free();
+    }
+}

src/SampSharp.OpenMp.Entities/Systems/SafeEventHandlerRegistration.cs → src/SampSharp.OpenMp.Entities/Hosting/SafeEventHandlerRegistration.cs

@@ -2,7 +2,7 @@
 
 namespace SampSharp.Entities;
 
-internal sealed class SafeEventHandlerRegistration<TComponent, TEventHandler>(SampSharpEnvironment environment, TEventHandler handler, Func<TComponent, IEventDispatcher<TEventHandler>> dispatcherProvider) : IDisposable
+internal sealed class SafeEventHandlerRegistration<TComponent, TEventHandler>(SafeComponentHandle<TComponent> component, TEventHandler handler, Func<TComponent, IEventDispatcher<TEventHandler>> dispatcherProvider) : IDisposable
     where TComponent : unmanaged, IComponent.IManagedInterface
     where TEventHandler : class, IEventHandler<TEventHandler>
 {
@@ -17,8 +17,6 @@ public void Dispose()
 
         _disposed = true;
 
-        var component = environment.Components.QueryComponent<TComponent>();
-
         if (!component.HasValue)
         {
             TEventHandler.Marshaller.Marshal(handler).Free();

src/SampSharp.OpenMp.Entities/Hosting/SafeEventHandlerSampSharpEnvironmentExtensions.cs

@@ -25,7 +25,7 @@ public static class SafeEventHandlerSampSharpEnvironmentExtensions
             ArgumentNullException.ThrowIfNull(dispatcherProvider);
             ArgumentNullException.ThrowIfNull(handler);
 
-            var component = environment.Components.QueryComponent<TComponent>();
+            var component = environment.SafeComponentHandleProvider.Get<TComponent>();
 
             if (!component.HasValue)
             {
@@ -44,7 +44,7 @@ public static class SafeEventHandlerSampSharpEnvironmentExtensions
                 return null;
             }
 
-            return new SafeEventHandlerRegistration<TComponent, TEventHandler>(environment, handler, dispatcherProvider);
+            return new SafeEventHandlerRegistration<TComponent, TEventHandler>(component, handler, dispatcherProvider);
         }
 
         /// <summary>
@@ -61,7 +61,7 @@ public IDisposable AddEventHandler<TComponent, TEventHandler>(Func<TComponent, I
             where TComponent : unmanaged, IComponent.IManagedInterface
             where TEventHandler : class, IEventHandler<TEventHandler>
         {
-            var registration = TryAddEventHandler(environment, dispatcherProvider, handler, priority);
+            var registration = environment.TryAddEventHandler(dispatcherProvider, handler, priority);
             if (registration is null)
             {
                 throw new InvalidOperationException("Failed to add event handler.");

src/SampSharp.OpenMp.Entities/Hosting/SampSharpEnvironment.cs

@@ -13,4 +13,5 @@ namespace SampSharp.Entities;
 /// <param name="EntryAssembly">The assembly which was configured to launch in open.mp. Used to discover game mode classes and other application types.</param>
 /// <param name="Core">The <see cref="ICore" /> interface for the open.mp server. Provides access to core server functionality and extensions.</param>
 /// <param name="Components">The <see cref="IComponentList" /> of open.mp. Manages all game components (players, vehicles, objects, etc.) accessible on the server.</param>
-public record SampSharpEnvironment(Assembly EntryAssembly, ICore Core, IComponentList Components);
+/// <param name="SafeComponentHandleProvider">A provider of safe handles of open.mp components.</param>
+public record SampSharpEnvironment(Assembly EntryAssembly, ICore Core, IComponentList Components, ISafeComponentHandleProvider SafeComponentHandleProvider);