Просмотр исходного кода

feat: Add concurrency control to group ratio management with mutexes

CaIon 11 месяцев назад
Родитель
Сommit
fb9f595044
1 измененных файлов с 17 добавлено и 0 удалено
  1. 17 0
      setting/group_ratio.go

+ 17 - 0
setting/group_ratio.go

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