From 3e3b8230ac138cc3279ccc65fd32b8121a66880f Mon Sep 17 00:00:00 2001
From: longkeyy <longkeyy@gmail.com>
Date: Sun, 9 Feb 2025 11:02:45 +0800
Subject: [PATCH] fix: add read/write locks for ModelRatio and GroupRatio to
 prevent concurrent map read/write issues (#2067)

---
 relay/billing/ratio/group.go |  6 ++++++
 relay/billing/ratio/model.go | 11 +++++++++--
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/relay/billing/ratio/group.go b/relay/billing/ratio/group.go
index 8e9c5b73..b7f62e6e 100644
--- a/relay/billing/ratio/group.go
+++ b/relay/billing/ratio/group.go
@@ -3,8 +3,10 @@ package ratio
 import (
 	"encoding/json"
 	"github.com/songquanpeng/one-api/common/logger"
+	"sync"
 )
 
+var groupRatioLock sync.RWMutex
 var GroupRatio = map[string]float64{
 	"default": 1,
 	"vip":     1,
@@ -20,11 +22,15 @@ func GroupRatio2JSONString() string {
 }
 
 func UpdateGroupRatioByJSONString(jsonStr string) error {
+	groupRatioLock.Lock()
+	defer groupRatioLock.Unlock()
 	GroupRatio = make(map[string]float64)
 	return json.Unmarshal([]byte(jsonStr), &GroupRatio)
 }
 
 func GetGroupRatio(name string) float64 {
+	groupRatioLock.RLock()
+	defer groupRatioLock.RUnlock()
 	ratio, ok := GroupRatio[name]
 	if !ok {
 		logger.SysError("group ratio not found: " + name)
diff --git a/relay/billing/ratio/model.go b/relay/billing/ratio/model.go
index d4deff8b..a0d57e26 100644
--- a/relay/billing/ratio/model.go
+++ b/relay/billing/ratio/model.go
@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"strings"
+	"sync"
 
 	"github.com/songquanpeng/one-api/common/logger"
 )
@@ -15,6 +16,8 @@ const (
 	RMB       = USD / USD2RMB
 )
 
+var modelRatioLock sync.RWMutex
+
 // ModelRatio
 // https://platform.openai.com/docs/models/model-endpoint-compatibility
 // 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-3-haiku-20240307":    0.25 / 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-5-sonnet-20240620": 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,
 	// https://cloud.baidu.com/doc/WENXINWORKSHOP/s/hlrk4akp7
 	"ERNIE-4.0-8K":       0.120 * RMB,
@@ -417,11 +420,15 @@ func ModelRatio2JSONString() string {
 }
 
 func UpdateModelRatioByJSONString(jsonStr string) error {
+	modelRatioLock.Lock()
+	defer modelRatioLock.Unlock()
 	ModelRatio = make(map[string]float64)
 	return json.Unmarshal([]byte(jsonStr), &ModelRatio)
 }
 
 func GetModelRatio(name string, channelType int) float64 {
+	modelRatioLock.RLock()
+	defer modelRatioLock.RUnlock()
 	if strings.HasPrefix(name, "qwen-") && strings.HasSuffix(name, "-internet") {
 		name = strings.TrimSuffix(name, "-internet")
 	}