get_proof function stucks

[NikhilSharmaWe]

Description

When I attempt to call get_proof to retrieve a proof for a specific entry, the execution hangs indefinitely on the first call itself — it doesn’t return or panic, it just gets stuck.

Reproduce

use parking_lot::RwLock;
use qmdb::config::Config;
use qmdb::def::{DEFAULT_ENTRY_SIZE, IN_BLOCK_IDX_BITS, OP_CREATE, OP_WRITE, OP_DELETE, SHARD_COUNT};
use qmdb::entryfile::EntryBz;
use qmdb::merkletree::proof::check_proof;
use qmdb::tasks::TasksManager;
use qmdb::test_helper::SimpleTask;
use qmdb::utils::byte0_to_shard_id;
use qmdb::utils::changeset::ChangeSet;
use qmdb::utils::hasher;
use qmdb::{AdsCore, AdsWrap, ADS};
use std::sync::Arc;

fn main() {
    println!("=== QMDB Proof Generation Demo ===\n");

    // Initialize QMDB
    let ads_dir = "proof_demo_db";
    let config = Config::from_dir(ads_dir);
    AdsCore::init_dir(&config);
    let mut ads = AdsWrap::new(&config);

    println!("1. Creating initial entries...");
    create_initial_entries(&mut ads);

    println!("\n2. Testing proof generation for existing entries...");
    test_existing_entry_proofs(&ads);

    println!("\n3. Testing proof generation for non-existent entries...");
    test_non_existent_entry_proofs(&ads);

    println!("\n4. Creating more entries and testing proofs...");
    add_more_entries(&mut ads);
    test_proof_verification(&ads);

    println!("\n5. Testing proofs after updates...");
    update_entries(&mut ads);
    test_proofs_after_updates(&ads);

    println!("\n=== Demo completed successfully! ===");
}

/// Create initial entries across different shards
fn create_initial_entries(ads: &mut AdsWrap<SimpleTask>) {
    let mut changeset = ChangeSet::new();
    
    // Create entries with known patterns for easier testing
    let test_entries = [
        ([0u8, 1, 2], b"value_001"),
        ([50u8, 1, 2], b"value_050"), 
        ([100u8, 1, 2], b"value_100"),
        ([150u8, 1, 2], b"value_150"),
        ([200u8, 1, 2], b"value_200"),
    ];

    for (key_prefix, value) in &test_entries {
        let mut key = [0u8; 32];
        key[0] = key_prefix[0];
        key[1] = key_prefix[1]; 
        key[2] = key_prefix[2];
        
        let key_hash = hasher::hash(&key);
        let shard_id = byte0_to_shard_id(key_hash[0]) as u8;
        
        changeset.add_op(OP_CREATE, shard_id, &key_hash, &key, *value, None);
        println!("  Created entry: key={:02x?} shard={} value={:?}", 
                 &key[0..3], shard_id, std::str::from_utf8(*value).unwrap_or("binary"));
    }
    
    changeset.sort();
    apply_changeset(ads, changeset, 1);
}

/// Test proof generation for entries that exist
fn test_existing_entry_proofs(ads: &AdsWrap<SimpleTask>) {
    let metadb = ads.get_metadb();
    let metadb_guard = metadb.read();
    
    // Test proofs for each shard
    for shard_id in 0..SHARD_COUNT {
        let next_sn = metadb_guard.get_next_serial_num(shard_id);
        let oldest_sn = metadb_guard.get_oldest_active_sn(shard_id);
        
        println!("  Shard {}: SN range [{}, {})", shard_id, oldest_sn, next_sn);
        
        if next_sn > oldest_sn {
            // Test proof for first active serial number
            println!("    Testing proof for SN {}...", oldest_sn);
            match ads.get_proof(shard_id, oldest_sn) {
                Ok(mut proof_path) => {
                    println!("    ✅ Proof generated successfully!");
                    println!("       Serial number: {}", proof_path.serial_num);
                    println!("       Entry hash: {:02x?}...", &proof_path.left_of_twig[0].self_hash[0..4]);
                    // println!("       Twig root: {:02x?}...", &proof_path.()[0..4]);
                    
                    // Verify the proof
                    match check_proof(&mut proof_path) {
                        Ok(_) => println!("    ✅ Proof verification passed!"),
                        Err(e) => println!("    ❌ Proof verification failed: {}", e),
                    }
                },
                Err(e) => {
                    println!("    ❌ Failed to generate proof: {}", e);
                }
            }
            
            // Test proof for a serial number in the middle (if range is large enough)
            if next_sn > oldest_sn + 10 {
                let mid_sn = oldest_sn + 5;
                println!("    Testing proof for mid-range SN {}...", mid_sn);
                match ads.get_proof(shard_id, mid_sn) {
                    Ok(mut proof_path) => {
                        println!("    ✅ Mid-range proof generated successfully!");
                        match check_proof(&mut proof_path) {
                            Ok(_) => println!("    ✅ Mid-range proof verification passed!"),
                            Err(e) => println!("    ❌ Mid-range proof verification failed: {}", e),
                        }
                    },
                    Err(e) => println!("    ❌ Failed to generate mid-range proof: {}", e),
                }
            }
        } else {
            println!("    No active entries in this shard");
        }
        println!();
    }
    
    drop(metadb_guard);
}

/// Test proof generation for serial numbers that don't exist
fn test_non_existent_entry_proofs(ads: &AdsWrap<SimpleTask>) {
    let metadb = ads.get_metadb();
    let metadb_guard = metadb.read();
    
    for shard_id in 0..4 { // Test first 4 shards
        let next_sn = metadb_guard.get_next_serial_num(shard_id);
        
        // Test serial number beyond current range
        let invalid_sn = next_sn + 1000;
        println!("  Testing invalid SN {} for shard {} (next_sn: {})...", 
                 invalid_sn, shard_id, next_sn);
        
        match ads.get_proof(shard_id, invalid_sn) {
            Ok(proof_path) => {
                println!("    ⚠️  Unexpectedly got proof for invalid SN: {}", proof_path.serial_num);
            },
            Err(e) => {
                println!("    ✅ Correctly failed for invalid SN: {}", e);
                // Check if it's the expected error
                if e.contains("twig_file is empty") || e.contains("get proof failed") {
                    println!("    ✅ Got expected error type");
                }
            }
        }
    }
    
    drop(metadb_guard);
}

/// Add more entries to test proof generation with more data
fn add_more_entries(ads: &mut AdsWrap<SimpleTask>) {
    let mut changeset = ChangeSet::new();
    
    // Add more entries with different patterns
    for i in 0..10 {
        let mut key = [0u8; 32];
        let mut value = [0u8; 16];
        
        key[0] = (i * 25) as u8;  // Spread across shards
        key[3] = i as u8;         // Make keys unique
        value[0] = (100 + i) as u8;
        
        let key_hash = hasher::hash(&key);
        let shard_id = byte0_to_shard_id(key_hash[0]) as u8;
        
        changeset.add_op(OP_CREATE, shard_id, &key_hash, &key, &value, None);
        println!("  Adding entry {}: key={:02x?} shard={}", i, &key[0..4], shard_id);
    }
    
    changeset.sort();
    apply_changeset(ads, changeset, 2);
}

/// Test proof verification with the new entries
fn test_proof_verification(ads: &AdsWrap<SimpleTask>) {
    let metadb = ads.get_metadb();
    let metadb_guard = metadb.read();
    
    let mut total_proofs_tested = 0;
    let mut successful_proofs = 0;
    
    for shard_id in 0..SHARD_COUNT {
        let next_sn = metadb_guard.get_next_serial_num(shard_id);
        let oldest_sn = metadb_guard.get_oldest_active_sn(shard_id);
        
        if next_sn > oldest_sn {
            // Test a few serial numbers from this shard
            let test_count = std::cmp::min(3, next_sn - oldest_sn);
            
            for i in 0..test_count {
                let test_sn = oldest_sn + i;
                total_proofs_tested += 1;
                
                match ads.get_proof(shard_id, test_sn) {
                    Ok(mut proof_path) => {
                        match check_proof(&mut proof_path) {
                            Ok(_) => {
                                successful_proofs += 1;
                                println!("  ✅ Shard {} SN {}: Proof OK", shard_id, test_sn);
                            },
                            Err(e) => {
                                println!("  ❌ Shard {} SN {}: Proof verification failed: {}", 
                                         shard_id, test_sn, e);
                            }
                        }
                    },
                    Err(e) => {
                        println!("  ❌ Shard {} SN {}: Proof generation failed: {}", 
                                 shard_id, test_sn, e);
                    }
                }
            }
        }
    }
    
    println!("  Summary: {}/{} proofs verified successfully", successful_proofs, total_proofs_tested);
    drop(metadb_guard);
}

/// Update some existing entries and test proofs
fn update_entries(ads: &mut AdsWrap<SimpleTask>) {
    let mut changeset = ChangeSet::new();
    
    // Update a few entries
    let update_keys = [
        [0u8, 1, 2],    // Should exist from initial creation
        [50u8, 1, 2],   // Should exist from initial creation
    ];
    
    for key_prefix in &update_keys {
        let mut key = [0u8; 32];
        key[0] = key_prefix[0];
        key[1] = key_prefix[1];
        key[2] = key_prefix[2];
        
        let new_value = b"updated_value";
        let key_hash = hasher::hash(&key);
        let shard_id = byte0_to_shard_id(key_hash[0]) as u8;
        
        changeset.add_op(OP_WRITE, shard_id, &key_hash, &key, new_value, None);
        println!("  Updating entry: key={:02x?} shard={}", &key[0..3], shard_id);
    }
    
    changeset.sort();
    apply_changeset(ads, changeset, 3);
}

/// Test proofs after updates to ensure they still work
fn test_proofs_after_updates(ads: &AdsWrap<SimpleTask>) {
    println!("  Testing proofs after updates...");
    
    let metadb = ads.get_metadb();
    let metadb_guard = metadb.read();
    
    for shard_id in 0..4 { // Test first 4 shards
        let next_sn = metadb_guard.get_next_serial_num(shard_id);
        let oldest_sn = metadb_guard.get_oldest_active_sn(shard_id);
        
        if next_sn > oldest_sn {
            // Test the most recent serial number (likely from updates)
            let recent_sn = next_sn - 1;
            
            match ads.get_proof(shard_id, recent_sn) {
                Ok(mut proof_path) => {
                    println!("    Shard {}: Generated proof for recent SN {}", shard_id, recent_sn);
                    
                    match check_proof(&mut proof_path) {
                        Ok(_) => println!("    ✅ Recent entry proof verified successfully"),
                        Err(e) => println!("    ❌ Recent entry proof verification failed: {}", e),
                    }
                },
                Err(e) => {
                    println!("    ❌ Failed to generate proof for recent SN {}: {}", recent_sn, e);
                }
            }
        }
    }
    
    drop(metadb_guard);
}

/// Helper function to apply a changeset to QMDB
fn apply_changeset(ads: &mut AdsWrap<SimpleTask>, changeset: ChangeSet, height: i64) {
    let task = SimpleTask::new(vec![changeset]);
    let task_list = vec![RwLock::new(Some(task))];
    let task_id = height << IN_BLOCK_IDX_BITS;
    
    ads.start_block(height, Arc::new(TasksManager::new(task_list, task_id)));
    let shared_ads = ads.get_shared();
    shared_ads.insert_extra_data(height, format!("block_{}", height));
    shared_ads.add_task(task_id);
    ads.flush();
    
    let curr_height = ads.get_metadb().read().get_curr_height();
    println!("  Applied changeset at height {}, current height: {}", height, curr_height);
}

Please inform me if I am using the function properly, and what could be the possible issue.

[rahulghangas]

@absolute0kelvin could you point us in the right direction

[rahulghangas]

cc @zouguangxian
if you could please shed some light on this

[NikhilSharmaWe]

get_proof working
It is handled by the FlushedShard type, which handles all the operations for a shard.

pub struct FlusherShard {
  
    // Channel for proof requests
    proof_req_sender: SyncSender<ProofReqElem>,     // Send proof requests
    proof_req_receiver: Receiver<ProofReqElem>,     // Receive proof requests
    
}

And flusher thread processes requests during its flush cycle at a specific point when the tree is in a consistent state. The handle_proof_req is called at a specific point in the flush cycle when it's safe to generate proofs.

So, if we want to use the get_proof method for Ads, then get_proof requests are handled specifically in the flush method after the changes have been made for a particular height. We cannot make the requests earlier or later, because the tree is not in a consistent state. After all, if we make it so, the approach is like that it might be the case that the changes for the next block are being done.

So, the only time I am seeing it is possible to use the get_proof method to make a proof request is to

• Make all the requests before the flush in a separate thread
• then call flush in the main thread
• the get_proof requests will be processed in the flush method through the handle_proof_req function

But this approach cannot be used in our case since we do not want to flush the changes in Ads.

We want to implement a memory overlay over QMDB, which will allow us to calculate the state root in-memory.
So, in our case, we only need to do the read operation over the Ads, so we are sure that the state is consistent, but we require direct access to the tree in order to handle the proof requests through Ads whenever we want.

The Ads do not provide direct access to the tree; the SeqAds do, but using SeqAds will be expensive since we will need to copy the entire state.

Is there a workaround in which we can enable the direct tree access through Ads in order to implement this or any other reasonable approach?