1 files changed, 31 insertions, 23 deletions

diff --git a/cmd/dmap/main.go b/cmd/dmap/main.go
index a8a52a2..8e980e8 100644
--- a/cmd/dmap/main.go
+++ b/cmd/dmap/main.go
@@ -3,18 +3,16 @@ package main
 import (
 	"context"
 	"flag"
+	"fmt"
 	"os"
-	"sync"
-
-	"net/http"
-	_ "net/http"
-	_ "net/http/pprof"
 
+	"github.com/mimecast/dtail/internal/cli"
 	"github.com/mimecast/dtail/internal/clients"
 	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/dlog"
 	"github.com/mimecast/dtail/internal/io/signal"
 	"github.com/mimecast/dtail/internal/omode"
+	"github.com/mimecast/dtail/internal/profiling"
 	"github.com/mimecast/dtail/internal/source"
 	"github.com/mimecast/dtail/internal/user"
 	"github.com/mimecast/dtail/internal/version"
@@ -23,61 +21,71 @@ import (
 // The evil begins here.
 func main() {
 	var displayVersion bool
+	var legacyAuthKeyPath string
 	var pprof string
+	var profileFlags profiling.Flags
 
 	args := config.Args{
-		Mode: omode.MapClient,
+		Mode:             omode.MapClient,
+		SSHAgentKeyIndex: -1,
 	}
 	userName := user.Name()
 
 	flag.BoolVar(&args.NoColor, "noColor", false, "Disable ANSII terminal colors")
+	flag.BoolVar(&args.NoAuthKey, "no-auth-key", false, "Disable auth-key fast reconnect feature")
+	flag.BoolVar(&args.LogPayload, "log-payload", false, "Also tee retrieved payload into the client log file (default: file keeps diagnostics only)")
 	flag.BoolVar(&args.Quiet, "quiet", false, "Quiet output mode")
+	flag.BoolVar(&args.InteractiveQuery, "interactive-query", false, "Enable interactive in-flight query control over supported sessions")
 	flag.BoolVar(&args.Plain, "plain", false, "Plain output mode")
 	flag.BoolVar(&args.TrustAllHosts, "trustAllHosts", false, "Trust all unknown host keys")
 	flag.BoolVar(&displayVersion, "version", false, "Display version")
 	flag.IntVar(&args.ConnectionsPerCPU, "cpc", config.DefaultConnectionsPerCPU,
 		"How many connections established per CPU core concurrently")
+	flag.IntVar(&args.SSHAgentKeyIndex, "agentKeyIndex", -1, "SSH agent key index to use (-1 for all keys)")
 	flag.IntVar(&args.SSHPort, "port", config.DefaultSSHPort, "SSH server port")
 	flag.IntVar(&args.Timeout, "timeout", 0, "Max time dtail server will collect data until disconnection")
 	flag.StringVar(&args.ConfigFile, "cfg", "", "Config file path")
+	flag.StringVar(&args.ControlTTYPath, "control-tty", "/dev/tty", "TTY device for interactive query control")
 	flag.StringVar(&args.Discovery, "discovery", "", "Server discovery method")
 	flag.StringVar(&args.LogDir, "logDir", "~/log", "Log dir")
 	flag.StringVar(&args.Logger, "logger", config.DefaultClientLogger, "Logger name")
 	flag.StringVar(&args.LogLevel, "logLevel", config.DefaultLogLevel, "Log level")
-	flag.StringVar(&args.SSHPrivateKeyFilePath, "key", "", "Path to private key")
+	cli.BindAuthKeyFlags(flag.CommandLine, &legacyAuthKeyPath, &args)
 	flag.StringVar(&args.QueryStr, "query", "", "Map reduce query")
 	flag.StringVar(&args.ServersStr, "servers", "", "Remote servers to connect")
 	flag.StringVar(&args.UserName, "user", userName, "Your system user name")
 	flag.StringVar(&args.What, "files", "", "File(s) to read")
 	flag.StringVar(&pprof, "pprof", "", "Start PProf server this address")
 
+	// Add profiling flags
+	profiling.AddFlags(&profileFlags)
+
 	flag.Parse()
+	if warning := cli.ApplyAuthKeyPathCompatibility(&args, legacyAuthKeyPath, cli.FlagWasSet("auth-key-path")); warning != "" {
+		fmt.Fprintln(os.Stderr, warning)
+	}
 	config.Setup(source.Client, &args, flag.Args())
 
 	if displayVersion {
-		version.PrintAndExit()
+		runtimeCfg := config.CurrentRuntime()
+		version.PrintAndExit(runtimeCfg.Client != nil && runtimeCfg.Client.TermColorsEnable)
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
-	var wg sync.WaitGroup
-	wg.Add(1)
-	dlog.Start(ctx, &wg, source.Client)
-
-	if pprof != "" {
-		dlog.Client.Info("Starting PProf", pprof)
-		go func() {
-			panic(http.ListenAndServe(pprof, nil))
-		}()
-	}
+	runtime := cli.NewClientRuntime(context.Background(), profileFlags, "dmap")
+	runtime.StartPProf(pprof)
+	runtime.LogStartupMetrics()
 
 	client, err := clients.NewMaprClient(args, clients.DefaultMode)
 	if err != nil {
+		runtime.Stop()
 		dlog.Client.FatalPanic(err)
 	}
 
-	status := client.Start(ctx, signal.InterruptCh(ctx))
-	cancel()
-
-	wg.Wait()
+	status := client.Start(
+		runtime.Context(),
+		signal.InterruptChWithCancel(runtime.Context(), runtime.Cancel),
+	)
+	runtime.LogShutdownMetrics()
+	runtime.Stop()
 	os.Exit(status)
 }