Track memory usage (#379)

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
 

cli/options.go

@@ -201,6 +201,7 @@ func RegisterDefaultOpts(reg *options.Registry) {
 		driver.OptIngestSampleSize,
 		csv.OptDelim,
 		csv.OptEmptyAsNull,
+		OptDebugTrackMemory,
 		progress.OptDebugSleep,
 	)
 }

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

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",
 		"",

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
+}