Added support for backend eviction

reneManqueros · Feb 28, 2023 · 14bfc38 · 14bfc38
1 parent a9ca647
commit 14bfc38
Show file tree

Hide file tree

Showing 4 changed files with 57 additions and 28 deletions.
diff --git a/cmd/balance.go b/cmd/balance.go
@@ -25,6 +25,7 @@ var balanceCmd = &cobra.Command{
 		isVerbose := false
 		timeout := 0
 		connections := 0
+		evict := 0
 		// Overrides
 		for _, v := range args {
 			argumentParts := strings.Split(v, "=")
@@ -50,17 +51,23 @@ var balanceCmd = &cobra.Command{
 				if argumentParts[0] == "verbose" && argumentParts[1] == "true" {
 					isVerbose = true
 				}
+				if argumentParts[0] == "evict" {
+					evict, _ = strconv.Atoi(argumentParts[1])
+				}
 			}
 		}
 
 		lb := models.LoadBalancer{
-			Network:     listenNetwork,
-			Timeout:     timeout,
-			Source:      listenAddress,
-			Mutex:       sync.Mutex{},
-			ConfigFile:  configFile,
-			IsVerbose:   isVerbose,
-			Connections: connections,
+			Network:           listenNetwork,
+			Timeout:           timeout,
+			Source:            listenAddress,
+			Mutex:             sync.Mutex{},
+			ConfigFile:        configFile,
+			IsVerbose:         isVerbose,
+			Connections:       connections,
+			EvictChances:      evict,
+			EvictableBackends: map[string]int{},
+			EvictableMutex:    sync.Mutex{},
 		}
 		lb.FromDisk()
 

diff --git a/go.mod b/go.mod
@@ -4,11 +4,11 @@ go 1.18
 
 require (
 	github.com/remeh/sizedwaitgroup v1.0.0
-	github.com/spf13/cobra v1.6.0
+	github.com/spf13/cobra v1.6.1
 	gopkg.in/yaml.v2 v2.4.0
 )
 
 require (
-	github.com/inconshreveable/mousetrap v1.0.1 // indirect
+	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
 )
diff --git a/go.sum b/go.sum
@@ -1,16 +1,12 @@
-github.com/cpuguy83/go-md2man/v2 v2.0.1/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
 github.com/cpuguy83/go-md2man/v2 v2.0.2/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
-github.com/inconshreveable/mousetrap v1.0.0 h1:Z8tu5sraLXCXIcARxBp/8cbvlwVa7Z1NHg9XEKhtSvM=
-github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8=
-github.com/inconshreveable/mousetrap v1.0.1 h1:U3uMjPSQEBMNp1lFxmllqCPM6P5u/Xq7Pgzkat/bFNc=
 github.com/inconshreveable/mousetrap v1.0.1/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
+github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
+github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
 github.com/remeh/sizedwaitgroup v1.0.0 h1:VNGGFwNo/R5+MJBf6yrsr110p0m4/OX4S3DCy7Kyl5E=
 github.com/remeh/sizedwaitgroup v1.0.0/go.mod h1:3j2R4OIe/SeS6YDhICBy22RWjJC5eNCJ1V+9+NVNYlo=
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
-github.com/spf13/cobra v1.4.0 h1:y+wJpx64xcgO1V+RcnwW0LEHxTKRi2ZDPSBjWnrg88Q=
-github.com/spf13/cobra v1.4.0/go.mod h1:Wo4iy3BUC+X2Fybo0PDqwJIv3dNRiZLHQymsfxlB84g=
-github.com/spf13/cobra v1.6.0 h1:42a0n6jwCot1pUmomAp4T7DeMD+20LFv4Q54pxLf2LI=
-github.com/spf13/cobra v1.6.0/go.mod h1:IOw/AERYS7UzyrGinqmz6HLUo219MORXGxhbaJUqzrY=
+github.com/spf13/cobra v1.6.1 h1:o94oiPyS4KD1mPy2fmcYYHHfCxLqYjJOhGsCHFZtEzA=
+github.com/spf13/cobra v1.6.1/go.mod h1:IOw/AERYS7UzyrGinqmz6HLUo219MORXGxhbaJUqzrY=
 github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=

diff --git a/models/loadbalancer.go b/models/loadbalancer.go
@@ -16,14 +16,17 @@ import (
 )
 
 type LoadBalancer struct {
-	Network     string
-	Source      string
-	ConfigFile  string
-	Backends    []Backend `yaml:"backends"`
-	Mutex       sync.Mutex
-	IsVerbose   bool
-	Timeout     int
-	Connections int
+	IsVerbose         bool
+	Timeout           int
+	Connections       int
+	Backends          []Backend `yaml:"backends"`
+	EvictChances      int
+	EvictableBackends map[string]int
+	EvictableMutex    sync.Mutex
+	Mutex             sync.Mutex
+	Network           string
+	Source            string
+	ConfigFile        string
 }
 
 func (l *LoadBalancer) GetBackend() (b Backend, err error) {
@@ -141,17 +144,40 @@ func (l *LoadBalancer) Listen() {
 			log.Println(fmt.Sprintf(`routing from %v through %s`, sourceConnection.LocalAddr(), destination))
 		}
 		swg.Add()
-		go l.handleRequest(sourceConnection, destination.Address, &swg)
+		go l.handleRequest(sourceConnection, destination, &swg)
 	}
 }
 
-func (l *LoadBalancer) handleRequest(sourceConnection net.Conn, destinationAddress string, tswg *sizedwaitgroup.SizedWaitGroup) {
+func (l *LoadBalancer) checkAndEvict(backend Backend) {
+	destinationAddress := backend.Address
+	evict := false
+	if l.EvictChances > 0 {
+		l.EvictableMutex.Lock()
+		if _, ok := l.EvictableBackends[destinationAddress]; ok {
+			l.EvictableBackends[destinationAddress]++
+		} else {
+			l.EvictableBackends[destinationAddress] = 1
+		}
+		if l.EvictableBackends[destinationAddress] > l.EvictChances {
+			evict = true
+			log.Println("evicting", destinationAddress)
+			delete(l.EvictableBackends, destinationAddress)
+		}
+		l.EvictableMutex.Unlock()
+	}
+	if evict == true {
+		l.Remove(backend)
+	}
+}
+
+func (l *LoadBalancer) handleRequest(sourceConnection net.Conn, backend Backend, tswg *sizedwaitgroup.SizedWaitGroup) {
 	var destinationConnection net.Conn
 	var err error
 
-	destinationConnection, err = net.DialTimeout(l.Network, destinationAddress, time.Duration(l.Timeout)*time.Millisecond)
+	destinationConnection, err = net.DialTimeout(l.Network, backend.Address, time.Duration(l.Timeout)*time.Millisecond)
 	if err != nil {
 		log.Println("handle request error", err)
+		l.checkAndEvict(backend)
 		return
 	}
 	wg := sync.WaitGroup{}