From b2a22eff04fbfd9801d865f8a7702d6f62dfac14 Mon Sep 17 00:00:00 2001
From: sschepens <sebastian.schepens@mercadolibre.com>
Date: Fri, 15 Jan 2021 09:32:49 -0300
Subject: [PATCH] Prevent race from prefetching (#4368)

Automatically submitted.
---
 plugin/cache/handler.go | 36 ++++++++++++++++--------------------
 1 file changed, 16 insertions(+), 20 deletions(-)

diff --git a/plugin/cache/handler.go b/plugin/cache/handler.go
index 350ec0bb0..2dd3c0646 100644
--- a/plugin/cache/handler.go
+++ b/plugin/cache/handler.go
@@ -39,39 +39,28 @@ func (c *Cache) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg)
 		ttl = i.ttl(now)
 	}
 	if i == nil {
-		if !do {
-			setDo(rc)
-		}
 		crr := &ResponseWriter{ResponseWriter: w, Cache: c, state: state, server: server, do: do}
-		return plugin.NextOrFailure(c.Name(), c.Next, ctx, crr, rc)
+		return c.doRefresh(ctx, state, crr)
 	}
 	if ttl < 0 {
 		servedStale.WithLabelValues(server).Inc()
 		// Adjust the time to get a 0 TTL in the reply built from a stale item.
 		now = now.Add(time.Duration(ttl) * time.Second)
-		addr := w.LocalAddr() // See https://github.com/coredns/coredns/issues/4271, unclear how, but pull this out of the goroutine, and get the address here.
-		go func() {
-			if !do {
-				setDo(rc)
-			}
-			crr := &ResponseWriter{Cache: c, state: state, server: server, prefetch: true, remoteAddr: addr, do: do}
-			plugin.NextOrFailure(c.Name(), c.Next, ctx, crr, rc)
-		}()
+		cw := newPrefetchResponseWriter(server, state, c)
+		go c.doPrefetch(ctx, state, cw, i, now)
+	} else if c.shouldPrefetch(i, now) {
+		cw := newPrefetchResponseWriter(server, state, c)
+		go c.doPrefetch(ctx, state, cw, i, now)
 	}
 	resp := i.toMsg(r, now, do)
 	w.WriteMsg(resp)
 
-	if c.shouldPrefetch(i, now) {
-		go c.doPrefetch(ctx, state, server, i, now)
-	}
 	return dns.RcodeSuccess, nil
 }
 
-func (c *Cache) doPrefetch(ctx context.Context, state request.Request, server string, i *item, now time.Time) {
-	cw := newPrefetchResponseWriter(server, state, c)
-
-	cachePrefetches.WithLabelValues(server).Inc()
-	plugin.NextOrFailure(c.Name(), c.Next, ctx, cw, state.Req)
+func (c *Cache) doPrefetch(ctx context.Context, state request.Request, cw *ResponseWriter, i *item, now time.Time) {
+	cachePrefetches.WithLabelValues(cw.server).Inc()
+	c.doRefresh(ctx, state, cw)
 
 	// When prefetching we loose the item i, and with it the frequency
 	// that we've gathered sofar. See we copy the frequencies info back
@@ -81,6 +70,13 @@ func (c *Cache) doPrefetch(ctx context.Context, state request.Request, server st
 	}
 }
 
+func (c *Cache) doRefresh(ctx context.Context, state request.Request, cw *ResponseWriter) (int, error) {
+	if !state.Do() {
+		setDo(state.Req)
+	}
+	return plugin.NextOrFailure(c.Name(), c.Next, ctx, cw, state.Req)
+}
+
 func (c *Cache) shouldPrefetch(i *item, now time.Time) bool {
 	if c.prefetch <= 0 {
 		return false