From 96ca09ee2106ecbd9b5d63566dcb6f96716552cd Mon Sep 17 00:00:00 2001
From: Tyler Treat <tyler.treat@apcera.com>
Date: Mon, 26 Jun 2017 10:50:23 -0500
Subject: [PATCH] Add --signal flag

---
 main.go                  | 37 +++++++++++++++++++---
 server/signal.go         | 66 ++++++++++++++++++++++++++++++++++++++++
 server/signal_windows.go |  8 +++++
 3 files changed, 107 insertions(+), 4 deletions(-)

diff --git a/main.go b/main.go
index 9920f1bf..d4c1e849 100644
--- a/main.go
+++ b/main.go
@@ -8,6 +8,8 @@ import (
 	"net"
 	"net/url"
 	"os"
+	"strconv"
+	"strings"
 
 	"github.com/nats-io/gnatsd/server"
 )
@@ -22,6 +24,7 @@ Server Options:
     -m, --http_port <port>           Use port for http monitoring
     -ms,--https_port <port>          Use port for https monitoring
     -c, --config <file>              Configuration file
+    -sl,--signal <signal>[=<pid>]    Send signal to gnatsd process (stop, quit, reopen, reload)
 
 Logging Options:
     -l, --log <file>                 File to redirect log output
@@ -67,10 +70,13 @@ func main() {
 	// Server Options
 	opts := &server.Options{}
 
-	var showVersion bool
-	var debugAndTrace bool
-	var configFile string
-	var showTLSHelp bool
+	var (
+		showVersion   bool
+		debugAndTrace bool
+		configFile    string
+		signal        string
+		showTLSHelp   bool
+	)
 
 	// Parse flags
 	flag.IntVar(&opts.Port, "port", 0, "Port to listen on.")
@@ -94,6 +100,8 @@ func main() {
 	flag.IntVar(&opts.HTTPSPort, "https_port", 0, "HTTPS Port for /varz, /connz endpoints.")
 	flag.StringVar(&configFile, "c", "", "Configuration file.")
 	flag.StringVar(&configFile, "config", "", "Configuration file.")
+	flag.StringVar(&signal, "sl", "", "Send signal to gnatsd process (stop, quit, reopen, reload)")
+	flag.StringVar(&signal, "signal", "", "Send signal to gnatsd process (stop, quit, reopen, reload)")
 	flag.StringVar(&opts.PidFile, "P", "", "File to store process pid.")
 	flag.StringVar(&opts.PidFile, "pid", "", "File to store process pid.")
 	flag.StringVar(&opts.LogFile, "l", "", "File to store logging output.")
@@ -151,6 +159,27 @@ func main() {
 	// Snapshot flag options.
 	server.FlagSnapshot = opts.Clone()
 
+	// Process signal.
+	if signal != "" {
+		var (
+			pid           = -1
+			commandAndPid = strings.Split(signal, "=")
+		)
+		if l := len(commandAndPid); l == 2 {
+			p, err := strconv.Atoi(commandAndPid[1])
+			if err != nil {
+				usage()
+			}
+			pid = p
+		} else if l > 2 {
+			usage()
+		}
+		if err := server.ProcessSignal(commandAndPid[0], pid); err != nil {
+			server.PrintAndDie(err.Error())
+		}
+		os.Exit(0)
+	}
+
 	// Parse config if given
 	if configFile != "" {
 		fileOpts, err := server.ProcessConfigFile(configFile)
diff --git a/server/signal.go b/server/signal.go
index fbd84ca7..a516677f 100644
--- a/server/signal.go
+++ b/server/signal.go
@@ -4,11 +4,18 @@
 package server
 
 import (
+	"errors"
+	"fmt"
 	"os"
+	"os/exec"
 	"os/signal"
+	"strconv"
+	"strings"
 	"syscall"
 )
 
+const processName = "gnatsd"
+
 // Signal Handling
 func (s *Server) handleSignals() {
 	if s.getOpts().NoSigs {
@@ -37,3 +44,62 @@ func (s *Server) handleSignals() {
 		}
 	}()
 }
+
+// ProcessSignal sends the given signal command to the given process. If pid is
+// -1, this will send the signal to the single running instance of gnatsd. If
+// multiple instances are running, it returns an error.
+func ProcessSignal(command string, pid int) (err error) {
+	if pid == -1 {
+		pids, err := resolvePids()
+		if err != nil {
+			return err
+		}
+		if len(pids) == 0 {
+			return errors.New("No gnatsd processes running")
+		}
+		if len(pids) > 1 {
+			errStr := "Multiple gnatsd processes running:\n"
+			prefix := ""
+			for _, p := range pids {
+				errStr += fmt.Sprintf("%s%d", prefix, p)
+				prefix = "\n"
+			}
+			return errors.New(errStr)
+		}
+		pid = pids[0]
+	}
+
+	switch command {
+	case "stop":
+		err = syscall.Kill(pid, syscall.SIGKILL)
+	case "quit":
+		err = syscall.Kill(pid, syscall.SIGINT)
+	case "reopen":
+		err = syscall.Kill(pid, syscall.SIGUSR1)
+	case "reload":
+		err = syscall.Kill(pid, syscall.SIGHUP)
+	default:
+		err = fmt.Errorf("unknown signal %q", command)
+	}
+	return
+}
+
+// resolvePids returns the pids for all running gnatsd processes.
+func resolvePids() ([]int, error) {
+	// If pgrep isn't available, this will just bail out and the user will be
+	// required to specify a pid.
+	output, _ := exec.Command("pgrep", processName).Output()
+	pidStrs := strings.Split(string(output), "\n")
+	pids := make([]int, 0, len(pidStrs))
+	for _, pidStr := range pidStrs {
+		if pidStr == "" {
+			continue
+		}
+		pid, err := strconv.Atoi(pidStr)
+		if err != nil {
+			return nil, errors.New("Unable to resolve pid")
+		}
+		pids = append(pids, pid)
+	}
+	return pids, nil
+}
diff --git a/server/signal_windows.go b/server/signal_windows.go
index fabe08c1..44cfc903 100644
--- a/server/signal_windows.go
+++ b/server/signal_windows.go
@@ -24,3 +24,11 @@ func (s *Server) handleSignals() {
 		}
 	}()
 }
+
+// ProcessSignal sends the given signal command to the given process. If pid is
+// -1, this will send the signal to the single running instance of gnatsd. If
+// multiple instances are running, it returns an error.
+func ProcessSignal(command string, pid int) error {
+	// TODO
+	return errors.New("TODO")
+}