-
Notifications
You must be signed in to change notification settings - Fork 51
/
key.go
106 lines (84 loc) 路 1.99 KB
/
key.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
package memdb
import (
"sync"
"time"
"github.com/bricks-cloud/bricksllm/internal/key"
"github.com/bricks-cloud/bricksllm/internal/logger"
)
type Storage interface {
GetAllKeys() ([]*key.ResponseKey, error)
GetUpdatedKeys(interval time.Duration) ([]*key.ResponseKey, error)
}
type MemDb struct {
external Storage
hashToKeys map[string]*key.ResponseKey
hashToKeysLock sync.RWMutex
done chan bool
interval time.Duration
lg logger.Logger
}
func NewMemDb(ex Storage, lg logger.Logger, interval time.Duration) (*MemDb, error) {
hashToKeys := map[string]*key.ResponseKey{}
keys, err := ex.GetAllKeys()
if err != nil {
return nil, err
}
for _, k := range keys {
hashToKeys[k.Key] = k
}
return &MemDb{
external: ex,
hashToKeys: hashToKeys,
lg: lg,
interval: interval,
done: make(chan bool),
}, nil
}
func (mdb *MemDb) GetKey(hash string) *key.ResponseKey {
k, ok := mdb.hashToKeys[hash]
if ok {
return k
}
return nil
}
func (mdb *MemDb) SetKey(k *key.ResponseKey) {
mdb.hashToKeysLock.RLock()
defer mdb.hashToKeysLock.RUnlock()
mdb.hashToKeys[k.Key] = k
}
func (mdb *MemDb) RemoveKey(k *key.ResponseKey) {
mdb.hashToKeysLock.RLock()
defer mdb.hashToKeysLock.RUnlock()
delete(mdb.hashToKeys, k.Key)
}
func (mdb *MemDb) Listen() {
ticker := time.NewTicker(mdb.interval)
mdb.lg.Info("memdb started listening for key updates")
go func() {
for {
select {
case <-mdb.done:
mdb.lg.Info("memdb stopped")
return
case <-ticker.C:
keys, err := mdb.external.GetUpdatedKeys(mdb.interval)
if err != nil {
mdb.lg.Debugf("memdb failed to update keys: %v", err)
continue
}
if len(keys) == 0 {
continue
}
mdb.lg.Debugf("memdb updated at %s", time.Now().UTC().String())
for _, k := range keys {
mdb.lg.Debugf("memdb updated a key: %s", k.KeyId)
mdb.SetKey(k)
}
}
}
}()
}
func (mdb *MemDb) Stop() {
mdb.lg.Infof("shutting down memdb...")
mdb.done <- true
}