Merge pull request #56 from Shopify/close-fds

[Experimental] Pitchfork::Info.close_all_fds!

lib/pitchfork/flock.rb

@@ -19,6 +19,10 @@ def at_fork
       nil
     end
 
+    def to_io
+      @file
+    end
+
     def unlink
       File.unlink(@file.path)
     rescue Errno::ENOENT

lib/pitchfork/http_server.rb

@@ -129,6 +129,7 @@ def initialize(app, options = {})
       @respawn = false
       @last_check = Pitchfork.time_now
       @promotion_lock = Flock.new("pitchfork-promotion")
+      Info.keep_io(@promotion_lock)
 
       options = options.dup
       @ready_pipe = options.delete(:ready_pipe)
@@ -180,6 +181,7 @@ def start(sync = true)
       # It's also used by newly spawned children to send their soft_signal pipe
       # to the master when they are spawned.
       @control_socket.replace(Pitchfork.socketpair)
+      Info.keep_ios(@control_socket)
       @master_pid = $$
 
       # setup signal handlers before writing pid file in case people get
@@ -264,6 +266,7 @@ def listen(address, opt = {}.merge(listener_opts[address] || {}))
           io = server_cast(io)
         end
         logger.info "listening on addr=#{sock_name(io)} fd=#{io.fileno}"
+        Info.keep_io(io)
         LISTENERS << io
         io
       rescue Errno::EADDRINUSE => err

lib/pitchfork/info.rb

@@ -6,10 +6,49 @@ module Pitchfork
   module Info
     @workers_count = 0
     @fork_safe = true
+    @kept_ios = ObjectSpace::WeakMap.new
 
     class << self
       attr_accessor :workers_count
 
+      def keep_io(io)
+        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_ios!
+        ignored_ios = [$stdin, $stdout, $stderr]
+
+        @kept_ios.each_value do |io_like|
+          ignored_ios << (io_like.is_a?(IO) ? io_like : io_like.to_io)
+        end
+
+        ObjectSpace.each_object(IO) do |io|
+          closed = begin
+            io.closed?
+          rescue IOError
+            true
+          end
+
+          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
+
       def fork_safe?
         @fork_safe
       end

lib/pitchfork/worker.rb

@@ -198,7 +198,7 @@ def close # :nodoc:
     end
 
     def create_socketpair!
-      @to_io, @master = Pitchfork.socketpair
+      @to_io, @master = Info.keep_ios(Pitchfork.socketpair)
     end
 
     def after_fork_in_child
@@ -208,6 +208,8 @@ 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/slow/test_server.rb

@@ -80,6 +80,7 @@ class WebServerStartTest < Pitchfork::Test
 
   def test_preload_app
     tmp = Tempfile.new('test_preload_app_config')
+    Pitchfork::Info.keep_io(tmp)
     ObjectSpace.undefine_finalizer(tmp)
     app = lambda { ||
       tmp.sysseek(0)

test/slow/test_signals.rb

@@ -25,8 +25,8 @@ def setup
     @bs = 1 * 1024 * 1024
     @count = 100
     @port = unused_port
-    @sock = Tempfile.new('pitchfork.sock')
-    @tmp = Tempfile.new('pitchfork.write')
+    @sock = Pitchfork::Info.keep_io(Tempfile.new('pitchfork.sock'))
+    @tmp = Pitchfork::Info.keep_io(Tempfile.new('pitchfork.write'))
     @tmp.sync = true
     File.unlink(@sock.path)
     File.unlink(@tmp.path)
@@ -75,7 +75,7 @@ def test_worker_dies_on_dead_master
   end
 
   def test_sleepy_kill
-    rd, wr = IO.pipe
+    rd, wr = Pitchfork::Info.keep_ios(IO.pipe)
     pid = fork {
       rd.close
       app = lambda { |env| wr.syswrite('.'); sleep; [ 200, {}, [] ] }

test/slow/test_upload.rb

@@ -143,7 +143,7 @@ def test_put_keepalive_truncates_small_overwrite
   end
 
   def test_put_excessive_overwrite_closed
-    tmp = Tempfile.new('overwrite_check')
+    tmp = Pitchfork::Info.keep_io(Tempfile.new('overwrite_check'))
     tmp.sync = true
     start_server(lambda { |env|
       nr = 0

test/unit/test_ccc.rb

@@ -6,9 +6,9 @@ def test_ccc_tcpi
     host = '127.0.0.1'
     port = unused_port
 
-    rd, wr = IO.pipe
-    sleep_pipe = IO.pipe
-    ready_read, ready_write = IO.pipe
+    rd, wr = Pitchfork::Info.keep_ios(IO.pipe)
+    sleep_pipe = Pitchfork::Info.keep_ios(IO.pipe)
+    ready_read, ready_write = Pitchfork::Info.keep_ios(IO.pipe)
 
     pid = fork do
       ready_read.close

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