fancy-dns/internal/engine/coalesce_test.go

package engine

import (
	"context"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/miekg/dns"

	"git.k8s-git-k6n.net/mats/fancy-dns/internal/config"
)

func waitInflight(t *testing.T, h *harness, n int64) {
	t.Helper()
	deadline := time.Now().Add(5 * time.Second)
	for h.eng.Inflight() < n {
		if time.Now().After(deadline) {
			t.Fatalf("only %d of %d queries reached the engine", h.eng.Inflight(), n)
		}
		time.Sleep(2 * time.Millisecond)
	}
}

func TestIdenticalMissesShareOneUpstreamExchange(t *testing.T) {
	h := newHarness(t, hopt{})
	release := h.up.hold()
	const n = 25
	results := make([]Response, n)
	var wg sync.WaitGroup
	for i := 0; i < n; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			results[i] = h.ask("burst.test.", dns.TypeA)
		}()
	}
	h.up.waitArrived(t)
	waitInflight(t, h, n)
	release()
	wg.Wait()

	if got := h.up.hitCount("burst.test."); got != 1 {
		t.Fatalf("upstream saw %d queries for %d identical misses, want 1", got, n)
	}
	seen := map[*dns.Msg]bool{}
	for i, r := range results {
		requireRcode(t, r, dns.RcodeSuccess)
		if len(r.Msg.Answer) != 1 {
			t.Fatalf("client %d got %v", i, r.Msg)
		}
		if seen[r.Msg] {
			t.Fatal("two clients were handed the same message; one's edits would reach the other")
		}
		seen[r.Msg] = true
	}
	// The answer was cached once, by the leader.
	h.ask("burst.test.", dns.TypeA)
	if got := h.up.hitCount("burst.test."); got != 1 {
		t.Fatalf("a later query went upstream again: %d", got)
	}
}

func TestDifferentQuestionsDoNotShare(t *testing.T) {
	h := newHarness(t, hopt{})
	release := h.up.hold()
	var wg sync.WaitGroup
	for _, qt := range []uint16{dns.TypeA, dns.TypeAAAA} {
		for _, name := range []string{"one.test.", "two.test."} {
			wg.Add(1)
			go func() { defer wg.Done(); h.ask(name, qt) }()
		}
	}
	waitInflight(t, h, 4)
	release()
	wg.Wait()
	if got := h.up.totalHits(); got != 4 {
		t.Fatalf("upstream saw %d queries, want 4 (distinct names and types must not share)", got)
	}
}

// TestWaitersSurviveTheLeaderLeaving: the client whose query started the
// exchange hanging up says nothing about the upstream, so the others must not
// be failed with its cancellation.
func TestWaitersSurviveTheLeaderLeaving(t *testing.T) {
	h := newHarness(t, hopt{})
	release := h.up.hold()
	leaderCtx, cancelLeader := context.WithCancel(context.Background())
	done := make(chan Response, 1)
	go func() { done <- h.eng.Handle(leaderCtx, h.request("leave.test.", dns.TypeA)) }()
	h.up.waitArrived(t)

	var follower atomic.Pointer[Response]
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		r := h.ask("leave.test.", dns.TypeA)
		follower.Store(&r)
	}()
	waitInflight(t, h, 2)
	cancelLeader()
	<-done
	release()
	wg.Wait()
	requireRcode(t, *follower.Load(), dns.RcodeSuccess)
}

func TestFlightGroupCopiesAndIsolatesKeys(t *testing.T) {
	var g flightGroup
	var calls atomic.Int32
	gate := make(chan struct{})
	mk := func() (*dns.Msg, string, error) {
		calls.Add(1)
		<-gate
		m := new(dns.Msg)
		m.SetQuestion("x.test.", dns.TypeA)
		return m, "up", nil
	}
	k := flightKey{name: "x.test.", qtype: dns.TypeA}
	msgs := make(chan *dns.Msg, 3)
	for i := 0; i < 3; i++ {
		go func() {
			m, _, _, _ := g.do(context.Background(), k, mk)
			msgs <- m
		}()
	}
	time.Sleep(50 * time.Millisecond)
	close(gate)
	a, b, c := <-msgs, <-msgs, <-msgs
	if calls.Load() != 1 {
		t.Fatalf("fn ran %d times", calls.Load())
	}
	if a == b || b == c || a == c {
		t.Fatal("callers share a message")
	}
	// A different key (another subnet, say) is its own exchange.
	k2 := k
	k2.ecs = "192.0.2.0/24/0"
	if _, _, shared, _ := g.do(context.Background(), k2, func() (*dns.Msg, string, error) { return new(dns.Msg), "up", nil }); shared {
		t.Fatal("a different ECS shared an exchange")
	}
}

func TestNonQueryOpcodesAreNotImplemented(t *testing.T) {
	h := newHarness(t, hopt{})
	for name, mut := range map[string]func(*dns.Msg){
		"update": func(m *dns.Msg) { m.Opcode = dns.OpcodeUpdate },
		"notify": func(m *dns.Msg) { m.Opcode = dns.OpcodeNotify },
		"status": func(m *dns.Msg) { m.Opcode = dns.OpcodeStatus; m.Question = nil },
	} {
		t.Run(name, func(t *testing.T) {
			req := h.request("zone.test.", dns.TypeSOA)
			mut(req.Msg)
			res := h.askWith(req)
			requireRcode(t, res, dns.RcodeNotImplemented)
			if n := h.up.totalHits(); n != 0 {
				t.Fatalf("%d queries reached the upstream", n)
			}
		})
	}
}

// TestPopularEntryIsRefreshedBeforeItExpires: a name that is being asked for
// is replaced in the cache before its TTL runs out, so no client is the one
// that finds it expired.
func TestPopularEntryIsRefreshedBeforeItExpires(t *testing.T) {
	h := newHarness(t, hopt{
		clock:      true,
		serveStale: time.Minute,
		cacheCfg:   func(c *config.Cache) { c.Prefetch = true },
	})
	h.ask("hot.test.", dns.TypeA) // miss, stored with a 60s TTL
	h.ask("hot.test.", dns.TypeA)
	h.ask("hot.test.", dns.TypeA)
	if got := h.up.hitCount("hot.test."); got != 1 {
		t.Fatalf("setup: %d upstream queries", got)
	}

	h.clock.Advance(30 * time.Second) // half way: nothing yet
	h.ask("hot.test.", dns.TypeA)
	time.Sleep(50 * time.Millisecond)
	if got := h.up.hitCount("hot.test."); got != 1 {
		t.Fatalf("refreshed too early: %d upstream queries", got)
	}

	h.clock.Advance(25 * time.Second) // 5s left of 60: under a tenth
	res := h.ask("hot.test.", dns.TypeA)
	if !res.CacheHit || res.CacheKind != "hit" {
		t.Fatalf("the query that triggers the refresh must still be served fresh: %+v", res)
	}
	deadline := time.Now().Add(5 * time.Second)
	for h.up.hitCount("hot.test.") < 2 {
		if time.Now().After(deadline) {
			t.Fatal("the entry was not refreshed ahead of its expiry")
		}
		time.Sleep(5 * time.Millisecond)
	}
}

// TestUnpopularEntryIsNotRefreshed: an entry nobody came back for is allowed
// to expire; refreshing everything would turn the cache into a poller.
func TestUnpopularEntryIsNotRefreshed(t *testing.T) {
	h := newHarness(t, hopt{
		clock:      true,
		serveStale: time.Minute,
		cacheCfg:   func(c *config.Cache) { c.Prefetch = true },
	})
	h.ask("once.test.", dns.TypeA)
	h.clock.Advance(57 * time.Second)
	h.ask("once.test.", dns.TypeA) // a single hit, near expiry
	time.Sleep(100 * time.Millisecond)
	if got := h.up.hitCount("once.test."); got != 1 {
		t.Fatalf("a once-asked name was refreshed: %d upstream queries", got)
	}
}