From 9578caeab0b302ad470aacae5613d2d0c85530ac Mon Sep 17 00:00:00 2001
From: Charlie Vieth <charlie.vieth@gmail.com>
Date: Mon, 23 Sep 2019 02:16:26 -0400
Subject: [PATCH] pool the transient maps used for Msg Pack, Truncate and Len
 (#1006)

This improves runtime by 20-40% and more importantly significantly
reduces memory usage and allocations.
---
 msg.go          | 30 ++++++++++++++++++++++++++----
 msg_truncate.go |  7 ++++++-
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/msg.go b/msg.go
index e04fb5d7..a2e051b6 100644
--- a/msg.go
+++ b/msg.go
@@ -754,13 +754,24 @@ func (dns *Msg) Pack() (msg []byte, err error) {
 	return dns.PackBuffer(nil)
 }
 
+var compressionPackPool = sync.Pool{
+	New: func() interface{} {
+		return make(map[string]uint16)
+	},
+}
+
 // PackBuffer packs a Msg, using the given buffer buf. If buf is too small a new buffer is allocated.
 func (dns *Msg) PackBuffer(buf []byte) (msg []byte, err error) {
 	// If this message can't be compressed, avoid filling the
 	// compression map and creating garbage.
 	if dns.Compress && dns.isCompressible() {
-		compression := make(map[string]uint16) // Compression pointer mappings.
-		return dns.packBufferWithCompressionMap(buf, compressionMap{int: compression}, true)
+		compression := compressionPackPool.Get().(map[string]uint16)
+		msg, err := dns.packBufferWithCompressionMap(buf, compressionMap{int: compression}, true)
+		for k := range compression {
+			delete(compression, k)
+		}
+		compressionPackPool.Put(compression)
+		return msg, err
 	}
 
 	return dns.packBufferWithCompressionMap(buf, compressionMap{}, false)
@@ -972,6 +983,12 @@ func (dns *Msg) isCompressible() bool {
 		len(dns.Ns) > 0 || len(dns.Extra) > 0
 }
 
+var compressionPool = sync.Pool{
+	New: func() interface{} {
+		return make(map[string]struct{})
+	},
+}
+
 // Len returns the message length when in (un)compressed wire format.
 // If dns.Compress is true compression it is taken into account. Len()
 // is provided to be a faster way to get the size of the resulting packet,
@@ -980,8 +997,13 @@ func (dns *Msg) Len() int {
 	// If this message can't be compressed, avoid filling the
 	// compression map and creating garbage.
 	if dns.Compress && dns.isCompressible() {
-		compression := make(map[string]struct{})
-		return msgLenWithCompressionMap(dns, compression)
+		compression := compressionPool.Get().(map[string]struct{})
+		n := msgLenWithCompressionMap(dns, compression)
+		for k := range compression {
+			delete(compression, k)
+		}
+		compressionPool.Put(compression)
+		return n
 	}
 
 	return msgLenWithCompressionMap(dns, nil)
diff --git a/msg_truncate.go b/msg_truncate.go
index 89d40757..a711f006 100644
--- a/msg_truncate.go
+++ b/msg_truncate.go
@@ -54,7 +54,7 @@ func (dns *Msg) Truncate(size int) {
 		size -= Len(edns0)
 	}
 
-	compression := make(map[string]struct{})
+	compression := compressionPool.Get().(map[string]struct{})
 
 	l = headerSize
 	for _, r := range dns.Question {
@@ -88,6 +88,11 @@ func (dns *Msg) Truncate(size int) {
 		// Add the OPT record back onto the additional section.
 		dns.Extra = append(dns.Extra, edns0)
 	}
+
+	for k := range compression {
+		delete(compression, k)
+	}
+	compressionPool.Put(compression)
 }
 
 func truncateLoop(rrs []RR, size, l int, compression map[string]struct{}) (int, int) {