plugin/forward: add startup health checks with start_failed (#8663)

* plugin/forward: add startup health checks with start_failed

Signed-off-by: houyuwushang <180804215+houyuwushang@users.noreply.github.com>

* test: separate CoreDNS and third-party imports

Signed-off-by: houyuwushang <180804215+houyuwushang@users.noreply.github.com>

---------

Signed-off-by: houyuwushang <180804215+houyuwushang@users.noreply.github.com>

plugin/forward/README.md

@@ -57,6 +57,7 @@ forward FROM TO... {
     tls_servername NAME
     policy random|round_robin|sequential
     health_check DURATION [no_rec] [domain FQDN]
+    start_failed
     max_concurrent MAX
     next RCODE_1 [RCODE_2] [RCODE_3...]
     failfast_all_unhealthy_upstreams
@@ -77,6 +78,12 @@ forward FROM TO... {
 * `max_fails` is the number of subsequent failed health checks that are needed before considering
   an upstream to be down. If 0, the upstream will never be marked as down (nor health checked).
   Default is 2.
+* `start_failed` starts upstreams in an unhealthy state and immediately runs health checks using
+  the configured `health_check` settings. Each upstream becomes healthy after its first successful
+  check. This option is disabled by default and cannot be combined with `max_fails 0`.
+  Checks run asynchronously; startup does not wait for them. While all upstreams are unhealthy,
+  requests follow the usual all-unhealthy policy. Combine with `failfast_all_unhealthy_upstreams`
+  to return SERVFAIL until an upstream passes a check instead of trying a random upstream.
 * `max_connect_attempts` caps the total number of upstream connect attempts
   performed for a single incoming DNS request. The default cap is twice the number of
   configured upstreams, allowing two complete passes when all upstreams are healthy.

plugin/forward/forward.go

@@ -43,6 +43,7 @@ type Forward struct {
 	proxies     []*proxyPkg.Proxy
 	p           Policy
 	hcInterval  time.Duration
+	startFailed bool
 
 	from    string
 	ignored []string

plugin/forward/setup.go

@@ -75,6 +75,9 @@ func setup(c *caddy.Controller) error {
 func (f *Forward) OnStartup() (err error) {
 	for _, p := range f.proxies {
 		p.Start(f.hcInterval)
+		if f.startFailed {
+			p.StartFailed()
+		}
 	}
 	return nil
 }
@@ -158,6 +161,9 @@ func parseStanza(c *caddy.Controller) (*Forward, error) {
 	if f.maxAge > 0 && f.maxAge < f.expire {
 		return f, fmt.Errorf("max_age (%s) must not be less than expire (%s)", f.maxAge, f.expire)
 	}
+	if f.startFailed && f.maxfails == 0 {
+		return f, fmt.Errorf("start_failed requires health checking; max_fails must be greater than zero")
+	}
 
 	// Reject HTTPS upstreams that include a path, the doh implementation default to /dns-query path.
 	for _, addr := range to {
@@ -290,6 +296,11 @@ func parseBlock(c *caddy.Controller, f *Forward) error {
 			return err
 		}
 		f.maxfails = uint32(n)
+	case "start_failed":
+		if c.NextArg() {
+			return c.ArgErr()
+		}
+		f.startFailed = true
 	case "max_connect_attempts":
 		if !c.NextArg() {
 			return c.ArgErr()

plugin/forward/startup_health_test.go

@@ -0,0 +1,193 @@
+package forward
+
+import (
+	"fmt"
+	"net"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/coredns/caddy"
+	"github.com/coredns/coredns/plugin/pkg/dnstest"
+	"github.com/coredns/coredns/plugin/test"
+
+	"github.com/miekg/dns"
+)
+
+func TestStartFailedParse(t *testing.T) {
+	for _, tc := range []struct {
+		name, options string
+		wantErr       bool
+	}{
+		{"enabled", "start_failed", false},
+		{"extra argument", "start_failed true", true},
+		{"health disabled before", "max_fails 0\nstart_failed", true},
+		{"health disabled after", "start_failed\nmax_fails 0", true},
+		{"maximum threshold", "start_failed\nmax_fails 4294967295", false},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			_, err := parseForward(caddy.NewTestController("dns", "forward . 127.0.0.1 {\n"+tc.options+"\n}"))
+			if (err != nil) != tc.wantErr {
+				t.Fatalf("parse error = %v, want error %v", err, tc.wantErr)
+			}
+		})
+	}
+}
+
+func TestStartFailedHealthAndRouting(t *testing.T) {
+	for _, tc := range []struct {
+		name, options string
+		failfast      bool
+	}{
+		{"fail fast", "failfast_all_unhealthy_upstreams", true},
+		{"random fallback", "", false},
+		{"maximum threshold", "max_fails 4294967295\nfailfast_all_unhealthy_upstreams", true},
+		{"TCP checks", "force_tcp\nfailfast_all_unhealthy_upstreams", true},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			primary := newStartupUpstream(t)
+			backup := newStartupUpstream(t)
+			input := fmt.Sprintf("forward . %s %s {\nstart_failed\npolicy sequential\nhealth_check 10ms no_rec domain probe.example.\n%s\n}", primary.addr, backup.addr, tc.options)
+			fs, err := parseForward(caddy.NewTestController("dns", input))
+			if err != nil {
+				t.Fatal(err)
+			}
+			f := fs[0]
+			if err := f.OnStartup(); err != nil {
+				t.Fatal(err)
+			}
+			t.Cleanup(func() { _ = f.OnShutdown() })
+			for i, u := range []*startupUpstream{primary, backup} {
+				select {
+				case q := <-u.probes:
+					if q.Question[0].Name != "probe.example." || q.Question[0].Qtype != dns.TypeNS || q.RecursionDesired {
+						t.Fatalf("unexpected startup probe: %s", q)
+					}
+				case <-time.After(5 * time.Second):
+					t.Fatal("no startup probe received")
+				}
+				if !f.proxies[i].Down(f.maxfails) {
+					t.Fatal("upstream is healthy before its probe returns")
+				}
+			}
+			query := func() int {
+				t.Helper()
+				w := dnstest.NewRecorder(&test.ResponseWriter{})
+				rcode, _ := f.ServeDNS(t.Context(), w, new(dns.Msg).SetQuestion("example.org.", dns.TypeA))
+				return rcode
+			}
+			wantCode := dns.RcodeSuccess
+			if tc.failfast {
+				wantCode = dns.RcodeServerFailure
+			}
+			if got := query(); got != wantCode {
+				t.Fatalf("all-down rcode = %d, want %d", got, wantCode)
+			}
+			beforePrimary, beforeBackup := primary.queries.Load(), backup.queries.Load()
+			wantQueries := int32(1)
+			if tc.failfast {
+				wantQueries = 0
+			}
+			if beforePrimary+beforeBackup != wantQueries {
+				t.Fatal("all-down policy forwarded an unexpected number of queries")
+			}
+
+			backup.release()
+			waitStartupHealthy(t, func() bool { return !f.proxies[1].Down(f.maxfails) })
+			if got := query(); got != dns.RcodeSuccess {
+				t.Fatalf("backup rcode = %d", got)
+			}
+			if primary.queries.Load() != beforePrimary || backup.queries.Load() != beforeBackup+1 {
+				t.Fatal("query did not select the healthy backup")
+			}
+
+			primary.release()
+			waitStartupHealthy(t, func() bool { return !f.proxies[0].Down(f.maxfails) })
+			if got := query(); got != dns.RcodeSuccess {
+				t.Fatalf("primary rcode = %d", got)
+			}
+			if primary.queries.Load() != beforePrimary+1 {
+				t.Fatal("recovered primary was not selected")
+			}
+		})
+	}
+}
+
+func TestStartupHealthDefaultUnchanged(t *testing.T) {
+	u := newStartupUpstream(t)
+	fs, err := parseForward(caddy.NewTestController("dns", "forward . "+u.addr))
+	if err != nil {
+		t.Fatal(err)
+	}
+	f := fs[0]
+	if err := f.OnStartup(); err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _ = f.OnShutdown() })
+	if f.proxies[0].Down(f.maxfails) {
+		t.Fatal("default upstream started unhealthy")
+	}
+	w := dnstest.NewRecorder(&test.ResponseWriter{})
+	if rcode, err := f.ServeDNS(t.Context(), w, new(dns.Msg).SetQuestion("example.org.", dns.TypeA)); err != nil || rcode != dns.RcodeSuccess {
+		t.Fatalf("query = %d, %v", rcode, err)
+	}
+	select {
+	case <-u.probes:
+		t.Fatal("default startup initiated a probe")
+	default:
+	}
+}
+
+type startupUpstream struct {
+	addr    string
+	release func()
+	probes  chan *dns.Msg
+	queries atomic.Int32
+}
+
+func newStartupUpstream(t *testing.T) *startupUpstream {
+	t.Helper()
+	gate := make(chan struct{})
+	u := &startupUpstream{probes: make(chan *dns.Msg, 16), release: sync.OnceFunc(func() { close(gate) })}
+	s := dnstest.NewMultipleServer(func(w dns.ResponseWriter, r *dns.Msg) {
+		m := new(dns.Msg).SetReply(r)
+		if r.Question[0].Qtype == dns.TypeNS {
+			select {
+			case u.probes <- r.Copy():
+			default:
+			}
+			<-gate
+			// Even REFUSED proves that the endpoint is reachable.
+			m.Rcode = dns.RcodeRefused
+		} else {
+			u.queries.Add(1)
+			m.Answer = []dns.RR{test.A("example.org. 30 IN A 192.0.2.1")}
+		}
+		_ = w.WriteMsg(m)
+	})
+	_, port, err := net.SplitHostPort(s.Addr)
+	if err != nil {
+		s.Close()
+		t.Fatal(err)
+	}
+	u.addr = net.JoinHostPort("127.0.0.1", port)
+	t.Cleanup(s.Close)
+	t.Cleanup(u.release)
+	return u
+}
+
+func waitStartupHealthy(t *testing.T, healthy func() bool) {
+	t.Helper()
+	deadline := time.NewTimer(5 * time.Second)
+	defer deadline.Stop()
+	tick := time.NewTicker(time.Millisecond)
+	defer tick.Stop()
+	for !healthy() {
+		select {
+		case <-tick.C:
+		case <-deadline.C:
+			t.Fatal("health probe did not recover upstream")
+		}
+	}
+}

plugin/pkg/proxy/doq_test.go

@@ -521,6 +521,7 @@ func TestDoQHealthCheck(t *testing.T) {
 	})
 
 	p := NewProxy("TestDoQHealth", server.addr(), transport.QUIC)
+	p.initiallyDown.Store(true)
 	p.SetTLSConfig(clientTLS)
 	defer p.Stop()
 	hc := p.GetHealthchecker()
@@ -529,6 +530,9 @@ func TestDoQHealthCheck(t *testing.T) {
 	if err := hc.Check(p); err != nil {
 		t.Fatalf("health check failed: %v", err)
 	}
+	if p.Down(2) {
+		t.Error("successful QUIC health check did not clear startup failure")
+	}
 	msg := <-query
 	if len(msg.Question) != 1 || msg.Question[0].Name != "health.example." || msg.Question[0].Qtype != dns.TypeNS {
 		t.Fatalf("unexpected health query: %v", msg.Question)

plugin/pkg/proxy/health.go

@@ -6,7 +6,6 @@ import (
 	"errors"
 	"net"
 	"net/http"
-	"sync/atomic"
 	"time"
 
 	"github.com/coredns/coredns/plugin/pkg/doh"
@@ -130,7 +129,7 @@ func (h *doqHc) Check(p *Proxy) error {
 		return err
 	}
 
-	atomic.StoreUint32(&p.fails, 0)
+	p.resetFails()
 	return nil
 }
 
@@ -219,7 +218,7 @@ func (h *dnsHc) Check(p *Proxy) error {
 		return err
 	}
 
-	atomic.StoreUint32(&p.fails, 0)
+	p.resetFails()
 	return nil
 }
 
@@ -293,7 +292,7 @@ func (h *dohHc) Check(p *Proxy) error {
 		return err
 	}
 
-	atomic.StoreUint32(&p.fails, 0)
+	p.resetFails()
 	return nil
 }
 

plugin/pkg/proxy/health_test.go

@@ -33,11 +33,15 @@ func TestHealth(t *testing.T) {
 	hc.SetWriteTimeout(10 * time.Millisecond)
 
 	p := NewProxy("TestHealth", s.Addr, transport.DNS)
+	p.initiallyDown.Store(true)
 	p.readTimeout = 10 * time.Millisecond
 	err := hc.Check(p)
 	if err != nil {
 		t.Errorf("check failed: %v", err)
 	}
+	if p.Down(2) {
+		t.Error("successful DNS health check did not clear startup failure")
+	}
 
 	time.Sleep(20 * time.Millisecond)
 	i1 := atomic.LoadUint32(&i)
@@ -110,11 +114,15 @@ func TestHealthHTTPS(t *testing.T) {
 	hc.SetWriteTimeout(1 * time.Second)
 
 	p := NewProxy("TestHealthHTTPS", s.URL, transport.HTTPS)
+	p.initiallyDown.Store(true)
 	p.readTimeout = 10 * time.Millisecond
 	err := hc.Check(p)
 	if err != nil {
 		t.Fatalf("check failed: %v", err)
 	}
+	if p.Down(2) {
+		t.Error("successful HTTPS health check did not clear startup failure")
+	}
 
 	time.Sleep(20 * time.Millisecond)
 	i1 := atomic.LoadUint32(&i)

plugin/pkg/proxy/proxy.go

@@ -16,6 +16,7 @@ import (
 // Proxy defines an upstream host.
 type Proxy struct {
 	fails         uint32
+	initiallyDown atomic.Bool
 	addr          string
 	proxyName     string
 
@@ -136,14 +137,28 @@ func (p *Proxy) Healthcheck() {
 	})
 }
 
-// Down returns true if this proxy is down, i.e. has *more* fails than maxfails.
+// StartFailed marks the proxy down until its first successful health check and
+// starts probing immediately. Start must be called first to set the probe interval.
+func (p *Proxy) StartFailed() {
+	p.initiallyDown.Store(true)
+	p.Healthcheck()
+}
+
+// Down reports whether the proxy is awaiting its first successful startup check
+// or has more failures than maxfails. A zero maxfails disables both checks.
 func (p *Proxy) Down(maxfails uint32) bool {
 	if maxfails == 0 {
 		return false
 	}
 
 	fails := atomic.LoadUint32(&p.fails)
-	return fails > maxfails
+	return p.initiallyDown.Load() || fails > maxfails
+}
+
+// resetFails makes the proxy available after a successful health check.
+func (p *Proxy) resetFails() {
+	atomic.StoreUint32(&p.fails, 0)
+	p.initiallyDown.Store(false)
 }
 
 // Stop stops health checking and closes the DoQ transport, when configured.