From 53f124c9ec30f12685c458e2924477010a2d0499 Mon Sep 17 00:00:00 2001 From: loveuer Date: Thu, 25 Jul 2024 17:31:57 +0800 Subject: [PATCH] feat: add middleware.cache --- go.mod | 4 +- go.sum | 4 +- internal/database/cache/client.go | 5 ++ internal/middleware/cache/cache.go | 129 +++++++++++++++++++++++++++++ internal/model/writer.go | 21 +++++ 5 files changed, 159 insertions(+), 4 deletions(-) create mode 100644 internal/middleware/cache/cache.go create mode 100644 internal/model/writer.go diff --git a/go.mod b/go.mod index a3e4e1f..5a0f7f0 100644 --- a/go.mod +++ b/go.mod @@ -8,11 +8,12 @@ require ( github.com/glebarez/sqlite v1.10.0 github.com/go-redis/redis/v8 v8.11.5 github.com/golang-jwt/jwt/v5 v5.2.0 + github.com/google/uuid v1.6.0 github.com/hashicorp/golang-lru/v2 v2.0.7 github.com/jackc/pgtype v1.12.0 github.com/jedib0t/go-pretty/v6 v6.5.9 github.com/loveuer/esgo2dump v0.3.3 - github.com/loveuer/nf v0.2.7 + github.com/loveuer/nf v0.2.8 github.com/loveuer/ngorm/v2 v2.1.1 github.com/rabbitmq/amqp091-go v1.10.0 github.com/samber/lo v1.39.0 @@ -35,7 +36,6 @@ require ( github.com/glebarez/go-sqlite v1.21.2 // indirect github.com/go-sql-driver/mysql v1.7.0 // indirect github.com/golang/protobuf v1.5.2 // indirect - github.com/google/uuid v1.6.0 // indirect github.com/jackc/chunkreader/v2 v2.0.1 // indirect github.com/jackc/pgconn v1.13.0 // indirect github.com/jackc/pgio v1.0.0 // indirect diff --git a/go.sum b/go.sum index f12a62e..b5f1edb 100644 --- a/go.sum +++ b/go.sum @@ -146,8 +146,8 @@ github.com/lib/pq v1.10.2 h1:AqzbZs4ZoCBp+GtejcpCpcxM3zlSMx29dXbUSeVtJb8= github.com/lib/pq v1.10.2/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o= github.com/loveuer/esgo2dump v0.3.3 h1:/AidoaFV7bDRyT1ycyBKs4XGmyVs2ShaUKrpEBiUWkM= github.com/loveuer/esgo2dump v0.3.3/go.mod h1:thZvfsO0kd7Ck3TA0jc9rRc4CuIa4Iuiq6tF3tCqXEY= -github.com/loveuer/nf v0.2.7 h1:p17Y2yvO6u1qnSvYawQsMzBb2Xw7EzkV87r6BZzhdfA= -github.com/loveuer/nf v0.2.7/go.mod h1:M6reF17/kJBis30H4DxR5hrtgo/oJL4AV4cBe4HzJLw= +github.com/loveuer/nf v0.2.8 h1:Qo0M748TglS6E5geh1LG0IBkrjLm+5yUs3II9l50tEQ= +github.com/loveuer/nf v0.2.8/go.mod h1:M6reF17/kJBis30H4DxR5hrtgo/oJL4AV4cBe4HzJLw= github.com/loveuer/ngorm/v2 v2.1.1 h1:v+ut5BjeSBFU87o800pI8Q3fXEOUAkvk5+btMw2oOEc= github.com/loveuer/ngorm/v2 v2.1.1/go.mod h1:BVhFGQsRMdcf08MtmwwRihwCR/x7wDd0Fzy8Xj+edM0= github.com/mattn/go-colorable v0.1.1/go.mod h1:FuOcm+DKB9mbwrcAfNl7/TZVBZ6rcnceauSikq3lYCQ= diff --git a/internal/database/cache/client.go b/internal/database/cache/client.go index cd544dc..16c941d 100644 --- a/internal/database/cache/client.go +++ b/internal/database/cache/client.go @@ -23,6 +23,11 @@ func handleValue(value any) ([]byte, error) { err error ) + switch value.(type) { + case []byte: + return value.([]byte), nil + } + if imp, ok := value.(encoded_value); ok { bs, err = imp.MarshalBinary() } else { diff --git a/internal/middleware/cache/cache.go b/internal/middleware/cache/cache.go new file mode 100644 index 0000000..1224ecf --- /dev/null +++ b/internal/middleware/cache/cache.go @@ -0,0 +1,129 @@ +package cache + +import ( + "bytes" + "encoding/json" + "errors" + "github.com/loveuer/nf" + "github.com/loveuer/nf/nft/log" + "io" + "net/http" + "strconv" + "time" + "ultone/internal/database/cache" + "ultone/internal/model" +) + +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.CustomResponseWriter{Body: bytes.NewBuffer(make([]byte, 0, 256)), ResponseWriter: c.Writer} + c.Writer = blw + + rerr := c.Next() + + resp, _ := io.ReadAll(blw.Body) + + 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 + } +} diff --git a/internal/model/writer.go b/internal/model/writer.go new file mode 100644 index 0000000..6eed927 --- /dev/null +++ b/internal/model/writer.go @@ -0,0 +1,21 @@ +package model + +import ( + "bytes" + "github.com/loveuer/nf" +) + +type CustomResponseWriter struct { + nf.ResponseWriter + Body *bytes.Buffer +} + +func (w CustomResponseWriter) Write(b []byte) (int, error) { + w.Body.Write(b) + return w.ResponseWriter.Write(b) +} + +func (w CustomResponseWriter) WriteString(s string) (int, error) { + w.Body.WriteString(s) + return w.ResponseWriter.WriteString(s) +}