Custom Strategies
A Freighter strategy is any struct that implements the Strategy
interface. This allows users to implement their own custom strategies which suit their needs. All they have to do is create a new struct which implements the interface and pass it to Freighter.
type Strategy interface {
Handle(r *http.Request, p *pool.ServerPool) *pool.Backend
}
Re-implementing Round-Robin
Let's try to implement our own version of the Round-Robin strategy! Round-Robin is one of the simplest load-balancing strategies to distribute load across a set of backends. All it does is forward incoming requests to backends turn wise.
First, initialize a new project and add freighter as a dependency.
go mod init github.com/my/repo
go get github.com/ShardulNalegave/freighter
Now create main.go
and add following code to it.
package main
import (
"net/url"
"time"
"github.com/ShardulNalegave/freighter"
"github.com/ShardulNalegave/freighter/pool"
)
func main() {
srv := freighter.NewFreighter(&freighter.Options{
URL: &url.URL{
Host: ":5000",
},
EnableConsoleLogging: true,
HealthCheckInterval: time.Second * 5,
Backends: []*pool.Backend{
pool.NewBackend(&url.URL{Host: ":8080", Scheme: "http"}, nil),
pool.NewBackend(&url.URL{Host: ":8081", Scheme: "http"}, nil),
},
// TODO: Add our new strategy
})
srv.ListenAndServe()
}
Now, all thats let is to implement our strategy.
type MyStrategy struct{
current uint64
}
Taking a look at the Strategy
interface, we see that Handle
method should be defined on MyStrategy
import (
"net/http"
"sync/atomic"
"github.com/ShardulNalegave/freighter/pool"
)
func (s *MyStrategy) Handle(r *http.Request, p *pool.ServerPool) *pool.Backend {
next := s.increment(p)
l := len(p.Backends) + int(s.current)
for i := next; i < l; i++ {
index := int(i) % len(p.Backends)
if p.Backends[index].IsAlive() {
if i != next {
atomic.StoreUint64(&s.current, uint64(index))
}
return p.Backends[index]
}
}
return nil
}
func (s *MyStrategy) increment(p *pool.ServerPool) int {
return int(atomic.AddUint64(&s.current, uint64(1)) % uint64(len(p.Backends)))
}
What we are doing here is:-
- When a request is received, get the next backend using the
current
index - Forward the request by using its reverse-proxy.
- Atomically increment
current
by 1. - If
current
is equal to number to backends then resetcurrent
to 0.
Finally, pass our new strategy to the Freighter instance.
Complete Code
package main
import (
"net/url"
"time"
"github.com/ShardulNalegave/freighter"
"github.com/ShardulNalegave/freighter/pool"
)
func main() {
srv := freighter.NewFreighter(&freighter.Options{
URL: &url.URL{
Host: ":5000",
},
EnableConsoleLogging: true,
HealthCheckInterval: time.Second * 5,
Strategy: &MyStrategy{}, // Our new strategy
Backends: []*pool.Backend{
pool.NewBackend(&url.URL{Host: ":8080", Scheme: "http"}, nil),
pool.NewBackend(&url.URL{Host: ":8081", Scheme: "http"}, nil),
},
})
srv.ListenAndServe()
}
type MyStrategy struct{
current uint64
}
func (s *MyStrategy) Handle(r *http.Request, p *pool.ServerPool) *pool.Backend {
next := s.increment(p)
l := len(p.Backends) + int(s.current)
for i := next; i < l; i++ {
index := int(i) % len(p.Backends)
if p.Backends[index].IsAlive() {
if i != next {
atomic.StoreUint64(&s.current, uint64(index))
}
return p.Backends[index]
}
}
return nil
}
func (s *MyStrategy) increment(p *pool.ServerPool) int {
return int(atomic.AddUint64(&s.current, uint64(1)) % uint64(len(p.Backends)))
}