Properly close inherited servers on fork

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