Fix panic when calling driver.Close concurrently (#626)

* Fix panic when calling driver.Close concurrently

Doing so might well still yield unexpected or undesired results, but at least
it shouldn't cause a panic.

* TestKit: close resources on disconnect & catch backend panics

* fixup! TestKit: close resources on disconnect & catch backend panics

* fixup! TestKit: close resources on disconnect & catch backend panics

---------

Co-authored-by: Stephen Cathcart <[email protected]>

Author: robsdedude

neo4j/driver.go

@@ -37,6 +37,7 @@ type Driver interface {
 	// or error describing the problem.
 	VerifyConnectivity() error
 	// Close the driver and all underlying connections
+	// This function may not be called while the driver is in use (i.e., concurrently).
 	Close() error
 	// IsEncrypted determines whether the driver communication with the server
 	// is encrypted. This is a static check. The function can also be called on

neo4j/internal/bolt/bolt3.go

@@ -151,6 +151,10 @@ func (b *bolt3) ServerName() string {
 	return b.serverName
 }
 
+func (b *bolt3) ConnId() string {
+	return b.connId
+}
+
 func (b *bolt3) ServerVersion() string {
 	return b.serverVersion
 }

neo4j/internal/bolt/bolt4.go

@@ -174,6 +174,10 @@ func (b *bolt4) ServerName() string {
 	return b.serverName
 }
 
+func (b *bolt4) ConnId() string {
+	return b.connId
+}
+
 func (b *bolt4) ServerVersion() string {
 	return b.serverVersion
 }

neo4j/internal/bolt/bolt5.go

@@ -185,6 +185,10 @@ func (b *bolt5) ServerName() string {
 	return b.serverName
 }
 
+func (b *bolt5) ConnId() string {
+	return b.connId
+}
+
 func (b *bolt5) ServerVersion() string {
 	return b.serverVersion
 }

neo4j/internal/db/connection.go

@@ -124,6 +124,8 @@ type Connection interface {
 	Bookmark() string
 	// ServerName returns the name of the remote server
 	ServerName() string
+	// ConnId returns the connection id as assigned by the server ("" if not available)
+	ConnId() string
 	// ServerVersion returns the server version on pattern Neo4j/1.2.3
 	ServerVersion() string
 	// IsAlive returns true if the connection is fully functional.

neo4j/internal/pool/pool.go

@@ -108,12 +108,22 @@ func (p *Pool) Close(ctx context.Context) {
 	p.queueMut.Unlock()
 	// Go through each server and close all connections to it
 	p.serversMut.Lock()
-	for n, s := range p.servers {
-		s.closeAll(ctx, p.closeConnection)
-		delete(p.servers, n)
+	pendingConnections := 0
+	for _, s := range p.servers {
+		s.startClosing(ctx, p.closeConnection)
+		pendingConnections += s.size()
 	}
 	p.serversMut.Unlock()
-	p.log.Infof(log.Pool, p.logId, "Closed")
+	if pendingConnections == 0 {
+		p.log.Infof(log.Pool, p.logId, "Closed")
+	} else {
+		p.log.Warnf(
+			log.Pool,
+			p.logId,
+			"Called close with %d in-flight connections (will be closed when work is done).",
+			pendingConnections,
+		)
+	}
 }
 
 // For testing
@@ -194,8 +204,8 @@ func (p *Pool) Borrow(
 	auth *idb.ReAuthToken,
 ) (idb.Connection, error) {
 	for {
-		if p.closed {
-			return nil, &errorutil.PoolClosed{}
+		if err := p.checkClosed(); err != nil {
+			return nil, err
 		}
 		serverNames := getServerNames()
 		if len(serverNames) == 0 {
@@ -294,6 +304,10 @@ func (p *Pool) tryBorrow(
 	var unlock = new(sync.Once)
 	defer unlock.Do(p.serversMut.Unlock)
 
+	if err := p.checkClosed(); err != nil {
+		return nil, err
+	}
+
 	srv := p.servers[serverName]
 	for {
 		if srv != nil {
@@ -351,6 +365,13 @@ func (p *Pool) tryBorrow(
 	return c, nil
 }
 
+func (p *Pool) checkClosed() error {
+	if p.closed {
+		return &errorutil.PoolClosed{}
+	}
+	return nil
+}
+
 func (p *Pool) unreg(ctx context.Context, serverName string, c idb.Connection, now time.Time) {
 	p.serversMut.Lock()
 	defer p.serversMut.Unlock()
@@ -396,13 +417,19 @@ func (p *Pool) closeConnection(ctx context.Context, c idb.Connection) {
 func (p *Pool) Return(ctx context.Context, c idb.Connection) {
 	if p.closed {
 		p.log.Warnf(log.Pool, p.logId, "Trying to return connection to closed pool")
-		return
 	}
 
 	// Get the name of the server that the connection belongs to.
 	serverName := c.ServerName()
 	isAlive := c.IsAlive()
-	p.log.Debugf(log.Pool, p.logId, "Returning connection to %s {alive:%t}", serverName, isAlive)
+	p.log.Debugf(
+		log.Pool,
+		p.logId,
+		"Returning connection %s to %s {alive:%t}",
+		c.ConnId(),
+		serverName,
+		isAlive,
+	)
 
 	// If the connection is dead, remove all other idle connections on the same server that older
 	// or of the same age as the dead connection, otherwise perform normal cleanup of old connections

neo4j/internal/pool/server.go

@@ -196,12 +196,6 @@ func (s *server) removeIdleOlderThan(ctx context.Context, now time.Time, maxAge
 	}
 }
 
-func (s *server) closeAll(ctx context.Context, close closeFunc) {
-	s.closeAndEmptyConnections(ctx, &s.idle, close)
-	// Closing the busy connections could mean here that we do close from another thread.
-	s.closeAndEmptyConnections(ctx, &s.busy, close)
-}
-
 func (s *server) executeForAllConnections(callback func(c db.Connection)) {
 	for item := s.busy.Front(); item != nil; item = item.Next() {
 		callback(item.Value.(db.Connection))

neo4j/internal/pool/server_test.go

@@ -133,16 +133,24 @@ func TestServer(ot *testing.T) {
 		}
 	})
 
-	ot.Run("closeAll clears all connections", func(t *testing.T) {
+	ot.Run("startClosing sets closing flag", func(t *testing.T) {
 		s := NewServer()
 		// Register and return three connections
 		_, _ = populateServer(s, time.Now(), 3, 3)
-		s.closeAll(context.Background(), closeConnection)
-		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())
+		s.startClosing(context.Background(), closeConnection)
+		testutil.AssertTrue(t, s.closing)
+	})
+
+	ot.Run("closing flag makes returnBusy close connections", func(t *testing.T) {
+		s := NewServer()
+		// Register and return three connections
+		_, busy_connections := populateServer(s, time.Now(), 0, 3)
+		s.startClosing(context.Background(), closeConnection)
+		for i, c := range busy_connections {
+			testutil.AssertFalse(t, c.Closed)
+			s.returnBusy(context.Background(), c, closeConnection)
+			testutil.AssertTrue(t, c.Closed)
+			testutil.AssertIntEqual(t, s.busy.Len(), 2-i)
 		}
 	})
 }

neo4j/internal/testutil/connfake.go

@@ -75,6 +75,7 @@ type ConnFake struct {
 	ReAuthHook              func(context.Context, *idb.ReAuthToken) error
 	SsrEnabled              bool
 	PinHomeDatabaseCallback func(context.Context, string)
+	Closed                  bool
 }
 
 func (c *ConnFake) Connect(
@@ -92,6 +93,10 @@ func (c *ConnFake) ServerName() string {
 	return c.Name
 }
 
+func (c *ConnFake) ConnId() string {
+	return "bolt-1"
+}
+
 func (c *ConnFake) IsAlive() bool {
 	return c.Alive
 }
@@ -108,6 +113,7 @@ func (c *ConnFake) ForceReset(context.Context) {
 }
 
 func (c *ConnFake) Close(ctx context.Context) {
+	c.Closed = true
 }
 
 func (c *ConnFake) Birthdate() time.Time {

testkit-backend/backend.go

@@ -190,10 +190,29 @@ func (b *backend) writeLineLocked(s string) error {
 
 // Reads and writes to the socket until it is closed
 func (b *backend) serve() {
+	defer b.close()
 	for b.process() {
 	}
 }
 
+func (b *backend) close() {
+	b.closed = true
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	defer cancel()
+	for k, tx := range b.explicitTransactions {
+		_ = tx.Close(ctx)
+		delete(b.explicitTransactions, k)
+	}
+	for k, session := range b.sessionStates {
+		_ = session.session.Close(ctx)
+		delete(b.sessionStates, k)
+	}
+	for k, driver := range b.drivers {
+		_ = driver.Close(ctx)
+		delete(b.drivers, k)
+	}
+}
+
 func (b *backend) setError(err error) string {
 	id := b.nextId()
 	b.recordedErrors[id] = err