diff --git a/lib/pitchfork/info.rb b/lib/pitchfork/info.rb index 85f20cf8..e44b2a69 100644 --- a/lib/pitchfork/info.rb +++ b/lib/pitchfork/info.rb @@ -12,7 +12,8 @@ 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 @@ -20,23 +21,25 @@ 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 diff --git a/lib/pitchfork/worker.rb b/lib/pitchfork/worker.rb index d2b1d341..5884bc57 100644 --- a/lib/pitchfork/worker.rb +++ b/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 diff --git a/test/unit/test_children.rb b/test/unit/test_children.rb index 34808f67..85b96df1 100644 --- a/test/unit/test_children.rb +++ b/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 diff --git a/test/unit/test_info.rb b/test/unit/test_info.rb new file mode 100644 index 00000000..4f9fb50d --- /dev/null +++ b/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