plugin/loadbalance: reduce roundRobin allocations for homogeneous address sets (#8375)

* perf(loadbalance): fast-path zero-allocation roundRobin for homogeneous record sets

Signed-off-by: Manuel Rüger <manuel@rueg.eu>

* plugin/loadbalance: copy before shuffling in the fast path

The fast path shuffled the caller's slice in place, which is not safe.
roundRobin must not modify its input: a backend may hand back a slice it
owns rather than one built for the response. plugin/file does exactly that
- Lookup returns elem.Type(qtype), which is the zone tree's own []dns.RR -
so an in-place shuffle reorders the zone itself, visible to every other
query and racing with the ones running concurrently.

Copy the records into a fresh slice and shuffle that instead. This is still
a single allocation rather than the four slices the partitioning path builds,
so most of the gain is kept:

    name          old time/op    new time/op    delta
    RoundRobin      353 ns         244 ns       -31%

    name          old alloc/op   new alloc/op   delta
    RoundRobin      118 B           54 B        -54%

    name          old allocs/op  new allocs/op  delta
    RoundRobin        5              4          -20%

Also reorder the type check so a response led by a CNAME is rejected on the
first record instead of scanning the whole answer section first.

TestRoundRobinDoesNotMutateInput pins the contract; it fails against the
in-place version.

Signed-off-by: Manuel Rüger <manuel@rueg.eu>

---------

Signed-off-by: Manuel Rüger <manuel@rueg.eu>

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{}

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
+			}
+		}
+	}
+}