diff --git a/go-socks-lb.yml.template b/go-socks-lb.yml.template index 80ec42b..3ee11ca 100644 --- a/go-socks-lb.yml.template +++ b/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 diff --git a/go-socks-lb/main.go b/go-socks-lb/main.go index 31a1a53..2f5f9f2 100644 --- a/go-socks-lb/main.go +++ b/go-socks-lb/main.go @@ -16,7 +16,9 @@ import ( ) const ( - URL = "http://connectivitycheck.gstatic.com/generate_204" + 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,8 +96,8 @@ type proxyConf struct { URL string `yaml:"url"` Weight int `yaml:"weight"` } `yaml:"proxy"` - Sticky bool `yaml:"sticky"` - CacheCleanInterval int `yaml:"cache-clean-interval"` + LoadBalanceMode string `yaml:"load-balance-mode"` + CacheCleanInterval int `yaml:"cache-clean-interval"` } func NewPM(confFp string) (*ProxyManager, error) { @@ -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 + 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 } - idx := pm.Chooser.Pick().(int) - pm.Cache[addr] = idx - // fmt.Println("addr", addr, "using:", idx) - 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)