From bdddae759a2e1f77306e1a54523a136a0b64c078 Mon Sep 17 00:00:00 2001
From: Julien Portalier <julien@portalier.com>
Date: Mon, 9 Sep 2024 22:12:05 +0200
Subject: [PATCH] Add methods to Crystal::EventLoop (#14977)

Add `#after_fork_before_exec` to allow an evloop to do some cleanup before exec (UNIX only).
Add `#remove(io)` to allow an evloop to free resources when the IO is closed in a GC finalizer.
---
 src/crystal/scheduler.cr                         | 6 ++++++
 src/crystal/system/event_loop.cr                 | 5 +++++
 src/crystal/system/event_loop/file_descriptor.cr | 8 ++++++++
 src/crystal/system/event_loop/socket.cr          | 8 ++++++++
 src/crystal/system/file_descriptor.cr            | 4 ++++
 src/crystal/system/socket.cr                     | 4 ++++
 src/crystal/system/unix/event_loop_libevent.cr   | 9 +++++++++
 src/crystal/system/unix/process.cr               | 3 +++
 src/crystal/system/wasi/event_loop.cr            | 6 ++++++
 src/crystal/system/win32/event_loop_iocp.cr      | 6 ++++++
 src/io/file_descriptor.cr                        | 1 +
 src/socket.cr                                    | 1 +
 12 files changed, 61 insertions(+)

diff --git a/src/crystal/scheduler.cr b/src/crystal/scheduler.cr
index d3634e9aea6a..bed98ef4d05b 100644
--- a/src/crystal/scheduler.cr
+++ b/src/crystal/scheduler.cr
@@ -24,6 +24,12 @@ class Crystal::Scheduler
     Thread.current.scheduler.@event_loop
   end
 
+  def self.event_loop?
+    if scheduler = Thread.current?.try(&.scheduler?)
+      scheduler.@event_loop
+    end
+  end
+
   def self.enqueue(fiber : Fiber) : Nil
     Crystal.trace :sched, "enqueue", fiber: fiber do
       thread = Thread.current
diff --git a/src/crystal/system/event_loop.cr b/src/crystal/system/event_loop.cr
index 46954e6034ff..fb1042b21f96 100644
--- a/src/crystal/system/event_loop.cr
+++ b/src/crystal/system/event_loop.cr
@@ -17,6 +17,11 @@ abstract class Crystal::EventLoop
     Crystal::Scheduler.event_loop
   end
 
+  @[AlwaysInline]
+  def self.current? : self?
+    Crystal::Scheduler.event_loop?
+  end
+
   # Runs the loop.
   #
   # Returns immediately if events are activable. Set `blocking` to false to
diff --git a/src/crystal/system/event_loop/file_descriptor.cr b/src/crystal/system/event_loop/file_descriptor.cr
index a041263609d9..5fb6cbb95cb0 100644
--- a/src/crystal/system/event_loop/file_descriptor.cr
+++ b/src/crystal/system/event_loop/file_descriptor.cr
@@ -19,5 +19,13 @@ abstract class Crystal::EventLoop
 
     # Closes the file descriptor resource.
     abstract def close(file_descriptor : Crystal::System::FileDescriptor) : Nil
+
+    # Removes the file descriptor from the event loop. Can be used to free up
+    # memory resources associated with the file descriptor, as well as removing
+    # the file descriptor from kernel data structures.
+    #
+    # Called by `::IO::FileDescriptor#finalize` before closing the file
+    # descriptor. Errors shall be silently ignored.
+    abstract def remove(file_descriptor : Crystal::System::FileDescriptor) : Nil
   end
 end
diff --git a/src/crystal/system/event_loop/socket.cr b/src/crystal/system/event_loop/socket.cr
index e6f35478b487..6309aed391e0 100644
--- a/src/crystal/system/event_loop/socket.cr
+++ b/src/crystal/system/event_loop/socket.cr
@@ -62,5 +62,13 @@ abstract class Crystal::EventLoop
 
     # Closes the socket.
     abstract def close(socket : ::Socket) : Nil
+
+    # Removes the socket from the event loop. Can be used to free up memory
+    # resources associated with the socket, as well as removing the socket from
+    # kernel data structures.
+    #
+    # Called by `::Socket#finalize` before closing the socket. Errors shall be
+    # silently ignored.
+    abstract def remove(socket : ::Socket) : Nil
   end
 end
diff --git a/src/crystal/system/file_descriptor.cr b/src/crystal/system/file_descriptor.cr
index 481e00982e25..03868bc07034 100644
--- a/src/crystal/system/file_descriptor.cr
+++ b/src/crystal/system/file_descriptor.cr
@@ -39,6 +39,10 @@ module Crystal::System::FileDescriptor
     event_loop.write(self, slice)
   end
 
+  private def event_loop? : Crystal::EventLoop::FileDescriptor?
+    Crystal::EventLoop.current?
+  end
+
   private def event_loop : Crystal::EventLoop::FileDescriptor
     Crystal::EventLoop.current
   end
diff --git a/src/crystal/system/socket.cr b/src/crystal/system/socket.cr
index 10f902e9f0c1..8d5e8c9afaf0 100644
--- a/src/crystal/system/socket.cr
+++ b/src/crystal/system/socket.cr
@@ -99,6 +99,10 @@ module Crystal::System::Socket
   # Also used in `Socket#finalize`
   # def socket_close
 
+  private def event_loop? : Crystal::EventLoop::Socket?
+    Crystal::EventLoop.current?
+  end
+
   private def event_loop : Crystal::EventLoop::Socket
     Crystal::EventLoop.current
   end
diff --git a/src/crystal/system/unix/event_loop_libevent.cr b/src/crystal/system/unix/event_loop_libevent.cr
index b67bad63ff2f..4594f07ffe66 100644
--- a/src/crystal/system/unix/event_loop_libevent.cr
+++ b/src/crystal/system/unix/event_loop_libevent.cr
@@ -4,6 +4,9 @@ require "./event_libevent"
 class Crystal::LibEvent::EventLoop < Crystal::EventLoop
   private getter(event_base) { Crystal::LibEvent::Event::Base.new }
 
+  def after_fork_before_exec : Nil
+  end
+
   {% unless flag?(:preview_mt) %}
     # Reinitializes the event loop after a fork.
     def after_fork : Nil
@@ -93,6 +96,9 @@ class Crystal::LibEvent::EventLoop < Crystal::EventLoop
     file_descriptor.evented_close
   end
 
+  def remove(file_descriptor : Crystal::System::FileDescriptor) : Nil
+  end
+
   def read(socket : ::Socket, slice : Bytes) : Int32
     evented_read(socket, "Error reading socket") do
       LibC.recv(socket.fd, slice, slice.size, 0).to_i32
@@ -186,6 +192,9 @@ class Crystal::LibEvent::EventLoop < Crystal::EventLoop
     socket.evented_close
   end
 
+  def remove(socket : ::Socket) : Nil
+  end
+
   def evented_read(target, errno_msg : String, &) : Int32
     loop do
       bytes_read = yield
diff --git a/src/crystal/system/unix/process.cr b/src/crystal/system/unix/process.cr
index 06b18aea7b1d..420030f8ba53 100644
--- a/src/crystal/system/unix/process.cr
+++ b/src/crystal/system/unix/process.cr
@@ -185,6 +185,9 @@ struct Crystal::System::Process
       # child:
       pid = nil
       if will_exec
+        # notify event loop
+        Crystal::EventLoop.current.after_fork_before_exec
+
         # reset signal handlers, then sigmask (inherited on exec):
         Crystal::System::Signal.after_fork_before_exec
         LibC.sigemptyset(pointerof(newmask))
diff --git a/src/crystal/system/wasi/event_loop.cr b/src/crystal/system/wasi/event_loop.cr
index ba657b917154..c804c4be27aa 100644
--- a/src/crystal/system/wasi/event_loop.cr
+++ b/src/crystal/system/wasi/event_loop.cr
@@ -53,6 +53,9 @@ class Crystal::Wasi::EventLoop < Crystal::EventLoop
     file_descriptor.evented_close
   end
 
+  def remove(file_descriptor : Crystal::System::FileDescriptor) : Nil
+  end
+
   def read(socket : ::Socket, slice : Bytes) : Int32
     evented_read(socket, "Error reading socket") do
       LibC.recv(socket.fd, slice, slice.size, 0).to_i32
@@ -85,6 +88,9 @@ class Crystal::Wasi::EventLoop < Crystal::EventLoop
     socket.evented_close
   end
 
+  def remove(socket : ::Socket) : Nil
+  end
+
   def evented_read(target, errno_msg : String, &) : Int32
     loop do
       bytes_read = yield
diff --git a/src/crystal/system/win32/event_loop_iocp.cr b/src/crystal/system/win32/event_loop_iocp.cr
index d1aae09b680a..d3cfaf98d853 100644
--- a/src/crystal/system/win32/event_loop_iocp.cr
+++ b/src/crystal/system/win32/event_loop_iocp.cr
@@ -161,6 +161,9 @@ class Crystal::IOCP::EventLoop < Crystal::EventLoop
     LibC.CancelIoEx(file_descriptor.windows_handle, nil) unless file_descriptor.system_blocking?
   end
 
+  def remove(file_descriptor : Crystal::System::FileDescriptor) : Nil
+  end
+
   private def wsa_buffer(bytes)
     wsabuf = LibC::WSABUF.new
     wsabuf.len = bytes.size
@@ -271,6 +274,9 @@ class Crystal::IOCP::EventLoop < Crystal::EventLoop
 
   def close(socket : ::Socket) : Nil
   end
+
+  def remove(socket : ::Socket) : Nil
+  end
 end
 
 class Crystal::IOCP::Event
diff --git a/src/io/file_descriptor.cr b/src/io/file_descriptor.cr
index 622229e43e00..a9b303b4b58c 100644
--- a/src/io/file_descriptor.cr
+++ b/src/io/file_descriptor.cr
@@ -255,6 +255,7 @@ class IO::FileDescriptor < IO
   def finalize
     return if closed? || !close_on_finalize?
 
+    event_loop?.try(&.remove(self))
     file_descriptor_close { } # ignore error
   end
 
diff --git a/src/socket.cr b/src/socket.cr
index 1d367f805343..e97deea9eb04 100644
--- a/src/socket.cr
+++ b/src/socket.cr
@@ -430,6 +430,7 @@ class Socket < IO
   def finalize
     return if closed?
 
+    event_loop?.try(&.remove(self))
     socket_close { } # ignore error
   end