From 896e1d978f989925a56b79478fa14253f39e4238 Mon Sep 17 00:00:00 2001
From: CaIon <1808837298@qq.com>
Date: Tue, 10 Jun 2025 18:55:21 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A7=20fix(token=5Fcounter):=20enhance?=
 =?UTF-8?q?=20token=20encoder=20caching=20and=20concurrency=20handling?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 service/token_counter.go | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/service/token_counter.go b/service/token_counter.go
index 636dbd94..d27bb5ea 100644
--- a/service/token_counter.go
+++ b/service/token_counter.go
@@ -14,12 +14,19 @@ import (
 	"one-api/dto"
 	relaycommon "one-api/relay/common"
 	"strings"
+	"sync"
 	"unicode/utf8"
 )
 
 // tokenEncoderMap won't grow after initialization
 var defaultTokenEncoder tokenizer.Codec
 
+// tokenEncoderMap is used to store token encoders for different models
+var tokenEncoderMap = make(map[string]tokenizer.Codec)
+
+// tokenEncoderMutex protects tokenEncoderMap for concurrent access
+var tokenEncoderMutex sync.RWMutex
+
 func InitTokenEncoders() {
 	common.SysLog("initializing token encoders")
 	defaultTokenEncoder = codec.NewCl100kBase()
@@ -27,10 +34,33 @@ func InitTokenEncoders() {
 }
 
 func getTokenEncoder(model string) tokenizer.Codec {
+	// First, try to get the encoder from cache with read lock
+	tokenEncoderMutex.RLock()
+	if encoder, exists := tokenEncoderMap[model]; exists {
+		tokenEncoderMutex.RUnlock()
+		return encoder
+	}
+	tokenEncoderMutex.RUnlock()
+
+	// If not in cache, create new encoder with write lock
+	tokenEncoderMutex.Lock()
+	defer tokenEncoderMutex.Unlock()
+
+	// Double-check if another goroutine already created the encoder
+	if encoder, exists := tokenEncoderMap[model]; exists {
+		return encoder
+	}
+
+	// Create new encoder
 	modelCodec, err := tokenizer.ForModel(tokenizer.Model(model))
 	if err != nil {
+		// Cache the default encoder for this model to avoid repeated failures
+		tokenEncoderMap[model] = defaultTokenEncoder
 		return defaultTokenEncoder
 	}
+
+	// Cache the new encoder
+	tokenEncoderMap[model] = modelCodec
 	return modelCodec
 }