xinyin025/one-api
1
0
Fork 0
mirror of https://github.com/songquanpeng/one-api.git synced 2025-09-17 01:06:37 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix: add read/write locks for ModelRatio and GroupRatio to prevent concurrent map read/write issues (#2067)
Some checks are pending
CI / Unit tests (push) Waiting to run
Details
CI / commit_lint (push) Waiting to run
Details

Browse Source
This commit is contained in:
longkeyy 2025-02-09 11:02:45 +08:00 committed by GitHub
parent 07808122a6
commit 3e3b8230ac
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 15 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
relay/billing/ratio/group.go
View File

@ -3,8 +3,10 @@ package ratio
import ( import (
"encoding/json" "encoding/json"
"github.com/songquanpeng/one-api/common/logger" "github.com/songquanpeng/one-api/common/logger"
"sync"
) )
var groupRatioLock sync.RWMutex
var GroupRatio = map[string]float64{ var GroupRatio = map[string]float64{
"default": 1, "default": 1,
"vip": 1, "vip": 1,
@ -20,11 +22,15 @@ func GroupRatio2JSONString() string {
} }
func UpdateGroupRatioByJSONString(jsonStr string) error { func UpdateGroupRatioByJSONString(jsonStr string) error {
groupRatioLock.Lock()
defer groupRatioLock.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 {
groupRatioLock.RLock()
defer groupRatioLock.RUnlock()
ratio, ok := GroupRatio[name] ratio, ok := GroupRatio[name]
if !ok { if !ok {
logger.SysError("group ratio not found: " + name) logger.SysError("group ratio not found: " + name)

11
relay/billing/ratio/model.go
View File

@ -4,6 +4,7 @@ import (
"encoding/json" "encoding/json"
"fmt" "fmt"
"strings" "strings"
"sync"
"github.com/songquanpeng/one-api/common/logger" "github.com/songquanpeng/one-api/common/logger"
) )
@ -15,6 +16,8 @@ const (
RMB = USD / USD2RMB RMB = USD / USD2RMB
) )
var modelRatioLock sync.RWMutex
// ModelRatio // ModelRatio
// https://platform.openai.com/docs/models/model-endpoint-compatibility // https://platform.openai.com/docs/models/model-endpoint-compatibility
// https://cloud.baidu.com/doc/WENXINWORKSHOP/s/Blfmc9dlf // https://cloud.baidu.com/doc/WENXINWORKSHOP/s/Blfmc9dlf
@ -88,11 +91,11 @@ var ModelRatio = map[string]float64{
"claude-2.1": 8.0 / 1000 * USD, "claude-2.1": 8.0 / 1000 * USD,
"claude-3-haiku-20240307": 0.25 / 1000 * USD, "claude-3-haiku-20240307": 0.25 / 1000 * USD,
"claude-3-5-haiku-20241022": 1.0 / 1000 * USD, "claude-3-5-haiku-20241022": 1.0 / 1000 * USD,
"claude-3-5-haiku-latest": 1.0 / 1000 * USD, "claude-3-5-haiku-latest": 1.0 / 1000 * USD,
"claude-3-sonnet-20240229": 3.0 / 1000 * USD, "claude-3-sonnet-20240229": 3.0 / 1000 * USD,
"claude-3-5-sonnet-20240620": 3.0 / 1000 * USD, "claude-3-5-sonnet-20240620": 3.0 / 1000 * USD,
"claude-3-5-sonnet-20241022": 3.0 / 1000 * USD, "claude-3-5-sonnet-20241022": 3.0 / 1000 * USD,
"claude-3-5-sonnet-latest" : 3.0 / 1000 * USD, "claude-3-5-sonnet-latest": 3.0 / 1000 * USD,
"claude-3-opus-20240229": 15.0 / 1000 * USD, "claude-3-opus-20240229": 15.0 / 1000 * USD,
// https://cloud.baidu.com/doc/WENXINWORKSHOP/s/hlrk4akp7 // https://cloud.baidu.com/doc/WENXINWORKSHOP/s/hlrk4akp7
"ERNIE-4.0-8K": 0.120 * RMB, "ERNIE-4.0-8K": 0.120 * RMB,
@ -417,11 +420,15 @@ func ModelRatio2JSONString() string {
} }
func UpdateModelRatioByJSONString(jsonStr string) error { func UpdateModelRatioByJSONString(jsonStr string) error {
modelRatioLock.Lock()
defer modelRatioLock.Unlock()
ModelRatio = make(map[string]float64) ModelRatio = make(map[string]float64)
return json.Unmarshal([]byte(jsonStr), &ModelRatio) return json.Unmarshal([]byte(jsonStr), &ModelRatio)
} }
func GetModelRatio(name string, channelType int) float64 { func GetModelRatio(name string, channelType int) float64 {
modelRatioLock.RLock()
defer modelRatioLock.RUnlock()
if strings.HasPrefix(name, "qwen-") && strings.HasSuffix(name, "-internet") { if strings.HasPrefix(name, "qwen-") && strings.HasSuffix(name, "-internet") {
name = strings.TrimSuffix(name, "-internet") name = strings.TrimSuffix(name, "-internet")
} }