feat: DTail fork — server/client feature development

Squashed development of the snonux/dtail fork's product code (internal/, cmd/)
since diverging from mimecast/dtail. Major areas:

- Read/output path: the former "turbo" channel-less path is now the single,
  default server-side read/output path for cat/grep/tail and MapReduce; the old
  channel-based path and its config/env toggles were removed.
- MapReduce: single aggregate implementation (server + serverless) fed directly
  by a processor pipeline, with input-exhausted finalization via the shutdown
  coordinator; high-concurrency and data-race fixes.
- Journal source reads (journal:unit.service) via journalctl, Linux-gated behind
  a journal-v1 capability.
- Auth-key fast reconnect: in-memory per-user public-key cache with TTL/max-keys,
  registered over an authenticated session (AUTHKEY), checked before
  authorized_keys.
- Interactive query reload (--interactive-query) with SESSION START/UPDATE
  generation boundaries and capability negotiation.
- Client-side deadlines: --timeout / --shutdownAfter as context deadlines;
  follow shutdown handling.
- Client logging: diagnostics-only daily log by default, opt-in payload tee via
  --log-payload.
- Numerous correctness fixes (buffer-pool double-recycle races, EOF-sentinel
  leaks, glob-expansion cap, TOCTOU in CSV parsing) with accompanying unit tests.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

8 files changed, 377 insertions, 122 deletions

diff --git a/cmd/dcat/main.go b/cmd/dcat/main.go
index a50be51..83b713c 100644
--- a/cmd/dcat/main.go
+++ b/cmd/dcat/main.go
@@ -3,17 +3,14 @@ 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/profiling"
 	"github.com/mimecast/dtail/internal/source"
 	"github.com/mimecast/dtail/internal/user"
 	"github.com/mimecast/dtail/internal/version"
@@ -23,56 +20,63 @@ import (
 func main() {
 	var args config.Args
 	var displayVersion bool
+	var legacyAuthKeyPath string
 	var pprof string
+	var profileFlags profiling.Flags
 
 	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.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.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, "dcat")
+	runtime.StartPProf(pprof)
+	runtime.LogStartupMetrics()
 
 	client, err := clients.NewCatClient(args)
 	if err != nil {
-		panic(err)
+		runtime.Stop()
+		fmt.Fprintf(os.Stderr, "unable to create dcat client: %v\n", err)
+		os.Exit(1)
 	}
 
-	status := client.Start(ctx, signal.InterruptCh(ctx))
-	cancel()
-
-	wg.Wait()
+	status := client.Start(runtime.Context(), signal.InterruptCh(runtime.Context()))
+	runtime.LogShutdownMetrics()
+	runtime.Stop()
 	os.Exit(status)
 }
diff --git a/cmd/dgrep/main.go b/cmd/dgrep/main.go
index 19f818b..1c2837b 100644
--- a/cmd/dgrep/main.go
+++ b/cmd/dgrep/main.go
@@ -3,17 +3,14 @@ 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/profiling"
 	"github.com/mimecast/dtail/internal/source"
 	"github.com/mimecast/dtail/internal/user"
 	"github.com/mimecast/dtail/internal/version"
@@ -24,12 +21,17 @@ func main() {
 	var args config.Args
 	var displayVersion bool
 	var grep string
+	var legacyAuthKeyPath string
 	var pprof string
+	var profileFlags profiling.Flags
 	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.RegexInvert, "invert", false, "Invert regex")
+	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")
@@ -38,13 +40,15 @@ func main() {
 	flag.IntVar(&args.LContext.AfterContext, "after", 0, "Print lines of trailing context after matching lines")
 	flag.IntVar(&args.LContext.BeforeContext, "before", 0, "Print lines of leading context before matching lines")
 	flag.IntVar(&args.LContext.MaxCount, "max", 0, "Stop reading file after NUM matching lines")
+	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.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.RegexStr, "regex", ".", "Regular expression")
 	flag.StringVar(&args.ServersStr, "servers", "", "Remote servers to connect")
 	flag.StringVar(&args.UserName, "user", userName, "Your system user name")
@@ -52,37 +56,38 @@ func main() {
 	flag.StringVar(&grep, "grep", "", "Alias for -regex")
 	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)
+	runtime := cli.NewClientRuntime(context.Background(), profileFlags, "dgrep")
 
 	if grep != "" {
 		args.RegexStr = grep
 	}
 
-	if pprof != "" {
-		dlog.Client.Info("Starting PProf", pprof)
-		go func() {
-			panic(http.ListenAndServe(pprof, nil))
-		}()
-	}
+	runtime.StartPProf(pprof)
+	runtime.LogStartupMetrics()
 
 	client, err := clients.NewGrepClient(args)
 	if err != nil {
-		panic(err)
+		runtime.Stop()
+		fmt.Fprintf(os.Stderr, "unable to create dgrep client: %v\n", err)
+		os.Exit(1)
 	}
 
-	status := client.Start(ctx, signal.InterruptCh(ctx))
-	cancel()
-
-	wg.Wait()
+	status := client.Start(runtime.Context(), signal.InterruptCh(runtime.Context()))
+	runtime.LogShutdownMetrics()
+	runtime.Stop()
 	os.Exit(status)
 }
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)
 }
diff --git a/cmd/dserver/main.go b/cmd/dserver/main.go
index 3377273..936edce 100644
--- a/cmd/dserver/main.go
+++ b/cmd/dserver/main.go
@@ -3,15 +3,13 @@ package main
 import (
 	"context"
 	"flag"
-	"net/http"
-	_ "net/http"
-	_ "net/http/pprof"
 	"os"
 	"os/signal"
 	"sync"
 	"syscall"
 	"time"
 
+	"github.com/mimecast/dtail/internal/cli"
 	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/dlog"
 	"github.com/mimecast/dtail/internal/server"
@@ -46,13 +44,31 @@ func main() {
 	config.Setup(source.Server, &args, flag.Args())
 
 	if displayVersion {
-		version.PrintAndExit()
+		runtimeCfg := config.CurrentRuntime()
+		version.PrintAndExit(runtimeCfg.Client != nil && runtimeCfg.Client.TermColorsEnable)
 	}
-	version.Print()
+	version.Print(false)
+
+	// rootCtx is always cancelled on exit to ensure the internal goroutine
+	// spawned by context.WithCancel is released. When -shutdownAfter is set,
+	// ctx is replaced by a child WithTimeout context whose own cancel is also
+	// deferred, preventing the lostcancel leak flagged by go vet.
+	rootCtx, rootCancel := context.WithCancel(context.Background())
+	defer rootCancel()
+
+	ctx := rootCtx
+	cancel := context.CancelFunc(rootCancel)
 
-	ctx, cancel := context.WithCancel(context.Background())
 	if shutdownAfter > 0 {
-		ctx, cancel = context.WithTimeout(ctx, time.Duration(shutdownAfter)*time.Second)
+		// Override ctx with a timeout-bounded child; defer its cancel so the
+		// timeout goroutine is always cleaned up regardless of code path.
+		var timeoutCancel context.CancelFunc
+		ctx, timeoutCancel = context.WithTimeout(rootCtx, time.Duration(shutdownAfter)*time.Second)
+		defer timeoutCancel()
+		// Callers that invoke cancel() (e.g. the signal handler and post-serve
+		// cleanup) should trigger the timeout cancel so the server shuts down
+		// promptly even before the deadline fires.
+		cancel = timeoutCancel
 	}
 
 	sigCh := make(chan os.Signal, 10)
@@ -70,16 +86,36 @@ func main() {
 	wg.Add(1)
 	dlog.Start(ctx, &wg, source.Server)
 
+	var pprofServer *cli.PProfServer
 	if pprof != "" {
-		dlog.Client.Info("Starting PProf", pprof)
-		go func() {
-			panic(http.ListenAndServe(pprof, nil))
-		}()
+		// Enable mutex and block profiling so the /debug/pprof/mutex and
+		// /debug/pprof/block endpoints actually contain samples. These rates
+		// are gated on --pprof so they cost nothing when profiling is off.
+		cli.EnableProfilingRates()
+
+		// Assign to the outer pprofServer with '=' (declaring pprofErr
+		// separately) so it is NOT shadowed: the graceful Shutdown below relies
+		// on the outer var being non-nil to actually stop the pprof server.
+		var pprofErr error
+		pprofServer, pprofErr = cli.NewPProfServer(pprof)
+		if pprofErr != nil {
+			dlog.Client.Error("Unable to start PProf", pprofErr)
+		} else {
+			dlog.Client.Info("Starting PProf", pprofServer.Address())
+			pprofServer.Start(nil)
+		}
 	}
 
-	serv := server.New()
+	serv := server.New(config.CurrentRuntime())
 	status := serv.Start(ctx)
 	cancel()
+	if pprofServer != nil {
+		shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
+		if err := pprofServer.Shutdown(shutdownCtx); err != nil {
+			dlog.Client.Error("Unable to stop PProf", err)
+		}
+		shutdownCancel()
+	}
 
 	wg.Wait()
 	os.Exit(status)
diff --git a/cmd/dtail-tools/main.go b/cmd/dtail-tools/main.go
new file mode 100644
index 0000000..2b96a56
--- /dev/null
+++ b/cmd/dtail-tools/main.go
@@ -0,0 +1,60 @@
+package main
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/mimecast/dtail/internal/tools/benchmark"
+	"github.com/mimecast/dtail/internal/tools/pgo"
+	"github.com/mimecast/dtail/internal/tools/profile"
+)
+
+func main() {
+	if len(os.Args) < 2 {
+		printUsage()
+		os.Exit(1)
+	}
+
+	command := os.Args[1]
+	
+	// Remove command from args for subcommand parsing
+	os.Args = append([]string{os.Args[0]}, os.Args[2:]...)
+
+	switch command {
+	case "profile":
+		if err := profile.Run(); err != nil {
+			fmt.Fprintf(os.Stderr, "Error: %v\n", err)
+			os.Exit(1)
+		}
+	case "benchmark":
+		if err := benchmark.Run(); err != nil {
+			fmt.Fprintf(os.Stderr, "Error: %v\n", err)
+			os.Exit(1)
+		}
+	case "pgo":
+		if err := pgo.Run(); err != nil {
+			fmt.Fprintf(os.Stderr, "Error: %v\n", err)
+			os.Exit(1)
+		}
+	case "help", "-h", "--help":
+		printUsage()
+	default:
+		fmt.Fprintf(os.Stderr, "Unknown command: %s\n", command)
+		printUsage()
+		os.Exit(1)
+	}
+}
+
+func printUsage() {
+	fmt.Println("dtail-tools - DTail performance analysis toolkit")
+	fmt.Println()
+	fmt.Println("Usage: dtail-tools <command> [options]")
+	fmt.Println()
+	fmt.Println("Commands:")
+	fmt.Println("  profile    Run profiling on dtail commands")
+	fmt.Println("  benchmark  Run benchmarks and manage baselines")
+	fmt.Println("  pgo        Profile-Guided Optimization for dtail commands")
+	fmt.Println("  help       Show this help message")
+	fmt.Println()
+	fmt.Println("Run 'dtail-tools <command> -h' for command-specific help")
+}
\ No newline at end of file
diff --git a/cmd/dtail/main.go b/cmd/dtail/main.go
index e18923a..0105709 100644
--- a/cmd/dtail/main.go
+++ b/cmd/dtail/main.go
@@ -4,19 +4,16 @@ import (
 	"context"
 	"flag"
 	"fmt"
-	"net/http"
-	_ "net/http"
-	_ "net/http/pprof"
 	"os"
-	"sync"
 	"time"
 
+	"github.com/mimecast/dtail/internal/cli"
 	"github.com/mimecast/dtail/internal/clients"
 	"github.com/mimecast/dtail/internal/color"
 	"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"
@@ -30,14 +27,19 @@ func main() {
 	var displayWideColorTable bool
 	var displayVersion bool
 	var grep string
+	var legacyAuthKeyPath string
 	var pprof string
 	var shutdownAfter int
+	var profileFlags profiling.Flags
 
 	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.RegexInvert, "invert", false, "Invert regex")
+	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(&checkHealth, "checkHealth", false, "Deprecated, flag will be removed soon")
@@ -49,15 +51,17 @@ func main() {
 	flag.IntVar(&args.LContext.AfterContext, "after", 0, "Print lines of trailing context after matching lines")
 	flag.IntVar(&args.LContext.BeforeContext, "before", 0, "Print lines of leading context before matching lines")
 	flag.IntVar(&args.LContext.MaxCount, "max", 0, "Stop reading file after NUM matching lines")
+	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.IntVar(&shutdownAfter, "shutdownAfter", 3600*24, "Shutdown after so many seconds")
 	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.RegexStr, "regex", ".", "Regular expression")
 	flag.StringVar(&args.ServersStr, "servers", "", "Remote servers to connect")
@@ -66,13 +70,20 @@ func main() {
 	flag.StringVar(&grep, "grep", "", "Alias for -regex")
 	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)
+	}
 	if grep != "" {
 		args.RegexStr = grep
 	}
 	config.Setup(source.Client, &args, flag.Args())
 	if displayVersion {
-		version.PrintAndExit()
+		runtimeCfg := config.CurrentRuntime()
+		version.PrintAndExit(runtimeCfg.Client != nil && runtimeCfg.Client.TermColorsEnable)
 	}
 	if !args.Plain {
 		if displayWideColorTable {
@@ -83,30 +94,26 @@ func main() {
 		}
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
-	if shutdownAfter > 0 {
-		// NEXT: This does not work (auto shutdown)
-		ctx, cancel = context.WithTimeout(ctx, time.Duration(shutdownAfter)*time.Second)
-		defer cancel()
-	}
+	baseCtx, timeoutCancel := applyClientDeadlines(context.Background(), shutdownAfter, args.Timeout)
 
-	var wg sync.WaitGroup
-	wg.Add(1)
-	dlog.Start(ctx, &wg, source.Client)
+	runtime := cli.NewClientRuntime(baseCtx, profileFlags, "dtail")
+	exitWithError := func(err error) {
+		runtime.Stop()
+		timeoutCancel()
+		fmt.Fprintf(os.Stderr, "unable to initialize dtail client: %v\n", err)
+		os.Exit(1)
+	}
 
 	if checkHealth {
 		fmt.Println("WARN: DTail health check has moved to separate binary dtailhealth" +
 			" - please adjust the monitoring scripts!")
-		cancel()
+		runtime.Stop()
+		timeoutCancel()
 		os.Exit(1)
 	}
 
-	if pprof != "" {
-		dlog.Client.Info("Starting PProf", pprof)
-		go func() {
-			panic(http.ListenAndServe(pprof, nil))
-		}()
-	}
+	runtime.StartPProf(pprof)
+	runtime.LogStartupMetrics()
 
 	var client clients.Client
 	var err error
@@ -115,17 +122,64 @@ func main() {
 	switch args.QueryStr {
 	case "":
 		if client, err = clients.NewTailClient(args); err != nil {
-			panic(err)
+			exitWithError(err)
 		}
 	default:
 		if client, err = clients.NewMaprClient(args, clients.DefaultMode); err != nil {
-			panic(err)
+			exitWithError(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()
+	timeoutCancel()
 	os.Exit(status)
 }
+
+// applyClientDeadlines wraps ctx with the earliest of two absolute deadlines:
+// the --shutdownAfter safety cap and the user's --timeout data-collection
+// window. Whichever elapses first cancels the returned context, which propagates
+// through the follow client's reconnect and per-connection read loops (both
+// select on ctx.Done()), so client.Start returns cleanly and the process exits
+// instead of auto-reconnecting for another cycle.
+//
+// Making --timeout a client-side deadline (rather than relying solely on the
+// server closing the read) is what fixes the historical hang: the server-side
+// read deadline fires at N seconds, but in tail+query mode the session stays
+// alive via the map/aggregate command, so the client used to treat the closed
+// read as a transient drop and reconnect indefinitely. A client-side deadline
+// makes --timeout behave consistently for follows with or without --query, which
+// matches the flag's help text ("Max time ... until disconnection").
+//
+// The two deadlines compose naturally (context deadlines nest, so the earlier
+// one wins), so they never conflict with each other. A timeout of 0 (unset)
+// contributes no deadline, preserving the previous behaviour. OS signals reach
+// the same context via signal.InterruptChWithCancel in main.
+func applyClientDeadlines(ctx context.Context, shutdownAfter, timeout int) (
+	context.Context, context.CancelFunc) {
+
+	var cancels []context.CancelFunc
+	addDeadline := func(seconds int) {
+		if seconds <= 0 {
+			return
+		}
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, time.Duration(seconds)*time.Second)
+		cancels = append(cancels, cancel)
+	}
+
+	addDeadline(shutdownAfter)
+	addDeadline(timeout)
+
+	return ctx, func() {
+		// Release timers in reverse (inner first) to avoid leaking the parent
+		// timer while an inner context still references it.
+		for i := len(cancels) - 1; i >= 0; i-- {
+			cancels[i]()
+		}
+	}
+}
diff --git a/cmd/dtail/main_test.go b/cmd/dtail/main_test.go
new file mode 100644
index 0000000..b7dbb77
--- /dev/null
+++ b/cmd/dtail/main_test.go
@@ -0,0 +1,66 @@
+package main
+
+import (
+	"context"
+	"testing"
+	"time"
+)
+
+// TestApplyClientDeadlines verifies that --timeout and --shutdownAfter are turned
+// into client-side context deadlines and that the earliest one wins, which is
+// what makes the dtail follow client exit instead of reconnecting after a
+// timeout-induced disconnect (task xu0).
+func TestApplyClientDeadlines(t *testing.T) {
+	tests := []struct {
+		name          string
+		shutdownAfter int
+		timeout       int
+		wantDeadline  bool
+		// wantWithin bounds the deadline from the parent context when a deadline
+		// is expected. It reflects the smaller of the two configured windows.
+		wantWithin time.Duration
+	}{
+		{name: "both unset: no deadline", shutdownAfter: 0, timeout: 0, wantDeadline: false},
+		{name: "only shutdownAfter", shutdownAfter: 5, timeout: 0, wantDeadline: true, wantWithin: 5 * time.Second},
+		{name: "only timeout", shutdownAfter: 0, timeout: 4, wantDeadline: true, wantWithin: 4 * time.Second},
+		{name: "timeout smaller wins", shutdownAfter: 30, timeout: 3, wantDeadline: true, wantWithin: 3 * time.Second},
+		{name: "shutdownAfter smaller wins", shutdownAfter: 2, timeout: 60, wantDeadline: true, wantWithin: 2 * time.Second},
+		{name: "negative values ignored", shutdownAfter: -1, timeout: -1, wantDeadline: false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			before := time.Now()
+			ctx, cancel := applyClientDeadlines(context.Background(), tt.shutdownAfter, tt.timeout)
+			defer cancel()
+
+			deadline, ok := ctx.Deadline()
+			if ok != tt.wantDeadline {
+				t.Fatalf("Deadline() ok = %v, want %v", ok, tt.wantDeadline)
+			}
+			if !tt.wantDeadline {
+				return
+			}
+
+			// The effective deadline must equal the smaller configured window
+			// (allow a small slack for scheduling).
+			gotWindow := deadline.Sub(before)
+			if gotWindow > tt.wantWithin+time.Second || gotWindow < tt.wantWithin-time.Second {
+				t.Fatalf("effective deadline window = %s, want ~%s", gotWindow, tt.wantWithin)
+			}
+		})
+	}
+}
+
+// TestApplyClientDeadlinesCancelFires ensures the returned cancel function
+// actually cancels the context (releasing timers) without panicking.
+func TestApplyClientDeadlinesCancelFires(t *testing.T) {
+	ctx, cancel := applyClientDeadlines(context.Background(), 3600, 10)
+	cancel()
+
+	select {
+	case <-ctx.Done():
+	default:
+		t.Fatal("context was not canceled after calling cancel()")
+	}
+}
diff --git a/cmd/dtailhealth/main.go b/cmd/dtailhealth/main.go
index a0ca84e..33e84db 100644
--- a/cmd/dtailhealth/main.go
+++ b/cmd/dtailhealth/main.go
@@ -3,13 +3,12 @@ package main
 import (
 	"context"
 	"flag"
+	"fmt"
 	"os"
 	"sync"
+	"time"
 
-	"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"
@@ -28,11 +27,12 @@ func main() {
 	flag.StringVar(&args.Logger, "logger", config.DefaultHealthCheckLogger, "Logger name")
 	flag.StringVar(&args.LogLevel, "logLevel", "none", "Log level")
 	flag.StringVar(&args.ServersStr, "server", "", "Remote server to connect")
+	flag.BoolVar(&args.NoAuthKey, "no-auth-key", false, "Disable auth-key fast reconnect feature")
 	flag.StringVar(&pprof, "pprof", "", "Start PProf server this address")
 	flag.Parse()
 
 	if displayVersion {
-		version.PrintAndExit()
+		version.PrintAndExit(false)
 	}
 
 	config.Setup(source.HealthCheck, &args, flag.Args())
@@ -43,13 +43,35 @@ func main() {
 	wg.Add(1)
 	dlog.Start(ctx, &wg, source.HealthCheck)
 
+	var pprofServer *cli.PProfServer
 	if pprof != "" {
-		dlog.Client.Info("Starting PProf", pprof)
-		go func() {
-			panic(http.ListenAndServe(pprof, nil))
-		}()
+		pprofServer, pprofErr := cli.NewPProfServer(pprof)
+		if pprofErr != nil {
+			dlog.Client.Error("Unable to start PProf", pprofErr)
+		} else {
+			dlog.Client.Info("Starting PProf", pprofServer.Address())
+			pprofServer.Start(nil)
+		}
+	}
+
+	healthClient, err := clients.NewHealthClient(args)
+	status := 0
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "CRITICAL: unable to create dtailhealth client: %v\n", err)
+		status = 2
+	} else {
+		status = healthClient.Start(ctx, signal.NoCh(ctx))
+	}
+
+	if pprofServer != nil {
+		shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
+		if err := pprofServer.Shutdown(shutdownCtx); err != nil {
+			dlog.Client.Error("Unable to stop PProf", err)
+		}
+		shutdownCancel()
 	}
 
-	healthClient, _ := clients.NewHealthClient(args)
-	os.Exit(healthClient.Start(ctx, signal.NoCh(ctx)))
+	cancel()
+	wg.Wait()
+	os.Exit(status)
 }