73 lines
1.5 KiB
Go
73 lines
1.5 KiB
Go
package templatecache
|
||
|
||
import (
|
||
"bytes"
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/patrickmn/go-cache"
|
||
"net/http"
|
||
"time"
|
||
)
|
||
|
||
var (
|
||
Cache *cache.Cache
|
||
)
|
||
|
||
func init() {
|
||
Cache = cache.New(5*time.Minute, 5*time.Minute)
|
||
}
|
||
|
||
// 自定义 ResponseWriter 用于捕获 HTML 输出
|
||
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) {
|
||
// 只缓存 GET 请求（HEAD 由 _web 层处理，_server 不会收到 HEAD）
|
||
if ctx.Request.Method != "GET" {
|
||
ctx.Next()
|
||
return
|
||
}
|
||
|
||
// 生成缓存键（包含路径和查询参数）
|
||
cacheKey := ctx.Request.URL.String()
|
||
var cacheJson []byte
|
||
if cacheData, exists := Cache.Get(cacheKey); exists {
|
||
if item, ok := cacheData.([]byte); ok {
|
||
cacheJson = item
|
||
}
|
||
}
|
||
// 检查缓存
|
||
if cacheJson != nil {
|
||
ctx.Data(http.StatusOK, "application/json; charset=utf-8", cacheJson)
|
||
ctx.Abort()
|
||
return
|
||
}
|
||
|
||
writer := &templateResponseWriter{
|
||
body: bytes.NewBuffer(nil),
|
||
ResponseWriter: ctx.Writer,
|
||
}
|
||
ctx.Writer = writer
|
||
|
||
ctx.Next()
|
||
|
||
// 只有成功的响应才缓存
|
||
status := ctx.Writer.Status()
|
||
if status == 200 {
|
||
cacheJsonData := writer.body.Bytes()
|
||
cacheDuration := duration
|
||
if cacheDuration == 0 {
|
||
cacheDuration = 1 * time.Minute
|
||
}
|
||
Cache.Set(cacheKey, cacheJsonData, cacheDuration)
|
||
}
|
||
}
|
||
}
|