Fix dns server shutdown

dns/server.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"fmt"
 	"net"
-	"sync"
 	"time"
 
 	"github.com/miekg/dns"
@@ -41,36 +40,26 @@ func (d *dnsServer) Run(ctx context.Context) error {
 
 	var g run.Group
 	{
-		var wg sync.WaitGroup
-		wg.Add(1)
 		udp := &dns.Server{
 			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",
 		}
 		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)
 		})
 	}

server/server_test.go

@@ -41,8 +41,12 @@ func Test_Server(t *testing.T) {
 		DnsAddr:   dnsAddr,
 	})
 	errch := make(chan error)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	go func() {
-		err := svr.Run(context.Background())
+		err := svr.Run(ctx)
 		if err != nil {
 			candy.Log().Error("error running server", zap.Error(err))
 		}
@@ -259,8 +263,12 @@ func Test_Server_Shutdown(t *testing.T) {
 		t.Run(c.Name, func(t *testing.T) {
 			errch := make(chan error)
 			srv := New(c.Config)
+
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
 			go func() {
-				err := srv.Run(context.Background())
+				err := srv.Run(ctx)
 				if err != nil {
 					candy.Log().Error("error running server", zap.Error(err))
 				}