From 94c0bf978ac376be6bae91681704758d1edf2396 Mon Sep 17 00:00:00 2001
From: Pierre Souchay <p.souchay@criteo.com>
Date: Tue, 17 Apr 2018 00:50:00 +0200
Subject: [PATCH] Perform a binary search to find optimal size of DNS responses

Will fix https://github.com/hashicorp/consul/issues/4036

Instead of removing one by one the entries, find the optimal
size using binary search.

For SRV records, with 5k nodes, duration of DNS lookups is
divided by 4 or more.
---
 agent/dns.go | 42 ++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 40 insertions(+), 2 deletions(-)

diff --git a/agent/dns.go b/agent/dns.go
index 98cb5a9b5b..e5abb05fb0 100644
--- a/agent/dns.go
+++ b/agent/dns.go
@@ -717,6 +717,32 @@ func syncExtra(index map[string]dns.RR, resp *dns.Msg) {
 	resp.Extra = extra
 }
 
+func dnsBinaryTruncate(resp *dns.Msg, maxSize int, index map[string]dns.RR, hasExtra bool) int {
+	originalAnswser := resp.Answer
+	startIndex := 0
+	endIndex := len(resp.Answer)
+	for endIndex-startIndex > 1 {
+		median := startIndex + (endIndex-startIndex)/2
+
+		resp.Answer = originalAnswser[:median]
+		resp.Extra = originalAnswser[:median]
+		if hasExtra {
+			syncExtra(index, resp)
+		}
+		aLen := resp.Len()
+		if aLen <= maxSize {
+			if maxSize-aLen < 10 {
+				// We are good, increasing will go out of bounds
+				return median
+			}
+			startIndex = median
+		} else {
+			endIndex = median
+		}
+	}
+	return startIndex
+}
+
 // trimTCPResponse limit the MaximumSize of messages to 64k as it is the limit
 // of DNS responses
 func (d *DNSServer) trimTCPResponse(req, resp *dns.Msg) (trimmed bool) {
@@ -752,7 +778,13 @@ func (d *DNSServer) trimTCPResponse(req, resp *dns.Msg) (trimmed bool) {
 	// This enforces the given limit on 64k, the max limit for DNS messages
 	for len(resp.Answer) > 0 && resp.Len() > maxSize {
 		truncated = true
-		resp.Answer = resp.Answer[:len(resp.Answer)-1]
+		// More than 100 bytes, find with a binary search
+		if resp.Len()-maxSize > 100 {
+			bestIndex := dnsBinaryTruncate(resp, maxSize, index, hasExtra)
+			resp.Answer = resp.Answer[:bestIndex]
+		} else {
+			resp.Answer = resp.Answer[:len(resp.Answer)-1]
+		}
 		if hasExtra {
 			syncExtra(index, resp)
 		}
@@ -809,7 +841,13 @@ func trimUDPResponse(req, resp *dns.Msg, udpAnswerLimit int) (trimmed bool) {
 	compress := resp.Compress
 	resp.Compress = false
 	for len(resp.Answer) > 0 && resp.Len() > maxSize {
-		resp.Answer = resp.Answer[:len(resp.Answer)-1]
+		// More than 100 bytes, find with a binary search
+		if resp.Len()-maxSize > 100 {
+			bestIndex := dnsBinaryTruncate(resp, maxSize, index, hasExtra)
+			resp.Answer = resp.Answer[:bestIndex]
+		} else {
+			resp.Answer = resp.Answer[:len(resp.Answer)-1]
+		}
 		if hasExtra {
 			syncExtra(index, resp)
 		}