Hold lock during entire prefetch to preserve order

2 files changed, 5 insertions, 5 deletions

diff --git a/cache/cache.go b/cache/cache.go
index f21af5f..5d41bbe 100644
--- a/cache/cache.go
+++ b/cache/cache.go
@@ -152,6 +152,8 @@ func (c *Cache) Reset() {
 func (c *Cache) prefetch() bool { return c.client != nil }
 
 func (c *Cache) refresh(key uint64, old *dns.Msg) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
 	q := old.Question[0]
 	msg := dns.Msg{}
 	msg.SetQuestion(q.Name, q.Qtype)
@@ -159,8 +161,6 @@ func (c *Cache) refresh(key uint64, old *dns.Msg) {
 	if err != nil {
 		return // Retry on next request
 	}
-	c.mu.Lock()
-	defer c.mu.Unlock()
 	if !c.set(key, r) {
 		c.evict(key)
 	}
diff --git a/cache/cache_test.go b/cache/cache_test.go
index 4f4aa90..7879d65 100644
--- a/cache/cache_test.go
+++ b/cache/cache_test.go
@@ -336,15 +336,15 @@ func TestCacheEvictAndUpdate(t *testing.T) {
 	copy := msg.Copy()
 	copy.Answer[0].(*dns.A).Hdr.Ttl = 0
 	exchanger.setAnswer(copy)
+	copy = msg.Copy()
+	copy.Answer[0].(*dns.A).Hdr.Ttl = 30
+	exchanger.setAnswer(copy)
 
 	// Advance time so that msg is now considered expired. Query to trigger prefetch
 	c.now = func() time.Time { return now.Add(61 * time.Second) }
 	c.Get(key)
 
 	// Query again, causing another prefetch with a non-zero TTL
-	copy = msg.Copy()
-	copy.Answer[0].(*dns.A).Hdr.Ttl = 30
-	exchanger.setAnswer(copy)
 	c.Get(key)
 
 	// Last query refreshes key