Spaces:
Running
Running
File size: 2,887 Bytes
b110593 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
// _ _
// __ _____ __ ___ ___ __ _| |_ ___
// \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
// \ V V / __/ (_| |\ V /| | (_| | || __/
// \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
//
// Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
//
// CONTACT: [email protected]
//
package memwatch
import (
"fmt"
"runtime/metrics"
"sync"
)
const (
B = 1
KiB = 1 << (10 * iota) // 2^10
MiB = 1 << (10 * iota) // 2^20
GiB = 1 << (10 * iota) // 2^30
TiB = 1 << (10 * iota) // 2^40
)
var ErrNotEnoughMemory = fmt.Errorf("not enough memory")
// Monitor allows making statements about the memory ratio used by the application
type Monitor struct {
metricsReader metricsReader
limitSetter limitSetter
maxRatio float64
// state
mu sync.Mutex
limit int64
used int64
}
// Refresh retrieves the current memory stats from the runtime and stores them
// in the local cache
func (m *Monitor) Refresh() {
m.obtainCurrentUsage()
m.updateLimit()
}
// we have no intentions of ever modifying the limit, but SetMemoryLimit with a
// negative value is the only way to read the limit from the runtime
type limitSetter func(size int64) int64
// NewMonitor creates a [Monitor] with the given metrics reader and target
// ratio
//
// Typically this would be called with LiveHeapReader and
// debug.SetMemoryLimit
func NewMonitor(metricsReader metricsReader, limitSetter limitSetter,
maxRatio float64,
) *Monitor {
return &Monitor{
metricsReader: metricsReader,
limitSetter: limitSetter,
maxRatio: maxRatio,
}
}
func (m *Monitor) CheckAlloc(sizeInBytes int64) error {
m.mu.Lock()
defer m.mu.Unlock()
if float64(m.used+sizeInBytes)/float64(m.limit) > m.maxRatio {
return ErrNotEnoughMemory
}
return nil
}
func (m *Monitor) Ratio() float64 {
m.mu.Lock()
defer m.mu.Unlock()
return float64(m.used) / float64(m.limit)
}
// obtainCurrentUsage obtains the most recent live heap from runtime/metrics
func (m *Monitor) obtainCurrentUsage() {
m.setUsed(m.metricsReader())
}
func LiveHeapReader() int64 {
const liveHeapBytesMetric = "/gc/heap/live:bytes"
sample := make([]metrics.Sample, 1)
sample[0].Name = liveHeapBytesMetric
metrics.Read(sample)
if sample[0].Value.Kind() == metrics.KindBad {
panic(fmt.Sprintf("metric %q no longer supported", liveHeapBytesMetric))
}
return int64(sample[0].Value.Uint64())
}
// setUsed is a thread-safe way way to set the current usage
func (m *Monitor) setUsed(used int64) {
m.mu.Lock()
defer m.mu.Unlock()
m.used = used
}
func (m *Monitor) updateLimit() {
m.mu.Lock()
defer m.mu.Unlock()
// setting a negative limit is the only way to obtain the current limit
m.limit = m.limitSetter(-1)
}
type metricsReader func() int64
|