dnsserver: Rely on dns.Server.ShutdownContext to gracefully stop (#7517)

core/dnsserver/server.go

@@ -6,7 +6,6 @@ import (
 	"fmt"
 	"maps"
 	"net"
-	"runtime"
 	"runtime/debug"
 	"strings"
 	"sync"
@@ -42,7 +41,6 @@ type Server struct {
 	m      sync.Mutex     // protects the servers
 
 	zones        map[string][]*Config // zones keyed by their address
-	dnsWg        sync.WaitGroup       // used to wait on outstanding connections
 	graceTimeout time.Duration        // the maximum duration of a graceful shutdown
 	trace        trace.Trace          // the trace plugin for the server
 	debug        bool                 // disable recover()
@@ -52,8 +50,8 @@ type Server struct {
 	tsigSecret map[string]string
 
 	// Ensure Stop is idempotent when invoked concurrently (e.g., during reload and SIGTERM).
-	stopOnce   sync.Once
-	wgDoneOnce sync.Once
+	stopOnce sync.Once
+	stopErr  error
 }
 
 // MetadataCollector is a plugin that can retrieve metadata functions from all metadata providing plugins
@@ -74,14 +72,6 @@ func NewServer(addr string, group []*Config) (*Server, error) {
 		tsigSecret:   make(map[string]string),
 	}
 
-	// We have to bound our wg with one increment
-	// to prevent a "race condition" that is hard-coded
-	// into sync.WaitGroup.Wait() - basically, an add
-	// with a positive delta must be guaranteed to
-	// occur before Wait() is called on the wg.
-	// In a way, this kind of acts as a safety barrier.
-	s.dnsWg.Add(1)
-
 	for _, site := range group {
 		if site.Debug {
 			s.debug = true
@@ -209,44 +199,36 @@ func (s *Server) ListenPacket() (net.PacketConn, error) {
 	return p, nil
 }
 
-// Stop stops the server. It blocks until the server is
-// totally stopped. On POSIX systems, it will wait for
-// connections to close (up to a max timeout of a few
-// seconds); on Windows it will close the listener
-// immediately.
+// Stop attempts to gracefully stop the server.
+// It waits until the server is stopped and its connections are closed,
+// up to a max timeout of a few seconds. If unsuccessful, an error is returned.
+//
 // This implements Caddy.Stopper interface.
-func (s *Server) Stop() (err error) {
-	var onceErr error
+func (s *Server) Stop() error {
 	s.stopOnce.Do(func() {
-		if runtime.GOOS != "windows" {
-			// force connections to close after timeout
-			done := make(chan struct{})
-			go func() {
-				// decrement our initial increment used as a barrier, but only once
-				s.wgDoneOnce.Do(func() { s.dnsWg.Done() })
-				s.dnsWg.Wait()
-				close(done)
-			}()
+		ctx, cancelCtx := context.WithTimeout(context.Background(), s.graceTimeout)
+		defer cancelCtx()
 
-			// Wait for remaining connections to finish or
-			// force them all to close after timeout
-			select {
-			case <-time.After(s.graceTimeout):
-			case <-done:
-			}
-		}
-
-		// Close the listener now; this stops the server without delay
+		var wg sync.WaitGroup
 		s.m.Lock()
 		for _, s1 := range s.server {
 			// We might not have started and initialized the full set of servers
-			if s1 != nil {
-				onceErr = s1.Shutdown()
+			if s1 == nil {
+				continue
 			}
+
+			wg.Add(1)
+			go func() {
+				s1.ShutdownContext(ctx)
+				wg.Done()
+			}()
 		}
 		s.m.Unlock()
+		wg.Wait()
+
+		s.stopErr = ctx.Err()
 	})
-	return onceErr
+	return s.stopErr
 }
 
 // Address together with Stop() implement caddy.GracefulServer.