From 4d8a8f95f1413fe52ff2ce9a90246cbf92b20541 Mon Sep 17 00:00:00 2001
From: chenyahui <cyhone@qq.com>
Date: Sat, 9 Oct 2021 20:13:23 +0800
Subject: [PATCH] feat: add OnShareSingleFlightCallback

---
 cache.go  |  6 ++++--
 option.go | 21 ++++++++++++++++++++-
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/cache.go b/cache.go
index ef8435f..c407723 100644
--- a/cache.go
+++ b/cache.go
@@ -34,8 +34,9 @@ func Cache(
 	opts ...Option,
 ) gin.HandlerFunc {
 	cfg := &Config{
-		logger:           Discard{},
-		hitCacheCallback: defaultHitCacheCallback,
+		logger:                    Discard{},
+		hitCacheCallback:          defaultHitCacheCallback,
+		shareSingleFlightCallback: defaultShareSingleFlightCallback,
 	}
 
 	for _, opt := range opts {
@@ -113,6 +114,7 @@ func Cache(
 
 		if !inFlight {
 			replyWithCache(c, cfg, rawRespCache.(*responseCache))
+			cfg.shareSingleFlightCallback(c)
 		}
 	}
 }
diff --git a/option.go b/option.go
index b9c751a..90f1135 100644
--- a/option.go
+++ b/option.go
@@ -15,6 +15,7 @@ type Config struct {
 	hitCacheCallback OnHitCacheCallback
 
 	singleFlightForgetTimeout time.Duration
+	shareSingleFlightCallback OnShareSingleFlightCallback
 }
 
 // Option represents the optional function.
@@ -29,6 +30,7 @@ func WithLogger(l Logger) Option {
 	}
 }
 
+// Logger define the logger interface
 type Logger interface {
 	Errorf(string, ...interface{})
 }
@@ -37,10 +39,12 @@ type Logger interface {
 type Discard struct {
 }
 
+// Errorf will output the log at error level
 func (l Discard) Errorf(string, ...interface{}) {
 
 }
 
+// WithCacheStrategyByRequest set up the custom strategy by per request
 func WithCacheStrategyByRequest(getGetCacheStrategyByRequest GetCacheStrategyByRequest) Option {
 	return func(c *Config) {
 		if getGetCacheStrategyByRequest != nil {
@@ -49,8 +53,11 @@ func WithCacheStrategyByRequest(getGetCacheStrategyByRequest GetCacheStrategyByR
 	}
 }
 
+// OnHitCacheCallback define the callback when use cache
 type OnHitCacheCallback func(c *gin.Context)
 
+var defaultHitCacheCallback = func(c *gin.Context) {}
+
 // WithOnHitCache will be called when cache hit.
 func WithOnHitCache(cb OnHitCacheCallback) Option {
 	return func(c *Config) {
@@ -60,7 +67,19 @@ func WithOnHitCache(cb OnHitCacheCallback) Option {
 	}
 }
 
-var defaultHitCacheCallback = func(c *gin.Context) {}
+// OnShareSingleFlightCallback define the callback when share the singleflight result
+type OnShareSingleFlightCallback func(c *gin.Context)
+
+var defaultShareSingleFlightCallback = func(c *gin.Context) {}
+
+// WithOnShareSingleFlight will be called when share the singleflight result
+func WithOnShareSingleFlight(cb OnShareSingleFlightCallback) Option {
+	return func(c *Config) {
+		if cb != nil {
+			c.shareSingleFlightCallback = cb
+		}
+	}
+}
 
 // WithSingleFlightForgetTimeout to reduce the impact of long tail requests. when request in the singleflight,
 // after the forget timeout, singleflight.Forget will be called