marcel/zoraxy
1
0
Fork 0
mirror of https://github.com/tobychui/zoraxy.git synced 2025-06-04 22:57:20 +02:00
Code Issues Packages Projects Releases Wiki Activity

⚡️ fallbackUpstreams with preserve index

Browse Source
This commit is contained in:
Kawin Viriyaprasopsook 2024-07-24 14:47:33 +07:00
parent d17de5c200
commit e53c3cf3c4
1 changed files with 15 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
src/mod/dynamicproxy/loadbalance/originPicker.go
View File

@ -107,15 +107,22 @@ func getRandomUpstreamByWeight(upstreams []*Upstream) (*Upstream, int, error) {
return upstreams[0], 0, nil return upstreams[0], 0, nil
} }
// Preserve the index with upstreams
type upstreamWithIndex struct {
Upstream *Upstream
Index int
}
// Calculate total weight for upstreams with weight > 0 // Calculate total weight for upstreams with weight > 0
totalWeight := 0 totalWeight := 0
fallbackUpstreams := make([]*Upstream, 0) fallbackUpstreams := make([]upstreamWithIndex, 0, len(upstreams))
for _, upstream := range upstreams { for index, upstream := range upstreams {
if upstream.Weight > 0 { if upstream.Weight > 0 {
totalWeight += upstream.Weight totalWeight += upstream.Weight
} else { } else {
fallbackUpstreams = append(fallbackUpstreams, upstream) // Collect fallback upstreams // Collect fallback upstreams
fallbackUpstreams = append(fallbackUpstreams, upstreamWithIndex{upstream, index})
} }
} }
@ -124,7 +131,7 @@ func getRandomUpstreamByWeight(upstreams []*Upstream) (*Upstream, int, error) {
if len(fallbackUpstreams) > 0 { if len(fallbackUpstreams) > 0 {
// Randomly select one of the fallback upstreams // Randomly select one of the fallback upstreams
index := rand.Intn(len(fallbackUpstreams)) index := rand.Intn(len(fallbackUpstreams))
return fallbackUpstreams[index], index, nil return fallbackUpstreams[index].Upstream, fallbackUpstreams[index].Index, nil
} }
// No upstreams available at all // No upstreams available at all
return nil, -1, errors.New("no valid upstream servers available") return nil, -1, errors.New("no valid upstream servers available")
@ -134,10 +141,11 @@ func getRandomUpstreamByWeight(upstreams []*Upstream) (*Upstream, int, error) {
randomWeight := rand.Intn(totalWeight) randomWeight := rand.Intn(totalWeight)
// Select an upstream based on the random weight // Select an upstream based on the random weight
for i, upstream := range upstreams { for index, upstream := range upstreams {
if upstream.Weight > 0 { // Only consider upstreams with weight > 0 if upstream.Weight > 0 { // Only consider upstreams with weight > 0
if randomWeight < upstream.Weight { if randomWeight < upstream.Weight {
return upstream, i, nil // Return the selected upstream and its index // Return the selected upstream and its index
return upstream, index, nil
} }
randomWeight -= upstream.Weight randomWeight -= upstream.Weight
} }
@ -146,7 +154,7 @@ func getRandomUpstreamByWeight(upstreams []*Upstream) (*Upstream, int, error) {
// If we reach here, it means we should return a fallback upstream if available // If we reach here, it means we should return a fallback upstream if available
if len(fallbackUpstreams) > 0 { if len(fallbackUpstreams) > 0 {
index := rand.Intn(len(fallbackUpstreams)) index := rand.Intn(len(fallbackUpstreams))
return fallbackUpstreams[index], index, nil return fallbackUpstreams[index].Upstream, fallbackUpstreams[index].Index, nil
} }
return nil, -1, errors.New("failed to pick an upstream origin server") return nil, -1, errors.New("failed to pick an upstream origin server")