Cleanup router when session is closed

neo4j/directrouter.go

@@ -34,3 +34,6 @@ func (r *directRouter) Writers(database string) ([]string, error) {
 
 func (r *directRouter) Invalidate(database string) {
 }
+
+func (r *directRouter) CleanUp() {
+}

neo4j/driver.go

@@ -177,6 +177,7 @@ type sessionRouter interface {
 	Readers(database string) ([]string, error)
 	Writers(database string) ([]string, error)
 	Invalidate(database string)
+	CleanUp()
 }
 
 type driver struct {

neo4j/internal/router/router.go

@@ -169,3 +169,16 @@ func (r *Router) Invalidate(database string) {
 		dbRouter.dueUnix = 0
 	}
 }
+
+func (r *Router) CleanUp() {
+	r.log.Debugf(r.logId, "Cleaning up")
+	now := r.now().Unix()
+	r.dbRoutersMut.Lock()
+	defer r.dbRoutersMut.Unlock()
+
+	for dbName, dbRouter := range r.dbRouters {
+		if now > dbRouter.dueUnix {
+			delete(r.dbRouters, dbName)
+		}
+	}
+}

neo4j/internal/router/router_test.go

@@ -184,7 +184,7 @@ func TestWritersFailAfterNRetries(t *testing.T) {
 		t.Error("Should have failed")
 	}
 	if writers != nil {
-		t.Error("Should'nt have any writers")
+		t.Error("Should not have any writers")
 	}
 	if numsleep <= 0 {
 		t.Error("Should have slept plenty")
@@ -231,3 +231,35 @@ func TestWritersRetriesWhenNoWriters(t *testing.T) {
 		t.Error("Should have slept once")
 	}
 }
+
+func TestCleanUp(t *testing.T) {
+	table := &db.RoutingTable{TimeToLive: 1}
+	pool := &poolFake{
+		borrow: func(names []string, cancel context.CancelFunc) (poolpackage.Connection, error) {
+			return &connFake{table: table}, nil
+		},
+	}
+	now := time.Now()
+	router := New("router", func() []string { return []string{} }, nil, pool, logger)
+	router.now = func() time.Time { return now }
+
+	router.Readers("db1")
+	router.Readers("db2")
+
+	// Should be a router for each requested database
+	if len(router.dbRouters) != 2 {
+		t.Fatal("Should be two routing tables, one for each database")
+	}
+
+	// Should not remove these since they still have time to live
+	router.CleanUp()
+	if len(router.dbRouters) != 2 {
+		t.Fatal("Should not have removed routing tables")
+	}
+
+	router.now = func() time.Time { return now.Add(1 * time.Minute) }
+	router.CleanUp()
+	if len(router.dbRouters) != 0 {
+		t.Fatal("Should have cleaned up")
+	}
+}

neo4j/session.go

@@ -471,6 +471,7 @@ func (s *session) Close() error {
 	// Schedule cleanups
 	go func() {
 		s.pool.CleanUp()
+		s.router.CleanUp()
 	}()
 	return nil
 }