From 4513a361be5a36cd87243082e30e9f4b18e96594 Mon Sep 17 00:00:00 2001
From: Neil O'Toole <neilotoole@apache.org>
Date: Mon, 29 Jan 2024 09:02:42 -0700
Subject: [PATCH] Refine mem stats (#380)

---
 cli/cli.go            | 10 +++++++--
 libsq/core/ioz/ioz.go | 49 ++++++++++++++++++++++++-------------------
 2 files changed, 35 insertions(+), 24 deletions(-)

diff --git a/cli/cli.go b/cli/cli.go
index 04568e54..66bba338 100644
--- a/cli/cli.go
+++ b/cli/cli.go
@@ -26,7 +26,9 @@ import (
 	"os"
 	"strings"
 	"sync"
+	"time"
 
+	"github.com/c2h5oh/datasize"
 	"github.com/spf13/cobra"
 
 	"github.com/neilotoole/sq/cli/buildinfo"
@@ -99,9 +101,13 @@ func ExecuteWith(ctx context.Context, ru *run.Run, args []string) error {
 
 	if freq := OptDebugTrackMemory.Get(options.FromContext(ctx)); freq > 0 {
 		// Debug setting to log peak memory usage on exit.
-		memTracker := ioz.StartPeakMemoryTracker(ctx, freq)
+		peakSys, peakAllocs, gcPause := ioz.StartMemStatsTracker(ctx, freq)
 		defer func() {
-			log.Info("Peak memory usage", "mem", memTracker.String(), "bytes", memTracker.Load())
+			log.Info("Peak memory stats",
+				"sys", datasize.ByteSize(peakSys.Load()).HR(),
+				"heap", datasize.ByteSize(peakAllocs.Load()).HR(),
+				"gc_pause", time.Duration(gcPause.Load()).String(),
+			)
 		}()
 	}
 
diff --git a/libsq/core/ioz/ioz.go b/libsq/core/ioz/ioz.go
index 332c5bf8..f4d008b8 100644
--- a/libsq/core/ioz/ioz.go
+++ b/libsq/core/ioz/ioz.go
@@ -19,7 +19,6 @@ import (
 
 	"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"
@@ -796,41 +795,47 @@ 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
-}
+// StartMemStatsTracker starts a goroutine that tracks memory stats, returning
+// the peak values of [runtime.MemStats.Sys], [runtime.MemStats.TotalAlloc] and
+// [runtime.MemStats.PauseTotalNs]. The goroutine sleeps for sampleFreq between
+// each sample and exits when ctx is done.
+//
+//nolint:revive // datarace
+func StartMemStatsTracker(ctx context.Context, sampleFreq time.Duration) (sys *atomic.Uint64,
+	allocs, gcPauseNs *atomic.Uint64,
+) {
+	sys = &atomic.Uint64{}
+	allocs = &atomic.Uint64{}
+	gcPauseNs = &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{}
+		var done bool
 		for {
 			runtime.ReadMemStats(stats)
-			peak = peakMem.Load()
-			if stats.Sys > peak {
-				peakMem.Store(stats.Sys)
+
+			if stats.Sys > sys.Load() {
+				sys.Store(stats.Sys)
+			}
+
+			allocs.Store(stats.TotalAlloc)
+			gcPauseNs.Store(stats.PauseTotalNs)
+
+			if done {
+				return
 			}
 			select {
 			case <-ctx.Done():
-				return
+				// We perform one more loop to ensure we capture
+				// the final stats before return.
+				done = true
 			case <-ticker.C:
 			}
 		}
 	}()
 
-	return peakMem
+	return sys, allocs, gcPauseNs
 }