From 2b5f214f2a094eddf0afcce315121bf1c5b08451 Mon Sep 17 00:00:00 2001
From: guochao <jeffguorg@gmail.com>
Date: Thu, 6 Feb 2025 00:33:45 +0800
Subject: [PATCH] continue with not modified if only all upstreams respond not
 modified

---
 server.go | 32 +++++++++++++++++++++++---------
 1 file changed, 23 insertions(+), 9 deletions(-)

diff --git a/server.go b/server.go
index 18a24cc..00f0d5e 100644
--- a/server.go
+++ b/server.go
@@ -416,15 +416,18 @@ func (server *Server) fastesUpstream(r *http.Request, lastModified time.Time) (r
 	returnLock := &sync.Mutex{}
 	upstreams := len(server.Upstreams)
 	cancelFuncs := make([]func(), upstreams)
-	selectedCh := make(chan int, 1)
-	selectedOnce := &sync.Once{}
+	updateCh := make(chan int, 1)
+	updateOnce := &sync.Once{}
+	notModifiedCh := make(chan int, 1)
+	notModifiedOnce := &sync.Once{}
 	wg := &sync.WaitGroup{}
 	wg.Add(len(server.Upstreams))
 
-	defer close(selectedCh)
+	defer close(updateCh)
 	for idx := range server.Upstreams {
 		idx := idx
 		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
 		cancelFuncs[idx] = cancel
 
 		logger := slog.With("upstreamIdx", idx)
@@ -446,7 +449,7 @@ func (server *Server) fastesUpstream(r *http.Request, lastModified time.Time) (r
 			if response == nil {
 				return
 			}
-			if response.StatusCode != http.StatusOK {
+			if response.StatusCode != http.StatusOK && response.StatusCode != http.StatusNotModified {
 				return
 			}
 			locked := returnLock.TryLock()
@@ -455,9 +458,17 @@ func (server *Server) fastesUpstream(r *http.Request, lastModified time.Time) (r
 			}
 			defer returnLock.Unlock()
 
-			selectedOnce.Do(func() {
+			if response.StatusCode == http.StatusNotModified {
+				notModifiedOnce.Do(func() {
+					resultResponse, resultCh, resultErr = response, ch, err
+					notModifiedCh <- idx
+				})
+				return
+			}
+
+			updateOnce.Do(func() {
 				resultResponse, resultCh, resultErr = response, ch, err
-				selectedCh <- idx
+				updateCh <- idx
 
 				for cancelIdx, cancel := range cancelFuncs {
 					if cancelIdx == idx {
@@ -474,11 +485,14 @@ func (server *Server) fastesUpstream(r *http.Request, lastModified time.Time) (r
 	wg.Wait()
 
 	resultIdx = -1
-	select {
-	case idx := <-selectedCh:
+
+	if idx, ok := <-updateCh; ok {
 		resultIdx = idx
 		slog.With("upstreamIdx", resultIdx).Debug("upstream selected")
-	default:
+	} else if idx, ok := <-notModifiedCh; ok {
+		resultIdx = idx
+		slog.With("upstreamIdx", resultIdx).Debug("upstream selected(not modified)")
+	} else {
 		slog.Debug("no valid upstream found")
 	}