uhash_prover/cpu/

parallel.rs

//! Multi-threaded CPU solver with persistent thread pool.
//!
//! Workers are created once and reused across batches, eliminating per-batch
//! thread creation overhead and 2 MB UniversalHash scratchpad allocations.

use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::mpsc::{sync_channel, Receiver, SyncSender};
use std::sync::{Arc, Mutex};

use anyhow::Result;
use uhash_core::{meets_difficulty, UniversalHash};

use crate::solver::{ProofResult, Solver};

/// A single unit of work dispatched to one worker.
struct Job {
    challenge: Arc<Vec<u8>>,
    start_nonce: u64,
    total_lanes: usize,
    tid: usize,
    thread_count: usize,
    /// `Some(difficulty)` → proof-search mode; `None` → benchmark mode.
    difficulty: Option<u32>,
    found: Arc<AtomicBool>,
    winner: Arc<Mutex<ProofResult>>,
    /// Shared counter for actual hashes computed across all workers.
    actual_hashes: Arc<AtomicUsize>,
}

/// Pool of persistent worker threads.
struct WorkerPool {
    thread_count: usize,
    /// One sender per worker — each worker blocks on its own receiver.
    job_txs: Vec<SyncSender<Option<Job>>>,
    /// Workers send `()` here after finishing a job; main collects N signals.
    done_rx: Receiver<()>,
}

impl WorkerPool {
    fn new(thread_count: usize) -> Self {
        // Capacity = thread_count so all workers can signal without blocking.
        let (done_tx, done_rx) = sync_channel::<()>(thread_count);
        let mut job_txs = Vec::with_capacity(thread_count);

        for _ in 0..thread_count {
            let (job_tx, job_rx) = sync_channel::<Option<Job>>(1);
            let done_tx = done_tx.clone();

            std::thread::spawn(move || {
                // Allocate the 2 MB scratchpad once for the lifetime of this thread.
                let mut hasher = UniversalHash::new();
                // Input buffer grown to the right size on first use and reused.
                let mut input = Vec::with_capacity(128);

                while let Ok(Some(job)) = job_rx.recv() {
                    let mut lane = job.tid;
                    let mut count = 0usize;

                    if let Some(d) = job.difficulty {
                        // Proof-search: check difficulty, stop early on found.
                        while lane < job.total_lanes && !job.found.load(Ordering::Relaxed) {
                            let nonce = job.start_nonce.saturating_add(lane as u64);
                            input.clear();
                            input.extend_from_slice(&job.challenge);
                            input.extend_from_slice(&nonce.to_le_bytes());
                            let hash = hasher.hash(&input);
                            count += 1;
                            if meets_difficulty(&hash, d) {
                                if !job.found.swap(true, Ordering::Relaxed) {
                                    let mut guard =
                                        job.winner.lock().expect("winner mutex poisoned");
                                    *guard = Some((nonce, hash));
                                }
                                break;
                            }
                            lane += job.thread_count;
                        }
                    } else {
                        // Benchmark: hash all assigned lanes, no difficulty check.
                        while lane < job.total_lanes {
                            let nonce = job.start_nonce.saturating_add(lane as u64);
                            input.clear();
                            input.extend_from_slice(&job.challenge);
                            input.extend_from_slice(&nonce.to_le_bytes());
                            let _ = hasher.hash(&input);
                            count += 1;
                            lane += job.thread_count;
                        }
                    }

                    job.actual_hashes.fetch_add(count, Ordering::Relaxed);
                    let _ = done_tx.send(());
                }
                // `None` received → thread exits cleanly.
            });

            job_txs.push(job_tx);
        }

        WorkerPool {
            thread_count,
            job_txs,
            done_rx,
        }
    }

    /// Send one job to every worker and block until all complete.
    /// Returns `(proof_result, actual_hashes_computed)`.
    fn dispatch(
        &mut self,
        challenge: Arc<Vec<u8>>,
        start_nonce: u64,
        total_lanes: usize,
        difficulty: Option<u32>,
    ) -> (ProofResult, usize) {
        let found = Arc::new(AtomicBool::new(false));
        let winner = Arc::new(Mutex::new(None::<(u64, [u8; 32])>));
        let actual_hashes = Arc::new(AtomicUsize::new(0));

        for (tid, tx) in self.job_txs.iter().enumerate() {
            let job = Job {
                challenge: Arc::clone(&challenge),
                start_nonce,
                total_lanes,
                tid,
                thread_count: self.thread_count,
                difficulty,
                found: Arc::clone(&found),
                winner: Arc::clone(&winner),
                actual_hashes: Arc::clone(&actual_hashes),
            };
            tx.send(Some(job))
                .expect("worker channel closed unexpectedly");
        }

        // Collect completion signal from every worker.
        for _ in 0..self.thread_count {
            self.done_rx.recv().expect("worker done signal lost");
        }

        let result = *winner.lock().expect("winner mutex poisoned");
        (result, actual_hashes.load(Ordering::Relaxed))
    }
}

impl Drop for WorkerPool {
    fn drop(&mut self) {
        // Signal each worker to exit by sending `None`.
        for tx in &self.job_txs {
            let _ = tx.send(None);
        }
        // Workers will exit after receiving None; OS reclaims threads.
    }
}

/// Multi-threaded CPU solver backed by a persistent thread pool.
///
/// Threads and their 2 MB scratchpads are allocated once on first use and
/// reused for every subsequent batch, eliminating per-batch overhead.
pub struct ParallelCpuSolver {
    threads: usize,
    pool: Option<WorkerPool>,
}

impl ParallelCpuSolver {
    /// Create a new parallel solver with the given thread count.
    /// If `threads` is 0, it will be resolved in `recommended_lanes`.
    pub fn new(threads: usize) -> Self {
        Self {
            threads,
            pool: None,
        }
    }

    fn effective_threads(&self) -> usize {
        if self.threads == 0 {
            std::thread::available_parallelism()
                .map(|n| n.get())
                .unwrap_or(1)
        } else {
            self.threads
        }
    }

    fn ensure_pool(&mut self) {
        if self.pool.is_none() {
            let n = self.effective_threads();
            self.pool = Some(WorkerPool::new(n));
        }
    }
}

impl Solver for ParallelCpuSolver {
    fn backend_name(&self) -> &'static str {
        "cpu"
    }

    fn recommended_lanes(&mut self, requested: usize) -> usize {
        if requested == 0 {
            self.effective_threads() * 1024
        } else {
            requested
        }
    }

    fn find_proof_batch(
        &mut self,
        header_without_nonce: &[u8],
        start_nonce: u64,
        lanes: usize,
        difficulty: u32,
    ) -> Result<(ProofResult, usize)> {
        self.ensure_pool();
        let challenge = Arc::new(header_without_nonce.to_vec());
        let (result, actual) =
            self.pool
                .as_mut()
                .unwrap()
                .dispatch(challenge, start_nonce, lanes, Some(difficulty));
        Ok((result, actual))
    }

    fn benchmark_hashes(
        &mut self,
        header_without_nonce: &[u8],
        start_nonce: u64,
        lanes: usize,
    ) -> Result<usize> {
        self.ensure_pool();
        let challenge = Arc::new(header_without_nonce.to_vec());
        let (_, actual) = self
            .pool
            .as_mut()
            .unwrap()
            .dispatch(challenge, start_nonce, lanes, None);
        Ok(actual)
    }
}