-
Notifications
You must be signed in to change notification settings - Fork 4
/
collection.go
112 lines (99 loc) · 2.33 KB
/
collection.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
package loaders
import (
"context"
"github.com/bcc-code/bcc-media-platform/backend/utils"
"time"
)
// Collection is a collection of loaders or other advanced structures
type Collection[K comparable, V any] struct {
values *utils.SyncMap[K, *entry[K, V]]
expiry time.Duration
janitor *janitor
}
// NewCollection of loaders or other structures
func NewCollection[K comparable, V any](expiration time.Duration) *Collection[K, V] {
col := &Collection[K, V]{
values: &utils.SyncMap[K, *entry[K, V]]{},
expiry: expiration,
janitor: newJanitor(context.Background(), time.Minute),
}
col.janitor.run(col.DeleteExpired)
return col
}
type entry[K comparable, V any] struct {
Key K
Value V
ExpiresAt time.Time
// A function that should run when this entry is deleted.
// For example canceling a context.
OnDelete func()
}
// Keys returns all keys
func (c Collection[K, V]) Keys() []K {
var keys []K
c.values.Range(func(key K, _ *entry[K, V]) bool {
keys = append(keys, key)
return true
})
return keys
}
// Set a key to specified value
func (c Collection[K, V]) Set(key K, value V, opts ...any) {
e := &entry[K, V]{
Key: key,
Value: value,
ExpiresAt: time.Now().Add(c.expiry),
}
for _, opt := range opts {
switch v := opt.(type) {
case onDelete:
e.OnDelete = v
}
}
c.values.Store(key, e)
}
// Get a value by the specified key
func (c Collection[K, V]) Get(key K) (value V, ok bool) {
e, ok := c.values.Load(key)
if !ok {
return value, false
}
if e.ExpiresAt.Before(time.Now()) {
c.values.Delete(key)
if e.OnDelete != nil {
e.OnDelete()
}
return value, false
}
e.ExpiresAt = time.Now().Add(c.expiry)
return e.Value, true
}
// Delete the specified key (&entry)
func (c Collection[K, V]) Delete(key K) {
e, ok := c.values.LoadAndDelete(key)
if !ok {
return
}
c.values.Delete(key)
if e.OnDelete != nil {
e.OnDelete()
}
}
// DeleteExpired entries
func (c Collection[K, V]) DeleteExpired() {
var deleteKeys []K
c.values.Range(func(key K, value *entry[K, V]) bool {
if value.ExpiresAt.Before(time.Now()) {
deleteKeys = append(deleteKeys, value.Key)
}
return true
})
for _, k := range deleteKeys {
c.Delete(k)
}
}
type onDelete func()
// WithOnDelete tells the collection to run this function when entry is deleted
func WithOnDelete(cb func()) any {
return onDelete(cb)
}