-
Notifications
You must be signed in to change notification settings - Fork 0
/
key_locker.go
74 lines (59 loc) · 1.21 KB
/
key_locker.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
package loader
import (
"sync"
)
type KeyLocker[Key comparable] interface {
Lock(key Key) (unlock func())
}
func newInMemoryKeyLocker[Key comparable]() KeyLocker[Key] {
return &InMemoryKeyLocker[Key]{
locks: map[Key]*inMemoryKeyLockerItem{},
}
}
type InMemoryKeyLocker[Key comparable] struct {
root sync.Mutex
locks map[Key]*inMemoryKeyLockerItem
}
type inMemoryKeyLockerItem struct {
ref int32
m sync.Mutex
}
// Lock implements KeyLocker
func (l *InMemoryKeyLocker[Key]) Lock(key Key) func() {
item := l.getItem(key)
item.m.Lock()
unlocked := false
return func() {
if unlocked {
return
}
item.m.Unlock()
l.releaseItem(key)
unlocked = true
}
}
func (l *InMemoryKeyLocker[Key]) getItem(key Key) *inMemoryKeyLockerItem {
l.root.Lock()
defer l.root.Unlock()
item, ok := l.locks[key]
if !ok {
item = &inMemoryKeyLockerItem{}
l.locks[key] = item
}
item.ref += 1
return item
}
func (l *InMemoryKeyLocker[Key]) releaseItem(key Key) {
l.root.Lock()
defer l.root.Unlock()
item, ok := l.locks[key]
if !ok {
return
}
item.ref -= 1
if item.ref <= 0 {
delete(l.locks, key)
}
}
var _ KeyLocker[int] = &InMemoryKeyLocker[int]{}
var _ KeyLocker[string] = &InMemoryKeyLocker[string]{}