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

✨ weighted random upstream

Browse Source
This commit is contained in:
Kawin Viriyaprasopsook 2024-07-23 08:50:10 +07:00
parent 97ff48ee70
commit d17de5c200
1 changed files with 30 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -102,44 +102,53 @@ func (m *RouteManager) getSessionHandler(r *http.Request, upstreams []*Upstream)
/* Functions related to random upstream picking */ /* Functions related to random upstream picking */
// Get a random upstream by the weights defined in Upstream struct, return the upstream, index value and any error // Get a random upstream by the weights defined in Upstream struct, return the upstream, index value and any error
func getRandomUpstreamByWeight(upstreams []*Upstream) (*Upstream, int, error) { func getRandomUpstreamByWeight(upstreams []*Upstream) (*Upstream, int, error) {
totalUpstreams := len(upstreams) // If there is only one upstream, return it
if totalUpstreams == 1 { if len(upstreams) == 1 {
return upstreams[0], 0, nil return upstreams[0], 0, nil
} }
// Calculate total weight for upstreams with weight > 0
totalWeight := 0 totalWeight := 0
fallbackUpstreams := make([]*Upstream, 0) // List of upstreams with weight 0 fallbackUpstreams := make([]*Upstream, 0)
fallbackUpstreamsOriginalID := make([]int, 0)
// Calculate total weight and gather fallbacks for _, upstream := range upstreams {
for ix, upstream := range upstreams { if upstream.Weight > 0 {
totalWeight += upstream.Weight totalWeight += upstream.Weight
if upstream.Weight == 0 { } else {
fallbackUpstreams = append(fallbackUpstreams, upstream) fallbackUpstreams = append(fallbackUpstreams, upstream) // Collect fallback upstreams
fallbackUpstreamsOriginalID = append(fallbackUpstreamsOriginalID, ix)
} }
} }
// If there are no upstreams with weight > 0, return a fallback upstream if available
if totalWeight == 0 { if totalWeight == 0 {
// If total weight is 0, fallback to a random upstream with weight 0 if len(fallbackUpstreams) > 0 {
if len(fallbackUpstreams) == 0 { // Randomly select one of the fallback upstreams
return nil, -1, errors.New("no valid upstream servers available") index := rand.Intn(len(fallbackUpstreams))
return fallbackUpstreams[index], index, nil
} }
upstreamID := rand.Intn(len(fallbackUpstreams)) // No upstreams available at all
return fallbackUpstreams[upstreamID], fallbackUpstreamsOriginalID[upstreamID], nil return nil, -1, errors.New("no valid upstream servers available")
} }
// Generate a random number in the range of total weight // Random weight between 0 and total weight
r := rand.Intn(totalWeight) randomWeight := rand.Intn(totalWeight)
// Select upstream based on random number // Select an upstream based on the random weight
for i, upstream := range upstreams { for i, upstream := range upstreams {
r -= upstream.Weight if upstream.Weight > 0 { // Only consider upstreams with weight > 0
if r < 0 { if randomWeight < upstream.Weight {
return upstream, i, nil return upstream, i, nil // Return the selected upstream and its index
}
randomWeight -= upstream.Weight
} }
} }
// If we reach here, it means we should return a fallback upstream if available
if len(fallbackUpstreams) > 0 {
index := rand.Intn(len(fallbackUpstreams))
return 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")
} }