chore: include a wg while serving conns

jeroenrinzema · Nov 13, 2023 · 5ab1b13 · 5ab1b13
1 parent 413e316
commit 5ab1b13
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/wire.go b/wire.go
@@ -92,7 +92,6 @@ func (srv *Server) Serve(listener net.Listener) error {
 	defer srv.logger.Info("closing server")
 
 	srv.logger.Info("serving incoming connections", slog.String("addr", listener.Addr().String()))
-
 	srv.wg.Add(1)
 
 	// NOTE: handle graceful shutdowns
@@ -112,7 +111,10 @@ func (srv *Server) Serve(listener net.Listener) error {
 			return err
 		}
 
+		srv.wg.Add(1)
+
 		go func() {
+			defer srv.wg.Done()
 			ctx := context.Background()
 			err = srv.serve(ctx, conn)
 			if err != nil {