130 lines
2.4 KiB
Go
130 lines
2.4 KiB
Go
package cache
|
|
|
|
import (
|
|
"encoding/json"
|
|
"errors"
|
|
"net/http"
|
|
"strconv"
|
|
"time"
|
|
|
|
"esway/internal/database/cache"
|
|
"esway/internal/model"
|
|
|
|
"github.com/loveuer/nf"
|
|
"github.com/loveuer/nf/nft/log"
|
|
)
|
|
|
|
var (
|
|
defaultKeyFn = func(c *nf.Ctx) string {
|
|
return c.Request.URL.String()
|
|
}
|
|
defaultTimeout = 3600
|
|
defaultPrefix = "midd:cache"
|
|
)
|
|
|
|
type Config struct {
|
|
// if return "" (won't cache)
|
|
KeyFn func(c *nf.Ctx) string
|
|
|
|
// cache timeout(seconds)
|
|
Timeout int
|
|
|
|
Prefix string
|
|
Refresh bool
|
|
}
|
|
|
|
type store struct {
|
|
Body []byte `json:"body"`
|
|
Header http.Header `json:"header"`
|
|
When int64 `json:"when"`
|
|
}
|
|
|
|
func New(cfgs ...Config) nf.HandlerFunc {
|
|
if cache.Client == nil {
|
|
log.Panic("[middleware.cache] database cache client is nil")
|
|
}
|
|
|
|
var cfg Config
|
|
if len(cfgs) > 0 {
|
|
cfg = cfgs[0]
|
|
}
|
|
|
|
if cfg.KeyFn == nil {
|
|
cfg.KeyFn = defaultKeyFn
|
|
}
|
|
|
|
if cfg.Timeout <= 0 {
|
|
cfg.Timeout = defaultTimeout
|
|
}
|
|
|
|
if cfg.Prefix == "" {
|
|
cfg.Prefix = defaultPrefix
|
|
}
|
|
|
|
return func(c *nf.Ctx) error {
|
|
var (
|
|
key string
|
|
err error
|
|
bs []byte
|
|
res = new(store)
|
|
)
|
|
|
|
if key = cfg.KeyFn(c); key == "" {
|
|
return c.Next()
|
|
}
|
|
|
|
key = cfg.Prefix + ":" + key
|
|
duration := time.Duration(cfg.Timeout) * time.Second
|
|
|
|
if cfg.Refresh {
|
|
if bs, err = cache.Client.GetEx(c.Context(), key, duration); err != nil {
|
|
if !errors.Is(err, cache.ErrorKeyNotFound) {
|
|
log.Warn("[middleware.cache] cache get err: %s", err.Error())
|
|
}
|
|
goto FromNext
|
|
}
|
|
} else {
|
|
if bs, err = cache.Client.Get(c.Context(), key); err != nil {
|
|
if !errors.Is(err, cache.ErrorKeyNotFound) {
|
|
log.Warn("[middleware.cache] cache get err: %s", err.Error())
|
|
}
|
|
goto FromNext
|
|
}
|
|
}
|
|
|
|
if err = json.Unmarshal(bs, res); err != nil {
|
|
log.Warn("[middleware.cache] cache data unamrshal err: %s", err.Error())
|
|
goto FromNext
|
|
}
|
|
|
|
for key := range res.Header {
|
|
for idx := range res.Header[key] {
|
|
c.SetHeader(key, res.Header[key][idx])
|
|
}
|
|
}
|
|
|
|
c.SetHeader("X-Nf-Cache-At", strconv.Itoa(int(res.When)))
|
|
|
|
_, err = c.Write(res.Body)
|
|
return err
|
|
|
|
FromNext:
|
|
|
|
blw := model.NewCopyWriter(c.Writer)
|
|
c.Writer = blw
|
|
|
|
rerr := c.Next()
|
|
|
|
resp := blw.Bytes()
|
|
|
|
data := &store{Body: resp, Header: blw.Header().Clone(), When: time.Now().UnixMilli()}
|
|
cbs, _ := json.Marshal(data)
|
|
|
|
if err = cache.Client.SetEx(c.Context(), key, cbs, duration); err != nil {
|
|
log.Warn("[middleware.cache] cache client setex err: %s", err.Error())
|
|
}
|
|
|
|
return rerr
|
|
}
|
|
}
|