Skip to content

Commit

Permalink
✨ replace cache components
Browse files Browse the repository at this point in the history
  • Loading branch information
v_frgfeng committed Feb 23, 2024
1 parent 2e941a7 commit f352d0a
Show file tree
Hide file tree
Showing 3 changed files with 97 additions and 21 deletions.
95 changes: 78 additions & 17 deletions api/copilot/copilot.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -25,7 +24,7 @@ import (
"time"
)

var TokenCache *bigcache.BigCache
var TokenCache *freecache.Cache
var TokenExpiredError = errors.New("token expired")
var Client = resty.New()

Expand All @@ -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()
}
Expand All @@ -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)
Expand All @@ -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)
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -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天前统计数据
Expand Down
7 changes: 5 additions & 2 deletions go.mod
View file
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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
Expand Down
16 changes: 14 additions & 2 deletions go.sum
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand All @@ -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=
Expand Down Expand Up @@ -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=
Expand All @@ -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=
Expand All @@ -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=
Expand Down Expand Up @@ -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=
Expand Down

0 comments on commit f352d0a

Please sign in to comment.