Pitchfork::Info.close_all_ios!

lib/pitchfork/info.rb

@@ -12,31 +12,33 @@ class << self
       attr_accessor :workers_count
 
       def keep_io(io)
-        @kept_ios[io] = io if io && !io.to_io.closed?
+        @kept_ios[io] = io
         io
       end
 
       def keep_ios(ios)
         ios.each { |io| keep_io(io) }
       end
 
-      def close_all_fds!
-        ignored_fds = [$stdin.to_i, $stdout.to_i, $stderr.to_i]
+      def close_all_ios!
+        ignored_ios = [$stdin, $stdout, $stderr]
+
         @kept_ios.each_value do |io_like|
-          if io = io_like&.to_io
-            ignored_fds << io.to_i unless io.closed?
-          end
+          ignored_ios << (io_like.is_a?(IO) ? io_like : io_like.to_io)
         end
 
-        all_fds = Dir.children("/dev/fd").map(&:to_i)
-        all_fds -= ignored_fds
-
-        all_fds.each do |fd|
-          IO.for_fd(fd).close
-        rescue ArgumentError
-          # RubyVM internal file descriptor, leave it alone
-        rescue Errno::EBADF
-          # Likely a race condition
+        ObjectSpace.each_object(IO) do |io|
+          if !io.closed? && io.autoclose? && !ignored_ios.include?(io)
+            if io.is_a?(TCPSocket)
+              # If we inherited a TCP Socket, calling #close directly could send FIN or RST.
+              # So we first reopen /dev/null to avoid that.
+              io.reopen(File::NULL)
+            end
+            begin
+              io.close
+            rescue Errno::EBADF
+            end
+          end
         end
       end
 

test/unit/test_info.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+require 'test_helper'
+
+class TestInfo < Pitchfork::Test
+  def test_close_all_ios_except_marked_ones
+    r, w = IO.pipe
+
+    Pitchfork::Info.keep_io(w)
+
+    pid = Process.fork do
+      Pitchfork::Info.close_all_ios!
+
+      w.write(Marshal.dump([
+        $stdin.closed?,
+        $stdout.closed?,
+        $stderr.closed?,
+        r.closed?,
+        w.closed?
+      ]))
+      Process.exit!(0)
+    end
+
+    _, status = Process.wait2(pid)
+    assert_predicate status, :success?
+
+    info = Marshal.load(r)
+
+    assert_equal([
+      false, # $stdin
+      false, # $stdout
+      false, # $stderr
+      true, # r
+      false, # w
+    ], info)
+  end
+end