From eca3d49077068a0d69487898c62f71b5cd480356 Mon Sep 17 00:00:00 2001
From: chenyahui <cyhone@qq.com>
Date: Thu, 3 Feb 2022 11:43:28 +0800
Subject: [PATCH] feat: add BeforeReplyWithCacheCallback option (#7)

---
 cache.go  | 23 +++++++++++------------
 option.go | 28 ++++++++++++++++++++++++++--
 2 files changed, 37 insertions(+), 14 deletions(-)

diff --git a/cache.go b/cache.go
index c407723..c25984f 100644
--- a/cache.go
+++ b/cache.go
@@ -33,11 +33,7 @@ func Cache(
 	defaultExpire time.Duration,
 	opts ...Option,
 ) gin.HandlerFunc {
-	cfg := &Config{
-		logger:                    Discard{},
-		hitCacheCallback:          defaultHitCacheCallback,
-		shareSingleFlightCallback: defaultShareSingleFlightCallback,
-	}
+	cfg := newConfig()
 
 	for _, opt := range opts {
 		opt(cfg)
@@ -67,7 +63,7 @@ func Cache(
 
 		// read cache first
 		{
-			respCache := &responseCache{}
+			respCache := &ResponseCache{}
 			err := cacheStore.Get(cacheKey, &respCache)
 			if err == nil {
 				replyWithCache(c, cfg, respCache)
@@ -99,7 +95,7 @@ func Cache(
 
 			inFlight = true
 
-			respCache := &responseCache{}
+			respCache := &ResponseCache{}
 			respCache.fillWithCacheWriter(cacheWriter)
 
 			// only cache 2xx response
@@ -113,7 +109,7 @@ func Cache(
 		})
 
 		if !inFlight {
-			replyWithCache(c, cfg, rawRespCache.(*responseCache))
+			replyWithCache(c, cfg, rawRespCache.(*ResponseCache))
 			cfg.shareSingleFlightCallback(c)
 		}
 	}
@@ -141,16 +137,17 @@ func CacheByRequestPath(defaultCacheStore persist.CacheStore, defaultExpire time
 }
 
 func init() {
-	gob.Register(&responseCache{})
+	gob.Register(&ResponseCache{})
 }
 
-type responseCache struct {
+// ResponseCache record the http response cache
+type ResponseCache struct {
 	Status int
 	Header http.Header
 	Data   []byte
 }
 
-func (c *responseCache) fillWithCacheWriter(cacheWriter *responseCacheWriter) {
+func (c *ResponseCache) fillWithCacheWriter(cacheWriter *responseCacheWriter) {
 	c.Status = cacheWriter.Status()
 	c.Data = cacheWriter.body.Bytes()
 	c.Header = cacheWriter.Header().Clone()
@@ -175,8 +172,10 @@ func (w *responseCacheWriter) WriteString(s string) (int, error) {
 func replyWithCache(
 	c *gin.Context,
 	cfg *Config,
-	respCache *responseCache,
+	respCache *ResponseCache,
 ) {
+	cfg.beforeReplyWithCacheCallback(c, respCache)
+
 	c.Writer.WriteHeader(respCache.Status)
 
 	for key, values := range respCache.Header {
diff --git a/option.go b/option.go
index 90f1135..6f08541 100644
--- a/option.go
+++ b/option.go
@@ -14,10 +14,21 @@ type Config struct {
 
 	hitCacheCallback OnHitCacheCallback
 
+	beforeReplyWithCacheCallback BeforeReplyWithCacheCallback
+
 	singleFlightForgetTimeout time.Duration
 	shareSingleFlightCallback OnShareSingleFlightCallback
 }
 
+func newConfig() *Config {
+	return &Config{
+		logger:                       Discard{},
+		hitCacheCallback:             defaultHitCacheCallback,
+		beforeReplyWithCacheCallback: defaultBeforeReplyWithCacheCallback,
+		shareSingleFlightCallback:    defaultShareSingleFlightCallback,
+	}
+}
+
 // Option represents the optional function.
 type Option func(c *Config)
 
@@ -67,6 +78,19 @@ func WithOnHitCache(cb OnHitCacheCallback) Option {
 	}
 }
 
+type BeforeReplyWithCacheCallback func(c *gin.Context, cache *ResponseCache)
+
+var defaultBeforeReplyWithCacheCallback = func(c *gin.Context, cache *ResponseCache) {}
+
+// WithBeforeReplyWithCache will be called before replying with cache.
+func WithBeforeReplyWithCache(cb BeforeReplyWithCacheCallback) Option {
+	return func(c *Config) {
+		if cb != nil {
+			c.beforeReplyWithCacheCallback = cb
+		}
+	}
+}
+
 // OnShareSingleFlightCallback define the callback when share the singleflight result
 type OnShareSingleFlightCallback func(c *gin.Context)
 
@@ -81,8 +105,8 @@ func WithOnShareSingleFlight(cb OnShareSingleFlightCallback) Option {
 	}
 }
 
-// WithSingleFlightForgetTimeout to reduce the impact of long tail requests. when request in the singleflight,
-// after the forget timeout, singleflight.Forget will be called
+// WithSingleFlightForgetTimeout to reduce the impact of long tail requests.
+// singleflight.Forget will be called after the timeout has reached for each backend request when timeout is greater than zero.
 func WithSingleFlightForgetTimeout(forgetTimeout time.Duration) Option {
 	return func(c *Config) {
 		if forgetTimeout > 0 {