chore: pointer to wg

command.go

@@ -69,6 +69,10 @@ func (srv *Server) consumeCommands(ctx context.Context, conn net.Conn, reader *b
 			return err
 		}
 
+		if srv.closing.Load() {
+			return nil
+		}
+
 		// NOTE: we increase the wait group by one in order to make sure that idle
 		// connections are not blocking a close.
 		srv.wg.Add(1)

wire.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"net"
 	"sync"
+	"sync/atomic"
 
 	"log/slog"
 
@@ -31,6 +32,7 @@ func ListenAndServe(address string, handler ParseFn) error {
 // NewServer constructs a new Postgres server using the given address and server options.
 func NewServer(parse ParseFn, options ...OptionFn) (*Server, error) {
 	srv := &Server{
+		wg:         &sync.WaitGroup{},
 		parse:      parse,
 		logger:     slog.Default(),
 		closer:     make(chan struct{}),
@@ -52,7 +54,8 @@ func NewServer(parse ParseFn, options ...OptionFn) (*Server, error) {
 
 // Server contains options for listening to an address.
 type Server struct {
-	wg              sync.WaitGroup
+	closing         atomic.Bool
+	wg              *sync.WaitGroup
 	logger          *slog.Logger
 	types           *pgtype.ConnInfo
 	Auth            AuthStrategy
@@ -165,6 +168,11 @@ func (srv *Server) serve(ctx context.Context, conn net.Conn) error {
 
 // Close gracefully closes the underlaying Postgres server.
 func (srv *Server) Close() error {
+	if srv.closing.Load() {
+		return nil
+	}
+
+	srv.closing.Store(true)
 	close(srv.closer)
 	srv.wg.Wait()
 	return nil