Properly close inherited servers on fork

I noticed some UNIXServer and AppServer accumulating with
reforking enabled.

The reason is that the `at_exit { stop }` make `Server` instances
immortal, so they are never GC, and their open socket accumulate and
are never closed.

That callback is useless anyway, as on exit the finalizer of these
objects will be invoked, and they'll cleanup after themselves.

CHANGELOG.md

@@ -4,6 +4,10 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## Unreleased
+
+- Fix `AppProfiler::Server` to no longer be immortal. This avoid leaking servers when respawning the server after fork (#93)
+
 ## [0.1.4] - 2023-09-19
 
 - Allow passing custom parameters instance to Middleware#call, to programatically trigger profiling (#87)

lib/app_profiler/server.rb

@@ -198,11 +198,26 @@ def stop
       end
 
       class UNIX < Transport
+        class << self
+          def unlink_socket(path, pid)
+            ->(_) do
+              if Process.pid == pid && File.exist?(path)
+                begin
+                  File.unlink(path)
+                rescue SystemCallError
+                  # Let not raise in a finalizer
+                end
+              end
+            end
+          end
+        end
+
         def start
           FileUtils.mkdir_p(PROFILER_TEMPFILE_PATH)
           @socket_file = File.join(PROFILER_TEMPFILE_PATH, "app-profiler-#{Process.pid}.sock")
           File.unlink(@socket_file) if File.exist?(@socket_file) && File.socket?(@socket_file)
           @socket = UNIXServer.new(@socket_file)
+          ObjectSpace.define_finalizer(self, self.class.unlink_socket(@socket_file, Process.pid))
         end
 
         def client
@@ -213,6 +228,10 @@ def stop
           File.unlink(@socket_file) if File.exist?(@socket_file) && File.socket?(@socket_file)
           @socket.close
         end
+
+        def abandon
+          @socket.close
+        end
       end
 
       class TCP < Transport
@@ -238,6 +257,11 @@ def stop
           @port_file.unlink
           @socket.close
         end
+
+        def abandon
+          @port_file.close # NB: Tempfile finalizer checks Process.pid to avoid unlinking inherited IOs.
+          @socket.close
+        end
       end
 
       def initialize(transport, logger)
@@ -257,13 +281,16 @@ def initialize(transport, logger)
           "[AppProfiler::Server] listening on addr=#{@transport.socket.addr}"
         )
         @pid = Process.pid
-        at_exit { stop }
       end
 
       def client
         @transport.client
       end
 
+      def join(...)
+        @listen_thread.join(...)
+      end
+
       def serve
         return unless @listen_thread.nil?
 
@@ -324,10 +351,12 @@ def serve
       end
 
       def stop
-        return unless @pid == Process.pid
-
         @listen_thread.kill
-        @transport.stop
+        if @pid == Process.pid
+          @transport.stop
+        else
+          @transport.abandon
+        end
       end
     end
 
@@ -371,7 +400,10 @@ def stop
       private
 
       def profile_server
-        @profile_server = nil if @pid != Process.pid
+        if @pid != Process.pid
+          @profile_server&.stop
+          @profile_server = nil
+        end
         @profile_server
       end
 

test/app_profiler/profile_server_test.rb

@@ -82,6 +82,47 @@ class UNIXServerTest < ProfileServerTestCase
       ensure
         assert_nil(Server.stop)
       end
+
+      Transport = AppProfiler::Server.const_get(:ProfileServer)::Transport
+      test "servers are abandoned on fork" do
+        server = AppProfiler::Server.start(@logger)
+        server.join(1)
+
+        open_servers = ObjectSpace.each_object(Transport).select { |t| !t.socket.closed? }
+        assert_equal(1, open_servers.size)
+
+        parent_socket = File.join(ProfileServer::PROFILER_TEMPFILE_PATH, "app-profiler-#{Process.pid}.sock")
+        assert_equal(true, File.exist?(parent_socket))
+
+        r, w = IO.pipe
+        pid = fork do
+          open_servers = ObjectSpace.each_object(Transport).select { |t| !t.socket.closed? }
+          w.write(Marshal.dump(open_servers.size))
+
+          AppProfiler::Server.start(@logger)
+
+          open_servers = ObjectSpace.each_object(Transport).select { |t| !t.socket.closed? }
+          w.write(Marshal.dump(open_servers.size))
+          Marshal.load(r)
+        end
+        assert_equal(1, Marshal.load(r))
+        assert_equal(1, Marshal.load(r))
+
+        child_socket = File.join(ProfileServer::PROFILER_TEMPFILE_PATH, "app-profiler-#{pid}.sock")
+        assert_equal(true, File.exist?(parent_socket))
+        assert_equal(true, File.exist?(child_socket))
+
+        w.write(Marshal.dump(:done))
+        w.close
+
+        _, status = Process.wait2(pid)
+        assert_predicate(status, :success?)
+
+        assert_equal(true, File.exist?(parent_socket))
+        assert_equal(false, File.exist?(child_socket))
+      ensure
+        server&.stop
+      end
     end
 
     class ProfileServerTest < ProfileServerTestCase