diff --git a/plugin/loadbalance/loadbalance.go b/plugin/loadbalance/loadbalance.go index 09a432121..9f7888e86 100644 --- a/plugin/loadbalance/loadbalance.go +++ b/plugin/loadbalance/loadbalance.go @@ -36,6 +36,32 @@ func randomShuffle(res *dns.Msg) *dns.Msg { } func roundRobin(in []dns.RR) []dns.RR { + if len(in) <= 1 { + return in + } + + // A response that is only addresses - the common case - needs shuffling but no + // partitioning, so it can be served with a single copy instead of four slices. + // The copy is not optional: in must not be modified, because a backend may hand + // us a slice it owns. plugin/file, for example, answers straight out of the zone + // tree, so shuffling in place would reorder the zone itself for every other + // query racing with this one. + if t := in[0].Header().Rrtype; t == dns.TypeA || t == dns.TypeAAAA { + allSame := true + for _, r := range in[1:] { + if r.Header().Rrtype != t { + allSame = false + break + } + } + if allSame { + out := make([]dns.RR, len(in)) + copy(out, in) + roundRobinShuffle(out) + return out + } + } + cname := []dns.RR{} address := []dns.RR{} mx := []dns.RR{} diff --git a/plugin/loadbalance/loadbalance_test.go b/plugin/loadbalance/loadbalance_test.go index cfa2463b3..e7834b001 100644 --- a/plugin/loadbalance/loadbalance_test.go +++ b/plugin/loadbalance/loadbalance_test.go @@ -201,3 +201,56 @@ func handler() plugin.Handler { return dns.RcodeSuccess, nil }) } + +func BenchmarkRoundRobin(b *testing.B) { + answer := []dns.RR{ + test.A("a.example.org. 300 IN A 10.0.0.1"), + test.A("a.example.org. 300 IN A 10.0.0.2"), + test.A("a.example.org. 300 IN A 10.0.0.3"), + } + b.ReportAllocs() + for b.Loop() { + _ = roundRobin(answer) + } +} + +// TestRoundRobinDoesNotMutateInput guards the contract that roundRobin leaves its +// input alone. Backends may return a slice they own - plugin/file answers directly +// out of the zone tree - so shuffling in place would corrupt shared state and race +// with concurrent queries. +func TestRoundRobinDoesNotMutateInput(t *testing.T) { + inputs := map[string][]dns.RR{ + "addresses only": { + test.A("a.example.org. 300 IN A 10.0.0.1"), + test.A("a.example.org. 300 IN A 10.0.0.2"), + test.A("a.example.org. 300 IN A 10.0.0.3"), + test.A("a.example.org. 300 IN A 10.0.0.4"), + }, + "mixed": { + test.CNAME("a.example.org. 300 IN CNAME b.example.org."), + test.A("b.example.org. 300 IN A 10.0.0.1"), + test.A("b.example.org. 300 IN A 10.0.0.2"), + test.MX("example.org. 300 IN MX 10 mx.example.org."), + }, + } + + for name, in := range inputs { + before := make([]dns.RR, len(in)) + copy(before, in) + + // Shuffle repeatedly: a single call may leave the order untouched by chance. + for range 50 { + out := roundRobin(in) + if len(out) != len(in) { + t.Fatalf("%s: got %d records back, want %d", name, len(out), len(in)) + } + } + + for i := range in { + if in[i] != before[i] { + t.Errorf("%s: roundRobin reordered its input at index %d", name, i) + break + } + } + } +}