From 9778be760bd07e632de4a07ee0643a40b4dc419d Mon Sep 17 00:00:00 2001
From: Alfonso Acosta <fons@weave.works>
Date: Fri, 23 Jun 2017 09:18:51 +0000
Subject: [PATCH] Use a global lock instead

---
 probe/endpoint/dns_snooper_linux_amd64.go | 36 ++++++++---------------
 1 file changed, 12 insertions(+), 24 deletions(-)

diff --git a/probe/endpoint/dns_snooper_linux_amd64.go b/probe/endpoint/dns_snooper_linux_amd64.go
index 0aaa1113d..3963b60bd 100644
--- a/probe/endpoint/dns_snooper_linux_amd64.go
+++ b/probe/endpoint/dns_snooper_linux_amd64.go
@@ -24,20 +24,15 @@ const (
 
 // DNSSnooper is a snopper of DNS queries
 type DNSSnooper struct {
-	stop                chan struct{}
-	pcapHandle          *pcap.Handle
+	stop       chan struct{}
+	pcapHandle *pcap.Handle
+	// gcache is goroutine-safe, but the
+	// values cached values aren't
+	reverseDNSMutex     sync.RWMutex
 	reverseDNSCache     gcache.Cache
 	decodingErrorCounts map[string]uint64 // for limiting
 }
 
-type domainSet struct {
-	// Not worth using a RMutex since we don't expect a lot of contention
-	// and they are considerably larger (8 bytes vs 24 bytes), which would
-	// bloat the cache
-	mutex   *sync.Mutex
-	domains map[string]struct{}
-}
-
 // NewDNSSnooper creates a new snooper of DNS queries
 func NewDNSSnooper() (*DNSSnooper, error) {
 	pcapHandle, err := newPcapHandle()
@@ -110,13 +105,11 @@ func (s *DNSSnooper) CachedNamesForIP(ip string) []string {
 	if err != nil {
 		return result
 	}
-
-	domainSet := domains.(domainSet)
-	domainSet.mutex.Lock()
-	for domain := range domainSet.domains {
+	s.reverseDNSMutex.RLock()
+	for domain := range domains.(map[string]struct{}) {
 		result = append(result, domain)
 	}
-	domainSet.mutex.Unlock()
+	s.reverseDNSMutex.RUnlock()
 
 	return result
 }
@@ -281,17 +274,12 @@ func (s *DNSSnooper) processDNSMessage(dns *layers.DNS) {
 	log.Debugf("DNSSnooper: caught DNS lookup: %s -> %v", newDomain, ips)
 	for ip := range ips {
 		if existingDomains, err := s.reverseDNSCache.Get(ip); err != nil {
-			singleton := domainSet{
-				domains: map[string]struct{}{newDomain: {}},
-				mutex:   &sync.Mutex{},
-			}
-			s.reverseDNSCache.Set(ip, singleton)
+			s.reverseDNSCache.Set(ip, map[string]struct{}{newDomain: {}})
 		} else {
 			// TODO: Be smarter about the expiration of entries with pre-existing associated domains
-			domainSet := existingDomains.(domainSet)
-			domainSet.mutex.Lock()
-			domainSet.domains[newDomain] = struct{}{}
-			domainSet.mutex.Unlock()
+			s.reverseDNSMutex.Lock()
+			existingDomains.(map[string]struct{})[newDomain] = struct{}{}
+			s.reverseDNSMutex.Unlock()
 		}
 	}
 }