Handle transaction function panics

Author: fbiville

neo4j/session.go

@@ -259,56 +259,9 @@ func (s *session) runRetriable(
 		DatabaseName:            s.databaseName,
 	}
 	for state.Continue() {
-		// Establish new connection
-		conn, err := s.getConnection(mode)
-		if err != nil {
-			state.OnFailure(conn, err, false)
-			continue
-		}
-
-		// Begin transaction
-		txHandle, err := conn.TxBegin(db.TxConfig{
-			Mode:      mode,
-			Bookmarks: s.bookmarks,
-			Timeout:   config.Timeout,
-			Meta:      config.Metadata,
-		})
-		if err != nil {
-			state.OnFailure(conn, err, false)
-			s.pool.Return(conn)
-			continue
+		if workResult := s.tryRun(&state, mode, &config, work); workResult != nil {
+			return workResult, nil
 		}
-
-		// Construct a transaction like thing for client to execute stuff on
-		// and invoke the client work function.
-		tx := retryableTransaction{conn: conn, fetchSize: s.fetchSize, txHandle: txHandle}
-		x, err := work(&tx)
-		// Evaluate the returned error from all the work for retryable, this means
-		// that client can mess up the error handling.
-		if err != nil {
-			// If the client returns a client specific error that means that
-			// client wants to rollback. We don't do an explicit rollback here
-			// but instead realy on pool invoking reset on the connection, that
-			// will do an implicit rollback.
-			state.OnFailure(conn, err, false)
-			s.pool.Return(conn)
-			continue
-		}
-
-		// Commit transaction
-		err = conn.TxCommit(txHandle)
-		if err != nil {
-			state.OnFailure(conn, err, true)
-			s.pool.Return(conn)
-			continue
-		}
-
-		// Collect bookmark and return connection to pool
-		s.retrieveBookmarks(conn)
-		s.pool.Return(conn)
-
-		// All well
-		return x, nil
 	}
 
 	// When retries has occured wrap the error, the last error is always added but
@@ -339,6 +292,48 @@ func (s *session) WriteTransaction(
 	return s.runRetriable(db.WriteMode, work, configurers...)
 }
 
+func (s *session) tryRun(state *retry.State, mode db.AccessMode, config *TransactionConfig, work TransactionWork) interface{} {
+	// Establish new connection
+	conn, err := s.getConnection(mode)
+	if err != nil {
+		state.OnFailure(conn, err, false)
+		return nil
+	}
+	defer s.pool.Return(conn)
+	txHandle, err := conn.TxBegin(db.TxConfig{
+		Mode:      mode,
+		Bookmarks: s.bookmarks,
+		Timeout:   config.Timeout,
+		Meta:      config.Metadata,
+	})
+	if err != nil {
+		state.OnFailure(conn, err, false)
+		return nil
+	}
+
+	tx := retryableTransaction{conn: conn, fetchSize: s.fetchSize, txHandle: txHandle}
+	x, err := work(&tx)
+	// Evaluate the returned error from all the work for retryable, this means
+	// that client can mess up the error handling.
+	if err != nil {
+		// If the client returns a client specific error that means that
+		// client wants to rollback. We don't do an explicit rollback here
+		// but instead rely on pool invoking reset on the connection, that
+		// will do an implicit rollback.
+		state.OnFailure(conn, err, false)
+		return nil
+	}
+
+	err = conn.TxCommit(txHandle)
+	if err != nil {
+		state.OnFailure(conn, err, true)
+		return nil
+	}
+
+	s.retrieveBookmarks(conn)
+	return x
+}
+
 func (s *session) getServers(ctx context.Context, mode db.AccessMode) ([]string, error) {
 	if mode == db.ReadMode {
 		return s.router.Readers(ctx, s.bookmarks, s.databaseName, s.boltLogger)

neo4j/session_test.go

@@ -129,6 +129,29 @@ func TestSession(st *testing.T) {
 			AssertSameType(t, err.(*ConnectivityError).inner, &retry.CommitFailedDeadError{})
 			assertCleanSessionState(t, sess)
 		})
+
+		rt.Run("tx function panic returns conn to pool and bubbles up", func(t *testing.T) {
+			connectionIsReturned := false
+			_, pool, newSession := createSession()
+			pool.BorrowConn = &ConnFake{Alive: true}
+			pool.ReturnHook = func() {
+				connectionIsReturned = true
+			}
+			transactionFunction := func(Transaction) (interface{}, error) {
+				panic("oopsie")
+			}
+
+			defer func() {
+				err := recover()
+				if err != "oopsie" {
+					t.Errorf("expected panic to bubble up")
+				}
+				if !connectionIsReturned {
+					t.Errorf("expected connection to return to pool")
+				}
+			}()
+			_, _ = newSession.WriteTransaction(transactionFunction)
+		})
 	})
 
 	st.Run("Bookmarking", func(bt *testing.T) {