ADR-026 Home Database Cache (#618)

* ADR-026 Home Database Cache

---------

Co-authored-by: Robsdedude <dev@rouvenbauer.de>
Co-authored-by: Robsdedude <rouven.bauer@neo4j.com>

Author: StephenCathcart

neo4j/directrouter.go

@@ -34,15 +34,15 @@ func (r *directRouter) InvalidateReader(string, string) {}
 
 func (r *directRouter) InvalidateServer(string) {}
 
-func (r *directRouter) GetOrUpdateReaders(context.Context, func(context.Context) ([]string, error), string, *db.ReAuthToken, log.BoltLogger) ([]string, error) {
+func (r *directRouter) GetOrUpdateReaders(context.Context, func(context.Context) ([]string, error), db.DatabaseSelection, *db.ReAuthToken, log.BoltLogger, func(string)) ([]string, error) {
 	return []string{r.address}, nil
 }
 
 func (r *directRouter) Readers(string) []string {
 	return []string{r.address}
 }
 
-func (r *directRouter) GetOrUpdateWriters(context.Context, func(context.Context) ([]string, error), string, *db.ReAuthToken, log.BoltLogger) ([]string, error) {
+func (r *directRouter) GetOrUpdateWriters(context.Context, func(context.Context) ([]string, error), db.DatabaseSelection, *db.ReAuthToken, log.BoltLogger, func(string)) ([]string, error) {
 	return []string{r.address}, nil
 }
 

neo4j/driver_with_context.go

@@ -21,6 +21,7 @@ package neo4j
 import (
 	"context"
 	"fmt"
+	"github.com/neo4j/neo4j-go-driver/v5/neo4j/internal/homedb"
 	"net/url"
 	"strings"
 	"sync"
@@ -219,8 +220,14 @@ func NewDriverWithContext(target string, auth auth.TokenManager, configurers ...
 	d.connector.RoutingContext = routingContext
 	d.connector.Config = d.config
 
+	// Create cache for home database
+	d.cache, err = homedb.NewCache(homedb.DefaultCacheMaxSize)
+	if err != nil {
+		return nil, err
+	}
+
 	// Let the pool use the same log ID as the driver to simplify log reading.
-	d.pool = pool.New(d.config, d.connector.Connect, d.log, d.logId)
+	d.pool = pool.New(d.config, d.connector.Connect, d.log, d.logId, d.cache)
 
 	if !routing {
 		d.router = &directRouter{address: address}
@@ -295,12 +302,12 @@ type sessionRouter interface {
 	// note: bookmarks are lazily supplied, only when a new routing table needs to be fetched
 	// this is needed because custom bookmark managers may provide bookmarks from external systems
 	// they should not be called when it is not needed (e.g. when a routing table is cached)
-	GetOrUpdateReaders(ctx context.Context, bookmarks func(context.Context) ([]string, error), database string, auth *idb.ReAuthToken, boltLogger log.BoltLogger) ([]string, error)
+	GetOrUpdateReaders(ctx context.Context, bookmarks func(context.Context) ([]string, error), dbSelection idb.DatabaseSelection, auth *idb.ReAuthToken, boltLogger log.BoltLogger, onRoutingTableUpdated func(string)) ([]string, error)
 	// Readers returns the list of servers that can serve reads on the requested database.
 	Readers(database string) []string
 	// GetOrUpdateWriters returns the list of servers that can serve writes on the requested database.
 	// note: bookmarks are lazily supplied, see Readers documentation to learn why
-	GetOrUpdateWriters(ctx context.Context, bookmarks func(context.Context) ([]string, error), database string, auth *idb.ReAuthToken, boltLogger log.BoltLogger) ([]string, error)
+	GetOrUpdateWriters(ctx context.Context, bookmarks func(context.Context) ([]string, error), dbSelection idb.DatabaseSelection, auth *idb.ReAuthToken, boltLogger log.BoltLogger, onRoutingTableUpdated func(string)) ([]string, error)
 	// Writers returns the list of servers that can serve writes on the requested database.
 	Writers(database string) []string
 	// GetNameOfDefaultDatabase returns the name of the default database for the specified user.
@@ -329,15 +336,14 @@ type driverWithContext struct {
 	// this is *not* used by default by user-created session (see NewSession)
 	executeQueryBookmarkManager BookmarkManager
 	auth                        auth.TokenManager
+	cache                       *homedb.Cache
 }
 
 func (d *driverWithContext) Target() url.URL {
 	return *d.target
 }
 
-// TODO 6.0: remove unused Context parameter
-
-func (d *driverWithContext) NewSession(_ context.Context, config SessionConfig) SessionWithContext {
+func (d *driverWithContext) NewSession(ctx context.Context, config SessionConfig) SessionWithContext {
 	if config.DatabaseName == "" {
 		config.DatabaseName = idb.DefaultDatabase
 	}
@@ -363,7 +369,7 @@ func (d *driverWithContext) NewSession(_ context.Context, config SessionConfig)
 		return &erroredSessionWithContext{
 			err: &UsageError{Message: "Trying to create session on closed driver"}}
 	}
-	return newSessionWithContext(d.config, config, d.router, d.pool, d.log, reAuthToken)
+	return newSessionWithContext(ctx, d.config, config, d.router, d.pool, d.cache, d.log, reAuthToken)
 }
 
 func (d *driverWithContext) VerifyConnectivity(ctx context.Context) error {

neo4j/driver_with_context_testkit.go

@@ -43,11 +43,16 @@ func ForceRoutingTableUpdate(d DriverWithContext, database string, bookmarks []s
 		FromSession: false,
 		ForceReAuth: false,
 	}
-	_, err := driver.router.GetOrUpdateReaders(ctx, getBookmarks, database, auth, logger)
+	dbSelection := idb.DatabaseSelection{Name: database}
+	_, err := driver.router.GetOrUpdateReaders(ctx, getBookmarks, dbSelection, auth, logger, func(db string) {
+		if dbSelection.Name == "" {
+			dbSelection.Name = db
+		}
+	})
 	if err != nil {
 		return errorutil.WrapError(err)
 	}
-	_, err = driver.router.GetOrUpdateWriters(ctx, getBookmarks, database, auth, logger)
+	_, err = driver.router.GetOrUpdateWriters(ctx, getBookmarks, dbSelection, auth, logger, nil)
 	return errorutil.WrapError(err)
 }
 

neo4j/internal/bolt/bolt3.go

@@ -872,3 +872,11 @@ func (b *bolt3) GetCurrentAuth() (auth.TokenManager, iauth.Token) {
 func (b *bolt3) Telemetry(telemetry.API, func()) {
 	// TELEMETRY not support by this protocol version, so we ignore it.
 }
+
+func (b *bolt3) SetPinHomeDatabaseCallback(func(context.Context, string)) {
+	// Home database not supported by this protocol version, so we ignore it.
+}
+
+func (b *bolt3) IsSsrEnabled() bool {
+	return false
+}

neo4j/internal/bolt/bolt4.go

@@ -985,6 +985,14 @@ func (b *bolt4) Telemetry(telemetry.API, func()) {
 	// TELEMETRY not support by this protocol version, so we ignore it.
 }
 
+func (b *bolt4) SetPinHomeDatabaseCallback(func(context.Context, string)) {
+	// Home database not supported by this protocol version, so we ignore it.
+}
+
+func (b *bolt4) IsSsrEnabled() bool {
+	return false
+}
+
 func (b *bolt4) helloResponseHandler(checkUtcPatch bool) responseHandler {
 	return b.expectedSuccessHandler(b.onHelloSuccess(checkUtcPatch))
 }

neo4j/internal/bolt/bolt5.go

@@ -50,6 +50,11 @@ const (
 // Default fetch size
 const bolt5FetchSize = 1000
 
+const (
+	telemetryEnabledHintName = "telemetry.enabled"
+	ssrEnabledHintName       = "ssr.enabled"
+)
+
 type internalTx5 struct {
 	mode               idb.AccessMode
 	bookmarks          []string
@@ -93,28 +98,30 @@ func (i *internalTx5) toMeta(logger log.Logger, logId string, version db.Protoco
 }
 
 type bolt5 struct {
-	state            int
-	txId             idb.TxHandle
-	streams          openstreams
-	conn             io.ReadWriteCloser
-	serverName       string
-	queue            messageQueue
-	connId           string
-	logId            string
-	serverVersion    string
-	bookmark         string // Last bookmark
-	birthDate        time.Time
-	log              log.Logger
-	databaseName     string
-	err              error // Last fatal error
-	minor            int
-	lastQid          int64 // Last seen qid
-	idleDate         time.Time
-	auth             map[string]any
-	authManager      auth.TokenManager
-	resetAuth        bool
-	errorListener    ConnectionErrorListener
-	telemetryEnabled bool
+	state                   int
+	txId                    idb.TxHandle
+	streams                 openstreams
+	conn                    io.ReadWriteCloser
+	serverName              string
+	queue                   messageQueue
+	connId                  string
+	logId                   string
+	serverVersion           string
+	bookmark                string // Last bookmark
+	birthDate               time.Time
+	log                     log.Logger
+	databaseName            string
+	err                     error // Last fatal error
+	minor                   int
+	lastQid                 int64 // Last seen qid
+	idleDate                time.Time
+	auth                    map[string]any
+	authManager             auth.TokenManager
+	resetAuth               bool
+	errorListener           ConnectionErrorListener
+	telemetryEnabled        bool
+	ssrEnabled              bool
+	pinHomeDatabaseCallback func(context.Context, string)
 }
 
 func NewBolt5(
@@ -322,7 +329,7 @@ func (b *bolt5) TxBegin(
 		notificationConfig: txConfig.NotificationConfig,
 	}
 
-	b.queue.appendBegin(tx.toMeta(b.log, b.logId, b.Version()), b.beginResponseHandler())
+	b.queue.appendBegin(tx.toMeta(b.log, b.logId, b.Version()), b.beginResponseHandler(ctx))
 	if syncMessages {
 		if b.queue.send(ctx); b.err != nil {
 			return 0, b.err
@@ -562,7 +569,7 @@ func (b *bolt5) run(ctx context.Context, cypher string, params map[string]any, r
 	fetchSize := b.normalizeFetchSize(rawFetchSize)
 	stream := &stream{fetchSize: fetchSize}
 	b.Version()
-	b.queue.appendRun(cypher, params, tx.toMeta(b.log, b.logId, b.Version()), b.runResponseHandler(stream))
+	b.queue.appendRun(cypher, params, tx.toMeta(b.log, b.logId, b.Version()), b.runResponseHandler(ctx, stream))
 	b.queue.appendPullN(fetchSize, b.pullResponseHandler(stream))
 	if b.queue.send(ctx); b.err != nil {
 		return nil, b.err
@@ -850,6 +857,14 @@ func (b *bolt5) SetBoltLogger(boltLogger log.BoltLogger) {
 	b.queue.setBoltLogger(boltLogger)
 }
 
+func (b *bolt5) SetPinHomeDatabaseCallback(callback func(context.Context, string)) {
+	b.pinHomeDatabaseCallback = callback
+}
+
+func (b *bolt5) IsSsrEnabled() bool {
+	return b.ssrEnabled
+}
+
 func (b *bolt5) ReAuth(ctx context.Context, auth *idb.ReAuthToken) error {
 	if b.minor == 0 {
 		return b.fallbackReAuth(ctx, auth)
@@ -998,12 +1013,19 @@ func (b *bolt5) routeResponseHandler(table **idb.RoutingTable) responseHandler {
 	})
 }
 
-func (b *bolt5) beginResponseHandler() responseHandler {
-	return b.expectedSuccessHandler(onSuccessNoOp)
+func (b *bolt5) beginResponseHandler(ctx context.Context) responseHandler {
+	return b.expectedSuccessHandler(func(beginSuccess *success) {
+		if b.pinHomeDatabaseCallback != nil && beginSuccess.db != "" {
+			b.pinHomeDatabaseCallback(ctx, beginSuccess.db)
+		}
+	})
 }
 
-func (b *bolt5) runResponseHandler(stream *stream) responseHandler {
+func (b *bolt5) runResponseHandler(ctx context.Context, stream *stream) responseHandler {
 	return b.expectedSuccessHandler(func(runSuccess *success) {
+		if b.pinHomeDatabaseCallback != nil && runSuccess.db != "" {
+			b.pinHomeDatabaseCallback(ctx, runSuccess.db)
+		}
 		stream.attached = true
 		stream.keys = runSuccess.fields
 		stream.qid = runSuccess.qid
@@ -1124,7 +1146,8 @@ func (b *bolt5) onHelloSuccess(helloSuccess *success) {
 	b.logId = connectionLogId
 	b.queue.setLogId(connectionLogId)
 	b.initializeReadTimeoutHint(helloSuccess.configurationHints)
-	b.initializeTelemetryHint(helloSuccess.configurationHints)
+	b.initializeTelemetryEnabledHint(helloSuccess.configurationHints)
+	b.initializeSsrEnabledHint(helloSuccess.configurationHints)
 }
 
 func (b *bolt5) onCommitSuccess(commitSuccess *success) {
@@ -1172,19 +1195,30 @@ func (b *bolt5) initializeReadTimeoutHint(hints map[string]any) {
 	b.queue.in.connReadTimeout = time.Duration(readTimeout) * time.Second
 }
 
-const readTelemetryHintName = "telemetry.enabled"
+func (b *bolt5) initializeTelemetryEnabledHint(hints map[string]any) {
+	telemetryEnabledHint, ok := hints[telemetryEnabledHintName]
+	if !ok {
+		return
+	}
+	telemetryEnabled, ok := telemetryEnabledHint.(bool)
+	if !ok {
+		b.log.Infof(log.Bolt5, b.logId, `invalid %q value: %v, ignoring hint. Only boolean values are accepted`, telemetryEnabledHintName, telemetryEnabledHint)
+		return
+	}
+	b.telemetryEnabled = telemetryEnabled
+}
 
-func (b *bolt5) initializeTelemetryHint(hints map[string]any) {
-	readTelemetryHint, ok := hints[readTelemetryHintName]
+func (b *bolt5) initializeSsrEnabledHint(hints map[string]any) {
+	ssrEnabledHint, ok := hints[ssrEnabledHintName]
 	if !ok {
 		return
 	}
-	readTelemetry, ok := readTelemetryHint.(bool)
+	ssrEnabled, ok := ssrEnabledHint.(bool)
 	if !ok {
-		b.log.Infof(log.Bolt5, b.logId, `invalid %q value: %v, ignoring hint. Only boolean values are accepted`, readTelemetryHintName, readTelemetryHint)
+		b.log.Infof(log.Bolt5, b.logId, `invalid %q value: %v, ignoring hint. Only boolean values are accepted`, ssrEnabledHintName, ssrEnabledHint)
 		return
 	}
-	b.telemetryEnabled = readTelemetry
+	b.ssrEnabled = ssrEnabled
 }
 
 func (b *bolt5) extractSummary(success *success, stream *stream) *db.Summary {

neo4j/internal/bolt/connect.go

@@ -47,7 +47,7 @@ func (p *protocolVersion) formatProtocol() string {
 // new manifest-style negotiation.
 var versions = [4]protocolVersion{
 	{major: 0xFF, minor: 0x01, back: 0x00}, // Bolt manifest marker
-	{major: 5, minor: 7, back: 7},
+	{major: 5, minor: 8, back: 8},
 	{major: 4, minor: 4, back: 2},
 	{major: 3, minor: 0, back: 0},
 }

neo4j/internal/db/connection.go

@@ -170,6 +170,11 @@ type Connection interface {
 	GetCurrentAuth() (auth.TokenManager, iauth.Token)
 	// Telemetry sends telemetry information about the API usage to the server.
 	Telemetry(api telemetry.API, onSuccess func())
+	// SetPinHomeDatabaseCallback registers a callback to update the session's cached home database.
+	// The callback is triggered on successful BEGIN or RUN responses containing a database name.
+	SetPinHomeDatabaseCallback(callback func(ctx context.Context, database string))
+	// IsSsrEnabled returns true if the connection supports Server-Side Routing.
+	IsSsrEnabled() bool
 }
 
 type RoutingTable struct {

neo4j/internal/db/database_selection.go

@@ -0,0 +1,24 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [https://neo4j.com]
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * 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
+ *
+ * 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 db
+
+// DatabaseSelection encapsulates the database name and whether it is guessed.
+type DatabaseSelection struct {
+	Name          string
+	IsHomeDbGuess bool
+}