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

darkweak · Dec 8, 2024 · b9354c8 · b9354c8
1 parent 0e1579f
commit b9354c8
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 21 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/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()]
 }