server kills subprocesses correctly on cancel

1 files changed, 3 insertions, 6 deletions

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)