plugin/timeouts: add maxtcpqueries option to bound queries per TCP/TLS connection (#8376)

core/dnsserver/config.go

@@ -127,6 +127,11 @@ type Config struct {
 	WriteTimeout time.Duration
 	IdleTimeout  time.Duration
 
+	// MaxTCPQueries defines the maximum number of queries served on a single TCP/TLS
+	// connection before it is closed. -1 means unlimited. This is nil if not specified,
+	// allowing for a default to be used.
+	MaxTCPQueries *int
+
 	// TSIG secrets, [name]key.
 	TsigSecret map[string]string
 

core/dnsserver/config_test.go

@@ -89,3 +89,16 @@ func TestAddPluginToAllServerBlocks(t *testing.T) {
 		t.Fatalf("second server block has %d plugins, want 1", got)
 	}
 }
+
+func TestPropagateConfigParamsMaxTCPQueries(t *testing.T) {
+	n := 128
+	first := &Config{MaxTCPQueries: &n}
+	first.firstConfigInBlock = first
+	second := &Config{firstConfigInBlock: first}
+
+	propagateConfigParams([]*Config{first, second})
+
+	if second.MaxTCPQueries == nil || *second.MaxTCPQueries != n {
+		t.Fatalf("expected MaxTCPQueries to propagate to second config as %d, got %v", n, second.MaxTCPQueries)
+	}
+}

core/dnsserver/register.go

@@ -274,6 +274,7 @@ func propagateConfigParams(configs []*Config) {
 		c.ReadTimeout = c.firstConfigInBlock.ReadTimeout
 		c.WriteTimeout = c.firstConfigInBlock.WriteTimeout
 		c.IdleTimeout = c.firstConfigInBlock.IdleTimeout
+		c.MaxTCPQueries = c.firstConfigInBlock.MaxTCPQueries
 		c.TsigSecret = c.firstConfigInBlock.TsigSecret
 
 		// Propagate HTTPRequestValidateFunc so that custom path validators work in

core/dnsserver/server.go

@@ -34,10 +34,11 @@ import (
 // the same address and the listener may be stopped for
 // graceful termination (POSIX only).
 type Server struct {
-	Addr         string        // Address we listen on
-	IdleTimeout  time.Duration // Idle timeout for connection-oriented transports
-	ReadTimeout  time.Duration // Read timeout for connection-oriented transports
-	WriteTimeout time.Duration // Write timeout for connection-oriented transports that support it
+	Addr          string        // Address we listen on
+	IdleTimeout   time.Duration // Idle timeout for connection-oriented transports
+	ReadTimeout   time.Duration // Read timeout for connection-oriented transports
+	WriteTimeout  time.Duration // Write timeout for connection-oriented transports that support it
+	MaxTCPQueries int           // Maximum number of queries served on a single TCP/TLS connection. -1 means unlimited.
 
 	connPolicy                    proxyproto.ConnPolicyFunc // Proxy Protocol connection policy function
 	udpSessionTrackingTTL         time.Duration             // TTL for UDP PPv2 session tracking (0 = disabled)
@@ -74,13 +75,14 @@ type MetadataCollector interface {
 // queries are blocked unless queries from enableChaos are loaded.
 func NewServer(addr string, group []*Config) (*Server, error) {
 	s := &Server{
-		Addr:         addr,
-		zones:        make(map[string][]*Config),
-		graceTimeout: 5 * time.Second,
-		IdleTimeout:  10 * time.Second,
-		ReadTimeout:  3 * time.Second,
-		WriteTimeout: 5 * time.Second,
-		tsigSecret:   make(map[string]string),
+		Addr:          addr,
+		zones:         make(map[string][]*Config),
+		graceTimeout:  5 * time.Second,
+		IdleTimeout:   10 * time.Second,
+		ReadTimeout:   3 * time.Second,
+		WriteTimeout:  5 * time.Second,
+		MaxTCPQueries: tcpMaxQueries,
+		tsigSecret:    make(map[string]string),
 	}
 
 	for _, site := range group {
@@ -103,6 +105,9 @@ func NewServer(addr string, group []*Config) (*Server, error) {
 		if site.IdleTimeout != 0 {
 			s.IdleTimeout = site.IdleTimeout
 		}
+		if site.MaxTCPQueries != nil {
+			s.MaxTCPQueries = *site.MaxTCPQueries
+		}
 
 		// copy tsig secrets
 		maps.Copy(s.tsigSecret, site.TsigSecret)
@@ -167,7 +172,7 @@ func (s *Server) Serve(l net.Listener) error {
 	s.server[tcp] = &dns.Server{Listener: l,
 		Net:           "tcp",
 		TsigSecret:    s.tsigSecret,
-		MaxTCPQueries: tcpMaxQueries,
+		MaxTCPQueries: s.MaxTCPQueries,
 		ReadTimeout:   s.ReadTimeout,
 		WriteTimeout:  s.WriteTimeout,
 		IdleTimeout: func() time.Duration {

core/dnsserver/server_test.go

@@ -200,6 +200,64 @@ func TestGracefulStopTimeout_Internal(t *testing.T) {
 	}
 }
 
+// TestMaxTCPQueriesBoundary proves the user-visible behavior of MaxTCPQueries:
+// a persistent TCP connection may serve exactly the configured number of
+// queries before the server closes it.
+func TestMaxTCPQueriesBoundary(t *testing.T) {
+	n := 2
+	config := testConfig("dns", test.ErrorHandler())
+	config.MaxTCPQueries = &n
+
+	s, err := NewServer("127.0.0.1:0", []*Config{config})
+	if err != nil {
+		t.Fatalf("NewServer failed: %v", err)
+	}
+	defer s.Stop()
+
+	l, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("net.Listen failed: %v", err)
+	}
+	defer l.Close()
+
+	go s.Serve(l)
+
+	conn, err := net.DialTimeout("tcp", l.Addr().String(), 2*time.Second)
+	if err != nil {
+		t.Fatalf("net.DialTimeout failed: %v", err)
+	}
+	defer conn.Close()
+
+	dnsConn := &dns.Conn{Conn: conn}
+
+	for i := range n {
+		m := new(dns.Msg)
+		m.SetQuestion("example.org.", dns.TypeA)
+
+		dnsConn.SetWriteDeadline(time.Now().Add(2 * time.Second))
+		if err := dnsConn.WriteMsg(m); err != nil {
+			t.Fatalf("query %d: WriteMsg failed: %v", i, err)
+		}
+
+		dnsConn.SetReadDeadline(time.Now().Add(2 * time.Second))
+		if _, err := dnsConn.ReadMsg(); err != nil {
+			t.Fatalf("query %d: ReadMsg failed: %v", i, err)
+		}
+	}
+
+	// The connection should be closed by the server after serving n queries;
+	// query n+1 must not succeed on the same connection.
+	m := new(dns.Msg)
+	m.SetQuestion("example.org.", dns.TypeA)
+	dnsConn.SetWriteDeadline(time.Now().Add(2 * time.Second))
+	if err := dnsConn.WriteMsg(m); err == nil {
+		dnsConn.SetReadDeadline(time.Now().Add(2 * time.Second))
+		if _, err := dnsConn.ReadMsg(); err == nil {
+			t.Fatal("expected query beyond MaxTCPQueries to fail on the same connection, but it succeeded")
+		}
+	}
+}
+
 func BenchmarkCoreServeDNS(b *testing.B) {
 	s, err := NewServer("127.0.0.1:53", []*Config{testConfig("dns", testPlugin{})})
 	if err != nil {

core/dnsserver/server_tls.go

@@ -55,7 +55,7 @@ func (s *ServerTLS) Serve(l net.Listener) error {
 	s.server[tcp] = &dns.Server{Listener: l,
 		Net:           "tcp-tls",
 		TsigSecret:    s.tsigSecret,
-		MaxTCPQueries: tlsMaxQueries,
+		MaxTCPQueries: s.MaxTCPQueries,
 		ReadTimeout:   s.ReadTimeout,
 		WriteTimeout:  s.WriteTimeout,
 		IdleTimeout: func() time.Duration {
@@ -102,7 +102,3 @@ func (s *ServerTLS) OnStartupComplete() {
 		fmt.Print(out)
 	}
 }
-
-const (
-	tlsMaxQueries = -1
-)

core/dnsserver/server_tls_test.go

@@ -54,3 +54,58 @@ func TestServerTLSSetsTsigSecret(t *testing.T) {
 		t.Fatalf("expected tsig secret %q, got %q", "abcd", got)
 	}
 }
+
+func TestServerSetsMaxTCPQueries(t *testing.T) {
+	n := 128
+
+	t.Run("default is unlimited", func(t *testing.T) {
+		server, err := NewServer("127.0.0.1:0", []*Config{testConfig("dns", testPlugin{})})
+		if err != nil {
+			t.Fatalf("NewServer() failed: %v", err)
+		}
+
+		if err := server.Serve(&stubListener{}); err == nil {
+			t.Fatal("expected Serve() to return from stub listener")
+		}
+
+		if got := server.server[tcp].MaxTCPQueries; got != -1 {
+			t.Fatalf("expected default MaxTCPQueries -1, got %d", got)
+		}
+	})
+
+	t.Run("configured value reaches the TCP server", func(t *testing.T) {
+		config := testConfig("dns", testPlugin{})
+		config.MaxTCPQueries = &n
+
+		server, err := NewServer("127.0.0.1:0", []*Config{config})
+		if err != nil {
+			t.Fatalf("NewServer() failed: %v", err)
+		}
+
+		if err := server.Serve(&stubListener{}); err == nil {
+			t.Fatal("expected Serve() to return from stub listener")
+		}
+
+		if got := server.server[tcp].MaxTCPQueries; got != n {
+			t.Fatalf("expected MaxTCPQueries %d, got %d", n, got)
+		}
+	})
+
+	t.Run("configured value reaches the TLS server", func(t *testing.T) {
+		config := testConfig("tls", testPlugin{})
+		config.MaxTCPQueries = &n
+
+		server, err := NewServerTLS("tls://127.0.0.1:0", []*Config{config})
+		if err != nil {
+			t.Fatalf("NewServerTLS() failed: %v", err)
+		}
+
+		if err := server.Serve(&stubListener{}); err == nil {
+			t.Fatal("expected Serve() to return from stub listener")
+		}
+
+		if got := server.server[tcp].MaxTCPQueries; got != n {
+			t.Fatalf("expected MaxTCPQueries %d, got %d", n, got)
+		}
+	})
+}

plugin/timeouts/README.md

@@ -2,7 +2,7 @@
 
 ## Name
 
-*timeouts* - allows you to configure the supported server read, write and idle timeouts for the TCP, TLS, DoH and DoQ servers.
+*timeouts* - allows you to configure the supported server read, write and idle timeouts for the TCP, TLS, DoH and DoQ servers, and the maximum number of queries served on a single TCP or TLS connection.
 
 ## Description
 
@@ -15,7 +15,8 @@ over HTTPS. Allowing a longer idle timeout helps performance and reduces issues
 with such routers.
 
 The *timeouts* "plugin" allows you to configure CoreDNS server read, write and
-idle timeouts.
+idle timeouts, and the maximum number of queries CoreDNS will serve on a single
+TCP or TLS connection before closing it.
 
 ## Syntax
 
@@ -24,13 +25,24 @@ timeouts {
 	read DURATION
 	write DURATION
 	idle DURATION
+	maxtcpqueries MAXIMUM
 }
 ~~~
 
 For any timeouts that are not provided, default values are used which may vary
-depending on the server type. At least one timeout must be specified otherwise
+depending on the server type. At least one option must be specified otherwise
 the entire timeouts block should be omitted.
 
+* `maxtcpqueries` sets the maximum number of queries served on a single TCP or
+  TLS connection before CoreDNS closes it. **MAXIMUM** must be a positive
+  integer, or `-1` to allow an unlimited number of queries per connection
+  (the default). Long-lived connections that serve an unlimited number of
+  queries can cause uneven load distribution across CoreDNS replicas that sit
+  behind a connection-based load balancer, since new queries keep reusing the
+  same connection instead of establishing a new one. Setting a bound, e.g.
+  `maxtcpqueries 128`, forces clients to periodically reconnect, which allows
+  the load balancer to redistribute load.
+
 The configured timeouts apply where the selected server transport supports
 them. TCP, TLS and DoH servers use the read, write and idle timeouts. DoQ
 servers use the read timeout to bound receiving a query on an opened QUIC
@@ -94,3 +106,16 @@ configured. The timeouts are only applied to the TCP side of the server.
 	forward . /etc/resolv.conf
 }
 ~~~
+
+Start a standard TCP/UDP server that closes a TCP connection after it has
+served 128 queries, to help spread load evenly across replicas sitting behind
+a connection-based load balancer.
+
+~~~
+. {
+	timeouts {
+		maxtcpqueries 128
+	}
+	forward . /etc/resolv.conf
+}
+~~~

plugin/timeouts/timeouts.go

@@ -1,6 +1,7 @@
 package timeouts
 
 import (
+	"strconv"
 	"time"
 
 	"github.com/coredns/caddy"
@@ -36,6 +37,19 @@ func parseTimeouts(c *caddy.Controller) error {
 				return c.ArgErr()
 			}
 
+			if block == "maxtcpqueries" {
+				n, err := strconv.Atoi(timeoutArgs[0])
+				if err != nil {
+					return c.Errf("invalid value for maxtcpqueries '%s': %v", timeoutArgs[0], err)
+				}
+				if n == 0 || n < -1 {
+					return c.Errf("maxtcpqueries provided '%d' needs to be -1 (unlimited) or a positive integer", n)
+				}
+				config.MaxTCPQueries = &n
+				b++
+				continue
+			}
+
 			timeout, err := durations.NewDurationFromArg(timeoutArgs[0])
 			if err != nil {
 				return c.Err(err.Error())

plugin/timeouts/timeouts_test.go

@@ -5,58 +5,81 @@ import (
 	"testing"
 
 	"github.com/coredns/caddy"
+	"github.com/coredns/coredns/core/dnsserver"
 )
 
 func TestTimeouts(t *testing.T) {
+	n128, nUnlimited := 128, -1
+
 	tests := []struct {
-		input              string
-		shouldErr          bool
-		expectedRoot       string // expected root, set to the controller. Empty for negative cases.
-		expectedErrContent string // substring from the expected error. Empty for positive cases.
+		input                 string
+		shouldErr             bool
+		expectedRoot          string // expected root, set to the controller. Empty for negative cases.
+		expectedErrContent    string // substring from the expected error. Empty for positive cases.
+		expectedMaxTCPQueries *int   // expected Config.MaxTCPQueries after setup. nil means left unset.
 	}{
 		// positive
 		{`timeouts {
 			read 30s
-		}`, false, "", ""},
+		}`, false, "", "", nil},
 		{`timeouts {
 			read 1m
 			write 2m
-		}`, false, "", ""},
+		}`, false, "", "", nil},
 		{` timeouts {
 			idle 1h
-		}`, false, "", ""},
+		}`, false, "", "", nil},
 		{`timeouts {
 			read 10
 			write 20
 			idle 60
-		}`, false, "", ""},
-		// negative
-		{`timeouts`, true, "", "block with no timeouts specified"},
+		}`, false, "", "", nil},
 		{`timeouts {
-		}`, true, "", "block with no timeouts specified"},
+			maxtcpqueries 128
+		}`, false, "", "", &n128},
+		{`timeouts {
+			maxtcpqueries -1
+		}`, false, "", "", &nUnlimited},
+		{`timeouts {
+			read 10s
+			maxtcpqueries 128
+		}`, false, "", "", &n128},
+		// negative
+		{`timeouts`, true, "", "block with no timeouts specified", nil},
+		{`timeouts {
+		}`, true, "", "block with no timeouts specified", nil},
 		{`timeouts {
 			read 10s
 			giraffe 30s
-		}`, true, "", "unknown option"},
+		}`, true, "", "unknown option", nil},
 		{`timeouts {
 			read 10s 20s
 			write 30s
-		}`, true, "", "Wrong argument"},
+		}`, true, "", "Wrong argument", nil},
 		{`timeouts {
 			write snake
-		}`, true, "", "failed to parse duration"},
+		}`, true, "", "failed to parse duration", nil},
 		{`timeouts {
 			idle 0s
-		}`, true, "", "needs to be between"},
+		}`, true, "", "needs to be between", nil},
 		{`timeouts {
 			read 48h
-		}`, true, "", "needs to be between"},
+		}`, true, "", "needs to be between", nil},
+		{`timeouts {
+			maxtcpqueries 0
+		}`, true, "", "needs to be -1", nil},
+		{`timeouts {
+			maxtcpqueries -2
+		}`, true, "", "needs to be -1", nil},
+		{`timeouts {
+			maxtcpqueries snake
+		}`, true, "", "invalid value for maxtcpqueries", nil},
 	}
 
 	for i, test := range tests {
 		c := caddy.NewTestController("dns", test.input)
 		err := setup(c)
-		//cfg := dnsserver.GetConfig(c)
+		cfg := dnsserver.GetConfig(c)
 
 		if test.shouldErr && err == nil {
 			t.Errorf("Test %d: Expected error but found %s for input %s", i, err, test.input)
@@ -70,6 +93,18 @@ func TestTimeouts(t *testing.T) {
 			if !strings.Contains(err.Error(), test.expectedErrContent) {
 				t.Errorf("Test %d: Expected error to contain: %v, found error: %v, input: %s", i, test.expectedErrContent, err, test.input)
 			}
+			continue
+		}
+
+		switch {
+		case test.expectedMaxTCPQueries == nil && cfg.MaxTCPQueries != nil:
+			t.Errorf("Test %d: Expected Config.MaxTCPQueries to remain unset for input %s, got %d", i, test.input, *cfg.MaxTCPQueries)
+		case test.expectedMaxTCPQueries != nil:
+			if cfg.MaxTCPQueries == nil {
+				t.Errorf("Test %d: Expected Config.MaxTCPQueries to be %d for input %s, got unset", i, *test.expectedMaxTCPQueries, test.input)
+			} else if *cfg.MaxTCPQueries != *test.expectedMaxTCPQueries {
+				t.Errorf("Test %d: Expected Config.MaxTCPQueries to be %d for input %s, got %d", i, *test.expectedMaxTCPQueries, test.input, *cfg.MaxTCPQueries)
+			}
 		}
 	}
 }