61 lines
1.4 KiB
Go
61 lines
1.4 KiB
Go
package internal
|
|
|
|
import (
|
|
"context"
|
|
"log"
|
|
"runtime"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
const MonitorInterval = 30 * time.Second
|
|
|
|
type MemoryMonitor struct {
|
|
peakAlloc atomic.Uint64
|
|
}
|
|
|
|
func NewMemoryMonitor() *MemoryMonitor {
|
|
return &MemoryMonitor{}
|
|
}
|
|
|
|
// Snapshot reads current heap allocation, updates the peak, and returns both.
|
|
func (m *MemoryMonitor) Snapshot() (currentAlloc, peakAlloc uint64) {
|
|
var ms runtime.MemStats
|
|
runtime.ReadMemStats(&ms)
|
|
currentAlloc = ms.HeapAlloc
|
|
for {
|
|
peak := m.peakAlloc.Load()
|
|
if currentAlloc <= peak {
|
|
return currentAlloc, peak
|
|
}
|
|
if m.peakAlloc.CompareAndSwap(peak, currentAlloc) {
|
|
return currentAlloc, currentAlloc
|
|
}
|
|
}
|
|
}
|
|
|
|
// Run starts a periodic monitor loop that calls logFn with each snapshot.
|
|
// It blocks until ctx is cancelled.
|
|
func (m *MemoryMonitor) Run(ctx context.Context, logFn func(currentAlloc, peakAlloc uint64)) {
|
|
ticker := time.NewTicker(MonitorInterval)
|
|
defer ticker.Stop()
|
|
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case <-ticker.C:
|
|
current, peak := m.Snapshot()
|
|
logFn(current, peak)
|
|
}
|
|
}
|
|
}
|
|
|
|
// DefaultLogFn returns a log function with the given prefix.
|
|
func DefaultLogFn(prefix string) func(currentAlloc, peakAlloc uint64) {
|
|
return func(currentAlloc, peakAlloc uint64) {
|
|
log.Printf("[%s] heap alloc: %s | peak heap alloc: %s",
|
|
prefix, FormatBytes(currentAlloc), FormatBytes(peakAlloc))
|
|
}
|
|
}
|