internal/worker/mem_monitor.go

// Copyright 2022 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package worker

import "time"

// memMonitor is used to probe memory consumption
// in a separate Go routine.
type memMonitor struct {
	stp chan struct{} // for stopping the monitor
	res chan uint64   // for communicating results
}

// newMemMonitor creates and starts new monitor that
// samples memory consumption.
//
// If threshold > 0, then when memory consumption reaches threshold,
// the monitor will call onThreshold and stop.
func newMemMonitor(threshold uint64, onThreshold func()) *memMonitor {
	m := &memMonitor{make(chan struct{}), make(chan uint64)}
	var max uint64
	go func() {
		for {
			select {
			case <-m.stp:
				m.res <- max
				return
			default:
				h := currHeapUsage()
				if h > max {
					max = h
				}
				if threshold > 0 && h > threshold {
					onThreshold()
					m.res <- max
					return
				}
				// We sample memory every 50ms.
				time.Sleep(50 * time.Millisecond)
			}
		}
	}()
	return m
}

// stop stops the monitor and returns the peak
// memory consumption in bytes.
func (m *memMonitor) stop() uint64 {
	if m == nil {
		return 0
	}
	close(m.stp)
	return <-m.res
}