Merge pull request #22 from owenthereal/remove_dns_tcp

Wait for udp/tcp server before shutting it down

Author: owenthereal

dns/server.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"sync"
 	"time"
 
 	"github.com/miekg/dns"
@@ -38,35 +39,43 @@ func (d *dnsServer) Run(ctx context.Context) error {
 		mux.HandleFunc(tld+".", d.handleDNS)
 	}
 
-	ctx, cancel := context.WithCancel(ctx)
 	var g run.Group
 	{
+		var wg sync.WaitGroup
+		wg.Add(1)
 		udp := &dns.Server{
-			Handler:   mux,
-			Addr:      d.cfg.Addr,
-			Net:       "udp",
-			ReusePort: true,
+			Handler: mux,
+			Addr:    d.cfg.Addr,
+			Net:     "udp",
 		}
 		g.Add(func() error {
+			wg.Done()
 			return udp.ListenAndServe()
 		}, func(err error) {
+			// Wait for udp server before shutting it down
+			wg.Wait()
 			_ = udp.ShutdownContext(ctx)
 		})
 	}
 	{
+		var wg sync.WaitGroup
+		wg.Add(1)
 		tcp := &dns.Server{
-			Handler:   mux,
-			Addr:      d.cfg.Addr,
-			Net:       "tcp",
-			ReusePort: true,
+			Handler: mux,
+			Addr:    d.cfg.Addr,
+			Net:     "tcp",
 		}
 		g.Add(func() error {
+			wg.Done()
 			return tcp.ListenAndServe()
 		}, func(err error) {
+			// Wait for tcp server before shutting it down
+			wg.Wait()
 			_ = tcp.ShutdownContext(ctx)
 		})
 	}
 	{
+		ctx, cancel := context.WithCancel(ctx)
 		g.Add(func() error {
 			<-ctx.Done()
 			return ctx.Err()

server/server_test.go

@@ -176,7 +176,7 @@ func Test_Server(t *testing.T) {
 		}
 
 		select {
-		case <-time.After(20 * time.Second):
+		case <-time.After(5 * time.Second):
 			t.Fatal("error wait time out")
 		case err := <-errch:
 			if want, got := fmt.Sprintf("host root %s was removed", hostRoot), err.Error(); want != got {
@@ -259,7 +259,7 @@ func Test_Server_Shutdown(t *testing.T) {
 			}()
 
 			select {
-			case <-time.After(5 * time.Second):
+			case <-time.After(10 * time.Second):
 				t.Fatal("error wait time out")
 			case err := <-errch:
 				if want, got := c.WantErrMsg, err.Error(); !strings.Contains(got, want) {