diff --git a/client/internal/dns/mgmt/export_test.go b/client/internal/dns/mgmt/export_test.go new file mode 100644 index 000000000..676b777ba --- /dev/null +++ b/client/internal/dns/mgmt/export_test.go @@ -0,0 +1,25 @@ +package mgmt + +import "time" + +// pendingCount returns the number of domains whose initial resolve is still in +// flight. Test-only: lets tests wait for background resolves kicked off by +// UpdateFromServerDomains. +func (m *Resolver) pendingCount() int { + m.mutex.RLock() + defer m.mutex.RUnlock() + return len(m.pending) +} + +// waitForPendingResolves blocks until all background initial resolves have +// settled, or the timeout elapses. Returns true if all settled. Test-only. +func (m *Resolver) waitForPendingResolves(timeout time.Duration) bool { + deadline := time.Now().Add(timeout) + for m.pendingCount() > 0 { + if time.Now().After(deadline) { + return false + } + time.Sleep(10 * time.Millisecond) + } + return true +} diff --git a/client/internal/dns/mgmt/mgmt.go b/client/internal/dns/mgmt/mgmt.go index 921664601..6964a9bcb 100644 --- a/client/internal/dns/mgmt/mgmt.go +++ b/client/internal/dns/mgmt/mgmt.go @@ -619,30 +619,6 @@ func (m *Resolver) clearPending(dnsName string) { m.mutex.Unlock() } -// pendingCount returns the number of domains whose initial resolve is still in -// flight. Primarily useful for tests that need to wait for background resolves -// kicked off by UpdateFromServerDomains. -func (m *Resolver) pendingCount() int { - m.mutex.RLock() - defer m.mutex.RUnlock() - return len(m.pending) -} - -// WaitForPendingResolves blocks until all background initial resolves kicked -// off by UpdateFromServerDomains have settled, or the timeout elapses. It -// returns true if all settled. Intended for tests and shutdown paths that need -// the cache populated synchronously. -func (m *Resolver) WaitForPendingResolves(timeout time.Duration) bool { - deadline := time.Now().Add(timeout) - for m.pendingCount() > 0 { - if time.Now().After(deadline) { - return false - } - time.Sleep(10 * time.Millisecond) - } - return true -} - // awaitPendingResolve waits for an in-flight initial resolve of dnsName to // finish, bounded by dnsTimeout. It joins the existing resolveGroup flight so // concurrent ServeDNS waiters share one resolution. Returns true if a record diff --git a/client/internal/dns/mgmt/mgmt_test.go b/client/internal/dns/mgmt/mgmt_test.go index 95c74a4a9..921666664 100644 --- a/client/internal/dns/mgmt/mgmt_test.go +++ b/client/internal/dns/mgmt/mgmt_test.go @@ -329,7 +329,7 @@ func TestResolver_ManagementDomainProtection(t *testing.T) { if err != nil { t.Logf("Server domains update failed: %v", err) } - resolver.WaitForPendingResolves(10 * time.Second) + resolver.waitForPendingResolves(10 * time.Second) finalDomains := resolver.GetCachedDomains() @@ -367,7 +367,7 @@ func TestResolver_EmptyUpdateDoesNotRemoveDomains(t *testing.T) { if err != nil { t.Skipf("Skipping test due to DNS resolution failure: %v", err) } - resolver.WaitForPendingResolves(10 * time.Second) + resolver.waitForPendingResolves(10 * time.Second) // Verify domains were added cachedDomains := resolver.GetCachedDomains() @@ -402,7 +402,7 @@ func TestResolver_PartialUpdateReplacesOnlyUpdatedTypes(t *testing.T) { if err != nil { t.Skipf("Skipping test due to DNS resolution failure: %v", err) } - resolver.WaitForPendingResolves(10 * time.Second) + resolver.waitForPendingResolves(10 * time.Second) assert.Len(t, resolver.GetCachedDomains(), 3) // Update with partial ServerDomains (only signal domain - this should replace signal but preserve stun/turn) @@ -413,7 +413,7 @@ func TestResolver_PartialUpdateReplacesOnlyUpdatedTypes(t *testing.T) { if err != nil { t.Skipf("Skipping test due to DNS resolution failure: %v", err) } - resolver.WaitForPendingResolves(10 * time.Second) + resolver.waitForPendingResolves(10 * time.Second) // Should remove only the old signal domain assert.Len(t, removedDomains, 1, "Should remove only the old signal domain") @@ -448,7 +448,7 @@ func TestResolver_PartialUpdateAddsNewTypePreservesExisting(t *testing.T) { if err != nil { t.Skipf("Skipping test due to DNS resolution failure: %v", err) } - resolver.WaitForPendingResolves(10 * time.Second) + resolver.waitForPendingResolves(10 * time.Second) assert.Len(t, resolver.GetCachedDomains(), 3) // Update with partial ServerDomains (only flow domain - flow is intentionally excluded from @@ -460,7 +460,7 @@ func TestResolver_PartialUpdateAddsNewTypePreservesExisting(t *testing.T) { if err != nil { t.Skipf("Skipping test due to DNS resolution failure: %v", err) } - resolver.WaitForPendingResolves(10 * time.Second) + resolver.waitForPendingResolves(10 * time.Second) assert.Len(t, removedDomains, 0, "Should not remove any domains when only flow domain is provided")