From ddb063551896be86a5acbfb5d3a8b760bf4ae33b Mon Sep 17 00:00:00 2001
From: Robsdedude <rouven.bauer@neo4j.com>
Date: Thu, 14 Dec 2023 11:23:22 +0100
Subject: [PATCH] Minor code clean-ups (#554)

---
 neo4j/internal/pool/pool.go   | 24 +++++++++++++++++-------
 neo4j/session_with_context.go |  4 +---
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/neo4j/internal/pool/pool.go b/neo4j/internal/pool/pool.go
index f7f8757d..acc0c5da 100644
--- a/neo4j/internal/pool/pool.go
+++ b/neo4j/internal/pool/pool.go
@@ -197,7 +197,14 @@ serverLoop:
 	return nil, nil
 }
 
-func (p *Pool) Borrow(ctx context.Context, getServerNames func() []string, wait bool, boltLogger log.BoltLogger, idlenessTimeout time.Duration, auth *idb.ReAuthToken) (idb.Connection, error) {
+func (p *Pool) Borrow(
+	ctx context.Context,
+	getServerNames func() []string,
+	wait bool,
+	boltLogger log.BoltLogger,
+	idlenessTimeout time.Duration,
+	auth *idb.ReAuthToken,
+) (idb.Connection, error) {
 	for {
 		if p.closed {
 			return nil, &errorutil.PoolClosed{}
@@ -279,10 +286,13 @@ func (p *Pool) Borrow(ctx context.Context, getServerNames func() []string, wait
 	}
 }
 
-func (p *Pool) tryBorrow(ctx context.Context, serverName string, boltLogger log.BoltLogger, idlenessTimeout time.Duration, auth *idb.ReAuthToken) (idb.Connection, error) {
-	// For now, lock complete servers map to avoid over connecting but with the downside
-	// that long connect times will block connects to other servers as well. To fix this
-	// we would need to add a pending connect to the server and lock per server.
+func (p *Pool) tryBorrow(
+	ctx context.Context,
+	serverName string,
+	boltLogger log.BoltLogger,
+	idlenessTimeout time.Duration,
+	auth *idb.ReAuthToken,
+) (idb.Connection, error) {
 	p.serversMut.Lock()
 	var unlock = new(sync.Once)
 	defer unlock.Do(p.serversMut.Unlock)
@@ -346,10 +356,10 @@ func (p *Pool) tryBorrow(ctx context.Context, serverName string, boltLogger log.
 func (p *Pool) unreg(ctx context.Context, serverName string, c idb.Connection, now time.Time) {
 	p.serversMut.Lock()
 	defer p.serversMut.Unlock()
-	p.unregUnlocked(ctx, serverName, c, now)
+	p.unregLocked(ctx, serverName, c, now)
 }
 
-func (p *Pool) unregUnlocked(ctx context.Context, serverName string, c idb.Connection, now time.Time) {
+func (p *Pool) unregLocked(ctx context.Context, serverName string, c idb.Connection, now time.Time) {
 	defer func() {
 		// Close connection in another thread to avoid potential long blocking operation during close.
 		go c.Close(ctx)
diff --git a/neo4j/session_with_context.go b/neo4j/session_with_context.go
index 7467db1f..18b42b22 100644
--- a/neo4j/session_with_context.go
+++ b/neo4j/session_with_context.go
@@ -524,9 +524,7 @@ func (s *sessionWithContext) getConnection(ctx context.Context, mode idb.AccessM
 	if timeout > 0 {
 		var cancel context.CancelFunc
 		ctx, cancel = context.WithTimeout(ctx, timeout)
-		if cancel != nil {
-			defer cancel()
-		}
+		defer cancel()
 		deadline, _ := ctx.Deadline()
 		s.log.Debugf(log.Session, s.logId, "connection acquisition timeout is %s, resolved deadline is: %s", timeout, deadline)
 	} else if deadline, ok := ctx.Deadline(); ok {