chore: lock handler using a RW mutex

command.go

@@ -71,10 +71,10 @@ func (srv *Server) consumeCommands(ctx context.Context, conn net.Conn, reader *b
 
 		// 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)
+		srv.mu.RLock()
 		srv.logger.Debug("incoming command", slog.Int("length", length), slog.String("type", string(t)))
 		err = srv.handleCommand(ctx, conn, t, reader, writer)
-		srv.wg.Done()
+		srv.mu.RUnlock()
 		if errors.Is(err, io.EOF) {
 			return nil
 		}

wire.go

@@ -52,7 +52,7 @@ func NewServer(parse ParseFn, options ...OptionFn) (*Server, error) {
 
 // Server contains options for listening to an address.
 type Server struct {
-	wg              sync.WaitGroup
+	mu              sync.RWMutex
 	logger          *slog.Logger
 	types           *pgtype.ConnInfo
 	Auth            AuthStrategy
@@ -91,11 +91,11 @@ func (srv *Server) Serve(listener net.Listener) error {
 
 	srv.logger.Info("serving incoming connections", slog.String("addr", listener.Addr().String()))
 
-	srv.wg.Add(1)
+	srv.mu.Add(1)
 
 	// NOTE: handle graceful shutdowns
 	go func() {
-		defer srv.wg.Done()
+		defer srv.mu.Done()
 		<-srv.closer
 
 		err := listener.Close()
@@ -166,6 +166,6 @@ func (srv *Server) serve(ctx context.Context, conn net.Conn) error {
 // Close gracefully closes the underlaying Postgres server.
 func (srv *Server) Close() error {
 	close(srv.closer)
-	srv.wg.Wait()
+	srv.mu.Lock()
 	return nil
 }