forked from polarsignals/frostdb
-
Notifications
You must be signed in to change notification settings - Fork 0
/
memory.go
93 lines (75 loc) 路 2.04 KB
/
memory.go
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
package query
import (
"runtime/debug"
"sync/atomic"
"github.com/apache/arrow/go/v15/arrow/memory"
"github.com/prometheus/client_golang/prometheus"
"github.com/prometheus/client_golang/prometheus/promauto"
)
const PanicMemoryLimit = "memory limit exceeded"
var _ memory.Allocator = (*LimitAllocator)(nil)
// LimitAllocator is a wrapper around a memory.Allocator that panics if the memory usage exceeds the defined limit.
type LimitAllocator struct {
limit int64
allocated *atomic.Int64
allocator memory.Allocator
reg prometheus.Registerer
}
type AllocatorOption func(*LimitAllocator)
func WithRegistry(reg prometheus.Registerer) AllocatorOption {
return func(a *LimitAllocator) {
a.reg = reg
}
}
func NewLimitAllocator(limit int64, allocator memory.Allocator, options ...AllocatorOption) *LimitAllocator {
l := &LimitAllocator{
limit: limit,
allocated: &atomic.Int64{},
allocator: allocator,
reg: prometheus.NewRegistry(),
}
for _, option := range options {
option(l)
}
promauto.With(l.reg).NewGaugeFunc(prometheus.GaugeOpts{
Name: "memory_allocated_bytes",
Help: "The total number of bytes allocated by the allocator.",
}, func() float64 {
return float64(l.allocated.Load())
})
return l
}
func (a *LimitAllocator) Allocate(size int) []byte {
for {
allocated := a.allocated.Load()
if allocated+int64(size) > a.limit {
panic(PanicMemoryLimit)
}
if a.allocated.CompareAndSwap(allocated, allocated+int64(size)) {
return a.allocator.Allocate(size)
}
}
}
func (a *LimitAllocator) Reallocate(size int, b []byte) []byte {
if len(b) == size {
return b
}
diff := int64(size - len(b))
for {
allocated := a.allocated.Load()
if allocated+diff > a.limit {
debug.PrintStack()
panic(PanicMemoryLimit)
}
if a.allocated.CompareAndSwap(allocated, allocated+diff) {
return a.allocator.Reallocate(size, b)
}
}
}
func (a *LimitAllocator) Free(b []byte) {
a.allocated.Add(-int64(len(b)))
a.allocator.Free(b)
}
func (a *LimitAllocator) Allocated() int {
return int(a.allocated.Load())
}