🔧 fix(token_counter): enhance token encoder caching and concurrency handling

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
 }