Pitchfork::Info.close_all_ios!

lib/pitchfork/info.rb

@@ -12,31 +12,40 @@ class << self
       attr_accessor :workers_count
 
       def keep_io(io)
-        @kept_ios[io] = io if io && !io.to_io.closed?
+        raise ArgumentError, "#{io.inspect} doesn't respond to :to_io" unless io.respond_to?(:to_io)
+        @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
+        ObjectSpace.each_object(IO) do |io|
+          closed = begin
+            io.closed?
+          rescue IOError
+            true
+          end
 
-        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
+          if !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
 

lib/pitchfork/worker.rb

@@ -208,6 +208,7 @@ def after_fork_in_child
     private
 
     def pipe=(socket)
+      raise ArgumentError, "pipe can't be nil" unless socket
       Info.keep_io(socket)
       @master = MessageSocket.new(socket)
     end

test/unit/test_children.rb

@@ -18,19 +18,21 @@ def test_register
     end
 
     def test_message_worker_spawned
+      pipe = IO.pipe.last
       worker = Worker.new(0)
       @children.register(worker)
       assert_predicate @children, :pending_workers?
 
-      @children.update(Message::WorkerSpawned.new(0, 42, 0, nil))
+      @children.update(Message::WorkerSpawned.new(0, 42, 0, pipe))
       refute_predicate @children, :pending_workers?
       assert_equal 42, worker.pid
       assert_equal [worker], @children.workers
     end
 
     def test_message_mold_spawned
+      pipe = IO.pipe.last
       assert_nil @children.mold
-      @children.update(Message::MoldSpawned.new(nil, 42, 1, nil))
+      @children.update(Message::MoldSpawned.new(nil, 42, 1, pipe))
 
       assert_nil @children.mold
       assert_equal 0, @children.molds.size
@@ -40,8 +42,9 @@ def test_message_mold_spawned
     end
 
     def test_message_mold_ready
+      pipe = IO.pipe.last
       assert_nil @children.mold
-      @children.update(Message::MoldSpawned.new(nil, 42, 1, nil))
+      @children.update(Message::MoldSpawned.new(nil, 42, 1, pipe))
       mold = @children.update(Message::MoldReady.new(nil, 42, 1))
 
       assert_equal mold, @children.mold
@@ -52,26 +55,28 @@ def test_message_mold_ready
     end
 
     def test_reap_worker
+      pipe = IO.pipe.last
       worker = Worker.new(0)
       @children.register(worker)
       assert_predicate @children, :pending_workers?
 
-      @children.update(Message::WorkerSpawned.new(0, 42, 0, nil))
+      @children.update(Message::WorkerSpawned.new(0, 42, 0, pipe))
 
       assert_equal worker, @children.reap(worker.pid)
       assert_nil @children.reap(worker.pid)
     end
 
     def test_reap_old_molds
+      pipe = IO.pipe.last
       assert_nil @children.mold
-      @children.update(Message::MoldSpawned.new(nil, 24, 0, nil))
+      @children.update(Message::MoldSpawned.new(nil, 24, 0, pipe))
       @children.update(Message::MoldReady.new(nil, 24, 0))
 
       first_mold = @children.mold
       refute_nil first_mold
       assert_equal 24, first_mold.pid
 
-      @children.update(Message::MoldSpawned.new(nil, 42, 1, nil))
+      @children.update(Message::MoldSpawned.new(nil, 42, 1, pipe))
       @children.update(Message::MoldReady.new(nil, 42, 1))
       second_mold = @children.mold
       refute_nil second_mold

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