From f352d0ac2c401630e2e7b5d847cc4c4cfeaf5f21 Mon Sep 17 00:00:00 2001 From: v_frgfeng Date: Fri, 23 Feb 2024 17:51:25 +0800 Subject: [PATCH] :sparkles: replace cache components --- api/copilot/copilot.go | 95 ++++++++++++++++++++++++++++++++++-------- go.mod | 7 +++- go.sum | 16 ++++++- 3 files changed, 97 insertions(+), 21 deletions(-) diff --git a/api/copilot/copilot.go b/api/copilot/copilot.go index 01e4f85..6f42726 100644 --- a/api/copilot/copilot.go +++ b/api/copilot/copilot.go @@ -2,10 +2,9 @@ package copilot import ( "bufio" - "context" "errors" "fmt" - "github.com/allegro/bigcache/v3" + "github.com/coocood/freecache" "github.com/dalefengs/chat-api-proxy/api/genai" "github.com/dalefengs/chat-api-proxy/global" "github.com/dalefengs/chat-api-proxy/model/common/response" @@ -25,7 +24,7 @@ import ( "time" ) -var TokenCache *bigcache.BigCache +var TokenCache *freecache.Cache var TokenExpiredError = errors.New("token expired") var Client = resty.New() @@ -40,12 +39,8 @@ const ( var AuthCount = make(map[string]map[StatisticsType]int) func init() { - var err error - TokenCache, err = bigcache.New(context.Background(), bigcache.DefaultConfig(23*time.Minute)) - if err != nil { - log.Println("init CopilotTokenCache error ", err.Error()) - panic(err) - } + cacheSize := 10 * 1024 * 1024 // 10M + TokenCache = freecache.NewCache(cacheSize) log.Println("init CopilotTokenCache success") go ClearAuthCount() } @@ -61,6 +56,14 @@ func (co *CopilotApi) TokenHandler(c *gin.Context) { response.FailWithOpenAIError(http.StatusUnauthorized, err.Error(), c) return } + + tokenRawCache, err := GetTokenRawInfoCache(token) + if err == nil && len(tokenRawCache) > 0 { + global.SugarLog.Infow("TokenHandler get token raw cache", "token", token) + c.JSON(http.StatusOK, tokenRawCache) + return + } + _, respMap, httpStatus, err := GetCopilotToken(token, false) c.Status(httpStatus) @@ -80,6 +83,14 @@ func (co *CopilotApi) CoTokenHandler(c *gin.Context) { response.FailWithOpenAIError(http.StatusUnauthorized, err.Error(), c) return } + + tokenRawCache, err := GetTokenRawInfoCache(token) + if err == nil && len(tokenRawCache) > 0 { + global.SugarLog.Infow("CoTokenHandler get token raw cache", "token", token) + c.JSON(http.StatusOK, tokenRawCache) + return + } + _, respMap, httpStatus, err := GetCopilotToken(token, true) c.Status(httpStatus) @@ -126,7 +137,7 @@ func (co *CopilotApi) CompletionsHandler(c *gin.Context) { err = CompletionsRequest(c, req, copilotToken) // 如果 token 过期，重新获取一次 token if errors.Is(err, TokenExpiredError) { - TokenCache.Delete(token) // 删除缓存 + TokenCache.Del([]byte(token)) // 删除缓存 global.SugarLog.Infow("CompletionsHandler token expired, try get new token", "token", token) coCopilotToken, _, _, coErr := GetCopilotToken(token, true) if coErr != nil { @@ -178,7 +189,7 @@ func (co *CopilotApi) CountHandler(c *gin.Context) { // GetCopilotTokenWithCache 先从缓存中获取 CopilotToken，如果缓存中没有，再从 CoCopilot 获取 func GetCopilotTokenWithCache(token string) (copilotToken string, err error) { - cacheToken, cacheErr := TokenCache.Get(token) + cacheToken, cacheErr := TokenCache.Get([]byte(token)) if cacheErr != nil { global.SugarLog.Infow("CompletionsHandler get cache err, Try http fetch token", "err", cacheErr.Error(), "token", token) var tokenErr error @@ -196,6 +207,24 @@ func GetCopilotTokenWithCache(token string) (copilotToken string, err error) { return } +// GetTokenRawInfoCache 获取 token 的原始信息缓存 +func GetTokenRawInfoCache(token string) (respMap map[string]any, err error) { + tokenRawCache, err := TokenCache.Get([]byte(token + "_raw")) + if err != nil { + return + } + if tokenRawCache == nil { + global.SugarLog.Infow("GetTokenRawInfoCache token raw cache is nil", "token", token) + return + } + err = jsoniter.Unmarshal(tokenRawCache, &respMap) + if err != nil { + global.SugarLog.Errorw("GetTokenRawInfoCache get token raw cache json unmarshal error", "err", err) + return + } + return +} + // CompletionsRequest 请求 Copilot CompletionsHandler 接口 func CompletionsRequest(c *gin.Context, req map[string]interface{}, copilotToken string) (err error) { completionsURL := global.Config.Copilot.CompletionsURL @@ -316,7 +345,7 @@ func GetCopilotToken(key string, isCo bool) (token string, data map[string]inter go func() { jsonData, _ := jsoniter.MarshalIndent(AuthCount, "", " ") - global.SugarLog.Infow("request statistics", "statistics", string(jsonData)) + global.SugarLog.Infof("request statistics \n %s", string(jsonData)) }() if err != nil { global.SugarLog.Errorw("GetCopilotToken request http error", "err", err, "url", global.Config.Copilot.CoTokenURL, "key", key) @@ -354,10 +383,44 @@ func GetCopilotToken(key string, isCo bool) (token string, data map[string]inter err = errors.New("response token is empty") return } - global.SugarLog.Infow("GetCopilotToken GetCopilotToken Success", "key", key) - cacheErr := TokenCache.Set(key, []byte(token)) + + expires := 700 + expiresAt, ok := data["expires_at"] + var expiresTime string + if ok { + switch expiresAt.(type) { + case int, int64: + expiresAtInt := expiresAt.(int64) + if expiresAtInt > 0 { + expires = int(expiresAtInt - time.Now().Unix()) + expiresTime = time.Unix(expiresAtInt, 0).Format("2006-01-02 15:04:05") + } + case float64: + expiresAtFloat := expiresAt.(float64) + if expiresAtFloat > 0 { + expires = int(int64(expiresAtFloat) - time.Now().Unix()) + expiresTime = time.Unix(int64(expiresAtFloat), 0).Format("2006-01-02 15:04:05") + } + default: + expiresTime = "expiresAt type error" + global.SugarLog.Errorw("GetCopilotToken expires_at type error", "expiresAt", expiresAt) + } + } + + global.SugarLog.Infow("GetCopilotToken HTTP Authorisation token success", "key", key, "expires", expires, "expiresTime", expiresTime) + + cacheErr := TokenCache.Set([]byte(key), []byte(token), expires) + if cacheErr != nil { + global.SugarLog.Errorw("GetCopilotToken set token cache err", "err", cacheErr) + } + rawTokenInfo, err := jsoniter.Marshal(data) + if err != nil { + global.SugarLog.Errorw("GetCopilotToken token raw info json marshal error", "err", err) + return + } + cacheErr = TokenCache.Set([]byte(key+"_raw"), rawTokenInfo, expires) if cacheErr != nil { - global.SugarLog.Errorw("GetCopilotToken set cache err", "err", cacheErr) + global.SugarLog.Errorw("GetCopilotToken set token raw cache err", "err", cacheErr) } return } @@ -391,8 +454,6 @@ func IncAuthCount(authType StatisticsType) { AuthCount[day][authType] = 0 } AuthCount[day][authType]++ - jsonData, _ := jsoniter.MarshalIndent(AuthCount, "", " ") - global.SugarLog.Infof("request statistics \n %s", string(jsonData)) } // ClearAuthCount 每天 0 点清空7天前统计数据 diff --git a/go.mod b/go.mod index 985b960..126cfc0 100644 --- a/go.mod +++ b/go.mod @@ -1,9 +1,11 @@ module github.com/dalefengs/chat-api-proxy -go 1.18 +go 1.21 + +toolchain go1.21.0 require ( - github.com/allegro/bigcache/v3 v3.1.0 + github.com/coocood/freecache v1.2.4 github.com/fsnotify/fsnotify v1.7.0 github.com/fvbock/endless v0.0.0-20170109170031-447134032cb6 github.com/gin-gonic/gin v1.9.1 @@ -25,6 +27,7 @@ require ( cloud.google.com/go/compute/metadata v0.2.3 // indirect cloud.google.com/go/longrunning v0.5.4 // indirect github.com/bytedance/sonic v1.9.1 // indirect + github.com/cespare/xxhash/v2 v2.2.0 // indirect github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 // indirect github.com/gabriel-vasile/mimetype v1.4.2 // indirect github.com/gin-contrib/sse v0.1.0 // indirect diff --git a/go.sum b/go.sum index 92332f9..6341b5b 100644 --- a/go.sum +++ b/go.sum @@ -8,27 +8,32 @@ cloud.google.com/go/compute/metadata v0.2.3/go.mod h1:VAV5nSsACxMJvgaAuX6Pk2Aawl cloud.google.com/go/longrunning v0.5.4 h1:w8xEcbZodnA2BbW6sVirkkoC+1gP8wS57EUUgGS0GVg= cloud.google.com/go/longrunning v0.5.4/go.mod h1:zqNVncI0BOP8ST6XQD1+VcvuShMmq7+xFSzOL++V0dI= github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= -github.com/allegro/bigcache/v3 v3.1.0 h1:H2Vp8VOvxcrB91o86fUSVJFqeuz8kpyyB02eH3bSzwk= -github.com/allegro/bigcache/v3 v3.1.0/go.mod h1:aPyh7jEvrog9zAwx5N7+JUQX5dZTSGpxF1LAR4dr35I= github.com/benbjohnson/clock v1.1.0 h1:Q92kusRqC1XV2MjkWETPvjJVqKetz1OzxZB7mHJLju8= github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA= github.com/bytedance/sonic v1.5.0/go.mod h1:ED5hyg4y6t3/9Ku1R6dU/4KyJ48DZ4jPhfY1O2AihPM= github.com/bytedance/sonic v1.9.1 h1:6iJ6NqdoxCDr6mbY8h18oSO+cShGSMRGCEo7F2h0x8s= github.com/bytedance/sonic v1.9.1/go.mod h1:i736AoUSYt75HyZLoJW9ERYxcy6eaN6h4BZXU064P/U= github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU= +github.com/cespare/xxhash/v2 v2.1.2/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44= +github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/chenzhuoyu/base64x v0.0.0-20211019084208-fb5309c8db06/go.mod h1:DH46F32mSOjUmXrMHnKwZdA8wcEefY7UVqBKYGjpdQY= github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 h1:qSGYFH7+jGhDF8vLC+iwCD4WpbV1EBDSzWkJODFLams= github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311/go.mod h1:b583jCggY9gE99b6G5LEC39OIiVsWj+R97kbl5odCEk= github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw= github.com/cncf/udpa/go v0.0.0-20191209042840-269d4d468f6f/go.mod h1:M8M6+tZqaGXZJjfX53e64911xZQV5JYwmTeXPW+k8Sc= +github.com/coocood/freecache v1.2.4 h1:UdR6Yz/X1HW4fZOuH0Z94KwG851GWOSknua5VUbb/5M= +github.com/coocood/freecache v1.2.4/go.mod h1:RBUWa/Cy+OHdfTGFEhEuE1pMCMX51Ncizj7rthiQ3vk= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= github.com/envoyproxy/go-control-plane v0.9.4/go.mod h1:6rpuAdCZL397s3pYoYcLgu1mIlRU8Am5FuJP05cCM98= github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c= github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8= +github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0= github.com/fsnotify/fsnotify v1.7.0 h1:8JEhPFa5W2WU7YfeZzPNqzMP6Lwt7L2715Ggo0nosvA= github.com/fsnotify/fsnotify v1.7.0/go.mod h1:40Bi/Hjc2AVfZrqy+aj+yEI+/bRxZnMJyTJwOpGvigM= github.com/fvbock/endless v0.0.0-20170109170031-447134032cb6 h1:6VSn3hB5U5GeA6kQw4TwWIWbOhtvR2hmbBJnTOtqTWc= @@ -40,6 +45,7 @@ github.com/gin-contrib/sse v0.1.0/go.mod h1:RHrZQHXnP2xjPF+u1gW/2HnVO7nvIa9PG3Gm github.com/gin-gonic/gin v1.9.1 h1:4idEAncQnU5cB7BeOkPtxjfCSye0AAm1R0RVIqJ+Jmg= github.com/gin-gonic/gin v1.9.1/go.mod h1:hPrL7YrpYKXt5YId3A/Tnip5kqbEAP+KLuI3SUcPTeU= github.com/go-playground/assert/v2 v2.2.0 h1:JvknZsQTYeFEAhQwI4qEt9cyV5ONwRHC+lYKSsYSR8s= +github.com/go-playground/assert/v2 v2.2.0/go.mod h1:VDjEfimB/XKnb+ZQfWdccd7VUvScMdVu0Titje2rxJ4= github.com/go-playground/locales v0.14.1 h1:EWaQ/wswjilfKLTECiXz7Rh+3BjFhfDFKv/oXslEjJA= github.com/go-playground/locales v0.14.1/go.mod h1:hxrqLVvrK65+Rwrd5Fc6F2O76J/NuW9t0sjnWqG1slY= github.com/go-playground/universal-translator v0.18.1 h1:Bcnm0ZwsGyWbCzImXv+pAJnYK9S473LQFuzCbDbfSFY= @@ -78,6 +84,7 @@ github.com/google/go-cmp v0.5.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/ github.com/google/go-cmp v0.5.3/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= +github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= github.com/google/s2a-go v0.1.7 h1:60BLSyTrOV4/haCDW4zb1guZItoSq8foHCXrAnjBo/o= github.com/google/s2a-go v0.1.7/go.mod h1:50CgR4k1jNlWBu4UfS4AcfhVe1r6pdZPygJ3R8F0Qdw= @@ -99,9 +106,11 @@ github.com/klauspost/cpuid/v2 v2.2.4 h1:acbojRNwl3o09bUq+yDCtZFc1aiwaAAxtcn8YkZX github.com/klauspost/cpuid/v2 v2.2.4/go.mod h1:RVVoqg1df56z8g3pUjL/3lE5UfnlrJX8tyFgg4nqhuY= github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/leodido/go-urn v1.2.4 h1:XlAE/cm/ms7TE/VMVoduSpNBoyc2dOxHs5MZSwAN63Q= github.com/leodido/go-urn v1.2.4/go.mod h1:7ZrI8mTSeBSHl/UaRyKQW1qZeMgak41ANeCNaVckg+4= github.com/lestrrat-go/envload v0.0.0-20180220234015-a3eb8ddeffcc h1:RKf14vYWi2ttpEmkA4aQ3j4u9dStX2t4M8UM6qqNsG8= @@ -128,8 +137,10 @@ github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA= github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZVejAe8= +github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs= github.com/sagikazarmark/locafero v0.4.0 h1:HApY1R9zGo4DBgr7dqsTH/JJxLTTsOt7u6keLGt6kNQ= github.com/sagikazarmark/locafero v0.4.0/go.mod h1:Pe1W6UlPYUk/+wc/6KFhbORCfqzgYEpgQ3O5fPuL3H4= github.com/sagikazarmark/slog-shim v0.1.0 h1:diDBnUNK9N/354PgrxMywXnAwEr1QZcOr6gto+ugjYE= @@ -309,6 +320,7 @@ google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqw gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 h1:YR8cESwS4TdDjEe65xsg0ogRM/Nc3DYOhEAlW+xobZo= +gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/ini.v1 v1.67.0 h1:Dgnx+6+nfE+IfzjUEISNeydPJh9AXNNsWbGP9KzCsOA= gopkg.in/ini.v1 v1.67.0/go.mod h1:pNLf8WUiyNEtQjuu5G5vTm06TEv9tsIgeAvK8hOrP4k= gopkg.in/yaml.v2 v2.2.8 h1:obN1ZagJSUGI0Ek/LBmuj4SNLPfIny3KsKFopxRdj10=