-
-
Notifications
You must be signed in to change notification settings - Fork 3.9k
/
safemap.go
114 lines (102 loc) · 2.34 KB
/
safemap.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
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
package collection
import "sync"
const (
copyThreshold = 1000
maxDeletion = 10000
)
// SafeMap provides a map alternative to avoid memory leak.
// This implementation is not needed until issue below fixed.
// https://github.com/golang/go/issues/20135
type SafeMap struct {
lock sync.RWMutex
deletionOld int
deletionNew int
dirtyOld map[any]any
dirtyNew map[any]any
}
// NewSafeMap returns a SafeMap.
func NewSafeMap() *SafeMap {
return &SafeMap{
dirtyOld: make(map[any]any),
dirtyNew: make(map[any]any),
}
}
// Del deletes the value with the given key from m.
func (m *SafeMap) Del(key any) {
m.lock.Lock()
if _, ok := m.dirtyOld[key]; ok {
delete(m.dirtyOld, key)
m.deletionOld++
} else if _, ok := m.dirtyNew[key]; ok {
delete(m.dirtyNew, key)
m.deletionNew++
}
if m.deletionOld >= maxDeletion && len(m.dirtyOld) < copyThreshold {
for k, v := range m.dirtyOld {
m.dirtyNew[k] = v
}
m.dirtyOld = m.dirtyNew
m.deletionOld = m.deletionNew
m.dirtyNew = make(map[any]any)
m.deletionNew = 0
}
if m.deletionNew >= maxDeletion && len(m.dirtyNew) < copyThreshold {
for k, v := range m.dirtyNew {
m.dirtyOld[k] = v
}
m.dirtyNew = make(map[any]any)
m.deletionNew = 0
}
m.lock.Unlock()
}
// Get gets the value with the given key from m.
func (m *SafeMap) Get(key any) (any, bool) {
m.lock.RLock()
defer m.lock.RUnlock()
if val, ok := m.dirtyOld[key]; ok {
return val, true
}
val, ok := m.dirtyNew[key]
return val, ok
}
// Range calls f sequentially for each key and value present in the map.
// If f returns false, range stops the iteration.
func (m *SafeMap) Range(f func(key, val any) bool) {
m.lock.RLock()
defer m.lock.RUnlock()
for k, v := range m.dirtyOld {
if !f(k, v) {
return
}
}
for k, v := range m.dirtyNew {
if !f(k, v) {
return
}
}
}
// Set sets the value into m with the given key.
func (m *SafeMap) Set(key, value any) {
m.lock.Lock()
if m.deletionOld <= maxDeletion {
if _, ok := m.dirtyNew[key]; ok {
delete(m.dirtyNew, key)
m.deletionNew++
}
m.dirtyOld[key] = value
} else {
if _, ok := m.dirtyOld[key]; ok {
delete(m.dirtyOld, key)
m.deletionOld++
}
m.dirtyNew[key] = value
}
m.lock.Unlock()
}
// Size returns the size of m.
func (m *SafeMap) Size() int {
m.lock.RLock()
size := len(m.dirtyOld) + len(m.dirtyNew)
m.lock.RUnlock()
return size
}