diff --git a/cli/cli.go b/cli/cli.go index ab192558..04568e54 100644 --- a/cli/cli.go +++ b/cli/cli.go @@ -32,6 +32,7 @@ import ( "github.com/neilotoole/sq/cli/buildinfo" "github.com/neilotoole/sq/cli/flag" "github.com/neilotoole/sq/cli/run" + "github.com/neilotoole/sq/libsq/core/ioz" "github.com/neilotoole/sq/libsq/core/lg" "github.com/neilotoole/sq/libsq/core/lg/lga" "github.com/neilotoole/sq/libsq/core/options" @@ -85,6 +86,7 @@ func ExecuteWith(ctx context.Context, ru *run.Run, args []string) error { _ = ru.LogCloser() } }() + ctx = options.NewContext(ctx, options.Merge(options.FromContext(ctx), ru.Config.Options)) log := lg.FromContext(ctx) log.Info("EXECUTE", "args", strings.Join(args, " ")) @@ -95,6 +97,14 @@ func ExecuteWith(ctx context.Context, ru *run.Run, args []string) error { ctx = run.NewContext(ctx, ru) + if freq := OptDebugTrackMemory.Get(options.FromContext(ctx)); freq > 0 { + // Debug setting to log peak memory usage on exit. + memTracker := ioz.StartPeakMemoryTracker(ctx, freq) + defer func() { + log.Info("Peak memory usage", "mem", memTracker.String(), "bytes", memTracker.Load()) + }() + } + rootCmd := newCommandTree(ru) var err error diff --git a/cli/options.go b/cli/options.go index 261b9371..617a59e9 100644 --- a/cli/options.go +++ b/cli/options.go @@ -201,6 +201,7 @@ func RegisterDefaultOpts(reg *options.Registry) { driver.OptIngestSampleSize, csv.OptDelim, csv.OptEmptyAsNull, + OptDebugTrackMemory, progress.OptDebugSleep, ) } diff --git a/cli/options_test.go b/cli/options_test.go index abe406e5..46b9b7d9 100644 --- a/cli/options_test.go +++ b/cli/options_test.go @@ -19,7 +19,7 @@ func TestRegisterDefaultOpts(t *testing.T) { log.Debug("options.Registry (after)", "reg", reg) keys := reg.Keys() - require.Len(t, keys, 50) + require.Len(t, keys, 51) for _, opt := range reg.Opts() { opt := opt diff --git a/cli/output.go b/cli/output.go index 78675b42..81fc5429 100644 --- a/cli/output.go +++ b/cli/output.go @@ -120,6 +120,16 @@ command, sq falls back to "text". Available formats: `Delay before showing a progress bar.`, ) + OptDebugTrackMemory = options.NewDuration( + "debug.stats.frequency", + "", + 0, + 0, + "Memory usage sampling interval.", + `Memory usage sampling interval. If non-zero, peak memory usage is periodically +sampled, and reported on exit. If zero, memory usage sampling is disabled.`, + ) + OptCompact = options.NewBool( "compact", "", diff --git a/libsq/core/ioz/ioz.go b/libsq/core/ioz/ioz.go index 84140448..332c5bf8 100644 --- a/libsq/core/ioz/ioz.go +++ b/libsq/core/ioz/ioz.go @@ -11,12 +11,15 @@ import ( mrand "math/rand" "os" "path/filepath" + "runtime" "strings" "sync" + "sync/atomic" "time" "github.com/a8m/tree" "github.com/a8m/tree/ostree" + "github.com/c2h5oh/datasize" yaml "github.com/goccy/go-yaml" "github.com/neilotoole/sq/libsq/core/errz" @@ -792,3 +795,42 @@ func countNonDirs(entries []os.DirEntry) (count int) { } return count } + +// PeakMemory is an [atomic.Uint64] that tracks the peak memory usage. +type PeakMemory struct { + atomic.Uint64 +} + +// String returns a human-friendly representation. +func (p *PeakMemory) String() string { + v := p.Load() + return datasize.ByteSize(v).HR() +} + +// StartPeakMemoryTracker starts a goroutine that tracks the peak memory usage, +// per [runtime.MemStats.Sys] and [runtime.ReadMemStats]. The goroutine sleeps +// for sampleFreq between each sample and exits when ctx is done. +func StartPeakMemoryTracker(ctx context.Context, sampleFreq time.Duration) *PeakMemory { + peakMem := &PeakMemory{} + go func() { + ticker := time.NewTicker(sampleFreq) + defer ticker.Stop() + + var peak uint64 + stats := &runtime.MemStats{} + for { + runtime.ReadMemStats(stats) + peak = peakMem.Load() + if stats.Sys > peak { + peakMem.Store(stats.Sys) + } + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + } + }() + + return peakMem +}