117 lines
2.2 KiB
Go
117 lines
2.2 KiB
Go
package templatecache
|
|
|
|
import (
|
|
"bytes"
|
|
"net/http"
|
|
"strings"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/patrickmn/go-cache"
|
|
)
|
|
|
|
const (
|
|
maxCacheItems = 2000
|
|
maxCacheBytes = 64 << 20
|
|
)
|
|
|
|
var (
|
|
Cache *cache.Cache
|
|
cacheBytes int64
|
|
)
|
|
|
|
// 章节详情、资讯详情等 ID 空间过大，禁止进进程缓存。
|
|
var highCardinalityContains = []string{
|
|
"/chapter/info",
|
|
"/consultinginfo/",
|
|
}
|
|
|
|
func init() {
|
|
Cache = cache.New(5*time.Minute, 5*time.Minute)
|
|
Cache.OnEvicted(func(_ string, v interface{}) {
|
|
if b, ok := v.([]byte); ok {
|
|
atomic.AddInt64(&cacheBytes, -int64(len(b)))
|
|
}
|
|
})
|
|
}
|
|
|
|
func shouldBypassCache(path string) bool {
|
|
if path == "" {
|
|
return false
|
|
}
|
|
if i := strings.IndexByte(path, '?'); i >= 0 {
|
|
path = path[:i]
|
|
}
|
|
for _, part := range highCardinalityContains {
|
|
if strings.Contains(path, part) {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func cacheBudgetAllows(n int) bool {
|
|
if n <= 0 {
|
|
return false
|
|
}
|
|
if Cache.ItemCount() >= maxCacheItems {
|
|
return false
|
|
}
|
|
return atomic.LoadInt64(&cacheBytes)+int64(n) <= maxCacheBytes
|
|
}
|
|
|
|
type templateResponseWriter struct {
|
|
gin.ResponseWriter
|
|
body *bytes.Buffer
|
|
}
|
|
|
|
func (w *templateResponseWriter) Write(b []byte) (int, error) {
|
|
w.body.Write(b)
|
|
return w.ResponseWriter.Write(b)
|
|
}
|
|
|
|
func WithCache(duration time.Duration) gin.HandlerFunc {
|
|
return func(ctx *gin.Context) {
|
|
if ctx.Request.Method != "GET" {
|
|
ctx.Next()
|
|
return
|
|
}
|
|
if shouldBypassCache(ctx.Request.URL.Path) {
|
|
ctx.Next()
|
|
return
|
|
}
|
|
|
|
cacheKey := ctx.Request.URL.String()
|
|
if cacheData, exists := Cache.Get(cacheKey); exists {
|
|
if item, ok := cacheData.([]byte); ok && len(item) > 0 {
|
|
ctx.Data(http.StatusOK, "application/json; charset=utf-8", item)
|
|
ctx.Abort()
|
|
return
|
|
}
|
|
}
|
|
|
|
writer := &templateResponseWriter{
|
|
body: bytes.NewBuffer(nil),
|
|
ResponseWriter: ctx.Writer,
|
|
}
|
|
ctx.Writer = writer
|
|
ctx.Next()
|
|
|
|
if ctx.Writer.Status() != 200 {
|
|
return
|
|
}
|
|
raw := writer.body.Bytes()
|
|
if !cacheBudgetAllows(len(raw)) {
|
|
return
|
|
}
|
|
copied := append([]byte(nil), raw...)
|
|
atomic.AddInt64(&cacheBytes, int64(len(copied)))
|
|
cacheDuration := duration
|
|
if cacheDuration == 0 {
|
|
cacheDuration = time.Minute
|
|
}
|
|
Cache.Set(cacheKey, copied, cacheDuration)
|
|
}
|
|
}
|