rust_multi_thread_walkdir

rust

use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::Arc;
use std::thread;
use std::time::{Duration, Instant};

use crossbeam::channel::{self, Receiver, Sender};

fn main() -> anyhow::Result<()> {
    let begin = Instant::now();
    // let root = PathBuf::from("/Users/fri3nd/");
    let root = PathBuf::from("/Users/fri3nd");
    let (finish_tx, finish_rx) = channel::bounded(4);
    let (task_tx, task_rx) = channel::unbounded();
    let numa = Arc::new(AtomicUsize::new(0));
    let size = Arc::new(AtomicU64::new(0));

    for id in 1..=32 {
        let ftx = finish_tx.clone();
        let tx = task_tx.clone();
        let rx = task_rx.clone();
        let num = numa.clone();
        let size = size.clone();

        thread::spawn(move || {
            walk_dir(id, tx, rx, num, size)?;
            ftx.send(())?;
            Ok::<_, anyhow::Error>(())
        });
    }

    drop(finish_tx);
    task_tx.send(root)?;
    for _ in finish_rx.iter() {}
    drop(task_tx);

    let duration = begin.elapsed();
    dbg!(duration);
    dbg!(numa);
    dbg!(size);

    // thread::sleep(Duration::from_secs(1000));
    Ok(())
}

fn walk_dir(
    id: usize,
    tx: Sender<PathBuf>,
    rx: Receiver<PathBuf>,
    num: Arc<AtomicUsize>,
    size: Arc<AtomicU64>,
) -> anyhow::Result<()> {
    loop {
        if let Ok(it) = rx.recv_timeout(Duration::from_millis(5000)) {
            println!("{} receive task {:?}", id, it);
            let subdir = inner_dir_read(it, num.clone(), size.clone());
            for sub in subdir {
                let _ = tx.send(sub);
            }
        }

        if rx.is_empty() {
            break;
        }
    }

    // println!("{} finished", id);
    Ok(())
}

fn inner_dir_read(path: PathBuf, num: Arc<AtomicUsize>, size: Arc<AtomicU64>) -> Vec<PathBuf> {
    let mut subdirs = vec![];
    if let Ok(children) = std::fs::read_dir(path) {
        for child in children {
            if let Ok(direntry) = child {
                if let Ok(metadata) = direntry.metadata() {
                    let subp = direntry.path();
                    if metadata.is_dir() {
                        num.fetch_add(1, Ordering::Release);
                        subdirs.push(subp);
                    } else {
                        num.fetch_add(1, Ordering::Release);
                        size.fetch_add(metadata.len(), Ordering::Release);
                        // println!("{:?}", subp);
                    }
                }
            }
        }
    }

    subdirs
}