feat(core): upstream select it's own storage (#582)

* feat(core): upstream select it's own storage * fix(ci): vendor and use sync.Mutex instead of reference
darkweak · Dec 22, 2024 · 7649435 · 7649435
1 parent 0e1579f
commit 7649435
Show file tree

Hide file tree

Showing 9 changed files with 58 additions and 34 deletions.
diff --git a/pkg/middleware/middleware.go b/pkg/middleware/middleware.go
@@ -358,26 +358,51 @@ func (s *SouinBaseHandler) Store(
 						hn := strings.Split(hname, ":")
 						vhs.Set(hn[0], rq.Header.Get(hn[0]))
 					}
-					for _, storer := range s.Storers {
-						wg.Add(1)
-						go func(currentStorer types.Storer, currentRes http.Response) {
-							defer wg.Done()
-							if currentStorer.SetMultiLevel(
-								cachedKey,
-								variedKey,
-								response,
-								vhs,
-								currentRes.Header.Get("Etag"), ma,
-								variedKey,
-							) == nil {
-								s.Configuration.GetLogger().Debugf("Stored the key %s in the %s provider", variedKey, currentStorer.Name())
-								currentRes.Request = rq
-							} else {
-								mu.Lock()
-								fails = append(fails, fmt.Sprintf("; detail=%s-INSERTION-ERROR", currentStorer.Name()))
-								mu.Unlock()
+					if upstreamStorerTarget := res.Header.Get("X-Souin-Storer"); upstreamStorerTarget != "" {
+						res.Header.Del("X-Souin-Storer")
+
+						var overridedStorer types.Storer
+						for _, storer := range s.Storers {
+							if strings.Contains(strings.ToLower(storer.Name()), strings.ToLower(upstreamStorerTarget)) {
+								overridedStorer = storer
 							}
-						}(storer, res)
+						}
+
+						if overridedStorer.SetMultiLevel(
+							cachedKey,
+							variedKey,
+							response,
+							vhs,
+							res.Header.Get("Etag"), ma,
+							variedKey,
+						) == nil {
+							s.Configuration.GetLogger().Debugf("Stored the key %s in the %s provider", variedKey, overridedStorer.Name())
+							res.Request = rq
+						} else {
+							fails = append(fails, fmt.Sprintf("; detail=%s-INSERTION-ERROR", overridedStorer.Name()))
+						}
+					} else {
+						for _, storer := range s.Storers {
+							wg.Add(1)
+							go func(currentStorer types.Storer, currentRes http.Response) {
+								defer wg.Done()
+								if currentStorer.SetMultiLevel(
+									cachedKey,
+									variedKey,
+									response,
+									vhs,
+									currentRes.Header.Get("Etag"), ma,
+									variedKey,
+								) == nil {
+									s.Configuration.GetLogger().Debugf("Stored the key %s in the %s provider", variedKey, currentStorer.Name())
+									currentRes.Request = rq
+								} else {
+									mu.Lock()
+									fails = append(fails, fmt.Sprintf("; detail=%s-INSERTION-ERROR", currentStorer.Name()))
+									mu.Unlock()
+								}
+							}(storer, res)
+						}
 					}
 
 					wg.Wait()

diff --git a/pkg/middleware/writer.go b/pkg/middleware/writer.go
@@ -24,7 +24,7 @@ func NewCustomWriter(rq *http.Request, rw http.ResponseWriter, b *bytes.Buffer)
 		Req:        rq,
 		Rw:         rw,
 		Headers:    http.Header{},
-		mutex:      &sync.Mutex{},
+		mutex:      sync.Mutex{},
 	}
 }
 
@@ -35,7 +35,7 @@ type CustomWriter struct {
 	Req         *http.Request
 	Headers     http.Header
 	headersSent bool
-	mutex       *sync.Mutex
+	mutex       sync.Mutex
 	statusCode  int
 }
 

diff --git a/pkg/surrogate/providers/common.go b/pkg/surrogate/providers/common.go
@@ -104,7 +104,7 @@ type baseStorage struct {
 	dynamic    bool
 	keepStale  bool
 	logger     core.Logger
-	mu         *sync.Mutex
+	mu         sync.Mutex
 	duration   time.Duration
 }
 
@@ -159,7 +159,6 @@ func (s *baseStorage) init(config configurationtypes.AbstractConfigurationInterf
 	s.dynamic = config.GetDefaultCache().GetCDN().Dynamic
 	s.logger = config.GetLogger()
 	s.keysRegexp = keysRegexp
-	s.mu = &sync.Mutex{}
 	s.duration = storageToInfiniteTTLMap[s.Storage.Name()]
 }
 

diff --git a/pkg/surrogate/providers/common_test.go b/pkg/surrogate/providers/common_test.go
@@ -30,7 +30,7 @@ func mockCommonProvider() *baseStorage {
 			Keys:       make(map[string]configurationtypes.SurrogateKeys),
 			keysRegexp: make(map[string]keysRegexpInner),
 			dynamic:    true,
-			mu:         &sync.Mutex{},
+			mu:         sync.Mutex{},
 			logger:     zap.NewNop().Sugar(),
 		},
 	}

diff --git a/plugins/traefik/override/middleware/writer.go b/plugins/traefik/override/middleware/writer.go
@@ -23,7 +23,7 @@ func NewCustomWriter(rq *http.Request, rw http.ResponseWriter, b *bytes.Buffer)
 		Req:        rq,
 		Rw:         rw,
 		Headers:    http.Header{},
-		mutex:      &sync.Mutex{},
+		mutex:      sync.Mutex{},
 	}
 }
 
@@ -34,7 +34,7 @@ type CustomWriter struct {
 	Req         *http.Request
 	Headers     http.Header
 	headersSent bool
-	mutex       *sync.Mutex
+	mutex       sync.Mutex
 	statusCode  int
 }
 

diff --git a/plugins/traefik/override/surrogate/providers/common.go b/plugins/traefik/override/surrogate/providers/common.go
@@ -79,7 +79,7 @@ type baseStorage struct {
 	keysRegexp map[string]keysRegexpInner
 	dynamic    bool
 	keepStale  bool
-	mu         *sync.Mutex
+	mu         sync.Mutex
 	duration   time.Duration
 }
 
@@ -116,7 +116,7 @@ func (s *baseStorage) init(config configurationtypes.AbstractConfigurationInterf
 
 	s.dynamic = config.GetDefaultCache().GetCDN().Dynamic
 	s.keysRegexp = keysRegexp
-	s.mu = &sync.Mutex{}
+	s.mu = sync.Mutex{}
 	s.duration = storageToInfiniteTTLMap[s.Storage.Name()]
 }
 

diff --git a/plugins/traefik/vendor/github.com/darkweak/souin/pkg/middleware/writer.go b/plugins/traefik/vendor/github.com/darkweak/souin/pkg/middleware/writer.go
diff --git a/plugins/traefik/vendor/github.com/darkweak/souin/pkg/surrogate/providers/common.go b/plugins/traefik/vendor/github.com/darkweak/souin/pkg/surrogate/providers/common.go
diff --git a/plugins/tyk/override/cache/surrogate/providers/common.go b/plugins/tyk/override/cache/surrogate/providers/common.go
@@ -72,7 +72,7 @@ type baseStorage struct {
 	dynamic    bool
 	keepStale  bool
 	logger     *zap.Logger
-	mu         *sync.Mutex
+	mu         sync.Mutex
 }
 
 func (s *baseStorage) init(config configurationtypes.AbstractConfigurationInterface) {
@@ -103,7 +103,7 @@ func (s *baseStorage) init(config configurationtypes.AbstractConfigurationInterf
 	s.dynamic = config.GetDefaultCache().GetCDN().Dynamic
 	s.logger = config.GetLogger()
 	s.keysRegexp = keysRegexp
-	s.mu = &sync.Mutex{}
+	s.mu = sync.Mutex{}
 }
 
 func (s *baseStorage) storeTag(tag string, cacheKey string, re *regexp.Regexp) {