refactor: polish handling of fail in transaction (#18497)

* refactor: clear txn buffer at once when txn fail.

* fix: http handler is no longer sticky if fail to start query.

Author: youngsofun

src/query/service/src/servers/http/v1/query/http_query.rs

@@ -101,8 +101,14 @@ impl HttpQueryRequest {
             } else {
                 s.txn_state.clone()
             };
+            let need_sticky = match &s.internal {
+                Some(internal) => internal.has_temp_table,
+                None => false,
+            };
             HttpSessionConf {
                 txn_state,
+                need_sticky,
+                need_keep_alive: need_sticky,
                 ..s.clone()
             }
         });
@@ -748,7 +754,7 @@ impl HttpQuery {
         );
 
         if is_stopped
-            && txn_state != TxnState::AutoCommit
+            && txn_state == TxnState::Active
             && !self.is_txn_mgr_saved.load(Ordering::Relaxed)
             && self
                 .is_txn_mgr_saved
@@ -764,7 +770,7 @@ impl HttpQuery {
                 .await;
         }
 
-        let need_sticky = txn_state != TxnState::AutoCommit || has_temp_table;
+        let need_sticky = txn_state == TxnState::Active || has_temp_table;
         let need_keep_alive = need_sticky;
 
         Ok(HttpSessionConf {

src/query/storages/common/session/src/transaction.rs

@@ -167,7 +167,7 @@ impl TxnManager {
 
     pub fn set_fail(&mut self) {
         if let TxnState::Active = self.state {
-            self.state = TxnState::Fail;
+            self.force_set_fail()
         }
     }
 
@@ -177,6 +177,8 @@ impl TxnManager {
 
     pub fn force_set_fail(&mut self) {
         self.state = TxnState::Fail;
+        self.txn_buffer.clear();
+        // keep the txn_id until commit/abort for tracing
     }
 
     pub fn is_fail(&self) -> bool {

tests/suites/1_stateful/09_http_handler/09_0008_forward.py

@@ -28,7 +28,12 @@ def do_query(query, port=8000, session=None, node_id=None, wait=100):
     return response
 
 
-def test_txn():
+def get_txn_state(resp):
+    return (resp.get("state") == "Succeeded",
+    resp.get("session").get("need_sticky"),
+    resp.get("session").get("txn_state"))
+
+def test_txn_success():
     resp = do_query("create or replace table t1(a int)").json()
     assert not (resp.get("session").get("need_sticky")), resp
 
@@ -37,21 +42,54 @@ def test_txn():
     node_id = resp.get("node_id")
     session = resp.get("session")
 
-    # can not find txn state in node 2
-    resp = do_query("insert into t1 values (1)", port=8002, session=session).json()
-    assert resp.get("state") == "Failed", resp.text
-
     # forward to node 1
     resp = do_query(
         "insert into t1 values (2)", port=8002, session=session, node_id=node_id
     ).json()
-    assert resp.get("state") == "Succeeded", resp
-    assert resp.get("session").get("need_sticky"), resp
+    assert get_txn_state(resp) == (True, True, "Active"), resp
 
     # return need_sticky = false after commit
     resp = do_query("commit").json()
+    assert get_txn_state(resp) == (True, False, "AutoCommit"), resp
+
+
+def test_txn_fail():
+    resp = do_query("create or replace table t1(a int)").json()
     assert not (resp.get("session").get("need_sticky")), resp
 
+    resp = do_query("begin").json()
+    assert resp.get("session").get("need_sticky"), resp
+    node_id = resp.get("node_id")
+    session = resp.get("session")
+
+    # fail
+    resp = do_query(
+        "select 1/0",  session=session
+    ).json()
+    assert get_txn_state(resp) == (False, False, "Fail"), resp
+
+    # fail because wrong node
+    resp = do_query(
+        "select 1", port=8002, session=session
+    ).json()
+    session = resp.get("session")
+    assert get_txn_state(resp) == (False, False, "Fail"), resp
+
+    # keep fail state until commit/abort
+    resp = do_query(
+        "select 1", session=session, node_id=node_id
+    ).json()
+    assert get_txn_state(resp) == (False, False, "Fail"), resp
+    session = resp.get("session")
+
+    # return need_sticky = false after commit
+    resp = do_query("commit", session=session).json()
+    assert get_txn_state(resp) == (True, False, "AutoCommit"), resp
+
+    # return need_sticky = false after abort
+    resp = do_query("abort", session=session).json()
+    assert get_txn_state(resp) == (True, False, "AutoCommit"), resp
+
 
 def test_query():
     """each query is sticky"""
@@ -104,7 +142,8 @@ def main():
         return
 
     test_query()
-    test_txn()
+    test_txn_success()
+    test_txn_fail()
 
 
 if __name__ == "__main__":