Ecosystem/ai-gateway
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: select channel without database (#158)

Browse Source
This commit is contained in:
JustSong 2023-06-21 17:04:18 +08:00
parent 1932c56ea8
commit ba54c71948
1 changed files with 32 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

44
model/cache.go
View File

@ -2,8 +2,11 @@ package model
import ( import (
"encoding/json" "encoding/json"
"errors"
"fmt" "fmt"
"math/rand"
"one-api/common" "one-api/common"
"strings"
"sync" "sync"
"time" "time"
) )
@ -57,18 +60,15 @@ func CacheGetUserGroup(id int) (group string, err error) {
return group, err return group, err
} }
var channelId2channel map[int]*Channel
var channelSyncLock sync.RWMutex
var group2model2channels map[string]map[string][]*Channel var group2model2channels map[string]map[string][]*Channel
var channelSyncLock sync.RWMutex
func InitChannelCache() { func InitChannelCache() {
channelSyncLock.Lock() newChannelId2channel := make(map[int]*Channel)
defer channelSyncLock.Unlock()
channelId2channel = make(map[int]*Channel)
var channels []*Channel var channels []*Channel
DB.Find(&channels) DB.Find(&channels)
for _, channel := range channels { for _, channel := range channels {
channelId2channel[channel.Id] = channel newChannelId2channel[channel.Id] = channel
} }
var abilities []*Ability var abilities []*Ability
DB.Find(&abilities) DB.Find(&abilities)
@ -76,11 +76,26 @@ func InitChannelCache() {
for _, ability := range abilities { for _, ability := range abilities {
groups[ability.Group] = true groups[ability.Group] = true
} }
group2model2channels = make(map[string]map[string][]*Channel) newGroup2model2channels := make(map[string]map[string][]*Channel)
for group := range groups { for group := range groups {
group2model2channels[group] = make(map[string][]*Channel) newGroup2model2channels[group] = make(map[string][]*Channel)
// TODO: implement this
} }
for _, channel := range channels {
groups := strings.Split(channel.Group, ",")
for _, group := range groups {
models := strings.Split(channel.Models, ",")
for _, model := range models {
if _, ok := newGroup2model2channels[group][model]; !ok {
newGroup2model2channels[group][model] = make([]*Channel, 0)
}
newGroup2model2channels[group][model] = append(newGroup2model2channels[group][model], channel)
}
}
}
channelSyncLock.Lock()
group2model2channels = newGroup2model2channels
channelSyncLock.Unlock()
common.SysLog("Channels synced from database")
} }
func SyncChannelCache(frequency int) { func SyncChannelCache(frequency int) {
@ -95,7 +110,12 @@ func CacheGetRandomSatisfiedChannel(group string, model string) (*Channel, error
if !common.RedisEnabled { if !common.RedisEnabled {
return GetRandomSatisfiedChannel(group, model) return GetRandomSatisfiedChannel(group, model)
} }
return GetRandomSatisfiedChannel(group, model) channelSyncLock.RLock()
// TODO: implement this defer channelSyncLock.RUnlock()
return nil, nil channels := group2model2channels[group][model]
if len(channels) == 0 {
return nil, errors.New("channel not found")
}
idx := rand.Intn(len(channels))
return channels[idx], nil
} }