add cached shuffle mode to try to use as many different connecitons

go-socks-lb.yml.template

@@ -2,5 +2,5 @@ proxy:{% for idx in running_idx %}
   - url: "socks5://192.168.122.128:108{{idx}}"
     weight: 5
 {% end %}
-sticky: false
+load-balance-mode: "cached-shuffle"
 cache-clean-interval: 0

go-socks-lb/main.go

@@ -17,6 +17,8 @@ import (
 
 const (
 	URL           = "http://connectivitycheck.gstatic.com/generate_204"
+	cachedShuffle = "cached-shuffle"
+	sticky        = "sticky"
 )
 
 func proxyTestStatusCode(proxyURL string, URL string, StatusCode int) bool {
@@ -84,7 +86,7 @@ type proxyInst struct {
 type ProxyManager struct {
 	Proxys             []proxyInst
 	Cache              map[string]int
-	Sticky             bool
+	LoadBalanceMode    string
 	CacheCleanInterval int
 	Chooser            wr.Chooser
 	mux                sync.Mutex
@@ -94,7 +96,7 @@ type proxyConf struct {
 		URL    string `yaml:"url"`
 		Weight int    `yaml:"weight"`
 	} `yaml:"proxy"`
-	Sticky             bool `yaml:"sticky"`
+	LoadBalanceMode    string `yaml:"load-balance-mode"`
 	CacheCleanInterval int    `yaml:"cache-clean-interval"`
 }
 
@@ -120,7 +122,7 @@ func NewPM(confFp string) (*ProxyManager, error) {
 		pm.Proxys = append(pm.Proxys, pi)
 		chooseArr = append(chooseArr, wr.Choice{Item: idx, Weight: uint(pi.Weight)})
 	}
-	pm.Sticky = cfg.Sticky
+	pm.LoadBalanceMode = cfg.LoadBalanceMode
 	pm.CacheCleanInterval = cfg.CacheCleanInterval
 	pm.Cache = make(map[string]int)
 	pm.Chooser = wr.NewChooser(chooseArr...)
@@ -132,17 +134,31 @@ func (pm *ProxyManager) Get(addr string) string {
 	pm.mux.Lock()
 	defer pm.mux.Unlock()
 	// fmt.Println(pm.Cache)
-	if pm.Sticky {
+	if pm.LoadBalanceMode == sticky {
 		idx, ok := pm.Cache[addr]
 		if ok {
 			// fmt.Println("match addr", addr, "using:", idx)
 			return pm.Proxys[idx].URL
-		}
-	}
+		} else {
 			idx := pm.Chooser.Pick().(int)
 			pm.Cache[addr] = idx
-	// fmt.Println("addr", addr, "using:", idx)
 			return pm.Proxys[idx].URL
+		}
+	} else if pm.LoadBalanceMode == cachedShuffle {
+		idx, ok := pm.Cache[addr]
+		if ok {
+			pm.Cache[addr] = (idx + 1) % (len(pm.Proxys))
+			return pm.Proxys[pm.Cache[addr]].URL
+		} else {
+			idx := pm.Chooser.Pick().(int)
+			pm.Cache[addr] = idx
+			return pm.Proxys[idx].URL
+		}
+	} else {
+		idx := pm.Chooser.Pick().(int)
+		return pm.Proxys[idx].URL
+	}
+	//return ""
 }
 
 func (pm *ProxyManager) ClearCache() {
@@ -190,10 +206,13 @@ func main() {
 	if err != nil {
 		return
 	}
-	if pm.Sticky {
-		fmt.Printf("Sticky with interval %ds.\n", pm.CacheCleanInterval)
+	fmt.Printf("Load balance mode: ")
+	if pm.LoadBalanceMode == sticky {
+		fmt.Printf("sticky with cache clearing interval: %ds.\n", pm.CacheCleanInterval)
+	} else if pm.LoadBalanceMode == cachedShuffle {
+		fmt.Println("cached shuffling.")
 	} else {
-		fmt.Println("Randomize every connection.")
+		fmt.Println("randomize every connection.")
 	}
 	conf := &Config{PM: pm}
 	server, err := New(conf)