From 6bc7474d87941f97f0d769145b338411da0b619e Mon Sep 17 00:00:00 2001
From: Jeroen Rinzema <jeroen@rinzema.dev>
Date: Fri, 1 Sep 2023 17:57:07 +0200
Subject: [PATCH] chore: pointer to wg

---
 command.go |  4 ++++
 wire.go    | 10 +++++++++-
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/command.go b/command.go
index 67ca2f2..07d8b7d 100644
--- a/command.go
+++ b/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)
diff --git a/wire.go b/wire.go
index 632bed0..1a776f7 100644
--- a/wire.go
+++ b/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