+ Add ability to inherit server sockets on restart

bin/puma

@@ -7,10 +7,4 @@ require 'puma/cli'
 
 cli = Puma::CLI.new ARGV
 
-begin
 cli.run
-rescue => e
-  raise e if $DEBUG
-  STDERR.puts e.message
-  exit 1
-end

lib/puma/app/status.rb

@@ -30,7 +30,8 @@ module Puma
 
         when "/restart"
           if @cli and @cli.restart_on_stop!
-            @server.stop
+            @server.begin_restart
+
             return [200, {}, ['{ "status": "ok" }']]
           else
             return [200, {}, ['{ "status": "not configured" }']]

lib/puma/cli.rb

@@ -33,9 +33,26 @@ module Puma
       @restart = false
       @temp_status_path = nil
 
+      @listeners = []
+
       setup_options
 
       generate_restart_data
+
+      @inherited_fds = {}
+      remove = []
+
+      ENV.each do |k,v|
+        if k =~ /PUMA_INHERIT_\d+/
+          fd, url = v.split(":", 2)
+          @inherited_fds[url] = fd.to_i
+          remove << k
+        end
+      end
+
+      remove.each do |k|
+        ENV.delete k
+      end
     end
 
     def restart_on_stop!
@@ -62,6 +79,8 @@ module Puma
 
       @restart_dir ||= Dir.pwd
 
+      @original_argv = ARGV.dup
+
       if defined? Rubinius::OS_ARGV
         @restart_argv = Rubinius::OS_ARGV
       else
@@ -72,20 +91,47 @@ module Puma
         # picked up in PATH.
         #
         if File.exists?($0)
-          @restart_argv = [Gem.ruby, $0] + ARGV
+          arg0 = [Gem.ruby, $0]
         else
-          @restart_argv = [Gem.ruby, "-S", $0] + ARGV
+          arg0 = [Gem.ruby, "-S", $0]
         end
+
+        @restart_argv = arg0 + ARGV
       end
     end
 
     def restart!
+      @options[:on_restart].each do |blk|
+        blk.call self
+      end
+
       if IS_JRUBY
+        @listeners.each_with_index do |(str,io),i|
+          io.close
+
+          # We have to unlink a unix socket path that's not being used
+          uri = URI.parse str
+          if uri.scheme == "unix"
+            path = "#{uri.host}#{uri.path}"
+            File.unlink path
+          end
+        end
+
         require 'puma/jruby_restart'
         JRubyRestart.chdir_exec(@restart_dir, Gem.ruby, *@restart_argv)
       else
+        @listeners.each_with_index do |(l,io),i|
+          ENV["PUMA_INHERIT_#{i}"] = "#{io.to_i}:#{l}"
+        end
+
+        if cmd = @options[:restart_cmd]
+          argv = cmd.split(' ') + @original_argv
+        else
+          argv = @restart_argv
+        end
+
         Dir.chdir @restart_dir
-        Kernel.exec(*@restart_argv)
+        Kernel.exec(*argv)
       end
     end
 
@@ -163,6 +209,11 @@ module Puma
           end
         end
 
+        o.on "--restart-cmd CMD",
+             "The puma command to run during a hot restart",
+             "Default: inferred" do |cmd|
+          @options[:restart_cmd] = cmd
+        end
       end
 
       @parser.banner = "puma <options> <rackup file>"
@@ -190,7 +241,10 @@ module Puma
       if path = @options[:state]
         state = { "pid" => Process.pid }
 
-        state["config"] = @config
+        cfg = @config.dup
+        cfg.options.delete :on_restart
+
+        state["config"] = cfg
 
         File.open(path, "w") do |f|
           f.write state.to_yaml
@@ -237,9 +291,20 @@ module Puma
         uri = URI.parse str
         case uri.scheme
         when "tcp"
+          if fd = @inherited_fds.delete(str)
+            log "* Inherited #{str}"
+            io = server.inherit_tcp_listener uri.host, uri.port, fd
+          else
             log "* Listening on #{str}"
-          server.add_tcp_listener uri.host, uri.port
+            io = server.add_tcp_listener uri.host, uri.port
+          end
+
+          @listeners << [str, io]
         when "unix"
+          if fd = @inherited_fds.delete(str)
+            log "* Inherited #{str}"
+            io = server.inherit_unix_listener uri.path, fd
+          else
             log "* Listening on #{str}"
             path = "#{uri.host}#{uri.path}"
 
@@ -253,9 +318,11 @@ module Puma
               end
             end
 
-          server.add_unix_listener path, umask
+            io = server.add_unix_listener path, umask
+          end
+
+          @listeners << [str, io]
         when "ssl"
-          log "* Listening on #{str}"
           params = Rack::Utils.parse_query uri.query
           require 'openssl'
 
@@ -274,12 +341,38 @@ module Puma
 
           ctx.verify_mode = OpenSSL::SSL::VERIFY_NONE
 
-          server.add_ssl_listener uri.host, uri.port, ctx
+          if fd = @inherited_fds.delete(str)
+            log "* Inherited #{str}"
+            io = server.inherited_ssl_listener fd, ctx
+          else
+            log "* Listening on #{str}"
+            io = server.add_ssl_listener uri.host, uri.port, ctx
+          end
+
+          @listeners << [str, io]
         else
           error "Invalid URI: #{str}"
         end
       end
 
+      # If we inherited fds but didn't use them (because of a
+      # configuration change), then be sure to close them.
+      @inherited_fds.each do |str, fd|
+        log "* Closing unused inherited connection: #{str}"
+
+        begin
+          IO.for_fd(fd).close
+        rescue SystemCallError
+        end
+
+        # We have to unlink a unix socket path that's not being used
+        uri = URI.parse str
+        if uri.scheme == "unix"
+          path = "#{uri.host}#{uri.path}"
+          File.unlink path
+        end
+      end
+
       @server = server
 
       if str = @options[:control_url]
@@ -314,6 +407,11 @@ module Puma
         @status = status
       end
 
+      Signal.trap "SIGUSR2" do
+        @restart = true
+        server.begin_restart
+      end
+
       log "Use Ctrl-C to stop"
 
       begin

lib/puma/configuration.rb

@@ -8,10 +8,15 @@ module Puma
     def initialize(options)
       @options = options
       @options[:binds] ||= []
+      @options[:on_restart] ||= []
     end
 
     attr_reader :options
 
+    def initialize_copy(other)
+      @options = @options.dup
+    end
+
     def load
       if path = @options[:config_file]
         DSL.new(@options)._load_from path
@@ -150,6 +155,15 @@ module Puma
         @options[:binds] << url
       end
 
+      # Code to run before doing a restart. This code should
+      # close logfiles, database connections, etc.
+      #
+      # This can be called multiple times to add code each time.
+      #
+      def on_restart(&blk)
+        @options[:on_restart] << blk
+      end
+
       # Store the pid of the server in the file at +path+.
       def pidfile(path)
         @options[:pidfile] = path

lib/puma/const.rb

@@ -133,6 +133,7 @@ module Puma
 
     STOP_COMMAND = "?".freeze
     HALT_COMMAND = "!".freeze
+    RESTART_COMMAND = "R".freeze
 
     RACK_INPUT = "rack.input".freeze
     RACK_URL_SCHEME = "rack.url_scheme".freeze

lib/puma/server.rb

@@ -109,6 +109,13 @@ module Puma
       end
       s.listen backlog
       @ios << s
+      s
+    end
+
+    def inherit_tcp_listener(host, port, fd)
+      s = TCPServer.for_fd(fd)
+      @ios << s
+      s
     end
 
     def add_ssl_listener(host, port, ctx, optimize_for_latency=true, backlog=1024)
@@ -119,6 +126,13 @@ module Puma
       s.listen backlog
       @proto_env[HTTPS_KEY] = HTTPS
       @ios << OpenSSL::SSL::SSLServer.new(s, ctx)
+      s
+    end
+
+    def inherited_ssl_listener(fd, ctx)
+      s = TCPServer.for_fd(fd)
+      @ios << OpenSSL::SSL::SSLServer.new(s, ctx)
+      s
     end
 
     # Tell the server to listen on +path+ as a UNIX domain socket.
@@ -131,10 +145,22 @@ module Puma
 
       begin
         old_mask = File.umask(umask)
-        @ios << UNIXServer.new(path)
+        s = UNIXServer.new(path)
+        @ios << s
       ensure
         File.umask old_mask
       end
+
+      s
+    end
+
+    def inherit_unix_listener(path, fd)
+      @unix_paths << path
+
+      s = UNIXServer.for_fd fd
+      @ios << s
+
+      s
     end
 
     def backlog
@@ -187,10 +213,12 @@ module Puma
 
           graceful_shutdown if @status == :stop
         ensure
+          unless @status == :restart
             @ios.each { |i| i.close }
             @unix_paths.each { |i| File.unlink i }
           end
         end
+      end
 
       return @thread
     end
@@ -206,6 +234,9 @@ module Puma
       when HALT_COMMAND
         @status = :halt
         return true
+      when RESTART_COMMAND
+        @status = :restart
+        return true
       end
 
       return false
@@ -585,5 +616,10 @@ module Puma
 
       @thread.join if @thread && sync
     end
+
+    def begin_restart
+      @persistent_wakeup.close
+      @notify << RESTART_COMMAND
+    end
   end
 end