Skip to content

Commit ddb6ba8

Browse files
will-bitlightlabswill-bitlightlabs
committed
feat: enhance block rollback process with comprehensive cleanup
Signed-off-by: will-bitlightlabs <[email protected]>
1 parent c74c36a commit ddb6ba8

File tree

1 file changed

+136
-9
lines changed

1 file changed

+136
-9
lines changed

‎src/blocks.rs‎

Lines changed: 136 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1671,6 +1671,9 @@ impl BlockProcessor {
16711671
let mut total_txs_removed = 0;
16721672
let mut total_utxos_restored = 0;
16731673
let mut total_utxos_removed = 0;
1674+
let mut total_spk_entries_cleaned = 0;
1675+
let mut total_inputs_cleaned = 0;
1676+
let mut total_outs_refs_cleaned = 0;
16741677

16751678
let block_spends_table = db
16761679
.open_table(TABLE_BLOCK_SPENDS)
@@ -1696,6 +1699,26 @@ impl BlockProcessor {
16961699
.open_table(TABLE_BLOCK_HEIGHTS)
16971700
.map_err(|e| BlockProcError::Custom(format!("Block heights table error: {}", e)))?;
16981701

1702+
// Open the SPKs table for script pubkey cleanup during rollback
1703+
let mut spks_table = db
1704+
.open_table(TABLE_SPKS)
1705+
.map_err(|e| BlockProcError::Custom(format!("SPKs table error: {}", e)))?;
1706+
1707+
// Open the inputs table to clean up input references
1708+
let mut inputs_table = db
1709+
.open_table(TABLE_INPUTS)
1710+
.map_err(|e| BlockProcError::Custom(format!("Inputs table error: {}", e)))?;
1711+
1712+
// Open the outs table to clean up spending relationships
1713+
let mut outs_table = db
1714+
.open_table(TABLE_OUTS)
1715+
.map_err(|e| BlockProcError::Custom(format!("Outs table error: {}", e)))?;
1716+
1717+
// Open tx_blocks table to clean up transaction-block associations
1718+
let mut tx_blocks_table = db
1719+
.open_table(TABLE_TX_BLOCKS)
1720+
.map_err(|e| BlockProcError::Custom(format!("TX-Blocks table error: {}", e)))?;
1721+
16991722
// Iterate through blocks to roll back (should be in descending height order)
17001723
for &(height, block_id) in blocks {
17011724
log::info!(
@@ -1708,6 +1731,9 @@ impl BlockProcessor {
17081731
let mut block_utxos_restored = 0;
17091732
let mut block_utxos_removed = 0;
17101733
let mut block_txs_removed = 0;
1734+
let mut block_spk_entries_cleaned = 0;
1735+
let mut block_inputs_cleaned = 0;
1736+
let mut block_outs_refs_cleaned = 0;
17111737

17121738
// 1. Restore UTXOs spent in this block
17131739
if let Some(spends_record) = block_spends_table
@@ -1750,27 +1776,119 @@ impl BlockProcessor {
17501776
.map_err(|e| BlockProcError::Custom(format!("Tx lookup error: {}", e)))?
17511777
{
17521778
let tx = tx_record.value();
1753-
let num_outputs = tx.as_ref().outputs.len();
1779+
let outputs = tx.as_ref().outputs.as_slice();
1780+
let num_outputs = outputs.len();
17541781
block_utxos_removed += num_outputs;
17551782
total_utxos_removed += num_outputs;
17561783

1757-
for vout in 0..num_outputs {
1758-
utxos_table.remove(&(txno, vout as u32)).map_err(|e| {
1784+
// Get the number of inputs for this transaction
1785+
let inputs_count = tx.as_ref().inputs.len();
1786+
1787+
for (vout_idx, output) in outputs.iter().enumerate() {
1788+
// Remove UTXOs
1789+
utxos_table.remove(&(txno, vout_idx as u32)).map_err(|e| {
17591790
BlockProcError::Custom(format!("UTXO removal error: {}", e))
17601791
})?;
17611792

17621793
log::debug!(
17631794
target: NAME,
17641795
"Removed UTXO: txno={}, vout={}",
17651796
txno,
1766-
vout
1797+
vout_idx
1798+
);
1799+
1800+
// 4. Clean up script pubkey index for this transaction output
1801+
let script = &output.script_pubkey;
1802+
if !script.is_empty() {
1803+
let txnos = spks_table
1804+
.get(script.as_slice())
1805+
.map_err(|e| {
1806+
BlockProcError::Custom(format!("SPKs lookup error: {}", e))
1807+
})?
1808+
.map(|t| t.value().to_vec());
1809+
1810+
if let Some(mut txnos) = txnos {
1811+
// Remove this transaction from the list
1812+
if let Some(pos) = txnos.iter().position(|&t| t == txno) {
1813+
txnos.remove(pos);
1814+
block_spk_entries_cleaned += 1;
1815+
1816+
// If the list is not empty, update it; otherwise, remove
1817+
// the entry
1818+
if !txnos.is_empty() {
1819+
spks_table.insert(script.as_slice(), txnos).map_err(
1820+
|e| {
1821+
BlockProcError::Custom(format!(
1822+
"SPKs update error: {}",
1823+
e
1824+
))
1825+
},
1826+
)?;
1827+
} else {
1828+
spks_table.remove(script.as_slice()).map_err(|e| {
1829+
BlockProcError::Custom(format!(
1830+
"SPKs removal error: {}",
1831+
e
1832+
))
1833+
})?;
1834+
}
1835+
1836+
log::debug!(
1837+
target: NAME,
1838+
"Cleaned up SPK index for txno={}, vout={}",
1839+
txno,
1840+
vout_idx
1841+
);
1842+
}
1843+
}
1844+
}
1845+
}
1846+
1847+
// 5. Clean up inputs table for this transaction
1848+
for input_idx in 0..inputs_count {
1849+
if inputs_table
1850+
.remove(&(txno, input_idx as u32))
1851+
.map_err(|e| {
1852+
BlockProcError::Custom(format!("Inputs removal error: {}", e))
1853+
})?
1854+
.is_some()
1855+
{
1856+
block_inputs_cleaned += 1;
1857+
log::debug!(
1858+
target: NAME,
1859+
"Removed input reference: txno={}, input_idx={}",
1860+
txno,
1861+
input_idx
1862+
);
1863+
}
1864+
}
1865+
total_inputs_cleaned += block_inputs_cleaned;
1866+
1867+
// 6. Clean up this transaction from spending relationships
1868+
if outs_table
1869+
.remove(txno)
1870+
.map_err(|e| {
1871+
BlockProcError::Custom(format!("Outs lookup error: {}", e))
1872+
})?
1873+
.is_some()
1874+
{
1875+
block_outs_refs_cleaned += 1;
1876+
log::debug!(
1877+
target: NAME,
1878+
"Removed spending relationship for txno={}",
1879+
txno
17671880
);
17681881
}
17691882
}
1883+
1884+
// 7. Remove transaction-block association
1885+
tx_blocks_table.remove(txno).map_err(|e| {
1886+
BlockProcError::Custom(format!("TX-Blocks removal error: {}", e))
1887+
})?;
17701888
}
17711889
}
17721890

1773-
// 4. Remove this block from the heights tables
1891+
// 8. Remove this block from the heights tables
17741892
heights_table
17751893
.remove(height)
17761894
.map_err(|e| BlockProcError::Custom(format!("Heights removal error: {}", e)))?;
@@ -1786,23 +1904,32 @@ impl BlockProcessor {
17861904
block_id
17871905
);
17881906

1907+
total_spk_entries_cleaned += block_spk_entries_cleaned;
1908+
total_outs_refs_cleaned += block_outs_refs_cleaned;
1909+
17891910
log::info!(
17901911
target: NAME,
1791-
"Block rollback stats for height {}: removed {} transactions, restored {} UTXOs, removed {} UTXOs",
1912+
"Block rollback stats for height {}: removed {} transactions, restored {} UTXOs, removed {} UTXOs, cleaned {} SPK entries, {} input refs, {} output refs",
17921913
height,
17931914
block_txs_removed,
17941915
block_utxos_restored,
1795-
block_utxos_removed
1916+
block_utxos_removed,
1917+
block_spk_entries_cleaned,
1918+
block_inputs_cleaned,
1919+
block_outs_refs_cleaned
17961920
);
17971921
}
17981922

17991923
log::info!(
18001924
target: NAME,
1801-
"Successfully rolled back {} blocks: removed {} transactions, restored {} UTXOs, removed {} UTXOs",
1925+
"Successfully rolled back {} blocks: removed {} transactions, restored {} UTXOs, removed {} UTXOs, cleaned {} SPK entries, {} input refs, {} output refs",
18021926
blocks.len(),
18031927
total_txs_removed,
18041928
total_utxos_restored,
1805-
total_utxos_removed
1929+
total_utxos_removed,
1930+
total_spk_entries_cleaned,
1931+
total_inputs_cleaned,
1932+
total_outs_refs_cleaned
18061933
);
18071934

18081935
Ok(())

0 commit comments

Comments
 (0)