diff --git a/cache_test.go b/cache_test.go new file mode 100644 index 0000000..f000c5b --- /dev/null +++ b/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) + } + } +} diff --git a/go.mod b/go.mod index f23d65b..ef3cdfc 100644 --- a/go.mod +++ b/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 diff --git a/resolver.go b/resolver.go index a9bfe36..106fee4 100644 --- a/resolver.go +++ b/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 {