-
Notifications
You must be signed in to change notification settings - Fork 1
/
lockable_map.go
125 lines (112 loc) · 2.9 KB
/
lockable_map.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
115
116
117
118
119
120
121
122
123
124
125
package lockable
import (
"sync"
)
// Map implements a map that can acquire key-specific locks.
//
// The zero value is not ready for use. Refer to [NewMap] to create a ready-to-use instance.
//
// Map has an interface that's deliberately similar to [sync.Map], but uses a combination of RWMutex/map to simulate it.
//
// [MutexMap] can be used instead if you want to use sync.Map internally.
// Refer to [sync.Map]'s document for use cases.
//
// Ex. usage:
//
// func main() {
// lockableMap := lockable.NewMap[string, int]()
//
//
// // This will only lock access the "potato" key
// // Keys do not need to exist prior to acquiring the key lock
// lockableMap.LockKey("potato")
// defer lockableMap.UnlockKey("potato")
//
// // Do async stuff....
//
// lockableMap.Store("potato", 10)
// }
//
// Refer to [Lockable] for more detailed exemples of locking.
type Map[T comparable, V any] struct {
internalMap map[T]V
mu *sync.RWMutex
Lockable[T]
}
// MutexMap acts the same as [Map] except it uses a sync.Map.
//
// The zero value is not ready for use. Refer to [NewMutexMap] to create a ready-to-use instance.
//
// Ex. usage:
// func main() {
// lockableMap := lockable.NewMutexMap[string]()
//
//
// // This will only lock access the "potato" key
// // Keys do not need to exist prior to acquiring the key lock
// lockableMap.LockKey("potato")
// defer lockableMap.UnlockKey("potato")
//
// // Do async stuff....
//
// lockableMap.Store("potato", 10)
// }
//
// Refer to [Lockable] for more detailed exemples of locking
type MutexMap[T comparable] struct {
sync.Map
Lockable[T]
}
// NewMap creates a ready-to-use Map instance.
//
// Refer to [Map] for usage.
func NewMap[T comparable, V any]() Map[T, V] {
return Map[T, V]{
internalMap: map[T]V{},
mu: &sync.RWMutex{},
Lockable: Lockable[T]{
locks: map[T]*versionedMutex{},
locksMu: &sync.Mutex{},
},
}
}
// NewMutexMap creates a ready-to-use MutexMap instance.
//
// Refer to [MutexMap] for usage.
func NewMutexMap[T comparable]() MutexMap[T] {
return MutexMap[T]{
Lockable: Lockable[T]{
locks: map[T]*versionedMutex{},
locksMu: &sync.Mutex{},
},
}
}
// Load effectively serves the same purpose as [sync.Map.Load]
func (m Map[T, V]) Load(key T) (V, bool) {
m.mu.RLock()
defer m.mu.RUnlock()
v, ok := m.internalMap[key]
return v, ok
}
// Store effectively serves the same purpose as [sync.Map.Store]
func (m Map[T, V]) Store(key T, value V) {
m.mu.Lock()
defer m.mu.Unlock()
m.internalMap[key] = value
}
// Delete effectively serves the same purpose as [sync.Map.Delete]
func (m Map[T, V]) Delete(key T) {
m.mu.Lock()
defer m.mu.Unlock()
delete(m.internalMap, key)
}
// Range effectively serves the same purpose as [sync.Map.Range]
func (m Map[T, V]) Range(fn func(key T, value V) bool) {
m.mu.RLock()
defer m.mu.RUnlock()
for k, v := range m.internalMap {
if keepGoing := fn(k, v); !keepGoing {
break
}
}
}