feat: Add concurrency control to group ratio management with mutexes
This commit is contained in:
@@ -4,6 +4,7 @@ import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"one-api/common"
|
||||
"sync"
|
||||
)
|
||||
|
||||
var groupRatio = map[string]float64{
|
||||
@@ -11,8 +12,12 @@ var groupRatio = map[string]float64{
|
||||
"vip": 1,
|
||||
"svip": 1,
|
||||
}
|
||||
var groupRatioMutex sync.RWMutex
|
||||
|
||||
func GetGroupRatioCopy() map[string]float64 {
|
||||
groupRatioMutex.RLock()
|
||||
defer groupRatioMutex.RUnlock()
|
||||
|
||||
groupRatioCopy := make(map[string]float64)
|
||||
for k, v := range groupRatio {
|
||||
groupRatioCopy[k] = v
|
||||
@@ -21,11 +26,17 @@ func GetGroupRatioCopy() map[string]float64 {
|
||||
}
|
||||
|
||||
func ContainsGroupRatio(name string) bool {
|
||||
groupRatioMutex.RLock()
|
||||
defer groupRatioMutex.RUnlock()
|
||||
|
||||
_, ok := groupRatio[name]
|
||||
return ok
|
||||
}
|
||||
|
||||
func GroupRatio2JSONString() string {
|
||||
groupRatioMutex.RLock()
|
||||
defer groupRatioMutex.RUnlock()
|
||||
|
||||
jsonBytes, err := json.Marshal(groupRatio)
|
||||
if err != nil {
|
||||
common.SysError("error marshalling model ratio: " + err.Error())
|
||||
@@ -34,11 +45,17 @@ func GroupRatio2JSONString() string {
|
||||
}
|
||||
|
||||
func UpdateGroupRatioByJSONString(jsonStr string) error {
|
||||
groupRatioMutex.Lock()
|
||||
defer groupRatioMutex.Unlock()
|
||||
|
||||
groupRatio = make(map[string]float64)
|
||||
return json.Unmarshal([]byte(jsonStr), &groupRatio)
|
||||
}
|
||||
|
||||
func GetGroupRatio(name string) float64 {
|
||||
groupRatioMutex.RLock()
|
||||
defer groupRatioMutex.RUnlock()
|
||||
|
||||
ratio, ok := groupRatio[name]
|
||||
if !ok {
|
||||
common.SysError("group ratio not found: " + name)
|
||||
|
||||
Reference in New Issue
Block a user