From 849951be1d1a7ee9f9302006ccb187bf5b4e36f3 Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Wed, 22 Jul 2026 23:51:18 +0300 Subject: =?UTF-8?q?feat:=20DTail=20fork=20=E2=80=94=20server/client=20feat?= =?UTF-8?q?ure=20development?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- internal/cli/authkeyflags.go | 43 + internal/cli/authkeyflags_test.go | 94 ++ internal/cli/pprof.go | 119 ++ internal/cli/pprof_test.go | 41 + internal/cli/runtime.go | 104 ++ internal/cli/runtime_test.go | 71 ++ internal/clients/baseclient.go | 213 +++- internal/clients/baseclient_retry_test.go | 280 +++++ internal/clients/catclient.go | 24 +- internal/clients/client_benchmark_test.go | 136 +++ internal/clients/connectors/connector.go | 18 + internal/clients/connectors/serverconnection.go | 248 +++- .../clients/connectors/serverconnection_test.go | 880 ++++++++++++++ internal/clients/connectors/serverless.go | 213 +++- internal/clients/connectors/sessiontransport.go | 213 ++++ internal/clients/grepclient.go | 24 +- internal/clients/handlers/basehandler.go | 319 ++++- internal/clients/handlers/basehandler_test.go | 356 ++++++ internal/clients/handlers/clienthandler.go | 15 +- internal/clients/handlers/handler.go | 6 + internal/clients/handlers/healthhandler.go | 13 +- internal/clients/handlers/maprhandler.go | 55 +- internal/clients/handlers/maprhandler_test.go | 278 +++++ internal/clients/healthclient.go | 22 +- internal/clients/interactive_control.go | 425 +++++++ internal/clients/interactive_control_test.go | 641 ++++++++++ internal/clients/maker.go | 8 + internal/clients/maprclient.go | 269 ++-- internal/clients/maprclient_test.go | 153 +++ internal/clients/query_regex.go | 30 + internal/clients/runtime_boundary.go | 219 ++++ internal/clients/runtime_boundary_test.go | 89 ++ internal/clients/session_spec.go | 14 + internal/clients/session_spec_test.go | 147 +++ internal/clients/stats.go | 28 +- internal/clients/tailclient.go | 25 +- internal/color/brush/brush.go | 30 +- internal/color/brush/brush_test.go | 90 ++ internal/color/color.go | 6 +- internal/config/args.go | 47 +- internal/config/args_test.go | 95 ++ internal/config/client.go | 41 +- internal/config/common.go | 3 + internal/config/config.go | 13 +- internal/config/env.go | 10 + internal/config/initializer.go | 87 +- internal/config/initializer_test.go | 260 ++++ internal/config/runtime.go | 18 + internal/config/server.go | 89 +- internal/ctxutil/sleep.go | 29 + internal/ctxutil/sleep_test.go | 32 + internal/discovery/comma.go | 15 +- internal/discovery/discovery.go | 34 +- internal/discovery/discovery_test.go | 176 +++ internal/io/dlog/dlog.go | 226 ++-- internal/io/dlog/dlog_test.go | 54 + internal/io/dlog/loggers/file.go | 122 +- internal/io/dlog/loggers/file_test.go | 183 +++ internal/io/dlog/loggers/fout.go | 68 +- internal/io/dlog/loggers/fout_test.go | 199 +++ internal/io/dlog/loggers/stdout.go | 86 +- internal/io/dlog/loggers/stdout_test.go | 190 +++ internal/io/dlog/rawlog_test.go | 76 ++ internal/io/dlog/rotation.go | 26 +- internal/io/dlog/rotation_test.go | 61 + internal/io/fs/catfile.go | 14 +- internal/io/fs/filereader.go | 8 +- internal/io/fs/permissions/permission.go | 1 - internal/io/fs/permissions/permission_linuxacl.go | 1 - internal/io/fs/permissions/permission_test.go | 1 - internal/io/fs/readfile.go | 238 +--- internal/io/fs/readfile_nozstd.go | 16 + internal/io/fs/readfile_processor.go | 359 ++++++ internal/io/fs/readfile_processor_optimized.go | 430 +++++++ internal/io/fs/readfile_processor_test.go | 869 +++++++++++++ internal/io/fs/readfile_zstd.go | 20 + internal/io/fs/readfilelcontext.go | 209 ---- internal/io/fs/rootedpath.go | 96 ++ internal/io/fs/rootedpath_test.go | 57 + internal/io/fs/tailfile.go | 14 +- internal/io/fs/validatedreadtarget.go | 148 +++ internal/io/fs/validatedreadtarget_test.go | 218 ++++ internal/io/journal/filter.go | 253 ++++ internal/io/journal/reader.go | 303 +++++ internal/io/journal/reader_test.go | 754 ++++++++++++ internal/io/journal/reader_unsupported.go | 45 + internal/io/journal/testhelper/mock.go | 397 ++++++ internal/io/journal/testhelper/mock_test.go | 247 ++++ internal/io/line/line.go | 18 +- internal/io/line/processor.go | 22 + internal/io/pool/bytesbuffer.go | 4 +- internal/io/pool/scanner_pool.go | 85 ++ internal/io/signal/signal.go | 45 + internal/mapr/aggregateset.go | 22 +- internal/mapr/client/aggregate.go | 73 +- internal/mapr/client/aggregate_test.go | 96 ++ internal/mapr/client/session_state.go | 95 ++ internal/mapr/client/session_state_test.go | 81 ++ internal/mapr/funcs/function.go | 64 +- internal/mapr/funcs/function_test.go | 114 +- internal/mapr/globalgroupset.go | 11 +- internal/mapr/globalgroupset_test.go | 95 ++ internal/mapr/groupset.go | 133 +- internal/mapr/groupset_avg_nan_test.go | 79 ++ internal/mapr/groupset_ordering_test.go | 125 ++ internal/mapr/groupset_percentage_test.go | 143 +++ internal/mapr/groupsetresult.go | 108 +- internal/mapr/groupsetresult_renderer_test.go | 113 ++ internal/mapr/logformat/csv.go | 97 +- internal/mapr/logformat/csv_test.go | 170 ++- internal/mapr/logformat/custom1.go | 5 +- internal/mapr/logformat/custom2.go | 5 +- internal/mapr/logformat/default.go | 239 +++- internal/mapr/logformat/default_benchmark_test.go | 44 + internal/mapr/logformat/default_test.go | 42 +- internal/mapr/logformat/delimited.go | 12 + internal/mapr/logformat/generic.go | 15 +- internal/mapr/logformat/generickv.go | 36 +- internal/mapr/logformat/mimecast.go | 5 +- internal/mapr/logformat/parser.go | 123 +- internal/mapr/logformat/parser_test.go | 69 ++ internal/mapr/logformat/variables.go | 107 ++ internal/mapr/logformat/variables_test.go | 145 +++ internal/mapr/parserfieldplan.go | 81 ++ internal/mapr/parserfieldplan_test.go | 32 + internal/mapr/query.go | 29 +- internal/mapr/query_test.go | 83 ++ internal/mapr/queryvariables.go | 92 ++ internal/mapr/queryvariables_test.go | 78 ++ internal/mapr/result_renderer.go | 34 + internal/mapr/safe_aggregateset.go | 72 ++ internal/mapr/safe_aggregateset_test.go | 147 +++ internal/mapr/selectcondition.go | 6 + internal/mapr/server/aggregate.go | 620 +++++++--- internal/mapr/server/aggregate_test.go | 714 +++++++++++ internal/mapr/server/groupkey.go | 31 + internal/mapr/server/parsername.go | 10 + internal/mapr/server/parsername_test.go | 62 + internal/mapr/token.go | 35 +- internal/mapr/token_test.go | 77 ++ internal/profiling/README.md | 320 +++++ internal/profiling/flags.go | 38 + internal/profiling/profiler.go | 227 ++++ internal/profiling/profiler_test.go | 269 ++++ internal/protocol/capabilities.go | 12 + internal/protocol/protocol.go | 6 + internal/protocol/session.go | 10 + internal/regex/bench_test.go | 111 ++ internal/regex/regex.go | 128 +- internal/regex/regex_literal_test.go | 226 ++++ internal/regex/regex_test.go | 23 + internal/server/auth_test.go | 50 + internal/server/continuous.go | 61 +- internal/server/continuous_test.go | 202 +++ internal/server/handlers/authkeycommand_test.go | 117 ++ internal/server/handlers/basehandler.go | 547 +++++++-- internal/server/handlers/basehandler_read_test.go | 295 +++++ internal/server/handlers/commandcancel_test.go | 143 +++ internal/server/handlers/commandtimeout_test.go | 150 +++ internal/server/handlers/framesize_test.go | 202 +++ internal/server/handlers/generation_output.go | 75 ++ internal/server/handlers/generation_output_test.go | 160 +++ internal/server/handlers/healthhandler.go | 28 +- internal/server/handlers/line_writer.go | 777 ++++++++++++ internal/server/handlers/line_writer_alloc_test.go | 54 + internal/server/handlers/line_writer_test.go | 1282 ++++++++++++++++++++ internal/server/handlers/lineprocessor.go | 177 +++ internal/server/handlers/lineprocessor_test.go | 82 ++ internal/server/handlers/mapcommand.go | 19 +- .../server/handlers/mapcommand_completion_test.go | 337 +++++ internal/server/handlers/mapcommand_race_test.go | 64 + internal/server/handlers/output_manager.go | 518 ++++++++ .../server/handlers/output_manager_race_test.go | 437 +++++++ internal/server/handlers/protocol_codec.go | 87 ++ internal/server/handlers/protocol_codec_test.go | 32 + internal/server/handlers/protocol_formatter.go | 40 + internal/server/handlers/readcommand.go | 513 +++++++- .../handlers/readcommand_cancellation_test.go | 65 + .../handlers/readcommand_epoch_order_test.go | 127 ++ .../server/handlers/readcommand_glob_cap_test.go | 214 ++++ .../server/handlers/readcommand_journal_test.go | 271 +++++ .../server/handlers/readcommand_semaphore_test.go | 108 ++ .../handlers/readcommand_sendmessage_test.go | 252 ++++ internal/server/handlers/readcommand_server.go | 250 ++++ internal/server/handlers/serverhandler.go | 247 +++- internal/server/handlers/sessioncommand.go | 258 ++++ internal/server/handlers/sessioncommand_test.go | 568 +++++++++ internal/server/handlers/shutdown_coordinator.go | 108 ++ internal/server/scheduler.go | 19 +- internal/server/server.go | 329 +++-- internal/server/stats.go | 63 +- internal/server/stats_test.go | 206 ++++ internal/session/spec.go | 183 +++ internal/session/spec_test.go | 117 ++ internal/ssh/client/authmethods.go | 178 ++- internal/ssh/client/authmethods_test.go | 194 +++ internal/ssh/client/customkeycallback.go | 6 +- internal/ssh/client/hostkeycallback.go | 7 +- internal/ssh/client/knownhostscallback.go | 205 +++- internal/ssh/client/knownhostscallback_test.go | 305 +++++ internal/ssh/client/simplecallback.go | 5 +- internal/ssh/server/authkeystore.go | 186 +++ internal/ssh/server/authkeystore_test.go | 178 +++ internal/ssh/server/hostkey.go | 70 +- internal/ssh/server/hostkey_test.go | 37 + internal/ssh/server/publickeycallback.go | 159 ++- internal/ssh/server/publickeycallback_test.go | 297 +++++ internal/ssh/ssh.go | 145 ++- internal/ssh/ssh_agent_test.go | 152 +++ internal/ssh/ssh_test.go | 113 ++ internal/tools/benchmark/benchmark.go | 385 ++++++ internal/tools/common/data_generator.go | 268 ++++ internal/tools/common/utils.go | 213 ++++ internal/tools/pgo/pgo.go | 1219 +++++++++++++++++++ internal/tools/pgo/pgo_test.go | 132 ++ internal/tools/profile/analyze.go | 221 ++++ internal/tools/profile/profile.go | 367 ++++++ internal/tools/profile/profile_test.go | 30 + internal/user/server/user.go | 122 +- internal/user/server/user_test.go | 203 ++++ internal/version/version.go | 15 +- 221 files changed, 32442 insertions(+), 1914 deletions(-) create mode 100644 internal/cli/authkeyflags.go create mode 100644 internal/cli/authkeyflags_test.go create mode 100644 internal/cli/pprof.go create mode 100644 internal/cli/pprof_test.go create mode 100644 internal/cli/runtime.go create mode 100644 internal/cli/runtime_test.go create mode 100644 internal/clients/baseclient_retry_test.go create mode 100644 internal/clients/client_benchmark_test.go create mode 100644 internal/clients/connectors/serverconnection_test.go create mode 100644 internal/clients/connectors/sessiontransport.go create mode 100644 internal/clients/handlers/basehandler_test.go create mode 100644 internal/clients/handlers/maprhandler_test.go create mode 100644 internal/clients/interactive_control.go create mode 100644 internal/clients/interactive_control_test.go create mode 100644 internal/clients/maprclient_test.go create mode 100644 internal/clients/query_regex.go create mode 100644 internal/clients/runtime_boundary.go create mode 100644 internal/clients/runtime_boundary_test.go create mode 100644 internal/clients/session_spec.go create mode 100644 internal/clients/session_spec_test.go create mode 100644 internal/color/brush/brush_test.go create mode 100644 internal/config/args_test.go create mode 100644 internal/config/initializer_test.go create mode 100644 internal/config/runtime.go create mode 100644 internal/ctxutil/sleep.go create mode 100644 internal/ctxutil/sleep_test.go create mode 100644 internal/discovery/discovery_test.go create mode 100644 internal/io/dlog/dlog_test.go create mode 100644 internal/io/dlog/loggers/file_test.go create mode 100644 internal/io/dlog/loggers/fout_test.go create mode 100644 internal/io/dlog/loggers/stdout_test.go create mode 100644 internal/io/dlog/rawlog_test.go create mode 100644 internal/io/dlog/rotation_test.go create mode 100644 internal/io/fs/readfile_nozstd.go create mode 100644 internal/io/fs/readfile_processor.go create mode 100644 internal/io/fs/readfile_processor_optimized.go create mode 100644 internal/io/fs/readfile_processor_test.go create mode 100644 internal/io/fs/readfile_zstd.go delete mode 100644 internal/io/fs/readfilelcontext.go create mode 100644 internal/io/fs/rootedpath.go create mode 100644 internal/io/fs/rootedpath_test.go create mode 100644 internal/io/fs/validatedreadtarget.go create mode 100644 internal/io/fs/validatedreadtarget_test.go create mode 100644 internal/io/journal/filter.go create mode 100644 internal/io/journal/reader.go create mode 100644 internal/io/journal/reader_test.go create mode 100644 internal/io/journal/reader_unsupported.go create mode 100644 internal/io/journal/testhelper/mock.go create mode 100644 internal/io/journal/testhelper/mock_test.go create mode 100644 internal/io/line/processor.go create mode 100644 internal/io/pool/scanner_pool.go create mode 100644 internal/mapr/client/aggregate_test.go create mode 100644 internal/mapr/client/session_state.go create mode 100644 internal/mapr/client/session_state_test.go create mode 100644 internal/mapr/globalgroupset_test.go create mode 100644 internal/mapr/groupset_avg_nan_test.go create mode 100644 internal/mapr/groupset_ordering_test.go create mode 100644 internal/mapr/groupset_percentage_test.go create mode 100644 internal/mapr/groupsetresult_renderer_test.go create mode 100644 internal/mapr/logformat/default_benchmark_test.go create mode 100644 internal/mapr/logformat/delimited.go create mode 100644 internal/mapr/logformat/parser_test.go create mode 100644 internal/mapr/logformat/variables.go create mode 100644 internal/mapr/logformat/variables_test.go create mode 100644 internal/mapr/parserfieldplan.go create mode 100644 internal/mapr/parserfieldplan_test.go create mode 100644 internal/mapr/queryvariables.go create mode 100644 internal/mapr/queryvariables_test.go create mode 100644 internal/mapr/result_renderer.go create mode 100644 internal/mapr/safe_aggregateset.go create mode 100644 internal/mapr/safe_aggregateset_test.go create mode 100644 internal/mapr/server/aggregate_test.go create mode 100644 internal/mapr/server/groupkey.go create mode 100644 internal/mapr/server/parsername.go create mode 100644 internal/mapr/server/parsername_test.go create mode 100644 internal/mapr/token_test.go create mode 100644 internal/profiling/README.md create mode 100644 internal/profiling/flags.go create mode 100644 internal/profiling/profiler.go create mode 100644 internal/profiling/profiler_test.go create mode 100644 internal/protocol/capabilities.go create mode 100644 internal/protocol/session.go create mode 100644 internal/regex/bench_test.go create mode 100644 internal/regex/regex_literal_test.go create mode 100644 internal/server/auth_test.go create mode 100644 internal/server/continuous_test.go create mode 100644 internal/server/handlers/authkeycommand_test.go create mode 100644 internal/server/handlers/basehandler_read_test.go create mode 100644 internal/server/handlers/commandcancel_test.go create mode 100644 internal/server/handlers/commandtimeout_test.go create mode 100644 internal/server/handlers/framesize_test.go create mode 100644 internal/server/handlers/generation_output.go create mode 100644 internal/server/handlers/generation_output_test.go create mode 100644 internal/server/handlers/line_writer.go create mode 100644 internal/server/handlers/line_writer_alloc_test.go create mode 100644 internal/server/handlers/line_writer_test.go create mode 100644 internal/server/handlers/lineprocessor.go create mode 100644 internal/server/handlers/lineprocessor_test.go create mode 100644 internal/server/handlers/mapcommand_completion_test.go create mode 100644 internal/server/handlers/mapcommand_race_test.go create mode 100644 internal/server/handlers/output_manager.go create mode 100644 internal/server/handlers/output_manager_race_test.go create mode 100644 internal/server/handlers/protocol_codec.go create mode 100644 internal/server/handlers/protocol_codec_test.go create mode 100644 internal/server/handlers/protocol_formatter.go create mode 100644 internal/server/handlers/readcommand_cancellation_test.go create mode 100644 internal/server/handlers/readcommand_epoch_order_test.go create mode 100644 internal/server/handlers/readcommand_glob_cap_test.go create mode 100644 internal/server/handlers/readcommand_journal_test.go create mode 100644 internal/server/handlers/readcommand_semaphore_test.go create mode 100644 internal/server/handlers/readcommand_sendmessage_test.go create mode 100644 internal/server/handlers/readcommand_server.go create mode 100644 internal/server/handlers/sessioncommand.go create mode 100644 internal/server/handlers/sessioncommand_test.go create mode 100644 internal/server/handlers/shutdown_coordinator.go create mode 100644 internal/server/stats_test.go create mode 100644 internal/session/spec.go create mode 100644 internal/session/spec_test.go create mode 100644 internal/ssh/client/authmethods_test.go create mode 100644 internal/ssh/client/knownhostscallback_test.go create mode 100644 internal/ssh/server/authkeystore.go create mode 100644 internal/ssh/server/authkeystore_test.go create mode 100644 internal/ssh/server/hostkey_test.go create mode 100644 internal/ssh/server/publickeycallback_test.go create mode 100644 internal/ssh/ssh_agent_test.go create mode 100644 internal/ssh/ssh_test.go create mode 100644 internal/tools/benchmark/benchmark.go create mode 100644 internal/tools/common/data_generator.go create mode 100644 internal/tools/common/utils.go create mode 100644 internal/tools/pgo/pgo.go create mode 100644 internal/tools/pgo/pgo_test.go create mode 100644 internal/tools/profile/analyze.go create mode 100644 internal/tools/profile/profile.go create mode 100644 internal/tools/profile/profile_test.go create mode 100644 internal/user/server/user_test.go (limited to 'internal') diff --git a/internal/cli/authkeyflags.go b/internal/cli/authkeyflags.go new file mode 100644 index 0000000..4c945d3 --- /dev/null +++ b/internal/cli/authkeyflags.go @@ -0,0 +1,43 @@ +package cli + +import ( + "flag" + + "github.com/mimecast/dtail/internal/config" +) + +const authKeyPathHelpText = "Path to auth key/private key (defaults to ~/.ssh/id_rsa via config)" + +// BindAuthKeyFlags registers the legacy and current auth-key flags. +func BindAuthKeyFlags(fs *flag.FlagSet, legacyKey *string, args *config.Args) { + fs.StringVar(legacyKey, "key", "", "Deprecated alias for -auth-key-path") + fs.StringVar(&args.SSHPrivateKeyFilePath, "auth-key-path", "", authKeyPathHelpText) +} + +// FlagWasSet reports whether the named flag was explicitly set on the command line. +func FlagWasSet(name string) bool { + var wasSet bool + flag.Visit(func(f *flag.Flag) { + if f.Name == name { + wasSet = true + } + }) + return wasSet +} + +// ApplyAuthKeyPathCompatibility copies the deprecated legacy key into args when +// the new flag was not explicitly set. If both were set, the new flag wins and +// a warning message is returned. +func ApplyAuthKeyPathCompatibility(args *config.Args, legacyKey string, authKeyPathSet bool) string { + if authKeyPathSet { + if legacyKey != "" { + return "WARN: -key is deprecated; ignoring it because -auth-key-path was also set" + } + return "" + } + + if legacyKey != "" { + args.SSHPrivateKeyFilePath = legacyKey + } + return "" +} diff --git a/internal/cli/authkeyflags_test.go b/internal/cli/authkeyflags_test.go new file mode 100644 index 0000000..64c76c3 --- /dev/null +++ b/internal/cli/authkeyflags_test.go @@ -0,0 +1,94 @@ +package cli + +import ( + "flag" + "testing" + + "github.com/mimecast/dtail/internal/config" +) + +func TestApplyAuthKeyPathCompatibilityLegacyOnly(t *testing.T) { + args := config.Args{} + + if warning := ApplyAuthKeyPathCompatibility(&args, "/tmp/legacy.pem", false); warning != "" { + t.Fatalf("unexpected warning: %q", warning) + } + + if got, want := args.SSHPrivateKeyFilePath, "/tmp/legacy.pem"; got != want { + t.Fatalf("unexpected auth key path: want %q got %q", want, got) + } +} + +func TestApplyAuthKeyPathCompatibilityPrefersExplicitAuthKeyPath(t *testing.T) { + tests := []struct { + name string + args []string + want string + warn bool + }{ + { + name: "legacy then auth", + args: []string{"-key", "/tmp/legacy.pem", "-auth-key-path", "/tmp/current.pem"}, + want: "/tmp/current.pem", + warn: true, + }, + { + name: "auth then legacy", + args: []string{"-auth-key-path", "/tmp/current.pem", "-key", "/tmp/legacy.pem"}, + want: "/tmp/current.pem", + warn: true, + }, + { + name: "explicit blank auth key path keeps blank", + args: []string{"-key", "/tmp/legacy.pem", "-auth-key-path="}, + want: "", + warn: true, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + fs := flag.NewFlagSet(t.Name(), flag.ContinueOnError) + var args config.Args + var legacyKey string + + BindAuthKeyFlags(fs, &legacyKey, &args) + if err := fs.Parse(tc.args); err != nil { + t.Fatalf("Parse failed: %v", err) + } + + var authKeyPathSet bool + fs.Visit(func(f *flag.Flag) { + if f.Name == "auth-key-path" { + authKeyPathSet = true + } + }) + + warning := ApplyAuthKeyPathCompatibility(&args, legacyKey, authKeyPathSet) + if gotWarn := warning != ""; gotWarn != tc.warn { + t.Fatalf("unexpected warning presence: want %v got %v (%q)", tc.warn, gotWarn, warning) + } + if got, want := args.SSHPrivateKeyFilePath, tc.want; got != want { + t.Fatalf("unexpected auth key path: want %q got %q", want, got) + } + }) + } +} + +func TestBindAuthKeyFlagsHelpText(t *testing.T) { + fs := flag.NewFlagSet(t.Name(), flag.ContinueOnError) + var args config.Args + var legacyKey string + + BindAuthKeyFlags(fs, &legacyKey, &args) + + if got, want := fs.Lookup("key").Usage, "Deprecated alias for -auth-key-path"; got != want { + t.Fatalf("unexpected legacy flag help: want %q got %q", want, got) + } + if got, want := fs.Lookup("auth-key-path").Usage, authKeyPathHelpText; got != want { + t.Fatalf("unexpected auth-key-path help: want %q got %q", want, got) + } + if got, want := fs.Lookup("auth-key-path").DefValue, ""; got != want { + t.Fatalf("unexpected auth-key-path default: want %q got %q", want, got) + } +} diff --git a/internal/cli/pprof.go b/internal/cli/pprof.go new file mode 100644 index 0000000..57cb38c --- /dev/null +++ b/internal/cli/pprof.go @@ -0,0 +1,119 @@ +package cli + +import ( + "context" + "errors" + "net" + "net/http" + "net/http/pprof" + "runtime" + "sync" + + "github.com/mimecast/dtail/internal/io/dlog" +) + +// Mutex and block profiling rates. The /debug/pprof/mutex and +// /debug/pprof/block endpoints return empty profiles unless these runtime +// rates are enabled first; they are off by default in the Go runtime. +// +// The values below were validated during task ss0: a mutex profile fraction of +// 5 (report ~1/5 of contention events) and a block profile rate of 100000ns +// (record blocking events longer than 100µs) were sufficient to measure +// outputManager mutex contention — 12.4ms total delay across three concurrent +// 100MB server-mode dcats, i.e. NOT a bottleneck; an idle follow session +// costs 0 CPU ticks/10s because the 1ms read poll only runs while direct output is +// active and the EOF-ack drops it back to 1s polling. Documented here so the +// locking design is not re-litigated. +const ( + mutexProfileFraction = 5 + blockProfileRateNanos = 100000 +) + +// PProfServer owns a dedicated pprof HTTP server lifecycle. +type PProfServer struct { + listener net.Listener + server *http.Server + done chan struct{} +} + +// NewPProfServer creates a pprof HTTP server bound to address. +func NewPProfServer(address string) (*PProfServer, error) { + listener, err := net.Listen("tcp", address) + if err != nil { + return nil, err + } + + return &PProfServer{ + listener: listener, + server: &http.Server{ + Handler: newPProfServeMux(), + }, + done: make(chan struct{}), + }, nil +} + +// EnableProfilingRates turns on mutex and block profiling collection so that +// the /debug/pprof/mutex and /debug/pprof/block endpoints actually contain +// samples. Without this the Go runtime keeps both rates at zero and those +// endpoints report empty profiles. Only call this when pprof is enabled so it +// costs nothing in the common (no --pprof) case. +func EnableProfilingRates() { + runtime.SetMutexProfileFraction(mutexProfileFraction) + runtime.SetBlockProfileRate(blockProfileRateNanos) +} + +// Address returns the bound pprof listener address. +func (s *PProfServer) Address() string { + if s == nil || s.listener == nil { + return "" + } + return s.listener.Addr().String() +} + +// Start serves the pprof HTTP endpoints until shutdown. +func (s *PProfServer) Start(wg *sync.WaitGroup) { + if s == nil { + return + } + + if wg != nil { + wg.Add(1) + } + + go func() { + if wg != nil { + defer wg.Done() + } + defer close(s.done) + + if err := s.server.Serve(s.listener); err != nil && !errors.Is(err, http.ErrServerClosed) { + dlog.Client.Error("PProf server exited", err) + } + }() +} + +// Shutdown stops the pprof HTTP server and waits for Serve to return. +func (s *PProfServer) Shutdown(ctx context.Context) error { + if s == nil { + return nil + } + + err := s.server.Shutdown(ctx) + <-s.done + return err +} + +func newPProfServeMux() *http.ServeMux { + mux := http.NewServeMux() + mux.HandleFunc("/debug/pprof/", pprof.Index) + mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline) + mux.HandleFunc("/debug/pprof/profile", pprof.Profile) + mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol) + mux.HandleFunc("/debug/pprof/trace", pprof.Trace) + + for _, name := range []string{"allocs", "block", "goroutine", "heap", "mutex", "threadcreate"} { + mux.Handle("/debug/pprof/"+name, pprof.Handler(name)) + } + + return mux +} diff --git a/internal/cli/pprof_test.go b/internal/cli/pprof_test.go new file mode 100644 index 0000000..aa1cd05 --- /dev/null +++ b/internal/cli/pprof_test.go @@ -0,0 +1,41 @@ +package cli + +import ( + "runtime" + "testing" +) + +// TestEnableProfilingRatesSetsMutexFraction verifies that EnableProfilingRates +// actually turns on mutex profiling. Without this the /debug/pprof/mutex +// endpoint reports an empty profile. runtime.SetMutexProfileFraction(-1) reads +// the current fraction without changing it, so it lets us assert the state. +func TestEnableProfilingRatesSetsMutexFraction(t *testing.T) { + // Save and restore global runtime state so this test does not leak into + // other tests in the package. The block profile rate has no getter, so we + // simply disable it again on cleanup. + prevMutex := runtime.SetMutexProfileFraction(-1) + t.Cleanup(func() { + runtime.SetMutexProfileFraction(prevMutex) + runtime.SetBlockProfileRate(0) + }) + + // Start from a known-disabled state. + runtime.SetMutexProfileFraction(0) + + EnableProfilingRates() + + if got := runtime.SetMutexProfileFraction(-1); got != mutexProfileFraction { + t.Errorf("mutex profile fraction = %d, want %d", got, mutexProfileFraction) + } +} + +// TestProfilingRateConstants pins the ss0-validated values so a change is a +// conscious decision rather than an accident. +func TestProfilingRateConstants(t *testing.T) { + if mutexProfileFraction != 5 { + t.Errorf("mutexProfileFraction = %d, want 5", mutexProfileFraction) + } + if blockProfileRateNanos != 100000 { + t.Errorf("blockProfileRateNanos = %d, want 100000", blockProfileRateNanos) + } +} diff --git a/internal/cli/runtime.go b/internal/cli/runtime.go new file mode 100644 index 0000000..950ed33 --- /dev/null +++ b/internal/cli/runtime.go @@ -0,0 +1,104 @@ +package cli + +import ( + "context" + "sync" + "time" + + "github.com/mimecast/dtail/internal/io/dlog" + "github.com/mimecast/dtail/internal/profiling" + "github.com/mimecast/dtail/internal/source" +) + +// ClientRuntime owns common client command runtime components. +type ClientRuntime struct { + ctx context.Context + cancel context.CancelFunc + wg sync.WaitGroup + pprofServer *PProfServer + profiler *profiling.Profiler + profileEnabled bool +} + +// NewClientRuntime starts logging and profiling for a client command. +func NewClientRuntime(parent context.Context, profileFlags profiling.Flags, profileName string) *ClientRuntime { + if parent == nil { + parent = context.Background() + } + ctx, cancel := context.WithCancel(parent) + runtime := &ClientRuntime{ + ctx: ctx, + cancel: cancel, + profiler: profiling.NewProfiler(profileFlags.ToConfig(profileName)), + profileEnabled: profileFlags.Enabled(), + } + + runtime.wg.Add(1) + dlog.Start(ctx, &runtime.wg, source.Client) + return runtime +} + +// Context returns the runtime context. +func (r *ClientRuntime) Context() context.Context { + return r.ctx +} + +// Cancel cancels the runtime context. +func (r *ClientRuntime) Cancel() { + r.cancel() +} + +// StartPProf starts the pprof server if an address is provided. +func (r *ClientRuntime) StartPProf(address string) { + if address == "" { + return + } + + r.stopPProf() + + server, err := NewPProfServer(address) + if err != nil { + dlog.Client.Error("Unable to start PProf", err) + return + } + + r.pprofServer = server + dlog.Client.Info("Starting PProf", server.Address()) + server.Start(&r.wg) +} + +// LogStartupMetrics logs startup profiling metrics when enabled. +func (r *ClientRuntime) LogStartupMetrics() { + if r.profileEnabled { + r.profiler.LogMetrics("startup") + } +} + +// LogShutdownMetrics logs shutdown profiling metrics when enabled. +func (r *ClientRuntime) LogShutdownMetrics() { + if r.profileEnabled { + r.profiler.LogMetrics("shutdown") + } +} + +// Stop stops profiling and logging runtime goroutines. +func (r *ClientRuntime) Stop() { + r.profiler.Stop() + r.stopPProf() + r.cancel() + r.wg.Wait() +} + +func (r *ClientRuntime) stopPProf() { + if r.pprofServer == nil { + return + } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + if err := r.pprofServer.Shutdown(ctx); err != nil { + dlog.Client.Error("Unable to stop PProf", err) + } + r.pprofServer = nil +} diff --git a/internal/cli/runtime_test.go b/internal/cli/runtime_test.go new file mode 100644 index 0000000..e024a4e --- /dev/null +++ b/internal/cli/runtime_test.go @@ -0,0 +1,71 @@ +package cli + +import ( + "context" + "net/http" + "testing" + "time" + + "github.com/mimecast/dtail/internal/io/dlog" + "github.com/mimecast/dtail/internal/profiling" +) + +func TestClientRuntimeStopShutsDownPProf(t *testing.T) { + prevClient := dlog.Client + dlog.Client = &dlog.DLog{} + t.Cleanup(func() { + dlog.Client = prevClient + }) + + ctx, cancel := context.WithCancel(context.Background()) + runtime := &ClientRuntime{ + ctx: ctx, + cancel: cancel, + profiler: profiling.NewProfiler(profiling.Config{}), + } + + runtime.StartPProf("127.0.0.1:0") + if runtime.pprofServer == nil { + t.Fatal("expected pprof server to start") + } + + url := "http://" + runtime.pprofServer.Address() + "/debug/pprof/" + waitForHTTPStatus(t, url, http.StatusOK) + + runtime.Stop() + waitForHTTPError(t, url) +} + +func waitForHTTPStatus(t *testing.T, url string, want int) { + t.Helper() + + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + resp, err := http.Get(url) + if err == nil { + resp.Body.Close() + if resp.StatusCode == want { + return + } + } + time.Sleep(10 * time.Millisecond) + } + + t.Fatalf("timed out waiting for %s to return %d", url, want) +} + +func waitForHTTPError(t *testing.T, url string) { + t.Helper() + + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + resp, err := http.Get(url) + if err != nil { + return + } + resp.Body.Close() + time.Sleep(10 * time.Millisecond) + } + + t.Fatalf("timed out waiting for %s to stop serving", url) +} diff --git a/internal/clients/baseclient.go b/internal/clients/baseclient.go index 013f2f2..165fbd5 100644 --- a/internal/clients/baseclient.go +++ b/internal/clients/baseclient.go @@ -2,6 +2,8 @@ package clients import ( "context" + "io" + "math/rand" "sync" "time" @@ -15,29 +17,52 @@ import ( gossh "golang.org/x/crypto/ssh" ) +const ( + initialRetryDelay = 2 * time.Second + maxRetryDelay = 60 * time.Second + retryJitterFactor = 0.2 // +/-20% jitter to avoid synchronized reconnect storms. +) + // This is the main client data structure. type baseClient struct { + mu *sync.RWMutex config.Args + runtime *clientRuntimeBoundary // To display client side stats stats *stats // We have one connection per remote server. connections []connectors.Connector // SSH auth methods to use to connect to the remote servers. sshAuthMethods []gossh.AuthMethod + // authCloser owns any ssh-agent connection acquired while building the + // auth methods; it must be closed once all SSH handshakes that consume + // sshAuthMethods have completed. + authCloser io.Closer // To deal with SSH host keys hostKeyCallback client.HostKeyCallback // Throttle how fast we initiate SSH connections concurrently throttleCh chan struct{} // Retry connection upon failure? retry bool + // The current connection-wide session specification. + sessionSpec SessionSpec // Connection maker helper. maker maker + // Optional factory override for retry/reconnect tests. + connectionFactory func(server string, authMethods []gossh.AuthMethod, + hostKeyCallback client.HostKeyCallback, sessionSpec SessionSpec, + interactive bool) connectors.Connector + // Optional sleep override for retry tests. + sleepFn func(context.Context, time.Duration) bool // Regex is the regular expresion object for line filtering Regex regex.Regex } func (c *baseClient) init() { dlog.Client.Debug("Initiating base client", c.Args.String()) + if c.runtime == nil { + c.runtime = newClientRuntimeBoundary(config.CurrentRuntime()) + } flag := regex.Default if c.Args.RegexInvert { @@ -52,25 +77,52 @@ func (c *baseClient) init() { if c.Args.Serverless { return } - c.sshAuthMethods, c.hostKeyCallback = client.InitSSHAuthMethods( + c.sshAuthMethods, c.hostKeyCallback, c.authCloser = client.InitSSHAuthMethods( c.Args.SSHAuthMethods, c.Args.SSHHostKeyCallback, c.Args.TrustAllHosts, - c.throttleCh, c.Args.SSHPrivateKeyFilePath) + c.Args.SSHPrivateKeyFilePath, c.Args.SSHAgentKeyIndex) } -func (c *baseClient) makeConnections(maker maker) { +func (c *baseClient) makeConnections(maker maker) error { c.maker = maker + if builder, ok := maker.(sessionSpecMaker); ok { + sessionSpec, err := builder.makeSessionSpec() + if err != nil { + dlog.Client.FatalPanic("unable to build session specification", err) + } + c.sessionSpec = sessionSpec + } - discoveryService := discovery.New(c.Discovery, c.ServersStr, discovery.Shuffle) + discoveryService, err := discovery.New(c.Discovery, c.ServersStr, discovery.Shuffle) + if err != nil { + return err + } for _, server := range discoveryService.ServerList() { c.connections = append(c.connections, c.makeConnection(server, c.sshAuthMethods, c.hostKeyCallback)) } - c.stats = newTailStats(len(c.connections)) + c.stats = newTailStats(len(c.connections), c.runtime.output, c.runtime.InterruptPause()) + return nil } func (c *baseClient) Start(ctx context.Context, statsCh <-chan string) (status int) { + if c.Args.InteractiveQuery { + return c.startInteractiveControl(ctx, statsCh) + } + return c.runConnections(ctx, statsCh) +} + +func (c *baseClient) runConnections(ctx context.Context, statsCh <-chan string) (status int) { dlog.Client.Trace("Starting base client") + // Release the ssh-agent connection (if any) once all handshakes and + // reconnect attempts that consume c.sshAuthMethods have finished. + if c.authCloser != nil { + defer func() { + if err := c.authCloser.Close(); err != nil { + dlog.Client.Debug("baseClient", "failed to close ssh-agent connection", err) + } + }() + } // Can be nil when serverless. if c.hostKeyCallback != nil { // Periodically check for unknown hosts, and ask the user whether to trust them or not. @@ -80,10 +132,11 @@ func (c *baseClient) Start(ctx context.Context, statsCh <-chan string) (status i go c.stats.Start(ctx, c.throttleCh, statsCh, c.Args.Quiet) var wg sync.WaitGroup - wg.Add(len(c.connections)) + connections := c.snapshotConnections() + wg.Add(len(connections)) var mutex sync.Mutex - for i, conn := range c.connections { + for i, conn := range connections { go func(i int, conn connectors.Connector) { defer wg.Done() connStatus := c.startConnection(ctx, i, conn) @@ -102,11 +155,14 @@ func (c *baseClient) Start(ctx context.Context, statsCh <-chan string) (status i func (c *baseClient) startConnection(ctx context.Context, i int, conn connectors.Connector) (status int) { + retryDelay := initialRetryDelay + retryRandom := newRetryRandom(i) + for { connCtx, cancel := context.WithCancel(ctx) - defer cancel() conn.Start(connCtx, cancel, c.throttleCh, c.stats.connectionsEstCh) + cancel() // Retrieve status code from handler (dtail client will exit with that status) status = conn.Handler().Status() @@ -122,20 +178,147 @@ func (c *baseClient) startConnection(ctx context.Context, i int, default: } - // Yes, we want to retry. - time.Sleep(time.Second * 2) - dlog.Client.Debug(conn.Server(), "Reconnecting") + // Yes, we want to retry with exponential backoff and jitter. + sleepDuration := jitterRetryDelay(retryDelay, retryRandom) + dlog.Client.Debug(conn.Server(), "Reconnecting", "backoff", sleepDuration) + if !c.sleepRetry(ctx, sleepDuration) { + return + } + + retryDelay = nextRetryDelay(retryDelay) conn = c.makeConnection(conn.Server(), c.sshAuthMethods, c.hostKeyCallback) - c.connections[i] = conn + c.replaceConnection(i, conn) + } +} + +func nextRetryDelay(current time.Duration) time.Duration { + if current <= 0 { + return initialRetryDelay + } + + next := current * 2 + if next > maxRetryDelay || next < current { + return maxRetryDelay + } + return next +} + +func jitterRetryDelay(base time.Duration, random *rand.Rand) time.Duration { + if base <= 0 || random == nil { + return base + } + + jitter := time.Duration(float64(base) * retryJitterFactor) + if jitter <= 0 { + return base + } + + minDelay := base - jitter + maxDelay := base + jitter + if maxDelay < minDelay { + return base + } + + return minDelay + time.Duration(random.Int63n(int64(maxDelay-minDelay+1))) +} + +func sleepWithContext(ctx context.Context, delay time.Duration) bool { + if delay <= 0 { + return true } + + timer := time.NewTimer(delay) + defer timer.Stop() + + select { + case <-ctx.Done(): + return false + case <-timer.C: + return true + } +} + +func newRetryRandom(seedOffset int) *rand.Rand { + return rand.New(rand.NewSource(time.Now().UnixNano() + int64(seedOffset))) } func (c *baseClient) makeConnection(server string, sshAuthMethods []gossh.AuthMethod, hostKeyCallback client.HostKeyCallback) connectors.Connector { - if c.Args.Serverless { + args, sessionSpec := c.snapshotConnectionState() + return c.makeConnectionWithState(server, sshAuthMethods, hostKeyCallback, args, sessionSpec) +} + +func (c *baseClient) makeConnectionWithState(server string, sshAuthMethods []gossh.AuthMethod, + hostKeyCallback client.HostKeyCallback, args config.Args, sessionSpec SessionSpec) connectors.Connector { + if c.connectionFactory != nil { + return c.connectionFactory(server, sshAuthMethods, hostKeyCallback, + sessionSpec, args.InteractiveQuery) + } + if args.Serverless { return connectors.NewServerless(c.UserName, c.maker.makeHandler(server), - c.maker.makeCommands()) + c.maker.makeCommands(), sessionSpec, args.InteractiveQuery, c.runtime) } return connectors.NewServerConnection(server, c.UserName, sshAuthMethods, - hostKeyCallback, c.maker.makeHandler(server), c.maker.makeCommands()) + hostKeyCallback, c.maker.makeHandler(server), c.maker.makeCommands(), + sessionSpec, args.InteractiveQuery, args.SSHPrivateKeyFilePath, + args.NoAuthKey, c.runtime) +} + +func (c *baseClient) sleepRetry(ctx context.Context, delay time.Duration) bool { + if c.sleepFn != nil { + return c.sleepFn(ctx, delay) + } + return sleepWithContext(ctx, delay) +} + +func (c *baseClient) snapshotConnectionState() (config.Args, SessionSpec) { + mu := c.stateMu() + mu.RLock() + defer mu.RUnlock() + + return c.Args, c.sessionSpec +} + +func (c *baseClient) snapshotMutableState() (config.Args, SessionSpec, []connectors.Connector) { + mu := c.stateMu() + mu.RLock() + defer mu.RUnlock() + + return c.Args, c.sessionSpec, append([]connectors.Connector(nil), c.connections...) +} + +func (c *baseClient) snapshotConnections() []connectors.Connector { + mu := c.stateMu() + mu.RLock() + defer mu.RUnlock() + + return append([]connectors.Connector(nil), c.connections...) +} + +func (c *baseClient) storeReloadState(args config.Args, spec SessionSpec) { + mu := c.stateMu() + mu.Lock() + defer mu.Unlock() + + c.Args = args + c.sessionSpec = spec +} + +func (c *baseClient) replaceConnection(i int, conn connectors.Connector) { + mu := c.stateMu() + mu.Lock() + defer mu.Unlock() + + c.connections[i] = conn +} + +func (c *baseClient) stateMu() *sync.RWMutex { + if c.mu == nil { + c.mu = newBaseClientMu() + } + return c.mu +} + +func newBaseClientMu() *sync.RWMutex { + return &sync.RWMutex{} } diff --git a/internal/clients/baseclient_retry_test.go b/internal/clients/baseclient_retry_test.go new file mode 100644 index 0000000..100f171 --- /dev/null +++ b/internal/clients/baseclient_retry_test.go @@ -0,0 +1,280 @@ +package clients + +import ( + "context" + "math/rand" + "sync/atomic" + "testing" + "time" + + "github.com/mimecast/dtail/internal/clients/connectors" + "github.com/mimecast/dtail/internal/clients/handlers" + "github.com/mimecast/dtail/internal/config" + "github.com/mimecast/dtail/internal/io/dlog" + "github.com/mimecast/dtail/internal/omode" + sshclient "github.com/mimecast/dtail/internal/ssh/client" + + gossh "golang.org/x/crypto/ssh" +) + +func TestNextRetryDelay(t *testing.T) { + tests := []struct { + name string + current time.Duration + want time.Duration + }{ + {name: "zero uses initial", current: 0, want: initialRetryDelay}, + {name: "doubles normally", current: 4 * time.Second, want: 8 * time.Second}, + {name: "caps at max", current: 40 * time.Second, want: maxRetryDelay}, + {name: "stays max at max", current: maxRetryDelay, want: maxRetryDelay}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := nextRetryDelay(tt.current); got != tt.want { + t.Fatalf("nextRetryDelay(%v) = %v, want %v", tt.current, got, tt.want) + } + }) + } +} + +func TestJitterRetryDelayWithinBounds(t *testing.T) { + base := 10 * time.Second + random := rand.New(rand.NewSource(1)) + + min := 8 * time.Second + max := 12 * time.Second + + for i := 0; i < 100; i++ { + got := jitterRetryDelay(base, random) + if got < min || got > max { + t.Fatalf("jitterRetryDelay() = %v, expected between %v and %v", got, min, max) + } + } +} + +func TestSleepWithContextCancellation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + start := time.Now() + if sleepWithContext(ctx, time.Second) { + t.Fatalf("sleepWithContext should stop when context is canceled") + } + + if time.Since(start) > 100*time.Millisecond { + t.Fatalf("sleepWithContext took too long to exit on canceled context") + } +} + +func TestStartConnectionReconnectsWithLatestSessionSpec(t *testing.T) { + originalLogger := dlog.Client + dlog.Client = &dlog.DLog{} + t.Cleanup(func() { + dlog.Client = originalLogger + }) + + first := &retryTestConnector{ + server: "srv1", + handler: &retryTestHandler{}, + } + second := &retryTestConnector{ + server: "srv1", + handler: &retryTestHandler{}, + } + + originalSpec := SessionSpec{ + Mode: omode.TailClient, + Files: []string{"/var/log/app.log"}, + Regex: "ERROR", + } + updatedSpec := SessionSpec{ + Mode: omode.TailClient, + Files: []string{"/var/log/next.log"}, + Regex: "WARN", + } + + sleepCalls := 0 + var capturedSpec SessionSpec + client := &baseClient{ + mu: newBaseClientMu(), + retry: true, + sessionSpec: originalSpec, + stats: &stats{ + connectionsEstCh: make(chan struct{}, 1), + }, + connections: []connectors.Connector{first}, + connectionFactory: func(server string, _ []gossh.AuthMethod, + _ sshclient.HostKeyCallback, sessionSpec SessionSpec, _ bool) connectors.Connector { + if server != "srv1" { + t.Fatalf("unexpected reconnect server %q", server) + } + capturedSpec = sessionSpec + return second + }, + } + client.sleepFn = func(context.Context, time.Duration) bool { + if sleepCalls == 0 { + sleepCalls++ + client.sessionSpec = updatedSpec + return true + } + return false + } + + status := client.startConnection(context.Background(), 0, first) + if status != 0 { + t.Fatalf("startConnection() status = %d, want 0", status) + } + if capturedSpec.Regex != updatedSpec.Regex || len(capturedSpec.Files) != 1 || capturedSpec.Files[0] != updatedSpec.Files[0] { + t.Fatalf("reconnect used stale session spec: got %#v want %#v", capturedSpec, updatedSpec) + } + if client.connections[0] != second { + t.Fatalf("expected retried connector to replace the original connection") + } +} + +func TestApplyInteractiveReloadConcurrentWithReconnect(t *testing.T) { + resetClientLogger(t) + + originalSpec := SessionSpec{ + Mode: omode.GrepClient, + Files: []string{"/var/log/app.log"}, + Regex: "ERROR", + } + nextSpec := SessionSpec{ + Mode: omode.GrepClient, + Files: []string{"/var/log/next.log"}, + Regex: "WARN", + } + originalArgs := config.Args{ + Mode: omode.GrepClient, + What: "/var/log/app.log", + RegexStr: "ERROR", + } + nextArgs := config.Args{ + Mode: omode.GrepClient, + What: "/var/log/next.log", + RegexStr: "WARN", + } + + var reconnects atomic.Int32 + client := &baseClient{ + mu: newBaseClientMu(), + Args: originalArgs, + retry: true, + sessionSpec: originalSpec, + stats: &stats{ + connectionsEstCh: make(chan struct{}, 1), + }, + connections: []connectors.Connector{ + newReloadRetryConnector("srv1", originalSpec), + }, + maker: &interactiveReloadMaker{}, + connectionFactory: func(server string, _ []gossh.AuthMethod, + _ sshclient.HostKeyCallback, sessionSpec SessionSpec, _ bool) connectors.Connector { + reconnects.Add(1) + return newReloadRetryConnector(server, sessionSpec) + }, + } + + client.sleepFn = func(context.Context, time.Duration) bool { + return reconnects.Load() < 200 + } + + done := make(chan int, 1) + go func() { + done <- client.startConnection(context.Background(), 0, client.connections[0]) + }() + + for i := 0; i < 200; i++ { + if err := client.applyInteractiveReload(nextArgs, nextSpec); err != nil { + t.Fatalf("applyInteractiveReload() error = %v", err) + } + if err := client.applyInteractiveReload(originalArgs, originalSpec); err != nil { + t.Fatalf("applyInteractiveReload() error = %v", err) + } + } + + if status := <-done; status != 0 { + t.Fatalf("startConnection() status = %d, want 0", status) + } +} + +func newReloadRetryConnector(server string, spec SessionSpec) *reloadRetryConnector { + return &reloadRetryConnector{ + interactiveReloadConnector: interactiveReloadConnector{ + server: server, + supported: true, + committedSpec: spec, + liveSpec: spec, + generation: 1, + }, + handler: &retryTestHandler{}, + } +} + +type reloadRetryConnector struct { + interactiveReloadConnector + handler handlers.Handler +} + +func (c *reloadRetryConnector) Handler() handlers.Handler { return c.handler } + +type retryTestConnector struct { + handler handlers.Handler + server string +} + +func (c *retryTestConnector) Start(context.Context, context.CancelFunc, chan struct{}, chan struct{}) { +} + +func (c *retryTestConnector) Server() string { return c.server } + +func (c *retryTestConnector) Handler() handlers.Handler { return c.handler } + +func (*retryTestConnector) SupportsQueryUpdates(time.Duration) bool { return false } + +func (*retryTestConnector) ApplySessionSpec(SessionSpec, time.Duration) error { return nil } + +func (*retryTestConnector) ApplySessionSpecWithGeneration(SessionSpec, uint64, time.Duration) error { + return nil +} + +func (*retryTestConnector) CommittedSession() (SessionSpec, uint64, bool) { + return SessionSpec{}, 0, false +} + +func (*retryTestConnector) RestoreCommittedSession(SessionSpec, uint64, bool) {} + +type retryTestHandler struct{} + +func (*retryTestHandler) Read([]byte) (int, error) { return 0, nil } + +func (*retryTestHandler) Write(p []byte) (int, error) { return len(p), nil } + +func (*retryTestHandler) Capabilities() []string { return nil } + +func (*retryTestHandler) HasCapability(string) bool { return false } + +func (*retryTestHandler) ReportServerError(string) {} + +func (*retryTestHandler) SendMessage(string) error { return nil } + +func (*retryTestHandler) Server() string { return "srv1" } + +func (*retryTestHandler) Status() int { return 0 } + +func (*retryTestHandler) Shutdown() {} + +func (*retryTestHandler) Done() <-chan struct{} { + done := make(chan struct{}) + close(done) + return done +} + +func (*retryTestHandler) WaitForCapabilities(time.Duration) bool { return false } + +func (*retryTestHandler) WaitForSessionAck(time.Duration) (handlers.SessionAck, bool) { + return handlers.SessionAck{}, false +} diff --git a/internal/clients/catclient.go b/internal/clients/catclient.go index bd65560..e2e247e 100644 --- a/internal/clients/catclient.go +++ b/internal/clients/catclient.go @@ -2,9 +2,7 @@ package clients import ( "errors" - "fmt" "runtime" - "strings" "github.com/mimecast/dtail/internal/clients/handlers" "github.com/mimecast/dtail/internal/config" @@ -26,14 +24,18 @@ func NewCatClient(args config.Args) (*CatClient, error) { c := CatClient{ baseClient: baseClient{ + mu: newBaseClientMu(), Args: args, throttleCh: make(chan struct{}, args.ConnectionsPerCPU*runtime.NumCPU()), retry: false, + runtime: newClientRuntimeBoundary(config.CurrentRuntime()), }, } c.init() - c.makeConnections(c) + if err := c.makeConnections(c); err != nil { + return nil, err + } return &c, nil } @@ -41,14 +43,18 @@ func (c CatClient) makeHandler(server string) handlers.Handler { return handlers.NewClientHandler(server) } +func (c CatClient) makeSessionSpec() (SessionSpec, error) { + return NewSessionSpec(c.Args), nil +} + func (c CatClient) makeCommands() (commands []string) { - regex, err := c.Regex.Serialize() + sessionSpec, err := c.makeSessionSpec() if err != nil { - dlog.Client.FatalPanic(err) + dlog.Client.FatalPanic("unable to build cat session spec", err) } - for _, file := range strings.Split(c.What, ",") { - commands = append(commands, fmt.Sprintf("%s:%s %s %s", - c.Mode.String(), c.Args.SerializeOptions(), file, regex)) + commands, err = sessionSpec.Commands() + if err != nil { + dlog.Client.FatalPanic("unable to build cat commands from session spec", err) } - return + return commands } diff --git a/internal/clients/client_benchmark_test.go b/internal/clients/client_benchmark_test.go new file mode 100644 index 0000000..9423274 --- /dev/null +++ b/internal/clients/client_benchmark_test.go @@ -0,0 +1,136 @@ +package clients + +import ( + "bytes" + "context" + "fmt" + "os" + "path/filepath" + "testing" + "time" + + "github.com/mimecast/dtail/internal/config" + "github.com/mimecast/dtail/internal/io/dlog" + "github.com/mimecast/dtail/internal/source" + "sync" +) + +func setupBenchmarkData(b *testing.B, lines int) string { + b.Helper() + + tmpDir := b.TempDir() + testFile := filepath.Join(tmpDir, "benchmark_data.log") + + f, err := os.Create(testFile) + if err != nil { + b.Fatalf("Failed to create test file: %v", err) + } + defer f.Close() + + // Create test data + for i := 0; i < lines; i++ { + line := fmt.Sprintf("INFO|1002-071143|1|test.go:%d|8|%d|7|0.21|471h0m21s|MAPREDUCE:STATS|currentConnections=%d|lifetimeConnections=%d|pattern=test-%d|data=%s\n", + i%100, i%50, i%10, i, i%5, "some-test-data-that-makes-the-line-longer") + f.WriteString(line) + } + + return testFile +} + +func BenchmarkDGrep(b *testing.B) { + benchmarkDGrepWithSize(b, 100000) // 100k lines +} + +// Benchmark with different file sizes +func BenchmarkDGrepSmallFile(b *testing.B) { + benchmarkDGrepWithSize(b, 1000) // 1k lines +} + +func BenchmarkDGrepMediumFile(b *testing.B) { + benchmarkDGrepWithSize(b, 50000) // 50k lines +} + +func BenchmarkDGrepLargeFile(b *testing.B) { + benchmarkDGrepWithSize(b, 500000) // 500k lines +} + +func benchmarkDGrepWithSize(b *testing.B, lines int) { + // Setup config. The direct-output read path is the only runtime path. + config.Server = &config.ServerConfig{ + MaxConcurrentCats: 10, + MaxConcurrentTails: 50, + MaxLineLength: 1024 * 1024, + } + + config.Common = &config.CommonConfig{ + Logger: "none", + LogLevel: "error", + } + + config.Client = &config.ClientConfig{ + TermColorsEnable: false, + } + + // Initialize logging + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + wg := &sync.WaitGroup{} + wg.Add(1) + dlog.Start(ctx, wg, source.Client) + + // Create test data + testFile := setupBenchmarkData(b, lines) + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + // Create grep client + args := config.Args{ + ServersStr: "serverless", + QueryStr: "", + What: testFile, + RegexStr: "pattern=test-1", + Serverless: true, + Plain: true, + } + + client, err := NewGrepClient(args) + if err != nil { + b.Fatalf("Failed to create grep client: %v", err) + } + + // Capture output + oldStdout := os.Stdout + r, w, _ := os.Pipe() + os.Stdout = w + + // Run grep + statusCh := make(chan int, 1) + go func() { + status := client.Start(ctx, nil) // nil for statsCh + statusCh <- status + }() + + // Wait for completion or timeout + select { + case status := <-statusCh: + if status != 0 { + b.Errorf("Grep failed with status: %d", status) + } + case <-time.After(30 * time.Second): + b.Error("Grep timed out") + } + + // Restore stdout + w.Close() + os.Stdout = oldStdout + + // Read captured output + var buf bytes.Buffer + buf.ReadFrom(r) + } + + // Report custom metrics + b.ReportMetric(float64(lines), "lines/op") + b.ReportMetric(float64(lines)/b.Elapsed().Seconds(), "lines/sec") +} diff --git a/internal/clients/connectors/connector.go b/internal/clients/connectors/connector.go index 3ab6a08..00de32e 100644 --- a/internal/clients/connectors/connector.go +++ b/internal/clients/connectors/connector.go @@ -2,8 +2,10 @@ package connectors import ( "context" + "time" "github.com/mimecast/dtail/internal/clients/handlers" + sessionspec "github.com/mimecast/dtail/internal/session" ) // Connector interface. @@ -14,4 +16,20 @@ type Connector interface { Server() string // Handler for the connection. Handler() handlers.Handler + // SupportsQueryUpdates reports whether the connected server advertised + // runtime query replacement support within the given timeout. + SupportsQueryUpdates(timeout time.Duration) bool + // ApplySessionSpec starts or updates the interactive session workload on an + // already connected server when query updates are supported. + ApplySessionSpec(spec sessionspec.Spec, timeout time.Duration) error + // ApplySessionSpecWithGeneration starts or updates the interactive session + // workload using the provided committed generation as the base for the + // session update command. + ApplySessionSpecWithGeneration(spec sessionspec.Spec, generation uint64, timeout time.Duration) error + // CommittedSession returns the last session spec and generation that the + // server acknowledged for this connection. + CommittedSession() (sessionspec.Spec, uint64, bool) + // RestoreCommittedSession resets the local committed session snapshot without + // advancing the generation. + RestoreCommittedSession(spec sessionspec.Spec, generation uint64, committed bool) } diff --git a/internal/clients/connectors/serverconnection.go b/internal/clients/connectors/serverconnection.go index 5c3d455..b98f815 100644 --- a/internal/clients/connectors/serverconnection.go +++ b/internal/clients/connectors/serverconnection.go @@ -2,20 +2,38 @@ package connectors import ( "context" + "encoding/base64" "fmt" "io" + "net" + "os" + "path/filepath" "strconv" "strings" + "sync" "time" "github.com/mimecast/dtail/internal/clients/handlers" - "github.com/mimecast/dtail/internal/config" "github.com/mimecast/dtail/internal/io/dlog" + "github.com/mimecast/dtail/internal/protocol" + sessionspec "github.com/mimecast/dtail/internal/session" "github.com/mimecast/dtail/internal/ssh/client" "golang.org/x/crypto/ssh" ) +// SSHSettings provides the connection settings needed by ServerConnection. +type SSHSettings interface { + SSHPort() int + SSHConnectTimeout() time.Duration +} + +const ( + defaultSSHConnectTimeout = 2 * time.Second + defaultSSHPort = 2222 + defaultCapabilityWait = 250 * time.Millisecond +) + // ServerConnection represents a connection to a single remote dtail server via // SSH protocol. type ServerConnection struct { @@ -28,30 +46,61 @@ type ServerConnection struct { config *ssh.ClientConfig handler handlers.Handler commands []string + sessionSpec sessionspec.Spec + sessionState committedSessionState + interactive bool + authKeyPath string + authKeyDisabled bool hostKeyCallback client.HostKeyCallback - throttlingDone bool + // throttleReleased ensures the throttle slot is returned to throttleCh + // exactly once, even if both the early-release path in handle() and the + // deferred cleanup in Start() execute concurrently or the same goroutine + // hits both paths. sync.Once is safe for concurrent callers whereas the + // previous bool guard was not synchronized. + throttleReleased sync.Once } +var _ Connector = (*ServerConnection)(nil) + // NewServerConnection returns a new DTail SSH server connection. func NewServerConnection(server string, userName string, authMethods []ssh.AuthMethod, hostKeyCallback client.HostKeyCallback, - handler handlers.Handler, commands []string) *ServerConnection { + handler handlers.Handler, commands []string, sessionSpec sessionspec.Spec, + interactive bool, authKeyPath string, authKeyDisabled bool, settings SSHSettings) *ServerConnection { dlog.Client.Debug(server, "Creating new connection", server, handler, commands) + sshConnectTimeout := defaultSSHConnectTimeout + defaultPort := defaultSSHPort + if settings != nil { + sshConnectTimeout = settings.SSHConnectTimeout() + defaultPort = settings.SSHPort() + } + if sshConnectTimeout <= 0 { + sshConnectTimeout = defaultSSHConnectTimeout + } + if defaultPort <= 0 { + defaultPort = defaultSSHPort + } + c := ServerConnection{ hostKeyCallback: hostKeyCallback, server: server, handler: handler, commands: commands, + sessionSpec: sessionSpec, + interactive: interactive, + authKeyPath: