async_fetcher/
checksum_system.rs

1// Copyright 2021-2022 System76 <info@system76.com>
2// SPDX-License-Identifier: MPL-2.0
3
4use crate::checksum::{Checksum, ChecksumError};
5use futures::prelude::*;
6use poolshark::{global::Pool, Poolable};
7use std::{path::Path, sync::Arc};
8
9/// Generates a stream of futures that validate checksums.
10///
11/// The caller can choose to distribute these futures across a thread pool.
12///
13/// ```no_test
14/// let mut stream = checksum_stream(checksums).map(tokio::spawn).buffered(8);
15/// while let Some((path, result)) = stream.next().await {
16///     eprintln!("{:?} checksum result: {:?}", path, result);
17/// }
18/// ```
19pub fn checksum_stream<I: Stream<Item = (Arc<Path>, Checksum)> + Send + Unpin + 'static>(
20    inputs: I,
21) -> impl Stream<Item = impl Future<Output = (Arc<Path>, Result<(), ChecksumError>)>> {
22    let buffer_pool = Arc::new(Pool::<Buf>::new(64, 1));
23
24    inputs.map(move |(dest, checksum)| {
25        let pool = buffer_pool.clone();
26
27        async {
28            tokio::task::spawn_blocking(move || {
29                let mut pooled_buf = pool.take();
30                let result = validate_checksum(&mut (&mut *pooled_buf).0[..], &dest, &checksum);
31                (dest, result)
32            })
33            .await
34            .unwrap()
35        }
36    })
37}
38
39/// Validates the checksum of a single file
40pub fn validate_checksum(
41    buf: &mut [u8],
42    dest: &Path,
43    checksum: &Checksum,
44) -> Result<(), ChecksumError> {
45    let error = match std::fs::File::open(&*dest) {
46        Ok(file) => match checksum.validate(file, buf) {
47            Ok(()) => return Ok(()),
48            Err(why) => why,
49        },
50        Err(why) => ChecksumError::from(why),
51    };
52
53    let _ = std::fs::remove_file(&*dest);
54    Err(error)
55}
56
57/// Poolable 8K buffer
58struct Buf(Box<[u8; 8 * 1024]>);
59
60impl Poolable for Buf {
61    fn empty() -> Self {
62        Self(Box::new([0u8; 8 * 1024]))
63    }
64
65    fn reset(&mut self) {
66        self.0.fill(0);
67    }
68
69    fn capacity(&self) -> usize {
70        1
71    }
72}
async_fetcher/checksum_system.rs

async_fetcher/
checksum_system.rs
