Add weighted-round-robin policy to loadbalance plugin (#5662)

* Add weighted-round-robin policy to loadbalance plugin

Signed-off-by: Gabor Dozsa <gabor.dozsa@ibm.com>
This commit is contained in:
Gabor Dozsa
2023-01-27 17:36:56 +01:00
committed by GitHub
parent bf7c2cf37b
commit 7da2cedaf0
8 changed files with 975 additions and 36 deletions

View File

@@ -5,11 +5,19 @@ import (
"github.com/miekg/dns"
)
// RoundRobinResponseWriter is a response writer that shuffles A, AAAA and MX records.
type RoundRobinResponseWriter struct{ dns.ResponseWriter }
const (
ramdomShufflePolicy = "round_robin"
weightedRoundRobinPolicy = "weighted"
)
// LoadBalanceResponseWriter is a response writer that shuffles A, AAAA and MX records.
type LoadBalanceResponseWriter struct {
dns.ResponseWriter
shuffle func(*dns.Msg) *dns.Msg
}
// WriteMsg implements the dns.ResponseWriter interface.
func (r *RoundRobinResponseWriter) WriteMsg(res *dns.Msg) error {
func (r *LoadBalanceResponseWriter) WriteMsg(res *dns.Msg) error {
if res.Rcode != dns.RcodeSuccess {
return r.ResponseWriter.WriteMsg(res)
}
@@ -18,11 +26,14 @@ func (r *RoundRobinResponseWriter) WriteMsg(res *dns.Msg) error {
return r.ResponseWriter.WriteMsg(res)
}
return r.ResponseWriter.WriteMsg(r.shuffle(res))
}
func randomShuffle(res *dns.Msg) *dns.Msg {
res.Answer = roundRobin(res.Answer)
res.Ns = roundRobin(res.Ns)
res.Extra = roundRobin(res.Extra)
return r.ResponseWriter.WriteMsg(res)
return res
}
func roundRobin(in []dns.RR) []dns.RR {
@@ -72,9 +83,9 @@ func roundRobinShuffle(records []dns.RR) {
}
// Write implements the dns.ResponseWriter interface.
func (r *RoundRobinResponseWriter) Write(buf []byte) (int, error) {
func (r *LoadBalanceResponseWriter) Write(buf []byte) (int, error) {
// Should we pack and unpack here to fiddle with the packet... Not likely.
log.Warning("RoundRobin called with Write: not shuffling records")
log.Warning("LoadBalance called with Write: not shuffling records")
n, err := r.ResponseWriter.Write(buf)
return n, err
}