From 005e28266ae8b4f9ee38aaf7c0aba5fd45ad5a06 Mon Sep 17 00:00:00 2001
From: Peilun Li <peilun.li@confluxnetwork.org>
Date: Tue, 15 Oct 2024 12:42:17 +0800
Subject: [PATCH] fix.

---
 common/append_merkle/src/lib.rs           |  4 +++
 common/append_merkle/src/node_manager.rs  |  3 ++-
 node/storage/src/log_store/log_manager.rs | 30 +++++++++++++++++------
 node/storage/src/log_store/tx_store.rs    |  3 +++
 tests/crash_test.py                       |  0
 5 files changed, 32 insertions(+), 8 deletions(-)
 mode change 100644 => 100755 tests/crash_test.py

diff --git a/common/append_merkle/src/lib.rs b/common/append_merkle/src/lib.rs
index d4292c1..86b550d 100644
--- a/common/append_merkle/src/lib.rs
+++ b/common/append_merkle/src/lib.rs
@@ -59,10 +59,12 @@ impl<E: HashElement, A: Algorithm<E>> AppendMerkleTree<E, A> {
                     },
                 );
             }
+            merkle.node_manager.commit();
             return merkle;
         }
         // Reconstruct the whole tree.
         merkle.recompute(0, 0, None);
+        merkle.node_manager.commit();
         // Commit the first version in memory.
         // TODO(zz): Check when the roots become available.
         merkle.commit(start_tx_seq);
@@ -116,6 +118,7 @@ impl<E: HashElement, A: Algorithm<E>> AppendMerkleTree<E, A> {
                     },
                 );
             }
+            merkle.node_manager.commit();
             merkle
         } else {
             let mut merkle = Self {
@@ -134,6 +137,7 @@ impl<E: HashElement, A: Algorithm<E>> AppendMerkleTree<E, A> {
             }
             // Reconstruct the whole tree.
             merkle.recompute(0, 0, None);
+            merkle.node_manager.commit();
             // Commit the first version in memory.
             merkle.commit(start_tx_seq);
             merkle
diff --git a/common/append_merkle/src/node_manager.rs b/common/append_merkle/src/node_manager.rs
index 005787e..3cdc54a 100644
--- a/common/append_merkle/src/node_manager.rs
+++ b/common/append_merkle/src/node_manager.rs
@@ -115,8 +115,9 @@ impl<E: HashElement> NodeManager<E> {
     }
 
     pub fn start_transaction(&mut self) {
-        if self.db_tx.is_none() {
+        if self.db_tx.is_some() {
             error!("start new tx before commit");
+            panic!("start new tx before commit");
         }
         self.db_tx = Some(self.db.start_transaction());
     }
diff --git a/node/storage/src/log_store/log_manager.rs b/node/storage/src/log_store/log_manager.rs
index ee694da..e83ccb0 100644
--- a/node/storage/src/log_store/log_manager.rs
+++ b/node/storage/src/log_store/log_manager.rs
@@ -651,8 +651,12 @@ impl LogManager {
                     .get_tx_by_seq_number(last_tx_seq)?
                     .expect("tx missing");
                 let current_len = pora_chunks_merkle.leaves();
-                let expected_len =
-                    sector_to_segment(last_tx.start_entry_index + last_tx.num_entries() as u64);
+                let expected_len = sector_to_segment(
+                    last_tx.start_entry_index
+                        + last_tx.num_entries() as u64
+                        + PORA_CHUNK_SIZE as u64
+                        - 1,
+                );
                 match expected_len.cmp(&(current_len)) {
                     Ordering::Less => {
                         bail!(
@@ -681,6 +685,8 @@ impl LogManager {
                             );
                             if current_len > expected_len {
                                 pora_chunks_merkle.revert_to_leaves(expected_len)?;
+                            } else {
+                                assert_eq!(current_len, expected_len);
                             }
                             start_tx_seq = Some(previous_tx.seq);
                         };
@@ -691,10 +697,20 @@ impl LogManager {
 
         let last_chunk_merkle = match start_tx_seq {
             Some(tx_seq) => {
-                tx_store.rebuild_last_chunk_merkle(pora_chunks_merkle.leaves(), tx_seq)?
+                let tx = tx_store.get_tx_by_seq_number(tx_seq)?.expect("tx missing");
+                if (tx.start_entry_index() + tx.num_entries() as u64) % PORA_CHUNK_SIZE as u64 == 0
+                {
+                    // The last chunk should be aligned, so it's empty.
+                    Merkle::new_with_depth(vec![], log2_pow2(PORA_CHUNK_SIZE) + 1, None)
+                } else {
+                    tx_store.rebuild_last_chunk_merkle(pora_chunks_merkle.leaves() - 1, tx_seq)?
+                }
             }
             // Initialize
-            None => Merkle::new_with_depth(vec![], 1, None),
+            None => {
+                pora_chunks_merkle.reset();
+                Merkle::new_with_depth(vec![], 1, None)
+            }
         };
 
         debug!(
@@ -704,10 +720,10 @@ impl LogManager {
             last_chunk_merkle.leaves(),
         );
         if last_chunk_merkle.leaves() != 0 {
-            pora_chunks_merkle.append(last_chunk_merkle.root());
-            // update the merkle root
-            pora_chunks_merkle.commit(start_tx_seq);
+            pora_chunks_merkle.update_last(last_chunk_merkle.root());
         }
+        // update the merkle root
+        pora_chunks_merkle.commit(start_tx_seq);
         let merkle = RwLock::new(MerkleManager {
             pora_chunks_merkle,
             last_chunk_merkle,
diff --git a/node/storage/src/log_store/tx_store.rs b/node/storage/src/log_store/tx_store.rs
index 39fb893..37d4b45 100644
--- a/node/storage/src/log_store/tx_store.rs
+++ b/node/storage/src/log_store/tx_store.rs
@@ -292,6 +292,9 @@ impl TransactionStore {
             match tx.start_entry_index.cmp(&last_chunk_start_index) {
                 cmp::Ordering::Greater => {
                     tx_list.push((tx_seq, tx.merkle_nodes));
+                    if tx.start_entry_index >= last_chunk_start_index + PORA_CHUNK_SIZE as u64 {
+                        break;
+                    }
                 }
                 cmp::Ordering::Equal => {
                     tx_list.push((tx_seq, tx.merkle_nodes));
diff --git a/tests/crash_test.py b/tests/crash_test.py
old mode 100644
new mode 100755