From b23870cb7b559d348fde5fb951becb84f65fc0a8 Mon Sep 17 00:00:00 2001
From: shao0222 <22172112+shao0222@users.noreply.github.com>
Date: Tue, 22 Aug 2023 19:09:55 +0800
Subject: [PATCH] fix: Ensure concurrency safety for tokenEncoder caching
 (close #459)

---
 controller/relay-utils.go | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/controller/relay-utils.go b/controller/relay-utils.go
index 5b3e0274..a99872bd 100644
--- a/controller/relay-utils.go
+++ b/controller/relay-utils.go
@@ -2,6 +2,7 @@ package controller
 
 import (
 	"fmt"
+	"sync"
 	"github.com/gin-gonic/gin"
 	"github.com/pkoukk/tiktoken-go"
 	"one-api/common"
@@ -10,8 +11,12 @@ import (
 var stopFinishReason = "stop"
 
 var tokenEncoderMap = map[string]*tiktoken.Tiktoken{}
+var tokenEncoderMapMux sync.Mutex
 
 func getTokenEncoder(model string) *tiktoken.Tiktoken {
+	tokenEncoderMapMux.Lock()
+	defer tokenEncoderMapMux.Unlock()
+
 	if tokenEncoder, ok := tokenEncoderMap[model]; ok {
 		return tokenEncoder
 	}