-
Notifications
You must be signed in to change notification settings - Fork 1.9k
/
usage_tracker.go
69 lines (57 loc) · 1.67 KB
/
usage_tracker.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
package csimanager
import (
"sync"
)
// volumeUsageTracker tracks the allocations that depend on a given volume
type volumeUsageTracker struct {
// state is a map of volumeUsageKey to a slice of allocation ids
state map[volumeUsageKey][]string
stateMu sync.Mutex
}
func newVolumeUsageTracker() *volumeUsageTracker {
return &volumeUsageTracker{
state: make(map[volumeUsageKey][]string),
}
}
type volumeUsageKey struct {
id string
usageOpts UsageOptions
}
func (v *volumeUsageTracker) allocsForKey(key volumeUsageKey) []string {
return v.state[key]
}
func (v *volumeUsageTracker) appendAlloc(key volumeUsageKey, allocID string) {
allocs := v.allocsForKey(key)
allocs = append(allocs, allocID)
v.state[key] = allocs
}
func (v *volumeUsageTracker) removeAlloc(key volumeUsageKey, needle string) {
allocs := v.allocsForKey(key)
var newAllocs []string
for _, allocID := range allocs {
if allocID != needle {
newAllocs = append(newAllocs, allocID)
}
}
if len(newAllocs) == 0 {
delete(v.state, key)
} else {
v.state[key] = newAllocs
}
}
func (v *volumeUsageTracker) Claim(allocID, volID string, usage *UsageOptions) {
v.stateMu.Lock()
defer v.stateMu.Unlock()
key := volumeUsageKey{id: volID, usageOpts: *usage}
v.appendAlloc(key, allocID)
}
// Free removes the allocation from the state list for the given alloc. If the
// alloc is the last allocation for the volume then it returns true.
func (v *volumeUsageTracker) Free(allocID, volID string, usage *UsageOptions) bool {
v.stateMu.Lock()
defer v.stateMu.Unlock()
key := volumeUsageKey{id: volID, usageOpts: *usage}
v.removeAlloc(key, allocID)
allocs := v.allocsForKey(key)
return len(allocs) == 0
}