Implement and use force-reset in session#run (Bolt v4)

Signed-off-by: Peter Wilhelmsson <peter.wilhelmsson@neo4j.com>

Author: fbiville

neo4j/db/connection.go

@@ -88,6 +88,7 @@ type Connection interface {
 	// Resets connection to same state as directly after a connect.
 	// Active streams will be discarded and the bookmark will be lost.
 	Reset()
+	ForceReset() error
 	// Closes the database connection as well as any underlying connection.
 	// The instance should not be used after being closed.
 	Close()

neo4j/internal/bolt/bolt3.go

@@ -735,3 +735,7 @@ func (b *bolt3) Close() {
 	b.conn.Close()
 	b.state = bolt3_dead
 }
+
+func (b *bolt3) ForceReset() error {
+	return nil
+}

neo4j/internal/bolt/bolt4.go

@@ -896,3 +896,17 @@ func (b *bolt4) Close() {
 func (b *bolt4) SelectDatabase(database string) {
 	b.databaseName = database
 }
+
+func (b *bolt4) ForceReset() error {
+	if b.state == bolt4_ready {
+		b.out.appendReset()
+		b.out.send(b.conn)
+		if b.err != nil {
+			return b.err
+		}
+		b.receiveMsg()
+		return b.err
+	}
+	b.Reset()
+	return b.err
+}

neo4j/internal/bolt/bolt4_test.go

@@ -491,6 +491,37 @@ func TestBolt4(ot *testing.T) {
 		bolt.Reset()
 	})
 
+	ot.Run("Forces reset in ready state", func(t *testing.T) {
+		bolt, cleanup := connectToServer(t, func(srv *bolt4server) {
+			srv.accept(4)
+			srv.waitForReset()
+			srv.sendSuccess(map[string]interface{}{})
+		})
+		defer cleanup()
+		defer bolt.Close()
+
+		err := bolt.ForceReset()
+		AssertNoError(t, err)
+		assertBoltState(t, bolt4_ready, bolt)
+	})
+
+	ot.Run("Forces reset while streaming", func(t *testing.T) {
+		bolt, cleanup := connectToServer(t, func(srv *bolt4server) {
+			srv.accept(4)
+			srv.serveRun(runResponse)
+			srv.waitForReset()
+			srv.sendSuccess(map[string]interface{}{})
+		})
+		defer cleanup()
+		defer bolt.Close()
+		_, err := bolt.Run(db.Command{Cypher: "MATCH (n) RETURN n"}, db.TxConfig{Mode: db.ReadMode})
+		AssertNoError(t, err)
+
+		err = bolt.ForceReset()
+		AssertNoError(t, err)
+		assertBoltState(t, bolt4_ready, bolt)
+	})
+
 	// Reset where state is ready
 
 	ot.Run("Buffer stream", func(t *testing.T) {

neo4j/internal/testutil/connfake.go

@@ -40,28 +40,29 @@ type RecordedTx struct {
 }
 
 type ConnFake struct {
-	Name          string
-	Version       string
-	Alive         bool
-	Birth         time.Time
-	Table         *db.RoutingTable
-	Err           error
-	Id            int
-	TxBeginHandle db.TxHandle
-	RunStream     db.StreamHandle
-	RunTxStream   db.StreamHandle
-	Nexts         []Next
-	Bookm         string
-	TxCommitErr   error
-	TxCommitHook  func()
-	TxRollbackErr error
-	ResetHook     func()
-	ConsumeSum    *db.Summary
-	ConsumeErr    error
-	ConsumeHook   func()
-	RecordedTxs   []RecordedTx // Appended to by Run/TxBegin
-	BufferErr     error
-	BufferHook    func()
+	Name           string
+	Version        string
+	Alive          bool
+	Birth          time.Time
+	Table          *db.RoutingTable
+	Err            error
+	Id             int
+	TxBeginHandle  db.TxHandle
+	RunStream      db.StreamHandle
+	RunTxStream    db.StreamHandle
+	Nexts          []Next
+	Bookm          string
+	TxCommitErr    error
+	TxCommitHook   func()
+	TxRollbackErr  error
+	ResetHook      func()
+	ConsumeSum     *db.Summary
+	ConsumeErr     error
+	ConsumeHook    func()
+	RecordedTxs    []RecordedTx // Appended to by Run/TxBegin
+	BufferErr      error
+	BufferHook     func()
+	ForceResetHook func() error
 }
 
 func (c *ConnFake) ServerName() string {
@@ -145,3 +146,10 @@ func (c *ConnFake) Next(streamHandle db.StreamHandle) (*db.Record, *db.Summary,
 	}
 	return next.Record, next.Summary, next.Err
 }
+
+func (c *ConnFake) ForceReset() error {
+	if c.ForceResetHook != nil {
+		return c.ForceResetHook()
+	}
+	return nil
+}

neo4j/session.go

@@ -397,9 +397,19 @@ func (s *session) Run(
 		c(&config)
 	}
 
-	conn, err := s.getConnection(s.defaultMode)
-	if err != nil {
-		return nil, err
+	var (
+		conn db.Connection
+		err  error
+	)
+	for  {
+		conn, err = s.getConnection(s.defaultMode)
+		if err != nil {
+			return nil, err
+		}
+		err = conn.ForceReset()
+		if err == nil {
+			break
+		}
 	}
 
 	stream, err := conn.Run(

neo4j/session_test.go

@@ -167,6 +167,41 @@ func TestSession(st *testing.T) {
 	})
 
 	st.Run("Run", func(bt *testing.T) {
+		bt.Run("Forces reset on acquired connection", func(t *testing.T) {
+			_, pool, sess := createSession()
+			forceResetCalls := 0
+			connection := &ConnFake{
+				ForceResetHook: func() error {
+					forceResetCalls++
+					return nil
+				},
+			}
+			pool.BorrowConn = connection
+
+			_, err := sess.Run("cypher", map[string]interface{}{})
+			AssertNoError(t, err)
+			AssertIntEqual(t, forceResetCalls, 1)
+		})
+
+		bt.Run("Picks connection from the pool until force-reset succeeds", func(t *testing.T) {
+			_, pool, sess := createSession()
+			forceResetCalls := 0
+			connection := &ConnFake{
+				ForceResetHook: func() error {
+					forceResetCalls++
+					if forceResetCalls == 1 {
+						return errors.New("force-reset failure")
+					}
+					return nil
+				},
+			}
+			pool.BorrowConn = connection
+
+			_, err := sess.Run("cypher", map[string]interface{}{})
+			AssertNoError(t, err)
+			AssertIntEqual(t, forceResetCalls, 2)
+		})
+
 		// Checks that chained Run results are buffered and that bookmarks are retrieved for
 		// those and that a Consume on the last result also gives the appropriate bookmark.
 		bt.Run("Chained and consume", func(t *testing.T) {