diff --git a/command.go b/command.go index 67ca2f2..5d976b9 100644 --- a/command.go +++ b/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 } diff --git a/wire.go b/wire.go index 632bed0..55e4533 100644 --- a/wire.go +++ b/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 }