server kills subprocesses correctly on cancel

6 files changed, 21 insertions, 54 deletions

diff --git a/cmd/dcat/main.go b/cmd/dcat/main.go
index 1ec945d..af46db1 100644
--- a/cmd/dcat/main.go
+++ b/cmd/dcat/main.go
@@ -9,7 +9,6 @@ import (
 	"github.com/mimecast/dtail/internal/color"
 	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/logger"
-	"github.com/mimecast/dtail/internal/pprof"
 	"github.com/mimecast/dtail/internal/user"
 	"github.com/mimecast/dtail/internal/version"
 )
@@ -23,7 +22,6 @@ func main() {
 	var displayVersion bool
 	var files string
 	var noColor bool
-	var pprofEnable bool
 	var serversStr string
 	var silentEnable bool
 	var sshPort int
@@ -34,7 +32,6 @@ func main() {
 	flag.BoolVar(&debugEnable, "debug", false, "Activate debug messages")
 	flag.BoolVar(&displayVersion, "version", false, "Display version")
 	flag.BoolVar(&noColor, "noColor", false, "Disable ANSII terminal colors")
-	flag.BoolVar(&pprofEnable, "pprofEnable", false, "Enable pprof server")
 	flag.BoolVar(&silentEnable, "silent", false, "Reduce output")
 	flag.BoolVar(&trustAllHosts, "trustAllHosts", false, "Auto trust all unknown host keys")
 	flag.IntVar(&connectionsPerCPU, "cpc", 10, "How many connections established per CPU core concurrently")
@@ -59,10 +56,6 @@ func main() {
 
 	logger.Start(ctx, serverEnable, debugEnable, silentEnable, silentEnable)
 
-	if pprofEnable || config.Common.PProfEnable {
-		pprof.Start()
-	}
-
 	args := clients.Args{
 		ConnectionsPerCPU: connectionsPerCPU,
 		ServersStr:        serversStr,
diff --git a/cmd/dgrep/main.go b/cmd/dgrep/main.go
index 74a501f..9393cf6 100644
--- a/cmd/dgrep/main.go
+++ b/cmd/dgrep/main.go
@@ -9,7 +9,6 @@ import (
 	"github.com/mimecast/dtail/internal/color"
 	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/logger"
-	"github.com/mimecast/dtail/internal/pprof"
 	"github.com/mimecast/dtail/internal/user"
 	"github.com/mimecast/dtail/internal/version"
 )
@@ -23,7 +22,6 @@ func main() {
 	var displayVersion bool
 	var files string
 	var noColor bool
-	var pprofEnable bool
 	var regex string
 	var serversStr string
 	var silentEnable bool
@@ -35,7 +33,6 @@ func main() {
 	flag.BoolVar(&debugEnable, "debug", false, "Activate debug messages")
 	flag.BoolVar(&displayVersion, "version", false, "Display version")
 	flag.BoolVar(&noColor, "noColor", false, "Disable ANSII terminal colors")
-	flag.BoolVar(&pprofEnable, "pprofEnable", false, "Enable pprof server")
 	flag.BoolVar(&silentEnable, "silent", false, "Reduce output")
 	flag.BoolVar(&trustAllHosts, "trustAllHosts", false, "Auto trust all unknown host keys")
 	flag.IntVar(&connectionsPerCPU, "cpc", 10, "How many connections established per CPU core concurrently")
@@ -60,10 +57,6 @@ func main() {
 	serverEnable := false
 	logger.Start(ctx, serverEnable, debugEnable, silentEnable, silentEnable)
 
-	if pprofEnable || config.Common.PProfEnable {
-		pprof.Start()
-	}
-
 	args := clients.Args{
 		ConnectionsPerCPU: connectionsPerCPU,
 		ServersStr:        serversStr,
diff --git a/cmd/dmap/main.go b/cmd/dmap/main.go
index f3f706a..07b7f2f 100644
--- a/cmd/dmap/main.go
+++ b/cmd/dmap/main.go
@@ -10,7 +10,6 @@ import (
 	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/logger"
 	"github.com/mimecast/dtail/internal/omode"
-	"github.com/mimecast/dtail/internal/pprof"
 	"github.com/mimecast/dtail/internal/user"
 	"github.com/mimecast/dtail/internal/version"
 )
@@ -24,7 +23,6 @@ func main() {
 	var displayVersion bool
 	var files string
 	var noColor bool
-	var pprofEnable bool
 	var queryStr string
 	var serversStr string
 	var silentEnable bool
@@ -36,7 +34,6 @@ func main() {
 	flag.BoolVar(&debugEnable, "debug", false, "Activate debug messages")
 	flag.BoolVar(&displayVersion, "version", false, "Display version")
 	flag.BoolVar(&noColor, "noColor", false, "Disable ANSII terminal colors")
-	flag.BoolVar(&pprofEnable, "pprofEnable", false, "Enable pprof server")
 	flag.BoolVar(&silentEnable, "silent", false, "Reduce output")
 	flag.BoolVar(&trustAllHosts, "trustAllHosts", false, "Auto trust all unknown host keys")
 	flag.IntVar(&connectionsPerCPU, "cpc", 10, "How many connections established per CPU core concurrently")
@@ -61,9 +58,6 @@ func main() {
 	serverEnable := false
 
 	logger.Start(ctx, serverEnable, debugEnable, silentEnable, silentEnable)
-	if pprofEnable || config.Common.PProfEnable {
-		pprof.Start()
-	}
 
 	args := clients.Args{
 		ConnectionsPerCPU: connectionsPerCPU,
diff --git a/cmd/drun/main.go b/cmd/drun/main.go
index 1b88b1b..e11900e 100644
--- a/cmd/drun/main.go
+++ b/cmd/drun/main.go
@@ -10,7 +10,6 @@ import (
 	"github.com/mimecast/dtail/internal/color"
 	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/logger"
-	"github.com/mimecast/dtail/internal/pprof"
 	"github.com/mimecast/dtail/internal/user"
 	"github.com/mimecast/dtail/internal/version"
 )
@@ -24,10 +23,10 @@ func main() {
 	var discovery string
 	var displayVersion bool
 	var noColor bool
-	var pprofEnable bool
 	var serversStr string
 	var silentEnable bool
 	var sshPort int
+	var timeout int
 	var trustAllHosts bool
 
 	userName := user.Name()
@@ -35,11 +34,11 @@ func main() {
 	flag.BoolVar(&debugEnable, "debug", false, "Activate debug messages")
 	flag.BoolVar(&displayVersion, "version", false, "Display version")
 	flag.BoolVar(&noColor, "noColor", false, "Disable ANSII terminal colors")
-	flag.BoolVar(&pprofEnable, "pprofEnable", false, "Enable pprof server")
 	flag.BoolVar(&silentEnable, "silent", false, "Reduce output")
 	flag.BoolVar(&trustAllHosts, "trustAllHosts", false, "Auto trust all unknown host keys")
 	flag.IntVar(&connectionsPerCPU, "cpc", 10, "How many connections established per CPU core concurrently")
 	flag.IntVar(&sshPort, "port", 2222, "SSH server port")
+	flag.IntVar(&timeout, "timeout", 0, "Command execution timeout")
 	flag.StringVar(&cfgFile, "cfg", "", "Config file path")
 	flag.StringVar(&command, "command", "", "Command to run")
 	flag.StringVar(&discovery, "discovery", "", "Server discovery method")
@@ -59,9 +58,6 @@ func main() {
 	serverEnable := false
 
 	logger.Start(ctx, serverEnable, debugEnable, silentEnable, silentEnable)
-	if pprofEnable || config.Common.PProfEnable {
-		pprof.Start()
-	}
 
 	args := clients.Args{
 		ConnectionsPerCPU: connectionsPerCPU,
@@ -70,6 +66,7 @@ func main() {
 		UserName:          userName,
 		What:              readCommand(command),
 		TrustAllHosts:     trustAllHosts,
+		Timeout:           timeout,
 	}
 
 	client, err := clients.NewRunClient(args)
diff --git a/cmd/dserver/main.go b/cmd/dserver/main.go
index 03834f1..94f1a87 100644
--- a/cmd/dserver/main.go
+++ b/cmd/dserver/main.go
@@ -4,12 +4,13 @@ import (
 	"context"
 	"flag"
 	"os"
+	"os/signal"
+	"syscall"
 	"time"
 
 	"github.com/mimecast/dtail/internal/color"
 	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/logger"
-	"github.com/mimecast/dtail/internal/pprof"
 	"github.com/mimecast/dtail/internal/server"
 	"github.com/mimecast/dtail/internal/user"
 	"github.com/mimecast/dtail/internal/version"
@@ -21,7 +22,6 @@ func main() {
 	var debugEnable bool
 	var displayVersion bool
 	var noColor bool
-	var pprofEnable bool
 	var shutdownAfter int
 	var sshPort int
 
@@ -30,7 +30,6 @@ func main() {
 	flag.BoolVar(&debugEnable, "debug", false, "Activate debug messages")
 	flag.BoolVar(&displayVersion, "version", false, "Display version")
 	flag.BoolVar(&noColor, "noColor", false, "Disable ANSII terminal colors")
-	flag.BoolVar(&pprofEnable, "pprofEnable", false, "Enable pprof server")
 	flag.IntVar(&shutdownAfter, "shutdownAfter", 0, "Automatically shutdown after so many seconds")
 	flag.IntVar(&sshPort, "port", 2222, "SSH server port")
 	flag.StringVar(&cfgFile, "cfg", "", "Config file path")
@@ -44,30 +43,28 @@ func main() {
 		version.PrintAndExit()
 	}
 
-	ctx := context.Background()
+	ctx, cancel := context.WithCancel(context.Background())
+	if shutdownAfter > 0 {
+		logger.Info("Enabling auto shutdown timer", shutdownAfter)
+		ctx, cancel = context.WithTimeout(ctx, time.Duration(shutdownAfter)*time.Second)
+	}
+
+	sigCh := make(chan os.Signal)
+	signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
+
+	go func() {
+		select {
+		case <-sigCh:
+			cancel()
+		case <-ctx.Done():
+		}
+	}()
 
 	serverEnable := true
 	silentEnable := false
 	nothingEnable := false
 	logger.Start(ctx, serverEnable, debugEnable, silentEnable, nothingEnable)
 
-	if shutdownAfter > 0 {
-		go func() {
-			defer os.Exit(1)
-
-			logger.Info("Enabling auto shutdown timer", shutdownAfter)
-			select {
-			case <-time.After(time.Duration(shutdownAfter) * time.Second):
-				logger.Info("Auto shutdown timer reached, shutting down now")
-			case <-ctx.Done():
-			}
-		}()
-	}
-
-	if pprofEnable || config.Common.PProfEnable {
-		pprof.Start()
-	}
-
 	serv := server.New()
 	status := serv.Start(ctx)
 	logger.Flush()
diff --git a/cmd/dtail/main.go b/cmd/dtail/main.go
index 76070ff..a243873 100644
--- a/cmd/dtail/main.go
+++ b/cmd/dtail/main.go
@@ -10,7 +10,6 @@ import (
 	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/logger"
 	"github.com/mimecast/dtail/internal/omode"
-	"github.com/mimecast/dtail/internal/pprof"
 	"github.com/mimecast/dtail/internal/user"
 	"github.com/mimecast/dtail/internal/version"
 )
@@ -25,7 +24,6 @@ func main() {
 	var displayVersion bool
 	var files string
 	var noColor bool
-	var pprofEnable bool
 	var queryStr string
 	var regex string
 	var serversStr string
@@ -39,7 +37,6 @@ func main() {
 	flag.BoolVar(&debugEnable, "debug", false, "Activate debug messages")
 	flag.BoolVar(&displayVersion, "version", false, "Display version")
 	flag.BoolVar(&noColor, "noColor", false, "Disable ANSII terminal colors")
-	flag.BoolVar(&pprofEnable, "pprofEnable", false, "Enable pprof server")
 	flag.BoolVar(&silentEnable, "silent", false, "Reduce output")
 	flag.BoolVar(&trustAllHosts, "trustAllHosts", false, "Auto trust all unknown host keys")
 	flag.IntVar(&connectionsPerCPU, "cpc", 10, "How many connections established per CPU core concurrently")
@@ -74,10 +71,6 @@ func main() {
 	}
 	logger.Start(ctx, serverEnable, debugEnable, silentEnable, silentEnable)
 
-	if pprofEnable || config.Common.PProfEnable {
-		pprof.Start()
-	}
-
 	args := clients.Args{
 		ConnectionsPerCPU: connectionsPerCPU,
 		ServersStr:        serversStr,