Contractless/src/verifications/verification_service.rs

use crate::common::types::Transaction;
use crate::rayon::ThreadPool;
use crate::rayon::ThreadPoolBuilder;
use crate::verifications::async_funcs::checks::block_balance::BlockBalanceTracker;
use crate::verifications::async_funcs::transactions::verify_transactions;
use crate::verifications::sync_funcs::transaction_verify_loop::COUNTER;
use crate::verifications::verification_types::{
    VerificationChunkJob, VerificationChunkResult, VerificationJob,
};
use crate::Arc;
use crate::AtomicBool;
use crate::AtomicOrdering;
use crate::Builder;
use crate::IntoParallelIterator;
use crate::OnceLock;
use crate::ParallelIterator;
use crate::Runtime;
use crate::{mpsc, oneshot, Mutex};

static VERIFICATION_SERVICE: OnceLock<VerificationService> = OnceLock::new();

// VerificationRequest wraps one verification job together with the reply channel that
// returns the combined results to the caller.
struct VerificationRequest {
    job: VerificationJob,
    reply_tx: oneshot::Sender<Result<Vec<String>, String>>,
}

// VerificationService owns the async queue that feeds verification work into the shared worker pool.
#[derive(Clone)]
pub struct VerificationService {
    job_tx: mpsc::Sender<VerificationRequest>,
}

impl VerificationService {
    pub fn start() -> Self {
        // The verification service keeps a dedicated rayon pool and Tokio runtime so
        // transaction verification can scale without blocking the rest of the node.
        let (job_tx, mut job_rx) = mpsc::channel::<VerificationRequest>(1);
        let worker_count = std::thread::available_parallelism()
            .map(|count| count.get())
            .unwrap_or(4);

        let rayon_pool = Arc::new(
            ThreadPoolBuilder::new()
                .num_threads(worker_count)
                .build()
                .expect("failed to create verification rayon pool"),
        );

        let verify_runtime = Arc::new(
            Builder::new_multi_thread()
                .worker_threads(worker_count)
                .enable_all()
                .build()
                .expect("failed to create verification runtime"),
        );

        tokio::spawn({
            let rayon_pool = rayon_pool.clone();
            let verify_runtime = verify_runtime.clone();
            async move {
                while let Some(request) = job_rx.recv().await {
                    let result =
                        Self::run_job(rayon_pool.clone(), verify_runtime.clone(), request.job)
                            .await;
                    let _ = request.reply_tx.send(result);
                }
            }
        });

        Self { job_tx }
    }

    pub async fn verify_block_transactions(
        &self,
        miner: String,
        transactions: Vec<Transaction>,
        db: crate::sled::Db,
        block_timestamp: u32,
    ) -> Result<Vec<String>, String> {
        let (reply_tx, reply_rx) = oneshot::channel();
        let request = VerificationRequest {
            job: VerificationJob {
                miner,
                transactions,
                db,
                block_timestamp,
            },
            reply_tx,
        };

        self.job_tx
            .send(request)
            .await
            .map_err(|_| "Verification service unavailable".to_string())?;

        reply_rx
            .await
            .map_err(|_| "Verification service response dropped".to_string())?
    }

    async fn run_job(
        rayon_pool: Arc<ThreadPool>,
        verify_runtime: Arc<Runtime>,
        job: VerificationJob,
    ) -> Result<Vec<String>, String> {
        // Transaction verification is split into chunks and processed in parallel, then
        // merged back into one result list for the caller.
        let chunk_size = 1000usize;
        let stop_flag = Arc::new(AtomicBool::new(false));
        let balance_tracker = Arc::new(Mutex::new(BlockBalanceTracker::new()));

        COUNTER.store(0, AtomicOrdering::SeqCst);

        let chunk_jobs = Self::split_transactions(job, chunk_size);

        let chunk_results: Vec<VerificationChunkResult> = rayon_pool.install(|| {
            chunk_jobs
                .into_par_iter()
                .map(|chunk_job| {
                    let stop_flag = stop_flag.clone();
                    let balance_tracker = balance_tracker.clone();
                    verify_runtime.block_on(async move {
                        match verify_transactions(
                            chunk_job.miner,
                            chunk_job.transactions,
                            &chunk_job.db,
                            stop_flag,
                            balance_tracker,
                            chunk_job.block_timestamp,
                        )
                        .await
                        {
                            Ok(results) => VerificationChunkResult {
                                results,
                                had_error: false,
                                error: None,
                            },
                            Err(error) => VerificationChunkResult {
                                results: Vec::new(),
                                had_error: true,
                                error: Some(error),
                            },
                        }
                    })
                })
                .collect()
        });

        if let Some(error) = chunk_results
            .iter()
            .find(|result| result.had_error)
            .and_then(|result| result.error.clone())
        {
            return Err(error);
        }

        let mut full_results = Vec::new();
        for chunk_result in chunk_results {
            full_results.extend(chunk_result.results);
        }

        Ok(full_results)
    }

    fn split_transactions(job: VerificationJob, chunk_size: usize) -> Vec<VerificationChunkJob> {
        let mut chunk_jobs = Vec::new();
        let VerificationJob {
            miner,
            transactions,
            db,
            block_timestamp,
        } = job;

        for chunk in transactions.chunks(chunk_size) {
            chunk_jobs.push(VerificationChunkJob {
                miner: miner.clone(),
                transactions: chunk.to_vec(),
                db: db.clone(),
                block_timestamp,
            });
        }

        chunk_jobs
    }
}

pub fn initialize_global_verification_service() -> VerificationService {
    // The global service is initialized once and then cloned anywhere the node needs it.
    let service = VerificationService::start();
    let _ = VERIFICATION_SERVICE.set(service.clone());
    service
}

pub fn global_verification_service() -> Option<VerificationService> {
    VERIFICATION_SERVICE.get().cloned()
}