From 14bfc38f12c3169c05d9603b50b6d0a6d5e498fd Mon Sep 17 00:00:00 2001 From: rene manqueros Date: Tue, 28 Feb 2023 15:51:28 -0600 Subject: [PATCH] Added support for backend eviction --- cmd/balance.go | 21 ++++++++++++------ go.mod | 4 ++-- go.sum | 12 ++++------- models/loadbalancer.go | 48 ++++++++++++++++++++++++++++++++---------- 4 files changed, 57 insertions(+), 28 deletions(-) diff --git a/cmd/balance.go b/cmd/balance.go index 2c786f7..b154e87 100644 --- 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 index e04458a..6374032 100644 --- 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 index 2e64667..d30b126 100644 --- 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 index b9003dd..9a7c901 100644 --- 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{}