tests: test dns middleware (#1451)

* dns: route dns tests through middleware

* run benchmarks in test pipeline to make sure they work

* require test-bench

Author: BeryJu

.github/workflows/benchmark.yml

@@ -28,6 +28,8 @@ jobs:
       - run: make test-env-start
       - run: make install-deps
       - run: make bench
+      - if: ${{ always() }}
+        run: make test-env-stop
       - uses: benchmark-action/github-action-benchmark@v1
         with:
           name: Gravity Benchmark

.github/workflows/test.yml

@@ -45,6 +45,26 @@ jobs:
         uses: codecov/test-results-action@v1
         with:
           token: ${{ secrets.CODECOV_TOKEN }}
+  test-bench:
+    name: Test (bench)
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-go@v5
+        with:
+          go-version-file: go.mod
+          cache: true
+      - uses: actions/setup-node@v4
+        with:
+          node-version-file: web/package.json
+          cache: "npm"
+          cache-dependency-path: web/package-lock.json
+      - run: make web-install gen-client-ts web-build
+      - run: make test-env-start
+      - run: make install-deps gen-clean gen-client-go
+      - run: make bench
+      - if: ${{ always() }}
+        run: make test-env-stop
   test-e2e:
     name: Test (e2e)
     runs-on: ubuntu-latest
@@ -73,6 +93,7 @@ jobs:
     needs:
       - test-unittest
       - test-e2e
+      - test-bench
     runs-on: ubuntu-latest
     steps:
       - uses: re-actors/alls-green@release/v1

pkg/roles/dns/dns_handler.go

@@ -31,7 +31,7 @@ func (r *Role) FindZone(fqdn string) *Zone {
 	return longestZone
 }
 
-func (ro *Role) Handler(w dns.ResponseWriter, r *dns.Msg) {
+func (ro *Role) handler(w dns.ResponseWriter, r *dns.Msg) {
 	span := sentry.StartTransaction(
 		context.TODO(),
 		"gravity.dns.request",

pkg/roles/dns/role.go

@@ -24,9 +24,11 @@ type Role struct {
 	ctx   context.Context
 	zones *watcher.Watcher[*Zone]
 
-	cfg     *RoleConfig
-	log     *zap.Logger
+	cfg *RoleConfig
+	log *zap.Logger
+
 	servers []*dns.Server
+	m       *dns.ServeMux
 }
 
 func init() {
@@ -41,6 +43,7 @@ func New(instance roles.Instance) *Role {
 		log:     instance.Log(),
 		i:       instance,
 		ctx:     instance.Context(),
+		m:       dns.NewServeMux(),
 	}
 	r.zones = watcher.New(
 		func(kv *mvccpb.KeyValue) (*Zone, error) {
@@ -106,12 +109,11 @@ func (r *Role) Start(ctx context.Context, config []byte) error {
 
 	r.zones.Start(r.ctx)
 
-	dnsMux := dns.NewServeMux()
-	dnsMux.HandleFunc(
+	r.m.HandleFunc(
 		types.DNSRootZone,
 		r.recoverMiddleware(
 			r.loggingMiddleware(
-				r.Handler,
+				r.handler,
 			),
 		),
 	)
@@ -136,19 +138,23 @@ func (r *Role) Start(ctx context.Context, config []byte) error {
 		{
 			Addr:    listen,
 			Net:     "udp",
-			Handler: dnsMux,
+			Handler: r.m,
 		},
 		{
 			Addr:    listen,
 			Net:     "tcp",
-			Handler: dnsMux,
+			Handler: r.m,
 		},
 	}
 	go srv(0)
 	go srv(1)
 	return nil
 }
 
+func (r *Role) Handler(w dns.ResponseWriter, req *dns.Msg) {
+	r.m.ServeDNS(w, req)
+}
+
 func (r *Role) Stop() {
 	r.zones.Stop()
 	for _, server := range r.servers {