Add more options to pumactl

README.md

@@ -118,10 +118,16 @@ You should place code to close global log files, redis connections, etc in this
 
 If you start puma with `-S some/path` then you can pass that same path to the `pumactl` program to control your server. For instance:
 
-    $ pumactl -S some/path restart
+    $ pumactl -S some/path command
+
+or 
+
+    $ pumactl -C url -T token command
 
 will cause the server to perform a restart. `pumactl` is a simple CLI frontend to the control/status app described above.
 
+Allowed commands: status, restart, halt, stop
+
 ## Managing multiple Pumas / init.d script 
 
 If you want an easy way to manage multiple scripts at once check [tools/jungle](https://github.com/puma/puma/tree/master/tools/jungle) for an init.d script.

lib/puma/control_cli.rb

@@ -1,139 +1,157 @@
 require 'optparse'
-
 require 'puma/const'
 require 'puma/configuration'
-
 require 'yaml'
 require 'uri'
-
 require 'socket'
-
 module Puma
   class ControlCLI
 
+    COMMANDS = %w{status restart stop halt}
+
+    def is_windows?
+      RUBY_PLATFORM =~ /(win|w)32$/ ? true : false
+    end
+
     def initialize(argv, stdout=STDOUT, stderr=STDERR)
-      @argv = argv
       @stdout = stdout
       @stderr = stderr
-      @path = nil
-      @config = nil
-    end
-
-    def setup_options
-      @parser = OptionParser.new do |o|
-        o.on "-S", "--state PATH", "Where the state file to use is" do |arg|
-          @path = arg
+      @options = {}
+
+      OptionParser.new do |option|
+        option.banner = "Usage: pumactl (-S status_file | -C url -T token) (#{COMMANDS.join("|")})"
+        option.on "-S", "--state PATH", "Where the state file to use is" do |arg|
+          @options[:status_path] = arg
+        end
+        option.on "-Q", "--quiet", "Not display messages" do |arg|
+          @options[:quiet_flag] = true
+        end
+        option.on "-P", "--pidfile PATH", "Pid file" do |arg|
+          @options[:pid_file] = arg
+        end
+        option.on "-C", "--control-url URL", "The bind url to use for the control server" do |arg|
+          @options[:control_url] = arg
         end
+        option.on "-T", "--control-token TOKEN", "The token to use as authentication for the control server" do |arg|
+          @options[:control_auth_token] = arg
+        end
+        option.on_tail("-H", "--help", "Show this message") do
+          @stdout.puts option
+          exit
+        end
+        option.on_tail("-V", "--version", "Show version") do
+          puts Const::PUMA_VERSION
+          exit
+        end
+      end.parse!(argv)
+
+      command = argv.shift
+      @options[:command] = command if command
+
+      # check present of command
+      unless @options[:command]
+        raise "Available commands: #{COMMANDS.join(", ")}"
+      end      
+      unless COMMANDS.include? @options[:command]
+        raise "Invalid command: #{@options[:command]}" 
       end
+
+    rescue => e
+      @stdout.puts e.message
+      exit 1
+    end
+
+    def message msg
+      @stdout.puts msg unless @options[:quiet_flag]
     end
 
-    def connect
-      if str = @config.options[:control_url]
-        uri = URI.parse str
-        case uri.scheme
-        when "tcp"
-          return TCPSocket.new uri.host, uri.port
-        when "unix"
-          path = "#{uri.host}#{uri.path}"
-          return UNIXSocket.new path
+    def prepare_configuration
+      if @options.has_key? :status_path
+        raise "Status file not found: #{@options[:status_path]} " unless File.exist? @options[:status_path]
+        status = YAML.load File.read(@options[:status_path])
+        if status.has_key? "config"
+          # get control_url
+          if status["config"].options.has_key?(:control_url)
+            @options[:control_url] = status["config"].options[:control_url]
+          end
+          # get control_auth_token
+          if status["config"].options.has_key?(:control_auth_token)
+            @options[:control_auth_token] = status["config"].options[:control_auth_token]
+          end
+          # get pid
+          @options[:pid] = status["pid"].to_i
         else
-          raise "Invalid URI: #{str}"
+          raise "Invalid status file: #{@options[:status_path]}"
         end
+      elsif @options.has_key? :pid_file
+        # get pid from pid_file
+        @options[:pid] = File.open(@options[:pid_file]).gets.to_i
       end
-
-      raise "No status address configured"
+
     end
 
-    def run
-      setup_options
-
-      @parser.order!(@argv) { |a| @parser.terminate a }
-
-      if @path
-        @state = YAML.load File.read(@path)
-        @config = @state['config']
-      end
-
-      cmd = @argv.shift
-
-      meth = "command_#{cmd}"
-
-      if respond_to?(meth)
-        __send__(meth)
+    def send_request
+      uri = URI.parse @options[:control_url]
+
+      # create server object by scheme
+      @server = case uri.scheme
+      when "tcp"
+        TCPSocket.new uri.host, uri.port
+      when "unix"
+        UNIXSocket.new "#{uri.host}#{uri.path}"
       else
-        raise "Unknown command: #{cmd}"
+        raise "Invalid scheme: #{uri.scheme}"
       end
-    end
-
-    def request(sock, url)
-      token = @config.options[:control_auth_token]
-      if token
-        url = "#{url}?token=#{token}"
-      end
-
-      sock << "GET #{url} HTTP/1.0\r\n\r\n"
-
-      rep = sock.read.split("\r\n")
-
-      m = %r!HTTP/1.\d (\d+)!.match(rep.first)
-      if m[1] == "403"
-        raise "Unauthorized access to server (wrong auth token)"
-      elsif m[1] != "200"
-        raise "Bad response code from server: #{m[1]}"
-      end
-
-      return rep.last
-    end
-
-    def command_pid
-      @stdout.puts "#{@state['pid']}"
-    end
-
-    def command_start
-      require 'puma/cli'
-
-      cli = Puma::CLI.new @argv, @stdout, @stderr
-      cli.run
-    end
-
-    def command_stop
-      sock = connect
-      body = request sock, "/stop"
-
-      if body != '{ "status": "ok" }'
-        raise "Invalid response: '#{body}'"
+
+      unless @options[:command] == "status"
+        url = "/#{@options[:command]}"
+        if @options.has_key?(:control_auth_token)
+          url = url + "?token=#{@options[:control_auth_token]}"
+        end
+        @server << "GET #{url} HTTP/1.0\r\n\r\n"
+        response = @server.read.split("\r\n")
+        (@http,@code,@message) = response.first.split(" ")
+        if @code == "403"
+          raise "Unauthorized access to server (wrong auth token)"
+        elsif @code != "200"
+          raise "Bad response from server: #{@code}"
+        end
+        message "Command #{@options[:command]} sent success"
       else
-        @stdout.puts "Requested stop from server"
+        message "Puma is started"
       end
+
+      @server.close
     end
 
-    def command_halt
-      sock = connect
-      body = request sock, "/halt"
-
-      if body != '{ "status": "ok" }'
-        raise "Invalid response: '#{body}'"
+    def send_signal
+      Process.getpgid(@options[:pid])
+      case @options[:command]
+      when "restart"
+        Process.kill("SIGUSR2", @options[:pid])
+      when "halt"
+        Process.kill("QUIT", @options[:pid])
+      when "stop"
+        Process.kill("SIGTERM", @options[:pid])
       else
-        @stdout.puts "Requested halt from server"
+        message "Puma is started"
+        return
       end
+      message "Command #{@options[:command]} sent success"
     end
 
-    def command_restart
-      sock = connect
-      body = request sock, "/restart"
-
-      if body != '{ "status": "ok" }'
-        raise "Invalid response: '#{body}'"
+    def run
+      prepare_configuration
+
+      if is_windows?
+        send_request
       else
-        @stdout.puts "Requested restart from server"
+        @options.has_key?(:control_url) ? send_request : send_signal
       end
-    end
-
-    def command_stats
-      sock = connect
-      body = request sock, "/stats"
 
-      @stdout.puts body
+    rescue => e
+      message e.message
+      exit 1
     end
   end
-end
+end