Enable -race in tests and fix race conditions (#141)

Author: JelteF

Makefile

@@ -39,7 +39,7 @@ lint: | $(GOLANGCI_LINT) ; $(info $(M) running golint…) @ ## Run the project l
 
 .PHONY: test
 test: ## Run all tests
-	$Q $(GO) test ./... -timeout 20s
+	$Q $(GO) test ./... -timeout 20s -race
 
 .PHONY: generate
 generate: | $(TOOLS) ; $(info $(M) running go generate…) @ ## Run go generate

command.go

@@ -131,13 +131,15 @@ func (srv *Session) consumeSingleCommand(ctx context.Context, reader *buffer.Rea
 		return err
 	}
 
+	// We hold closingMu for reading while checking closing + adding to the
+	// wait group, so that Close cannot finish wg.Wait before we are tracked.
+	srv.closingMu.RLock()
 	if srv.closing.Load() {
+		srv.closingMu.RUnlock()
 		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)
+	srv.closingMu.RUnlock()
 	srv.logger.Debug("<- incoming command", slog.Int("length", length), slog.String("type", t.String()))
 	err = srv.handleCommand(ctx, conn, t, reader, writer)
 	srv.wg.Done()

pkg/mock/client.go

@@ -139,4 +139,9 @@ func (client *Client) Close(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+
+	err = client.conn.Close()
+	if err != nil {
+		t.Fatal(err)
+	}
 }

wire.go

@@ -110,8 +110,21 @@ func NewServer(parse ParseFn, options ...OptionFn) (*Server, error) {
 
 // Server contains options for listening to an address.
 type Server struct {
-	closing          atomic.Bool
-	wg               sync.WaitGroup
+	closing atomic.Bool
+	// Used to make reading closing and calling wg.Add be a single atomic operation.
+	closingMu sync.RWMutex
+	// This WaitGroup tracks the number of connections that are actively
+	// serving a request. Idle connections are not counted. This is used during
+	// Shutdown and Close to wait for outstanding requests to finish.
+	//
+	// Additionally it also waits for the Serve loop to stop listening for new
+	// connections.
+	wg sync.WaitGroup
+	// Tracks total number of connection-serving goroutines, so tests can wait
+	// for them to exit before the test's `t` becomes invalid.
+	//
+	// Additionally it also waits for the Go routine
+	connWg           sync.WaitGroup
 	logger           *slog.Logger
 	Auth             AuthStrategy
 	BackendKeyData   BackendKeyDataFunc
@@ -149,21 +162,19 @@ func (srv *Server) ListenAndServe(address string) error {
 // preconfigured configurations. The given listener will be closed once the
 // server is gracefully closed.
 func (srv *Server) Serve(listener net.Listener) error {
-	// Early check to avoid logging and work if shutdown already started
+	srv.closingMu.RLock()
 	if srv.closing.Load() {
+		// Don't start serving if we're already shutting down or shut down.
+		srv.closingMu.RUnlock()
 		return nil
 	}
+	srv.wg.Add(1)
+	srv.closingMu.RUnlock()
 
-	srv.logger.Info("serving incoming connections", slog.String("addr", listener.Addr().String()))
-
-	// Double-check: if shutdown started between the check and Add, we must undo
-	if srv.closing.Load() {
-		return nil
-	}
+	defer srv.wg.Done()
 
-	srv.wg.Add(1)
-	go func() {
-		defer srv.wg.Done()
+	srv.logger.Info("serving incoming connections", slog.String("addr", listener.Addr().String()))
+	srv.wg.Go(func() {
 		<-srv.closer
 
 		srv.logger.Info("closing server")
@@ -172,7 +183,7 @@ func (srv *Server) Serve(listener net.Listener) error {
 		if err != nil {
 			srv.logger.Error("unexpected error while attempting to close the net listener", "err", err)
 		}
-	}()
+	})
 
 	for {
 		conn, err := listener.Accept()
@@ -190,7 +201,7 @@ func (srv *Server) Serve(listener net.Listener) error {
 			continue
 		}
 
-		go func() {
+		srv.connWg.Go(func() {
 			ctx := context.Background()
 			err = srv.serve(ctx, conn)
 			if err != nil {
@@ -200,7 +211,7 @@ func (srv *Server) Serve(listener net.Listener) error {
 					srv.logger.Error("an unexpected error got returned while serving a client connection", "err", err)
 				}
 			}
-		}()
+		})
 	}
 }
 
@@ -292,11 +303,14 @@ 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() {
+	srv.closingMu.Lock()
+	if !srv.closing.CompareAndSwap(false, true) {
+		srv.closingMu.Unlock()
+		srv.wg.Wait()
 		return nil
 	}
+	srv.closingMu.Unlock()
 
-	srv.closing.Store(true)
 	close(srv.closer)
 	srv.wg.Wait()
 	return nil
@@ -308,11 +322,14 @@ func (srv *Server) Close() error {
 // If the context has no deadline, the server's ShutdownTimeout is used.
 func (srv *Server) Shutdown(ctx context.Context) error {
 	// Check if already shutting down or shut down
-	if srv.closing.Load() {
+	srv.closingMu.Lock()
+	if !srv.closing.CompareAndSwap(false, true) {
 		// If already closing, just wait for existing shutdown to complete
+		srv.closingMu.Unlock()
 		srv.wg.Wait()
 		return nil
 	}
+	srv.closingMu.Unlock()
 
 	// Use the shorter of context deadline or server timeout
 	var shutdownCtx context.Context
@@ -329,13 +346,6 @@ func (srv *Server) Shutdown(ctx context.Context) error {
 	}
 	defer cancel()
 
-	// Atomically check and set closing state
-	if !srv.closing.CompareAndSwap(false, true) {
-		// Another goroutine beat us to it, just wait for shutdown to complete
-		srv.wg.Wait()
-		return nil
-	}
-
 	srv.logger.Info("starting graceful shutdown")
 
 	// Close the closer channel (we're the first/only one to get here)
@@ -358,6 +368,13 @@ func (srv *Server) Shutdown(ctx context.Context) error {
 	}
 }
 
+// Wait blocks until all connection-serving goroutines have finished. This is
+// intended to be called at the end of a test, so no go-routines are lingering.
+// This can block indefinitely if not all clients have disconnected.
+func (srv *Server) Wait() {
+	srv.connWg.Wait()
+}
+
 // isNormalConnectionClosure checks if an error represents a normal client connection closure
 // that shouldn't be logged as an error.
 func (srv *Server) isNormalConnectionClosure(err error) bool {

wire_test.go

@@ -37,6 +37,7 @@ func TListenAndServe(t *testing.T, server *Server) *net.TCPAddr {
 		if err != nil {
 			t.Fatal(err)
 		}
+		server.Wait()
 	})
 
 	go server.Serve(listener) //nolint:errcheck
@@ -1237,6 +1238,10 @@ func TListenAndServeWithoutCleanup(t *testing.T, server *Server) *net.TCPAddr {
 		t.Fatal(err)
 	}
 
+	t.Cleanup(func() {
+		server.Wait()
+	})
+
 	go server.Serve(listener) //nolint:errcheck
 	return listener.Addr().(*net.TCPAddr)
 }