ultone/internal/middleware/cache/cache.go

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
	}
}