Use routing per database from API level

Author: Peter Wilhelmsson

neo4j/directrouter.go

@@ -24,13 +24,13 @@ type directRouter struct {
 	server string
 }
 
-func (r *directRouter) Readers() ([]string, error) {
+func (r *directRouter) Readers(database string) ([]string, error) {
 	return []string{r.server}, nil
 }
 
-func (r *directRouter) Writers() ([]string, error) {
+func (r *directRouter) Writers(database string) ([]string, error) {
 	return []string{r.server}, nil
 }
 
-func (r *directRouter) Invalidate() {
+func (r *directRouter) Invalidate(database string) {
 }

neo4j/driver.go

@@ -174,9 +174,9 @@ func routingContextFromUrl(u *url.URL) (map[string]string, error) {
 }
 
 type sessionRouter interface {
-	Readers() ([]string, error)
-	Writers() ([]string, error)
-	Invalidate()
+	Readers(database string) ([]string, error)
+	Writers(database string) ([]string, error)
+	Invalidate(database string)
 }
 
 type driver struct {

neo4j/internal/router/readtable.go

@@ -28,7 +28,7 @@ import (
 
 // Tries to read routing table from any of the specified routers using new or existing connection
 // from the supplied pool.
-func readTable(ctx context.Context, pool Pool, routers []string, routerContext map[string]string) (*db.RoutingTable, error) {
+func readTable(ctx context.Context, pool Pool, database string, routers []string, routerContext map[string]string) (*db.RoutingTable, error) {
 	// Preserve last error to be returned, set a default for case of no routers
 	var err error = &ReadRoutingTableError{}
 
@@ -55,7 +55,7 @@ func readTable(ctx context.Context, pool Pool, routers []string, routerContext m
 		}
 
 		var table *db.RoutingTable
-		table, err = discovery.GetRoutingTable(db.DefaultDatabase, routerContext)
+		table, err = discovery.GetRoutingTable(database, routerContext)
 		if err == nil {
 			return table, nil
 		}

neo4j/internal/router/readtable_test.go

@@ -146,7 +146,7 @@ func TestReadTableTable(ot *testing.T) {
 		ot.Run(c.name, func(t *testing.T) {
 			ctx, cancel := context.WithCancel(context.Background())
 			c.pool.cancel = cancel
-			table, err := readTable(ctx, c.pool, c.routers, nil)
+			table, err := readTable(ctx, c.pool, "dbname", c.routers, nil)
 			c.assert(t, table, err)
 			if c.numReturns != len(c.pool.returned) {
 				t.Errorf("Expected %d returned connections but %d was returned", c.numReturns, len(c.pool.returned))

neo4j/internal/router/router.go

@@ -34,13 +34,17 @@ import (
 
 const missingWriterRetries = 100
 
+type databaseRouter struct {
+	dueUnix int64
+	table   *db.RoutingTable
+}
+
 // Thread safe
 type Router struct {
 	routerContext map[string]string
 	pool          Pool
-	table         *db.RoutingTable
-	dueUnix       int64
-	tableMut      sync.Mutex
+	dbRouters     map[string]*databaseRouter
+	dbRoutersMut  sync.Mutex
 	now           func() time.Time
 	sleep         func(time.Duration)
 	rootRouter    string
@@ -63,6 +67,7 @@ func New(rootRouter string, getRouters func() []string, routerContext map[string
 		getRouters:    getRouters,
 		routerContext: routerContext,
 		pool:          pool,
+		dbRouters:     make(map[string]*databaseRouter),
 		now:           time.Now,
 		sleep:         time.Sleep,
 		log:           logger,
@@ -72,54 +77,57 @@ func New(rootRouter string, getRouters func() []string, routerContext map[string
 	return r
 }
 
-func (r *Router) getTable() (*db.RoutingTable, error) {
-	r.tableMut.Lock()
-	defer r.tableMut.Unlock()
-
+func (r *Router) getTable(database string) (*db.RoutingTable, error) {
 	now := r.now()
 
-	if r.table != nil && now.Unix() < r.dueUnix {
-		return r.table, nil
+	r.dbRoutersMut.Lock()
+	defer r.dbRoutersMut.Unlock()
+
+	dbRouter := r.dbRouters[database]
+	if dbRouter != nil && now.Unix() < dbRouter.dueUnix {
+		return dbRouter.table, nil
 	}
 
 	var routers []string
-	if r.table != nil {
-		routers = r.table.Routers
+	if dbRouter != nil {
+		routers = dbRouter.table.Routers
 	}
 	if len(routers) == 0 {
 		routers = []string{r.rootRouter}
 	}
 
-	r.log.Infof(r.logId, "Reading routing table from any of %v", routers)
-	table, err := readTable(context.Background(), r.pool, routers, r.routerContext)
+	r.log.Infof(r.logId, "Reading routing table for '%s' from any of %v", database, routers)
+	table, err := readTable(context.Background(), r.pool, database, routers, r.routerContext)
 	if err != nil {
 		// Use hook to retrieve possibly different set of routers and retry
 		if r.getRouters != nil {
 			routers = r.getRouters()
-			table, err = readTable(context.Background(), r.pool, routers, r.routerContext)
+			table, err = readTable(context.Background(), r.pool, database, routers, r.routerContext)
 		}
 		if err != nil {
 			r.log.Error(r.logId, err)
 			return nil, err
 		}
 	}
-	r.table = table
-	r.dueUnix = now.Add(time.Duration(table.TimeToLive) * time.Second).Unix()
-	r.log.Debugf(r.logId, "New routing table, TTL %d", table.TimeToLive)
+	r.dbRouters[database] = &databaseRouter{
+		table:   table,
+		dueUnix: now.Add(time.Duration(table.TimeToLive) * time.Second).Unix(),
+	}
+	r.log.Debugf(r.logId, "New routing table for '%s', TTL %d", database, table.TimeToLive)
 
 	return table, nil
 }
 
-func (r *Router) Readers() ([]string, error) {
-	table, err := r.getTable()
+func (r *Router) Readers(database string) ([]string, error) {
+	table, err := r.getTable(database)
 	if err != nil {
 		return nil, err
 	}
 	return table.Readers, nil
 }
 
-func (r *Router) Writers() ([]string, error) {
-	table, err := r.getTable()
+func (r *Router) Writers(database string) ([]string, error) {
+	table, err := r.getTable(database)
 	if err != nil {
 		return nil, err
 	}
@@ -133,11 +141,8 @@ func (r *Router) Writers() ([]string, error) {
 		}
 		r.log.Debugf(r.logId, "Invalidating routing table, no writers")
 		r.sleep(100 * time.Millisecond)
-		r.tableMut.Lock()
-		// Reset due time to keep list of routers
-		r.dueUnix = 0
-		r.tableMut.Unlock()
-		table, err = r.getTable()
+		r.Invalidate(database)
+		table, err = r.getTable(database)
 		if err != nil {
 			return nil, err
 		}
@@ -153,11 +158,14 @@ func (r *Router) Context() map[string]string {
 	return r.routerContext
 }
 
-func (r *Router) Invalidate() {
-	r.tableMut.Lock()
-	defer r.tableMut.Unlock()
+func (r *Router) Invalidate(database string) {
+	r.log.Infof(r.logId, "Invalidating routing table for '%s'", database)
+	r.dbRoutersMut.Lock()
+	defer r.dbRoutersMut.Unlock()
 	// Reset due time to the 70s, this will make next access refresh the routing table using
 	// last set of routers instead of the original one.
-	r.dueUnix = 0
-	r.log.Infof(r.logId, "Invalidating routing table")
+	dbRouter := r.dbRouters[database]
+	if dbRouter != nil {
+		dbRouter.dueUnix = 0
+	}
 }

neo4j/internal/router/router_test.go

@@ -37,8 +37,6 @@ var logger = &log.ConsoleLogger{Errors: true, Infos: true, Warns: true}
 // Verifies that concurrent access works as expected relying on the race detector to
 // report supicious behavior.
 func TestMultithreading(t *testing.T) {
-	wg := sync.WaitGroup{}
-	wg.Add(2)
 
 	// Setup a router that needs to read the routing table essentially on every access to
 	// stress threading a bit more.
@@ -52,28 +50,34 @@ func TestMultithreading(t *testing.T) {
 	}
 	n := time.Now()
 	router := New("router", func() []string { return []string{} }, nil, pool, logger)
+	mut := sync.Mutex{}
 	router.now = func() time.Time {
+		// Need to lock here to make race detector happy
+		mut.Lock()
+		defer mut.Unlock()
 		n = n.Add(time.Duration(table.TimeToLive) * time.Second * 2)
 		return n
 	}
 
+	dbName := "dbname"
+	wg := sync.WaitGroup{}
+	wg.Add(2)
 	consumer := func() {
 		for i := 0; i < 30; i++ {
-			readers, err := router.Readers()
+			readers, err := router.Readers(dbName)
 			if len(readers) != 2 {
 				t.Error("Wrong number of readers")
 			}
 			if err != nil {
 				t.Error(err)
 			}
-			writers, err := router.Writers()
+			writers, err := router.Writers(dbName)
 			if len(writers) != 1 {
 				t.Error("Wrong number of writers")
 			}
 			if err != nil {
 				t.Error(err)
 			}
-
 		}
 		wg.Done()
 	}
@@ -106,27 +110,28 @@ func TestRespectsTimeToLiveAndInvalidate(t *testing.T) {
 	router.now = func() time.Time {
 		return n
 	}
+	dbName := "dbname"
 
 	// First access should trigger initial table read
-	router.Readers()
+	router.Readers(dbName)
 	assertNum(t, numfetch, 1, "Should have fetched initial")
 
 	// Second access with time set to same should not trigger a read
-	router.Readers()
+	router.Readers(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()
+	router.Readers(dbName)
 	assertNum(t, numfetch, 2, "Should have have fetched")
 
 	// Just another one to make sure we're cached
-	router.Readers()
+	router.Readers(dbName)
 	assertNum(t, numfetch, 2, "Should not have have fetched")
 
 	// Invalidate should force fetching
-	router.Invalidate()
-	router.Readers()
+	router.Invalidate(dbName)
+	router.Readers(dbName)
 	assertNum(t, numfetch, 3, "Should have have fetched")
 }
 
@@ -143,9 +148,10 @@ func TestUseGetRoutersHookWhenInitialRouterFails(t *testing.T) {
 	rootRouter := "rootRouter"
 	backupRouters := []string{"bup1", "bup2"}
 	router := New(rootRouter, func() []string { return backupRouters }, nil, pool, logger)
+	dbName := "dbname"
 
 	// Trigger read of routing table
-	router.Readers()
+	router.Readers(dbName)
 
 	expected := []string{rootRouter}
 	expected = append(expected, backupRouters...)
@@ -170,9 +176,10 @@ func TestWritersFailAfterNRetries(t *testing.T) {
 	router.sleep = func(time.Duration) {
 		numsleep++
 	}
+	dbName := "dbname"
 
 	// Should trigger a lot of retries to get a writer until it finally fails
-	writers, err := router.Writers()
+	writers, err := router.Writers(dbName)
 	if err == nil {
 		t.Error("Should have failed")
 	}
@@ -206,10 +213,11 @@ func TestWritersRetriesWhenNoWriters(t *testing.T) {
 	router.sleep = func(time.Duration) {
 		numsleep++
 	}
+	dbName := "dbname"
 
 	// Should trigger initial table read that contains no writers and a second table read
 	// that gets the writers
-	writers, err := router.Writers()
+	writers, err := router.Writers(dbName)
 	if err != nil {
 		t.Errorf("Got error: %s", err)
 	}

neo4j/session.go

@@ -325,7 +325,7 @@ func (s *session) runRetriable(
 			switch {
 			case e.IsRetriableCluster():
 				// Force routing tables to be updated before trying again
-				s.router.Invalidate()
+				s.router.Invalidate(s.databaseName)
 				maxClusterErrors--
 				if maxClusterErrors < 0 {
 					s.log.Errorf(s.logId, "Retriable transaction failed due to encountering too many cluster errors")
@@ -362,9 +362,9 @@ func (s *session) borrowConn(mode db.AccessMode) error {
 	var servers []string
 	var err error
 	if mode == db.ReadMode {
-		servers, err = s.router.Readers()
+		servers, err = s.router.Readers(s.databaseName)
 	} else {
-		servers, err = s.router.Writers()
+		servers, err = s.router.Writers(s.databaseName)
 	}
 	if err != nil {
 		return err