Respect connection context timout when retrieving routing table

Author: Peter Wilhelmsson

neo4j/directrouter.go

@@ -19,16 +19,18 @@
 
 package neo4j
 
+import "context"
+
 // A router implementation that never routes
 type directRouter struct {
 	address string
 }
 
-func (r *directRouter) Readers(database string) ([]string, error) {
+func (r *directRouter) Readers(ctx context.Context, database string) ([]string, error) {
 	return []string{r.address}, nil
 }
 
-func (r *directRouter) Writers(database string) ([]string, error) {
+func (r *directRouter) Writers(ctx context.Context, database string) ([]string, error) {
 	return []string{r.address}, nil
 }
 

neo4j/driver.go

@@ -21,6 +21,7 @@
 package neo4j
 
 import (
+	"context"
 	"fmt"
 	"net/url"
 	"strings"
@@ -223,8 +224,8 @@ func routingContextFromUrl(useRouting bool, u *url.URL) (map[string]string, erro
 }
 
 type sessionRouter interface {
-	Readers(database string) ([]string, error)
-	Writers(database string) ([]string, error)
+	Readers(ctx context.Context, database string) ([]string, error)
+	Writers(ctx context.Context, database string) ([]string, error)
 	Invalidate(database string)
 	CleanUp()
 }

neo4j/internal/pool/pool.go

@@ -249,15 +249,15 @@ func (p *Pool) Borrow(ctx context.Context, serverNames []string, wait bool) (db.
 	var err error
 	var conn db.Connection
 	for _, s := range penalties {
-		// Check if we have timed out
-		if timeOut() {
-			return nil, &PoolTimeout{servers: serverNames}
-		}
-
 		conn, err = p.tryBorrow(s.name)
 		if err == nil {
 			return conn, nil
 		}
+
+		// Check if we have timed out after failed borrow
+		if timeOut() {
+			return nil, &PoolTimeout{servers: serverNames}
+		}
 	}
 
 	// If there are no connections for any of the servers, there is no point in waiting for anything

neo4j/internal/router/router.go

@@ -72,7 +72,7 @@ func New(rootRouter string, getRouters func() []string, routerContext map[string
 	return r
 }
 
-func (r *Router) getTable(database string) (*db.RoutingTable, error) {
+func (r *Router) getTable(ctx context.Context, database string) (*db.RoutingTable, error) {
 	now := r.now()
 
 	r.dbRoutersMut.Lock()
@@ -92,20 +92,20 @@ func (r *Router) getTable(database string) (*db.RoutingTable, error) {
 	if dbRouter != nil && len(dbRouter.table.Routers) > 0 {
 		routers := dbRouter.table.Routers
 		r.log.Infof(log.Router, r.logId, "Reading routing table for '%s' from previously known routers: %v", database, routers)
-		table, err = readTable(context.Background(), r.pool, database, routers, r.routerContext)
+		table, err = readTable(ctx, r.pool, database, routers, r.routerContext)
 	}
 
 	// Try initial router if no routers or failed
-	if table == nil || err != nil {
+	if table == nil {
 		r.log.Infof(log.Router, r.logId, "Reading routing table from initial router: %s", r.rootRouter)
-		table, err = readTable(context.Background(), r.pool, database, []string{r.rootRouter}, r.routerContext)
+		table, err = readTable(ctx, r.pool, database, []string{r.rootRouter}, r.routerContext)
 	}
 
 	// Use hook to retrieve possibly different set of routers and retry
-	if err != nil && r.getRouters != nil {
+	if table == nil && r.getRouters != nil {
 		routers := r.getRouters()
 		r.log.Infof(log.Router, r.logId, "Reading routing table for '%s' from custom routers: %v", routers)
-		table, err = readTable(context.Background(), r.pool, database, routers, r.routerContext)
+		table, err = readTable(ctx, r.pool, database, routers, r.routerContext)
 	}
 
 	if err != nil {
@@ -130,8 +130,8 @@ func (r *Router) getTable(database string) (*db.RoutingTable, error) {
 	return table, nil
 }
 
-func (r *Router) Readers(database string) ([]string, error) {
-	table, err := r.getTable(database)
+func (r *Router) Readers(ctx context.Context, database string) ([]string, error) {
+	table, err := r.getTable(ctx, database)
 	if err != nil {
 		return nil, err
 	}
@@ -146,7 +146,7 @@ func (r *Router) Readers(database string) ([]string, error) {
 		r.log.Infof(log.Router, r.logId, "Invalidating routing table, no readers")
 		r.Invalidate(database)
 		r.sleep(100 * time.Millisecond)
-		table, err = r.getTable(database)
+		table, err = r.getTable(ctx, database)
 		if err != nil {
 			return nil, err
 		}
@@ -158,8 +158,8 @@ func (r *Router) Readers(database string) ([]string, error) {
 	return table.Readers, nil
 }
 
-func (r *Router) Writers(database string) ([]string, error) {
-	table, err := r.getTable(database)
+func (r *Router) Writers(ctx context.Context, database string) ([]string, error) {
+	table, err := r.getTable(ctx, database)
 	if err != nil {
 		return nil, err
 	}
@@ -174,7 +174,7 @@ func (r *Router) Writers(database string) ([]string, error) {
 		r.log.Infof(log.Router, r.logId, "Invalidating routing table, no writers")
 		r.Invalidate(database)
 		r.sleep(100 * time.Millisecond)
-		table, err = r.getTable(database)
+		table, err = r.getTable(ctx, database)
 		if err != nil {
 			return nil, err
 		}

neo4j/internal/router/router_test.go

@@ -65,14 +65,14 @@ func TestMultithreading(t *testing.T) {
 	wg.Add(2)
 	consumer := func() {
 		for i := 0; i < 30; i++ {
-			readers, err := router.Readers(dbName)
+			readers, err := router.Readers(context.Background(), dbName)
 			if len(readers) != 2 {
 				t.Error("Wrong number of readers")
 			}
 			if err != nil {
 				t.Error(err)
 			}
-			writers, err := router.Writers(dbName)
+			writers, err := router.Writers(context.Background(), dbName)
 			if len(writers) != 1 {
 				t.Error("Wrong number of writers")
 			}
@@ -114,25 +114,25 @@ func TestRespectsTimeToLiveAndInvalidate(t *testing.T) {
 	dbName := "dbname"
 
 	// First access should trigger initial table read
-	router.Readers(dbName)
+	router.Readers(context.Background(), dbName)
 	assertNum(t, numfetch, 1, "Should have fetched initial")
 
 	// Second access with time set to same should not trigger a read
-	router.Readers(dbName)
+	router.Readers(context.Background(), dbName)
 	assertNum(t, numfetch, 1, "Should not have have fetched")
 
 	// Third access with time passed table due should trigger fetch
 	n = n.Add(2 * time.Second)
-	router.Readers(dbName)
+	router.Readers(context.Background(), dbName)
 	assertNum(t, numfetch, 2, "Should have have fetched")
 
 	// Just another one to make sure we're cached
-	router.Readers(dbName)
+	router.Readers(context.Background(), dbName)
 	assertNum(t, numfetch, 2, "Should not have have fetched")
 
 	// Invalidate should force fetching
 	router.Invalidate(dbName)
-	router.Readers(dbName)
+	router.Readers(context.Background(), dbName)
 	assertNum(t, numfetch, 3, "Should have have fetched")
 }
 
@@ -160,13 +160,13 @@ func TestUsesRootRouterWhenPreviousRoutersFails(t *testing.T) {
 	dbName := "dbname"
 
 	// First access should trigger initial table read from root router
-	router.Readers(dbName)
+	router.Readers(context.Background(), dbName)
 	if borrows[0][0] != "rootRouter" {
 		t.Errorf("Should have connected to root upon first router request")
 	}
 	// Next access should go to otherRouter
 	n = n.Add(2 * time.Second)
-	router.Readers(dbName)
+	router.Readers(context.Background(), dbName)
 	if borrows[1][0] != "otherRouter" {
 		t.Errorf("Should have queried other router")
 	}
@@ -191,7 +191,7 @@ func TestUsesRootRouterWhenPreviousRoutersFails(t *testing.T) {
 		return &testutil.ConnFake{Table: &db.RoutingTable{TimeToLive: 1, Readers: []string{"aReader"}}}, nil
 	}
 	n = n.Add(2 * time.Second)
-	readers, err := router.Readers(dbName)
+	readers, err := router.Readers(context.Background(), dbName)
 	if err != nil {
 		t.Error(err)
 	}
@@ -219,7 +219,7 @@ func TestUseGetRoutersHookWhenInitialRouterFails(t *testing.T) {
 	dbName := "dbname"
 
 	// Trigger read of routing table
-	router.Readers(dbName)
+	router.Readers(context.Background(), dbName)
 
 	expected := []string{rootRouter}
 	expected = append(expected, backupRouters...)
@@ -247,7 +247,7 @@ func TestWritersFailAfterNRetries(t *testing.T) {
 	dbName := "dbname"
 
 	// Should trigger a lot of retries to get a writer until it finally fails
-	writers, err := router.Writers(dbName)
+	writers, err := router.Writers(context.Background(), dbName)
 	if err == nil {
 		t.Error("Should have failed")
 	}
@@ -285,7 +285,7 @@ func TestWritersRetriesWhenNoWriters(t *testing.T) {
 
 	// Should trigger initial table read that contains no writers and a second table read
 	// that gets the writers
-	writers, err := router.Writers(dbName)
+	writers, err := router.Writers(context.Background(), dbName)
 	if err != nil {
 		t.Errorf("Got error: %s", err)
 	}
@@ -323,7 +323,7 @@ func TestReadersRetriesWhenNoReaders(t *testing.T) {
 
 	// Should trigger initial table read that contains no readers and a second table read
 	// that gets the readers
-	readers, err := router.Readers(dbName)
+	readers, err := router.Readers(context.Background(), dbName)
 	if err != nil {
 		t.Errorf("Got error: %s", err)
 	}
@@ -349,8 +349,8 @@ func TestCleanUp(t *testing.T) {
 	router := New("router", func() []string { return []string{} }, nil, pool, logger, "routerid")
 	router.now = func() time.Time { return now }
 
-	router.Readers("db1")
-	router.Readers("db2")
+	router.Readers(context.Background(), "db1")
+	router.Readers(context.Background(), "db2")
 
 	// Should be a router for each requested database
 	if len(router.dbRouters) != 2 {

neo4j/internal/testutil/routerfake.go

@@ -19,6 +19,10 @@
 
 package testutil
 
+import (
+	"context"
+)
+
 type RouterFake struct {
 	Invalidated   bool
 	InvalidatedDb string
@@ -33,11 +37,11 @@ func (r *RouterFake) Invalidate(database string) {
 	r.Invalidated = true
 }
 
-func (r *RouterFake) Readers(database string) ([]string, error) {
+func (r *RouterFake) Readers(ctx context.Context, database string) ([]string, error) {
 	return r.ReadersRet, r.Err
 }
 
-func (r *RouterFake) Writers(database string) ([]string, error) {
+func (r *RouterFake) Writers(ctx context.Context, database string) ([]string, error) {
 	return r.WritersRet, r.Err
 }
 

neo4j/session.go

@@ -327,20 +327,15 @@ func (s *session) WriteTransaction(
 	return s.runRetriable(db.WriteMode, work, configurers...)
 }
 
-func (s *session) getServers(mode db.AccessMode) ([]string, error) {
+func (s *session) getServers(ctx context.Context, mode db.AccessMode) ([]string, error) {
 	if mode == db.ReadMode {
-		return s.router.Readers(s.databaseName)
+		return s.router.Readers(ctx, s.databaseName)
 	} else {
-		return s.router.Writers(s.databaseName)
+		return s.router.Writers(ctx, s.databaseName)
 	}
 }
 
 func (s *session) getConnection(mode db.AccessMode) (db.Connection, error) {
-	servers, err := s.getServers(mode)
-	if err != nil {
-		return nil, wrapError(err)
-	}
-
 	var ctx context.Context
 	if s.config.ConnectionAcquisitionTimeout > 0 {
 		var cancel context.CancelFunc
@@ -351,6 +346,11 @@ func (s *session) getConnection(mode db.AccessMode) (db.Connection, error) {
 	} else {
 		ctx = context.Background()
 	}
+	servers, err := s.getServers(ctx, mode)
+	if err != nil {
+		return nil, wrapError(err)
+	}
+
 	conn, err := s.pool.Borrow(ctx, servers, s.config.ConnectionAcquisitionTimeout != 0)
 	if err != nil {
 		return nil, wrapError(err)