Improve routing table fetching

neo4j/internal/bolt/chunker.go

@@ -8,20 +8,21 @@
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
  *
- *      https://www.apache.org/licenses/LICENSE-2.0
+ *     https://www.apache.org/licenses/LICENSE-2.0
  *
- *  Unless required by applicable law or agreed to in writing, software
- *  distributed under the License is distributed on an "AS IS" BASIS,
- *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- *  See the License for the specific language governing permissions and
- *  limitations under the License.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
  */
 
 package bolt
 
 import (
 	"context"
 	"encoding/binary"
+	"errors"
 	"github.com/neo4j/neo4j-go-driver/v5/neo4j/internal/errorutil"
 	rio "github.com/neo4j/neo4j-go-driver/v5/neo4j/internal/racing"
 	"io"
@@ -118,7 +119,7 @@ func processWriteError(err error, ctx context.Context) error {
 			Err:         err,
 		}
 	}
-	if err == context.Canceled {
+	if errors.Is(err, context.Canceled) {
 		return &errorutil.ConnectionWriteCanceled{
 			Err: err,
 		}

neo4j/internal/bolt/dechunker.go

@@ -8,20 +8,21 @@
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
  *
- *      https://www.apache.org/licenses/LICENSE-2.0
+ *     https://www.apache.org/licenses/LICENSE-2.0
  *
- *  Unless required by applicable law or agreed to in writing, software
- *  distributed under the License is distributed on an "AS IS" BASIS,
- *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- *  See the License for the specific language governing permissions and
- *  limitations under the License.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
  */
 
 package bolt
 
 import (
 	"context"
 	"encoding/binary"
+	"errors"
 	"github.com/neo4j/neo4j-go-driver/v5/neo4j/internal/errorutil"
 	rio "github.com/neo4j/neo4j-go-driver/v5/neo4j/internal/racing"
 	"net"
@@ -94,7 +95,7 @@ func processReadError(err error, ctx context.Context, readTimeout time.Duration)
 			Err:         err,
 		}
 	}
-	if err == context.Canceled {
+	if errors.Is(err, context.Canceled) {
 		return &errorutil.ConnectionReadCanceled{
 			Err: err,
 		}

neo4j/internal/errorutil/bolt.go

@@ -21,6 +21,7 @@ package errorutil
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	idb "github.com/neo4j/neo4j-go-driver/v5/neo4j/db"
 	"strings"
@@ -81,33 +82,35 @@ type timeout interface {
 }
 
 func IsTimeoutError(err error) bool {
-	if err == context.DeadlineExceeded {
+	if errors.Is(err, context.DeadlineExceeded) {
 		return true
 	}
 	timeoutErr, ok := err.(timeout)
 	return ok && timeoutErr.Timeout()
 }
 
 func IsFatalDuringDiscovery(err error) bool {
-	if _, ok := err.(*idb.FeatureNotSupportedError); ok {
+	var featureNotSupportedError *idb.FeatureNotSupportedError
+	if errors.As(err, &featureNotSupportedError) {
 		return true
 	}
-	if err, ok := err.(*idb.Neo4jError); ok {
-		if err.Code == "Neo.ClientError.Database.DatabaseNotFound" ||
-			err.Code == "Neo.ClientError.Transaction.InvalidBookmark" ||
-			err.Code == "Neo.ClientError.Transaction.InvalidBookmarkMixture" ||
-			err.Code == "Neo.ClientError.Statement.TypeError" ||
-			err.Code == "Neo.ClientError.Statement.ArgumentError" ||
-			err.Code == "Neo.ClientError.Request.Invalid" {
+	var neo4jErr *idb.Neo4jError
+	if errors.As(err, &neo4jErr) {
+		if neo4jErr.Code == "Neo.ClientError.Database.DatabaseNotFound" ||
+			neo4jErr.Code == "Neo.ClientError.Transaction.InvalidBookmark" ||
+			neo4jErr.Code == "Neo.ClientError.Transaction.InvalidBookmarkMixture" ||
+			neo4jErr.Code == "Neo.ClientError.Statement.TypeError" ||
+			neo4jErr.Code == "Neo.ClientError.Statement.ArgumentError" ||
+			neo4jErr.Code == "Neo.ClientError.Request.Invalid" {
 			return true
 		}
-		if strings.HasPrefix(err.Code, "Neo.ClientError.Security.") &&
-			err.Code != "Neo.ClientError.Security.AuthorizationExpired" {
+		if strings.HasPrefix(neo4jErr.Code, "Neo.ClientError.Security.") &&
+			neo4jErr.Code != "Neo.ClientError.Security.AuthorizationExpired" {
 			return true
 		}
 	}
-	if err == context.DeadlineExceeded ||
-		err == context.Canceled {
+	if errors.Is(err, context.DeadlineExceeded) ||
+		errors.Is(err, context.Canceled) {
 		return true
 	}
 	return false

neo4j/internal/pool/pool.go

@@ -204,12 +204,15 @@ serverLoop:
 	return nil, nil
 }
 
-func (p *Pool) Borrow(ctx context.Context, getServerNames func() []string, wait bool, boltLogger log.BoltLogger, idlenessThreshold time.Duration, auth *idb.ReAuthToken) (idb.Connection, error) {
+func (p *Pool) Borrow(ctx context.Context, getServerNames func(context.Context) ([]string, error), wait bool, boltLogger log.BoltLogger, idlenessThreshold time.Duration, auth *idb.ReAuthToken) (idb.Connection, error) {
 	for {
 		if p.closed {
 			return nil, &errorutil.PoolClosed{}
 		}
-		serverNames := getServerNames()
+		serverNames, err := getServerNames(ctx)
+		if err != nil {
+			return nil, err
+		}
 		if len(serverNames) == 0 {
 			return nil, &errorutil.PoolOutOfServers{}
 		}
@@ -221,8 +224,6 @@ func (p *Pool) Borrow(ctx context.Context, getServerNames func() []string, wait
 			return penalties[i].penalty < penalties[j].penalty
 		})
 
-		var err error
-
 		var conn idb.Connection
 		for _, s := range penalties {
 			conn, err = p.tryBorrow(ctx, s.name, boltLogger, idlenessThreshold, auth)

neo4j/internal/pool/pool_test.go

@@ -673,9 +673,9 @@ func deadConnectionAfterForceReset(name string, idleness time.Time) *ConnFake {
 	return result
 }
 
-func getServers(servers []string) func() []string {
-	return func() []string {
-		return servers
+func getServers(servers []string) func(context.Context) ([]string, error) {
+	return func(context.Context) ([]string, error) {
+		return servers, nil
 	}
 }
 

neo4j/internal/router/no_test.go

@@ -27,14 +27,17 @@ import (
 )
 
 type poolFake struct {
-	borrow   func(names []string, cancel context.CancelFunc, logger log.BoltLogger) (db.Connection, error)
+	borrow   func(ctx context.Context, names []string, cancel context.CancelFunc, logger log.BoltLogger) (db.Connection, error)
 	returned []db.Connection
 	cancel   context.CancelFunc
 }
 
-func (p *poolFake) Borrow(_ context.Context, getServers func() []string, _ bool, logger log.BoltLogger, _ time.Duration, _ *db.ReAuthToken) (db.Connection, error) {
-	servers := getServers()
-	return p.borrow(servers, p.cancel, logger)
+func (p *poolFake) Borrow(ctx context.Context, getServers func(context.Context) ([]string, error), _ bool, logger log.BoltLogger, _ time.Duration, _ *db.ReAuthToken) (db.Connection, error) {
+	servers, err := getServers(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return p.borrow(ctx, servers, p.cancel, logger)
 }
 
 func (p *poolFake) Return(_ context.Context, c db.Connection) {

neo4j/internal/router/readtable.go

@@ -49,10 +49,6 @@ func readTable(
 	for _, router := range routers {
 		var conn db.Connection
 		if conn, err = connectionPool.Borrow(ctx, getStaticServer(router), true, boltLogger, pool.DefaultLivenessCheckThreshold, auth); err != nil {
-			// Check if failed due to context timing out
-			if ctx.Err() != nil {
-				return nil, wrapError(router, ctx.Err())
-			}
 			if errorutil.IsFatalDuringDiscovery(err) {
 				return nil, err
 			}
@@ -75,8 +71,8 @@ func readTable(
 	return nil, err
 }
 
-func getStaticServer(server string) func() []string {
-	return func() []string {
-		return []string{server}
+func getStaticServer(server string) func(context.Context) ([]string, error) {
+	return func(context.Context) ([]string, error) {
+		return []string{server}, nil
 	}
 }

neo4j/internal/router/readtable_test.go

@@ -8,13 +8,13 @@
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
  *
- *      https://www.apache.org/licenses/LICENSE-2.0
+ *     https://www.apache.org/licenses/LICENSE-2.0
  *
- *  Unless required by applicable law or agreed to in writing, software
- *  distributed under the License is distributed on an "AS IS" BASIS,
- *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- *  See the License for the specific language governing permissions and
- *  limitations under the License.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
  */
 
 package router
@@ -27,6 +27,7 @@ import (
 	"github.com/neo4j/neo4j-go-driver/v5/neo4j/internal/errorutil"
 	"github.com/neo4j/neo4j-go-driver/v5/neo4j/log"
 	"testing"
+	"time"
 
 	"github.com/neo4j/neo4j-go-driver/v5/neo4j/db"
 	"github.com/neo4j/neo4j-go-driver/v5/neo4j/internal/testutil"
@@ -71,6 +72,18 @@ func TestReadTableTable(ot *testing.T) {
 		}
 	}
 
+	assertCancelledError := func(t *testing.T, err error) {
+		if !errors.Is(err, context.Canceled) {
+			t.Errorf("Error should be %T but was %T", context.Canceled, err)
+		}
+	}
+
+	assertDeadlineExceededError := func(t *testing.T, err error) {
+		if !errors.Is(err, context.DeadlineExceeded) {
+			t.Errorf("Error should be %T but was %T", context.DeadlineExceeded, err)
+		}
+	}
+
 	cases := []struct {
 		name       string
 		routers    []string
@@ -93,8 +106,9 @@ func TestReadTableTable(ot *testing.T) {
 			assert:    assertNoTable,
 			assertErr: assertRoutingTableError,
 			pool: &poolFake{
-				borrow: func(names []string, cancel context.CancelFunc,
-					_ log.BoltLogger) (idb.Connection, error) {
+				borrow: func(
+					ctx context.Context, names []string, cancel context.CancelFunc, _ log.BoltLogger,
+				) (idb.Connection, error) {
 					return nil, errors.New("borrow fail")
 				},
 			},
@@ -106,8 +120,9 @@ func TestReadTableTable(ot *testing.T) {
 			assert:    assertNoTable,
 			assertErr: assertNeo4jError,
 			pool: &poolFake{
-				borrow: func(names []string, cancel context.CancelFunc,
-					_ log.BoltLogger) (idb.Connection, error) {
+				borrow: func(
+					ctx context.Context, names []string, cancel context.CancelFunc, _ log.BoltLogger,
+				) (idb.Connection, error) {
 					return nil, &db.Neo4jError{Code: "Neo.ClientError.Security.Unauthorized"}
 				},
 			},
@@ -118,8 +133,9 @@ func TestReadTableTable(ot *testing.T) {
 			routers: standardRouters,
 			assert:  assertTable,
 			pool: &poolFake{
-				borrow: func(names []string, cancel context.CancelFunc,
-					_ log.BoltLogger) (idb.Connection, error) {
+				borrow: func(
+					ctx context.Context, names []string, cancel context.CancelFunc, _ log.BoltLogger,
+				) (idb.Connection, error) {
 					return &testutil.ConnFake{Table: &idb.RoutingTable{}}, nil
 				},
 			},
@@ -130,8 +146,9 @@ func TestReadTableTable(ot *testing.T) {
 			routers: standardRouters,
 			assert:  assertTable,
 			pool: &poolFake{
-				borrow: func(names []string, cancel context.CancelFunc,
-					_ log.BoltLogger) (idb.Connection, error) {
+				borrow: func(
+					ctx context.Context, names []string, cancel context.CancelFunc, _ log.BoltLogger,
+				) (idb.Connection, error) {
 					if names[0] == "router2" {
 						return &testutil.ConnFake{Table: &idb.
 							RoutingTable{}}, nil
@@ -147,8 +164,9 @@ func TestReadTableTable(ot *testing.T) {
 			assert:    assertNoTable,
 			assertErr: assertRoutingTableError,
 			pool: &poolFake{
-				borrow: func(names []string, cancel context.CancelFunc,
-					_ log.BoltLogger) (idb.Connection, error) {
+				borrow: func(
+					ctx context.Context, names []string, cancel context.CancelFunc, _ log.BoltLogger,
+				) (idb.Connection, error) {
 					return &testutil.ConnFake{Err: errors.New("GetRoutingTable fail")}, nil
 				},
 			},
@@ -158,17 +176,38 @@ func TestReadTableTable(ot *testing.T) {
 			name:    "Cancel context",
 			routers: standardRouters,
 			pool: &poolFake{
-				borrow: func(names []string, cancel context.CancelFunc,
-					_ log.BoltLogger) (idb.Connection, error) {
+				borrow: func(
+					ctx context.Context, names []string, cancel context.CancelFunc, _ log.BoltLogger,
+				) (idb.Connection, error) {
 					if names[0] == "router2" {
 						panic("Should not be called")
 					}
 					cancel()
-					return nil, errors.New("cancelled")
+					return nil, ctx.Err()
 				},
 			},
 			assert:     assertNoTable,
-			assertErr:  assertRoutingTableError,
+			assertErr:  assertCancelledError,
+			numReturns: 0,
+		},
+		{
+			name:    "Deadline exceeded context",
+			routers: standardRouters,
+			pool: &poolFake{
+				borrow: func(
+					ctx context.Context, names []string, _ context.CancelFunc, _ log.BoltLogger,
+				) (idb.Connection, error) {
+					if names[0] == "router2" {
+						panic("Should not be called")
+					}
+					ctx, cancel := context.WithDeadline(ctx, time.Now().Add(-1*time.Second))
+					err := ctx.Err()
+					cancel()
+					return nil, err
+				},
+			},
+			assert:     assertNoTable,
+			assertErr:  assertDeadlineExceededError,
 			numReturns: 0,
 		},
 	}