Fix race: pool kept connections after closure (#562)

Co-authored-by: Stephen Cathcart <stephen.cathcart@neotechnology.com>

neo4j/internal/pool/server.go

@@ -157,7 +157,7 @@ func (s *server) numBusy() int {
 	return s.busy.Len()
 }
 
-// Adds a db to busy list
+// Adds a connection to the busy list
 func (s *server) registerBusy(c db.Connection) {
 	// Update round-robin to indicate when this server was last used.
 	atomic.StoreUint32(&s.roundRobin, atomic.AddUint32(&sharedRoundRobin, 1))
@@ -197,9 +197,9 @@ func (s *server) removeIdleOlderThan(ctx context.Context, now time.Time, maxAge
 }
 
 func (s *server) closeAll(ctx context.Context) {
-	closeAndEmptyConnections(ctx, s.idle)
+	closeAndEmptyConnections(ctx, &s.idle)
 	// Closing the busy connections could mean here that we do close from another thread.
-	closeAndEmptyConnections(ctx, s.busy)
+	closeAndEmptyConnections(ctx, &s.busy)
 }
 
 func (s *server) executeForAllConnections(callback func(c db.Connection)) {
@@ -213,10 +213,10 @@ func (s *server) executeForAllConnections(callback func(c db.Connection)) {
 
 func (s *server) startClosing(ctx context.Context) {
 	s.closing = true
-	closeAndEmptyConnections(ctx, s.idle)
+	closeAndEmptyConnections(ctx, &s.idle)
 }
 
-func closeAndEmptyConnections(ctx context.Context, l list.List) {
+func closeAndEmptyConnections(ctx context.Context, l *list.List) {
 	for e := l.Front(); e != nil; e = e.Next() {
 		c := e.Value.(db.Connection)
 		go c.Close(ctx)

neo4j/internal/pool/server_test.go

@@ -86,13 +86,8 @@ func TestServer(ot *testing.T) {
 	ot.Run("removeIdleOlderThan", func(t *testing.T) {
 		s := NewServer()
 		// Register and return three connections
-		conns := make([]*testutil.ConnFake, 3)
 		now := time.Now()
-		for i := range conns {
-			c := &testutil.ConnFake{Birth: now, Alive: true}
-			conns[i] = c
-			registerIdle(s, c)
-		}
+		conns, _ := populateServer(s, now, 3, 0)
 
 		// Let the connection in the middle be too old
 		conns[1].Birth = now.Add(-20 * time.Second)
@@ -119,6 +114,32 @@ func TestServer(ot *testing.T) {
 		assertNilConnection(t, b1)
 		assertSize(t, s, 0)
 	})
+
+	ot.Run("startClosing clears idle connections", func(t *testing.T) {
+		s := NewServer()
+		// Register and return three connections
+		_, _ = populateServer(s, time.Now(), 3, 3)
+		s.startClosing(context.Background())
+		if s.idle.Len() != 0 {
+			t.Errorf("Expected 0 idle connections, found %d", s.idle.Len())
+		}
+		if s.busy.Len() != 3 {
+			t.Errorf("Expected 3 busy connections, found %d", s.busy.Len())
+		}
+	})
+
+	ot.Run("closeAll clears all connections", func(t *testing.T) {
+		s := NewServer()
+		// Register and return three connections
+		_, _ = populateServer(s, time.Now(), 3, 3)
+		s.closeAll(context.Background())
+		if s.idle.Len() != 0 {
+			t.Errorf("Expected 0 idle connections, found %d", s.idle.Len())
+		}
+		if s.busy.Len() != 0 {
+			t.Errorf("Expected 0 busy connections, found %d", s.busy.Len())
+		}
+	})
 }
 
 func TestServerPenalty(t *testing.T) {
@@ -286,6 +307,22 @@ func TestIdlenessThreshold(outer *testing.T) {
 	})
 }
 
+func populateServer(s *server, birth time.Time, idle, busy int) ([]*testutil.ConnFake, []*testutil.ConnFake) {
+	idleConns := make([]*testutil.ConnFake, idle)
+	for i := range idleConns {
+		c := &testutil.ConnFake{Birth: birth, Alive: true}
+		idleConns[i] = c
+		registerIdle(s, c)
+	}
+	busyConns := make([]*testutil.ConnFake, busy)
+	for i := range busyConns {
+		c := &testutil.ConnFake{Birth: birth, Alive: true}
+		busyConns[i] = c
+		s.registerBusy(c)
+	}
+	return idleConns, busyConns
+}
+
 func registerIdle(srv *server, connection db.Connection) {
 	srv.registerBusy(connection)
 	srv.returnBusy(context.Background(), connection)