refactor: generic rwmutex cache for ip and txt

cache_test.go

@@ -0,0 +1,141 @@
+package doh
+
+import (
+	"slices"
+	"sync"
+	"testing"
+	"testing/synctest"
+	"time"
+)
+
+// TestCacheExpiry checks that a fresh entry is a hit, an expired entry is a
+// miss, and reading an expired entry deletes it from the map. The fake clock
+// from synctest advances past the TTL without a real sleep.
+func TestCacheExpiry(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		c := newCache[int]()
+		c.set("k", 42, time.Second)
+
+		if v, ok := c.get("k"); !ok || v != 42 {
+			t.Fatalf("expected fresh hit (42, true), got (%d, %t)", v, ok)
+		}
+
+		time.Sleep(time.Second + time.Millisecond)
+
+		if v, ok := c.get("k"); ok || v != 0 {
+			t.Fatalf("expected miss after expiry (0, false), got (%d, %t)", v, ok)
+		}
+		if len(c.entries) != 0 {
+			t.Fatalf("expected expired entry to be deleted, %d remain", len(c.entries))
+		}
+	})
+}
+
+// TestCacheZeroTTL checks that a zero TTL is a no-op, so a disabled cache stays
+// empty.
+func TestCacheZeroTTL(t *testing.T) {
+	c := newCache[int]()
+	c.set("k", 42, 0)
+
+	// Check the map before any get. A get would report a miss for an entry that
+	// was stored and then evicted on read, so only the map state right after set
+	// proves that set stored nothing.
+	if len(c.entries) != 0 {
+		t.Fatalf("expected zero-TTL set to store nothing, got %d entries", len(c.entries))
+	}
+	if _, ok := c.get("k"); ok {
+		t.Fatal("expected zero-TTL set to be a no-op")
+	}
+}
+
+// TestCacheRefreshNotClobbered checks the guarantee that justifies the double-
+// checked locking in get: when a set refreshes an entry in the window after get
+// drops the read lock on an expired entry but before it takes the write lock,
+// get must return the refreshed value rather than delete it. The afterExpiredRead
+// seam drives the set into that exact window, so the result does not depend on
+// goroutine scheduling. Remove the re-check in get and this test fails.
+func TestCacheRefreshNotClobbered(t *testing.T) {
+	c := newCache[int]()
+	// Seed an already-expired entry so get's first read sees it as expired and
+	// heads for the write-locked delete. set cannot do this: it always stores a
+	// future expiry.
+	c.entries["k"] = cacheEntry[int]{val: 1, expire: time.Now().Add(-time.Second)}
+
+	c.afterExpiredRead = func() {
+		c.afterExpiredRead = nil // refresh exactly once
+		c.set("k", 2, time.Minute)
+	}
+
+	if v, ok := c.get("k"); !ok || v != 2 {
+		t.Fatalf("expected refreshed hit (2, true), got (%d, %t)", v, ok)
+	}
+	if len(c.entries) != 1 {
+		t.Fatalf("expected refreshed entry to survive, %d entries", len(c.entries))
+	}
+}
+
+// TestCacheDeletedDuringExpiredGet checks the other re-check branch: when
+// another get deletes the expired entry in the same window, the re-checking get
+// finds it gone and reports a clean miss instead of touching the map again.
+func TestCacheDeletedDuringExpiredGet(t *testing.T) {
+	c := newCache[int]()
+	c.entries["k"] = cacheEntry[int]{val: 1, expire: time.Now().Add(-time.Second)}
+
+	c.afterExpiredRead = func() {
+		c.afterExpiredRead = nil
+		c.mx.Lock()
+		delete(c.entries, "k")
+		c.mx.Unlock()
+	}
+
+	if v, ok := c.get("k"); ok || v != 0 {
+		t.Fatalf("expected miss after concurrent delete (0, false), got (%d, %t)", v, ok)
+	}
+	if len(c.entries) != 0 {
+		t.Fatalf("expected empty cache, %d entries", len(c.entries))
+	}
+}
+
+// TestCacheConcurrent hammers the cache from many goroutines with near-instant
+// TTLs so reads constantly hit the expired-entry delete path while writes and
+// other deletes run concurrently. Run with -race: it guards against the data
+// race that map deletion under a read lock would cause. The assertions add a
+// second check, that no read ever returns a value that was not stored and that
+// the map never holds a stray key or value.
+func TestCacheConcurrent(t *testing.T) {
+	c := newCache[int]()
+	keys := []string{"a", "b", "c", "d"}
+
+	const workers = 64
+	const iters = 1000
+
+	var wg sync.WaitGroup
+	for w := range workers {
+		wg.Go(func() {
+			for i := range iters {
+				k := keys[(w+i)%len(keys)]
+				if i%2 == 0 {
+					c.set(k, i, time.Nanosecond)
+				} else if v, ok := c.get(k); ok && (v < 0 || v >= iters) {
+					// A hit must return an iteration index some set call stored.
+					t.Errorf("get(%q) returned out-of-range value %d", k, v)
+				}
+			}
+		})
+	}
+	wg.Wait()
+
+	c.mx.RLock()
+	defer c.mx.RUnlock()
+	if len(c.entries) > len(keys) {
+		t.Errorf("cache holds %d entries, want at most %d", len(c.entries), len(keys))
+	}
+	for k, e := range c.entries {
+		if !slices.Contains(keys, k) {
+			t.Errorf("cache holds unexpected key %q", k)
+		}
+		if e.val < 0 || e.val >= iters {
+			t.Errorf("cache holds out-of-range value %d for key %q", e.val, k)
+		}
+	}
+}

go.mod

@@ -1,6 +1,6 @@
 module github.com/libp2p/go-doh-resolver
 
-go 1.24
+go 1.25
 
 require (
 	github.com/ipfs/go-log/v2 v2.5.1

resolver.go

@@ -15,23 +15,13 @@ import (
 )
 
 type Resolver struct {
-	mx  sync.Mutex
 	url string
 
-	// RR cache
-	ipCache     map[string]ipAddrEntry
-	txtCache    map[string]txtEntry
-	maxCacheTTL time.Duration
-}
-
-type ipAddrEntry struct {
-	ips    []net.IPAddr
-	expire time.Time
-}
+	// RR caches keyed by FQDN.
+	ipCache  *cache[[]net.IPAddr]
+	txtCache *cache[[]string]
 
-type txtEntry struct {
-	txt    []string
-	expire time.Time
+	maxCacheTTL time.Duration
 }
 
 type Option func(*Resolver) error
@@ -66,8 +56,8 @@ func NewResolver(url string, opts ...Option) (*Resolver, error) {
 
 	r := &Resolver{
 		url:         url,
-		ipCache:     make(map[string]ipAddrEntry),
-		txtCache:    make(map[string]txtEntry),
+		ipCache:     newCache[[]net.IPAddr](),
+		txtCache:    newCache[[]string](),
 		maxCacheTTL: time.Duration(math.MaxUint32) * time.Second,
 	}
 
@@ -106,7 +96,7 @@ func (r *Resolver) LookupIPAddr(ctx context.Context, domain string) (result []ne
 	}()
 
 	var ttl uint32
-	for i := 0; i < 2; i++ {
+	for range 2 {
 		r := <-resch
 		if r.err != nil {
 			return nil, r.err
@@ -139,64 +129,93 @@ func (r *Resolver) LookupTXT(ctx context.Context, domain string) ([]string, erro
 	return result, nil
 }
 
-func (r *Resolver) getCachedIPAddr(domain string) ([]net.IPAddr, bool) {
-	r.mx.Lock()
-	defer r.mx.Unlock()
+// cacheEntry is a cached value and the time it expires.
+type cacheEntry[V any] struct {
+	val    V
+	expire time.Time
+}
 
-	fqdn := dns.Fqdn(domain)
-	entry, ok := r.ipCache[fqdn]
+// cache is a TTL cache keyed by string, safe for concurrent use. A read that
+// hits a fresh entry takes only the read lock, so concurrent reads run in
+// parallel. Deleting an expired entry needs the write lock. Go cannot upgrade a
+// read lock to a write lock in place, so get drops the read lock, takes the
+// write lock, and re-checks the entry before it deletes.
+type cache[V any] struct {
+	mx      sync.RWMutex
+	entries map[string]cacheEntry[V]
+
+	// afterExpiredRead is a test seam. When non-nil, get calls it after dropping
+	// the read lock on an expired entry and before taking the write lock, so a
+	// test can drive a concurrent set or delete into that window deterministically.
+	// It is always nil in normal use.
+	afterExpiredRead func()
+}
+
+func newCache[V any]() *cache[V] {
+	return &cache[V]{entries: make(map[string]cacheEntry[V])}
+}
+
+// get returns the value stored under key, or the zero value and ok=false when
+// the key is absent or expired. It deletes an expired entry before returning.
+func (c *cache[V]) get(key string) (V, bool) {
+	c.mx.RLock()
+	entry, ok := c.entries[key]
+	c.mx.RUnlock()
+
+	var zero V
 	if !ok {
-		return nil, false
+		return zero, false
+	}
+	if !time.Now().After(entry.expire) {
+		return entry.val, true
 	}
 
-	if time.Now().After(entry.expire) {
-		delete(r.ipCache, fqdn)
-		return nil, false
+	// The entry is expired. Re-check it under the write lock before deleting: in
+	// the gap between dropping the read lock and taking the write lock, a
+	// concurrent set may have refreshed it, or another get may have deleted it.
+	if c.afterExpiredRead != nil {
+		c.afterExpiredRead()
 	}
+	c.mx.Lock()
+	defer c.mx.Unlock()
 
-	return entry.ips, true
+	entry, ok = c.entries[key]
+	if !ok {
+		return zero, false
+	}
+	if time.Now().After(entry.expire) {
+		delete(c.entries, key)
+		return zero, false
+	}
+	return entry.val, true
 }
 
-func (r *Resolver) cacheIPAddr(domain string, ips []net.IPAddr, ttl time.Duration) {
+// set stores val under key for the given TTL. A zero TTL stores nothing, so a
+// disabled cache stays empty.
+func (c *cache[V]) set(key string, val V, ttl time.Duration) {
 	if ttl == 0 {
 		return
 	}
 
-	r.mx.Lock()
-	defer r.mx.Unlock()
-
-	fqdn := dns.Fqdn(domain)
-	r.ipCache[fqdn] = ipAddrEntry{ips, time.Now().Add(ttl)}
+	c.mx.Lock()
+	defer c.mx.Unlock()
+	c.entries[key] = cacheEntry[V]{val: val, expire: time.Now().Add(ttl)}
 }
 
-func (r *Resolver) getCachedTXT(domain string) ([]string, bool) {
-	r.mx.Lock()
-	defer r.mx.Unlock()
-
-	fqdn := dns.Fqdn(domain)
-	entry, ok := r.txtCache[fqdn]
-	if !ok {
-		return nil, false
-	}
+func (r *Resolver) getCachedIPAddr(domain string) ([]net.IPAddr, bool) {
+	return r.ipCache.get(dns.Fqdn(domain))
+}
 
-	if time.Now().After(entry.expire) {
-		delete(r.txtCache, fqdn)
-		return nil, false
-	}
+func (r *Resolver) cacheIPAddr(domain string, ips []net.IPAddr, ttl time.Duration) {
+	r.ipCache.set(dns.Fqdn(domain), ips, ttl)
+}
 
-	return entry.txt, true
+func (r *Resolver) getCachedTXT(domain string) ([]string, bool) {
+	return r.txtCache.get(dns.Fqdn(domain))
 }
 
 func (r *Resolver) cacheTXT(domain string, txt []string, ttl time.Duration) {
-	if ttl == 0 {
-		return
-	}
-
-	r.mx.Lock()
-	defer r.mx.Unlock()
-
-	fqdn := dns.Fqdn(domain)
-	r.txtCache[fqdn] = txtEntry{txt, time.Now().Add(ttl)}
+	r.txtCache.set(dns.Fqdn(domain), txt, ttl)
 }
 
 func minTTL(a, b time.Duration) time.Duration {