can tail probe with a given timeout and then write a mapreduce result

3 files changed, 32 insertions, 0 deletions

diff --git a/cmd/dmap/main.go b/cmd/dmap/main.go
index 07b7f2f..dea4706 100644
--- a/cmd/dmap/main.go
+++ b/cmd/dmap/main.go
@@ -27,6 +27,7 @@ func main() {
 	var serversStr string
 	var silentEnable bool
 	var sshPort int
+	var timeout int
 	var trustAllHosts bool
 
 	userName := user.Name()
@@ -38,6 +39,7 @@ func main() {
 	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, "Max time dtail server will collect data until disconnection")
 	flag.StringVar(&cfgFile, "cfg", "", "Config file path")
 	flag.StringVar(&discovery, "discovery", "", "Server discovery method")
 	flag.StringVar(&files, "files", "", "File(s) to read")
@@ -67,6 +69,7 @@ func main() {
 		What:              files,
 		TrustAllHosts:     trustAllHosts,
 		Mode:              omode.MapClient,
+		Timeout:           timeout,
 	}
 
 	client, err := clients.NewMaprClient(args, queryStr)
diff --git a/cmd/dserver/main.go b/cmd/dserver/main.go
index 94f1a87..f9207d5 100644
--- a/cmd/dserver/main.go
+++ b/cmd/dserver/main.go
@@ -3,6 +3,10 @@ package main
 import (
 	"context"
 	"flag"
+	"fmt"
+	"net/http"
+	_ "net/http"
+	_ "net/http/pprof"
 	"os"
 	"os/signal"
 	"syscall"
@@ -22,6 +26,7 @@ func main() {
 	var debugEnable bool
 	var displayVersion bool
 	var noColor bool
+	var pprof int
 	var shutdownAfter int
 	var sshPort int
 
@@ -32,6 +37,7 @@ func main() {
 	flag.BoolVar(&noColor, "noColor", false, "Disable ANSII terminal colors")
 	flag.IntVar(&shutdownAfter, "shutdownAfter", 0, "Automatically shutdown after so many seconds")
 	flag.IntVar(&sshPort, "port", 2222, "SSH server port")
+	flag.IntVar(&pprof, "pprof", -1, "Start PProf server this port")
 	flag.StringVar(&cfgFile, "cfg", "", "Config file path")
 
 	flag.Parse()
@@ -65,6 +71,13 @@ func main() {
 	nothingEnable := false
 	logger.Start(ctx, serverEnable, debugEnable, silentEnable, nothingEnable)
 
+	if pprof > -1 {
+		// For debugging purposes only
+		pprofArgs := fmt.Sprintf("0.0.0.0:%d", pprof)
+		logger.Info("Starting PProf", pprofArgs)
+		go http.ListenAndServe(pprofArgs, nil)
+	}
+
 	serv := server.New()
 	status := serv.Start(ctx)
 	logger.Flush()
diff --git a/cmd/dtail/main.go b/cmd/dtail/main.go
index a243873..f062cdc 100644
--- a/cmd/dtail/main.go
+++ b/cmd/dtail/main.go
@@ -3,6 +3,10 @@ package main
 import (
 	"context"
 	"flag"
+	"fmt"
+	"net/http"
+	_ "net/http"
+	_ "net/http/pprof"
 	"os"
 
 	"github.com/mimecast/dtail/internal/clients"
@@ -21,6 +25,7 @@ func main() {
 	var connectionsPerCPU int
 	var debugEnable bool
 	var discovery string
+	var pprof int
 	var displayVersion bool
 	var files string
 	var noColor bool
@@ -29,6 +34,7 @@ func main() {
 	var serversStr string
 	var silentEnable bool
 	var sshPort int
+	var timeout int
 	var trustAllHosts bool
 
 	userName := user.Name()
@@ -41,6 +47,8 @@ func main() {
 	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, "Max time dtail server will collect data until disconnection")
+	flag.IntVar(&pprof, "pprof", -1, "Start PProf server this port")
 	flag.StringVar(&cfgFile, "cfg", "", "Config file path")
 	flag.StringVar(&discovery, "discovery", "", "Server discovery method")
 	flag.StringVar(&files, "files", "", "File(s) to read")
@@ -71,6 +79,13 @@ func main() {
 	}
 	logger.Start(ctx, serverEnable, debugEnable, silentEnable, silentEnable)
 
+	if pprof > -1 {
+		// For debugging purposes only
+		pprofArgs := fmt.Sprintf("0.0.0.0:%d", pprof)
+		logger.Info("Starting PProf", pprofArgs)
+		go http.ListenAndServe(pprofArgs, nil)
+	}
+
 	args := clients.Args{
 		ConnectionsPerCPU: connectionsPerCPU,
 		ServersStr:        serversStr,
@@ -80,6 +95,7 @@ func main() {
 		TrustAllHosts:     trustAllHosts,
 		Regex:             regex,
 		Mode:              omode.TailClient,
+		Timeout:           timeout,
 	}
 
 	var client clients.Client