-
Notifications
You must be signed in to change notification settings - Fork 54
/
custom_provider.go
140 lines (116 loc) 路 3.35 KB
/
custom_provider.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
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
package memdb
import (
"sync"
"time"
"github.com/bricks-cloud/bricksllm/internal/provider/custom"
"github.com/bricks-cloud/bricksllm/internal/stats"
"go.uber.org/zap"
)
type CustomProvidersStorage interface {
GetCustomProviders() ([]*custom.Provider, error)
GetUpdatedCustomProviders(updatedAt int64) ([]*custom.Provider, error)
}
type CustomProvidersMemDb struct {
external CustomProvidersStorage
lastUpdated int64
nameToProviders map[string]*custom.Provider
lock sync.RWMutex
done chan bool
interval time.Duration
log *zap.Logger
}
func NewCustomProvidersMemDb(ex CustomProvidersStorage, log *zap.Logger, interval time.Duration) (*CustomProvidersMemDb, error) {
nameToProviders := map[string]*custom.Provider{}
providers, err := ex.GetCustomProviders()
if err != nil {
return nil, err
}
numberOfProviders := 0
var latetest int64 = -1
for _, p := range providers {
nameToProviders[p.Provider] = p
numberOfProviders++
if p.UpdatedAt > latetest {
latetest = p.UpdatedAt
}
}
if numberOfProviders != 0 {
log.Sugar().Infof("custom provider settings memdb updated at %d with %d providers", latetest, numberOfProviders)
}
return &CustomProvidersMemDb{
external: ex,
nameToProviders: nameToProviders,
log: log,
lastUpdated: latetest,
interval: interval,
done: make(chan bool),
}, nil
}
func (mdb *CustomProvidersMemDb) GetProvider(name string) *custom.Provider {
provider, ok := mdb.nameToProviders[name]
if ok {
return provider
}
return nil
}
func (mdb *CustomProvidersMemDb) GetRouteConfig(name string, path string) *custom.RouteConfig {
provider, ok := mdb.nameToProviders[name]
if ok {
for _, rc := range provider.RouteConfigs {
if rc.Path == path {
return rc
}
}
}
return nil
}
func (mdb *CustomProvidersMemDb) SetProvider(provider *custom.Provider) {
mdb.lock.RLock()
defer mdb.lock.RUnlock()
mdb.nameToProviders[provider.Provider] = provider
}
func (mdb *CustomProvidersMemDb) Listen() {
ticker := time.NewTicker(mdb.interval)
mdb.log.Info("custom providers memdb started listening for provider updates")
go func() {
lastUpdated := mdb.lastUpdated
for {
select {
case <-mdb.done:
mdb.log.Info("memdb stopped")
return
case <-ticker.C:
providers, err := mdb.external.GetUpdatedCustomProviders(lastUpdated)
if err != nil {
stats.Incr("bricksllm.memdb.custom_proivders_memdb.listen.get_updated_custom_providers_error", nil, 1)
mdb.log.Sugar().Debugf("memdb failed to update keys: %v", err)
continue
}
if len(providers) == 0 {
continue
}
any := false
numberOfUpdated := 0
for _, provider := range providers {
if provider.UpdatedAt > lastUpdated {
lastUpdated = provider.UpdatedAt
}
existing := mdb.GetProvider(provider.Provider)
if existing == nil || provider.UpdatedAt > existing.UpdatedAt {
mdb.log.Sugar().Infof("custom providers memdb updated a provider: %s", provider.Id)
numberOfUpdated += 1
any = true
mdb.SetProvider(provider)
}
}
if any {
mdb.log.Sugar().Infof("custom providers memdb updated at %d with %d providers", lastUpdated, numberOfUpdated)
}
}
}
}()
}
func (mdb *CustomProvidersMemDb) Stop() {
mdb.log.Info("shutting down custom providers memdb...")
mdb.done <- true
}