6956988158

Committed 22 Nov 2023 11:52AM UTC coverage: 71.499% (+1.1%) from 70.425%

Build # 6956988158

Build Type

push

github

Committed by

ensc

Commit Message

CI: upload Cargo.lock as artifact

Run Details

1884 of 2635 relevant lines covered (71.5%)

383.9 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

57.36

/mod-proxy/src/cache.rs

use std::{os::unix::prelude::AsRawFd, time::Instant};
use std::sync::Arc;
use tokio::sync::RwLock;
use std::collections::HashMap;
// use chrono::{ NaiveDateTime, Utc };
use std::time::Duration;

use super::http;
use http::Time;

use crate::{ Result, Error };
use crate::util::pretty_dump_wrap as pretty;

lazy_static::lazy_static!{
    static ref CACHE: std::sync::RwLock<CacheImpl> = std::sync::RwLock::new(CacheImpl::new());
}

#[derive(Clone, Copy, Debug, Default)]
struct Stats {
    pub tm:                Duration,
}

impl Stats {
    pub async fn chunk(&mut self, response: &mut reqwest::Response) -> reqwest::Result<Option<bytes::Bytes>>
    {
        let start = std::time::Instant::now();
        let chunk = response.chunk().await;
        self.tm += start.elapsed();

        chunk
    }
}

impl crate::util::PrettyDump for Stats {
    fn pretty_dump(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!("{:.3}s", self.tm.as_secs_f32()))
    }
}

#[derive(Debug)]
enum State {
    None,

    Error(&'static str),

    Init {
        response:        reqwest::Response,
    },

    HaveMeta {
        response:        reqwest::Response,
        cache_info:        http::CacheInfo,
        file_size:        Option<u64>,
        stats:                Stats,
    },

    Downloading {
        response:        reqwest::Response,
        cache_info:        http::CacheInfo,
        file_size:        Option<u64>,
        file:                std::fs::File,
        file_pos:        u64,
        stats:                Stats,
    },

    Complete {
        cache_info:        http::CacheInfo,
        file:                std::fs::File,
        file_size:        u64,
    },

    Refresh {
        response:        reqwest::Response,
        cache_info:        http::CacheInfo,
        file:                std::fs::File,
        file_size:        u64,
    },
}

impl crate::util::PrettyDump for State {
    fn pretty_dump(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            State::None =>
                f.write_str("no state"),
            State::Error(e) =>
                f.write_fmt(format_args!("error {:?}", e)),
            State::Init { response } =>
                f.write_fmt(format_args!("INIT({})", pretty(response))),
            State::HaveMeta { response, cache_info, file_size, stats } =>
                f.write_fmt(format_args!("META({}, {}, {}, {})",
                                         pretty(response), pretty(cache_info),
                                         pretty(file_size), pretty(stats))),
            State::Downloading { response, cache_info, file_size, file, file_pos, stats } =>
                f.write_fmt(format_args!("DOWNLOADING({}, {}, {}, {}@{}, {})",
                                         pretty(response), pretty(cache_info),
                                         pretty(file_size), pretty(file),
                                         file_pos, pretty(stats))),
            State::Complete { cache_info, file, file_size } =>
                f.write_fmt(format_args!("COMPLETE({}, {}/{})",
                                         pretty(cache_info), pretty(file), file_size)),
            State::Refresh { response, cache_info, file, file_size } =>
                f.write_fmt(format_args!("REFRESH({},.{}, [{}/{}])",
                                         pretty(response), pretty(cache_info), pretty(file),
                                         file_size)),
        }
    }
}

impl State {
    pub fn take(&mut self, hint: &'static str) -> Self {
        std::mem::replace(self, State::Error(hint))
    }

    pub fn is_none(&self) -> bool {
        matches!(self, Self::None)
    }

    pub fn is_init(&self) -> bool {
        matches!(self, Self::Init { .. })
    }

    pub fn is_error(&self) -> bool {
        matches!(self, Self::Error(_))
    }

    pub fn is_refresh(&self) -> bool {
        matches!(self, Self::Refresh { .. })
    }

    pub fn is_have_meta(&self) -> bool {
        matches!(self, Self::HaveMeta { .. })
    }

    pub fn is_downloading(&self) -> bool {
        matches!(self, Self::Downloading { .. })
    }

    pub fn is_complete(&self) -> bool {
        matches!(self, Self::Complete { .. })
    }

    pub fn get_file_size(&self) -> Option<u64> {
        match self {
            Self::None |
            Self::Init { .. }        => None,

            Self::HaveMeta { file_size, .. }        => *file_size,
            Self::Downloading { file_size, .. }        => *file_size,

            Self::Complete { file_size, .. }        => Some(*file_size),
            Self::Refresh { file_size, .. }        => Some(*file_size),

            Self::Error(hint)        => panic!("get_file_size called in error state ({hint})"),
        }
    }

    pub fn get_cache_info(&self) -> Option<&http::CacheInfo> {
        match self {
            Self::None |
            Self::Error(_) |
            Self::Init { .. }        => None,

            Self::HaveMeta { cache_info, .. } |
            Self::Downloading { cache_info, .. } |
            Self::Complete { cache_info, .. } |
            Self::Refresh { cache_info, .. }        => Some(cache_info),
        }
    }

    fn read_file(file: &std::fs::File, ofs: u64, buf: &mut [u8], max: u64) -> Result<usize> {
        use nix::libc;

        assert!(max > ofs);

        let len = (buf.len() as u64).min(max - ofs) as usize;
        let buf_ptr = buf.as_mut_ptr() as *mut libc::c_void;

        // TODO: this would be nice, but does not work because we can not get
        // a mutable reference to 'file'
        //file.flush()?;

        let rc = unsafe { libc::pread(file.as_raw_fd(), buf_ptr, len, ofs as i64) };

        if rc < 0 {
            return Err(std::io::Error::last_os_error().into());
        }

        Ok(len)
    }

    pub fn read(&self, ofs: u64, buf: &mut [u8]) -> Result<Option<usize>> {
        match &self {
            State::Downloading { file, file_pos, .. } if ofs < *file_pos        => {
                Self::read_file(file, ofs, buf, *file_pos)
            },

            State::Complete { file, file_size, .. } if ofs < *file_size                => {
                Self::read_file(file, ofs, buf, *file_size)
            }

            State::Complete { file_size, .. } if ofs == *file_size        => Ok(0),

            State::Complete { file_size, .. } if ofs >= *file_size        =>
                Err(Error::Internal("file out-of-bound read")),

            _        => return Ok(None)
        }.map(Some)
    }

    pub fn is_outdated(&self, reftm: Instant, max_lt: Duration) -> bool {
        match self.get_cache_info() {
            None        => true,
            Some(info)        => info.is_outdated(reftm, max_lt),
        }
    }
}

#[derive(Debug)]
pub struct EntryData {
    pub key:                url::Url,
    state:                State,
    reftm:                Time,
}

impl std::fmt::Display for EntryData {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!("{}: reftm={}, state={}", self.key,
                    pretty(&self.reftm.local), pretty(&self.state)))
    }
}

impl EntryData {
    pub fn new(url: &url::Url) -> Self {
        Self {
            key:                url.clone(),
            state:                State::None,
            reftm:                Time::now(),
        }
    }

    pub fn is_complete(&self) -> bool {
        self.state.is_complete()
    }

    pub fn is_error(&self) -> bool {
        self.state.is_error()
    }

    pub fn is_running(&self) -> bool {
        self.state.is_have_meta() || self.state.is_downloading()
    }

    pub fn update_localtm(&mut self) {
        self.reftm = Time::now();
    }

    pub fn set_response(&mut self, response: reqwest::Response) {
        self.state = match self.state.take("set_respone") {
            State::None |
            State::Error(_)        => State::Init { response },

            State::Complete { cache_info, file, file_size } |
            State::Refresh { cache_info, file, file_size, .. } => State::Refresh {
                cache_info:        cache_info,
                file:                file,
                file_size:        file_size,
                response:        response,
            },

            s                        => panic!("unexpected state {s:?}"),
        }
    }

    pub fn is_outdated(&self, reftm: Instant, max_lt: Duration) -> bool {
        self.state.is_outdated(reftm, max_lt)
    }

    pub fn get_cache_info(&self) -> Option<&http::CacheInfo> {
        self.state.get_cache_info()
    }

    pub async fn fill_meta(&mut self) -> Result<()> {
        if !self.state.is_init() && !self.state.is_none() && !self.state.is_refresh() {
            return Ok(());
        }

        self.state = match self.state.take("fill_meta") {
            State::None                        => panic!("unexpected state"),

            State::Init{ response }        => {
                let hdrs = response.headers();

                State::HaveMeta {
                    cache_info:        http::CacheInfo::new(self.reftm, hdrs)?,
                    file_size:        response.content_length(),
                    response:        response,
                    stats:        Stats::default(),
                }
            },

            State::Refresh { file, file_size, response, cache_info }        => {
                let hdrs = response.headers();

                State::Complete {
                    cache_info:        cache_info.update(self.reftm, hdrs)?,
                    file:        file,
                    file_size:        file_size,
                }
            },

            _                                => unreachable!(),
        };

        Ok(())
    }

    fn signal_complete(&self, stats: Stats) {
        if let State::Complete { file_size, .. } = self.state {
            info!("downloaded {} with {} bytes in {}ms", self.key, file_size, stats.tm.as_millis());
        }
    }

    #[instrument(level = "trace")]
    pub async fn get_filesize(&mut self) -> Result<u64> {
        use std::io::Write;

        if let Some(sz) = self.state.get_file_size() {
            return Ok(sz);
        }

        match self.state.take("get_filesize") {
            State::HaveMeta { mut response, file_size: None, mut stats, cache_info }        => {
                let mut file = Cache::new_file()?;
                let mut pos = 0;


                while let Some(chunk) = stats.chunk(&mut response).await? {
                    pos += chunk.len() as u64;
                    file.write_all(&chunk)?;
                }

                self.state = State::Complete {
                    file:        file,
                    file_size:        pos,
                    cache_info:        cache_info,
                };

                self.signal_complete(stats);

                Ok(pos)
            },

            State::Downloading { mut response, mut file, file_pos, file_size: None, mut stats, cache_info } => {
                let mut pos = file_pos;

                while let Some(chunk) = stats.chunk(&mut response).await? {
                    pos += chunk.len() as u64;
                    file.write_all(&chunk)?;
                }

                self.state = State::Complete {
                    file:        file,
                    file_size:        pos,
                    cache_info:        cache_info,
                };

                self.signal_complete(stats);

                Ok(pos)
            }

            s                => panic!("unexpected state: {s:?}"),
        }
    }

    pub fn fill_request(&self, req: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
        match self.state.get_cache_info() {
            Some(info)        => info.fill_request(self.reftm.mono, req),
            None        => req,
        }
    }

    pub fn matches(&self, etag: Option<&str>) -> bool {
        let cache_info = self.state.get_cache_info();

        match cache_info.and_then(|c| c.not_after) {
            Some(t) if t < Time::now().mono                => return false,
            _                                                => {},
        }

        let self_etag = match cache_info {
            Some(c)        => c.etag.as_ref(),
            None        => None,
        };

        match (self_etag, etag) {
            (Some(a), Some(b)) if a == b                => {},
            (None, None)                                => {},
            _                                                => return false,
        }

        true
    }

    pub fn invalidate(&mut self)
    {
        match &self.state {
            State::Refresh { .. }        => self.state = State::None,
            State::Complete { .. }        => self.state = State::None,
            _                                => {},
        }
    }

    pub async fn read_some(&mut self, ofs: u64, buf: &mut [u8]) -> Result<usize>
    {
        use std::io::Write;

        trace!("state={:?}, ofs={}, #buf={}", self.state, ofs, buf.len());

        async fn fetch(response: &mut reqwest::Response, file: &mut std::fs::File,
                       buf: &mut [u8], stats: &mut Stats) -> Result<(usize, usize)> {
            match stats.chunk(response).await? {
                Some(data)        => {
                    let len = buf.len().min(data.len());

                    buf[0..len].clone_from_slice(&data.as_ref()[0..len]);
                    file.write_all(&data)?;

                    // TODO: it would be better to do this in State::read_file()
                    file.flush()?;

                    Ok((len, data.len()))
                },

                None                => Ok((0, 0))
            }
        }

        if self.state.is_init() {
            self.fill_meta().await?;
        }

        if let Some(sz) = self.state.read(ofs, buf)? {
            return Ok(sz);
        }

        match self.state.take("read_some") {
            State::HaveMeta { mut response, cache_info, file_size, mut stats }        => {
                let mut file = Cache::new_file()?;

                let res = fetch(&mut response, &mut file, buf, &mut stats).await?;

                self.state = match res {
                    (_, 0)        => State::Complete {
                        cache_info:        cache_info,
                        file:                file,
                        file_size:        0,
                    },

                    (_, sz)        => State::Downloading {
                        response:        response,
                        cache_info:        cache_info,
                        file_size:        file_size,
                        file:                file,
                        file_pos:        sz as u64,
                        stats:                stats,
                    }
                };

                self.signal_complete(stats);

                Ok(res.0)
            },

            // catched by 'self.state.read()' above
            State::Downloading { file_pos, .. } if ofs < file_pos        => unreachable!(),

            State::Downloading { mut response, cache_info, file_size, mut file, file_pos, mut stats } => {
                let res = fetch(&mut response, &mut file, buf, &mut stats).await?;

                self.state = match res {
                    (_, 0)        => State::Complete {
                        cache_info:        cache_info,
                        file:                file,
                        file_size:        file_pos,
                    },

                    (_, sz)        => State::Downloading {
                        response:        response,
                        cache_info:        cache_info,
                        file_size:        file_size,
                        file:                file,
                        file_pos:        file_pos + (sz as u64),
                        stats:                stats,
                    }
                };

                self.signal_complete(stats);

                Ok(res.0)
            }

            s                => panic!("unexpected state: {s:?}"),
        }

    }
}

pub type Entry = Arc<RwLock<EntryData>>;

struct CacheImpl {
    tmpdir:        std::path::PathBuf,
    entries:        HashMap<url::Url, Entry>,
    client:        Arc<reqwest::Client>,
    refcnt:        u32,

    abort_ch:        Option<tokio::sync::watch::Sender<()>>,
    gc:                Option<tokio::task::JoinHandle<()>>,
}

pub enum LookupResult {
    Found(Entry),
    Missing,
}

impl CacheImpl {
    fn new() -> Self {
        Self {
            tmpdir:        std::env::temp_dir(),
            entries:        HashMap::new(),
            client:        Arc::new(reqwest::Client::new()),
            abort_ch:        None,
            refcnt:        0,
            gc:                None,
        }
    }

    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    pub fn clear(&mut self) {
        self.entries.clear();
    }

    pub fn get_client(&self) -> Arc<reqwest::Client> {
        self.client.clone()
    }

    pub fn lookup_or_create(&mut self, key: &url::Url) -> Entry {
        match self.entries.get(key) {
            Some(v)        => v.clone(),
            None        => self.create(key),
        }
    }

    pub fn create(&mut self, key: &url::Url) -> Entry {
        Entry::new(RwLock::new(EntryData::new(key)))
    }

    pub fn replace(&mut self, key: &url::Url, entry: &Entry) {
        self.entries.insert(key.clone(), entry.clone());
    }

    pub fn remove(&mut self, key: &url::Url) {
        self.entries.remove(key);
    }

    /// Removes the `num` oldest cache entries
    pub fn gc_oldest(&mut self, mut num: usize) {
        if num == 0 {
            return;
        }

        let mut tmp = Vec::with_capacity(self.entries.len());

        for (key, e) in &self.entries {
            let entry = match e.try_read() {
                Ok(e)        => e,
                _        => continue,
            };

            tmp.push((key.clone(), entry.get_cache_info().map(|c| c.local_time)));
        }

        tmp.sort_by(|(_, tm_a), (_, tm_b)| tm_a.cmp(tm_b));

        let mut rm_cnt = 0;

        for (key, _) in tmp {
            if num == 0 {
                break;
            }

            debug!("gc: removing old {}", key);
            self.entries.remove(&key);
            num -= 1;
            rm_cnt += 1;
        }

        if rm_cnt > 0 {
            info!("gc: removed {} old entries", rm_cnt);
        }
    }

    /// Removes cache entries which are older than `max_lt`.
    ///
    /// Returns the number of remaining cache entries.
    pub fn gc_outdated(&mut self, max_lt: Duration) -> usize {
        let now = Instant::now();
        let mut outdated = Vec::new();
        let mut cnt = 0;

        for (key, e) in &self.entries {
            match e.try_read().map(|v| v.is_outdated(now, max_lt)) {
                Ok(true)        => outdated.push(key.clone()),
                _                => cnt += 1,
            }
        }

        let rm_cnt = outdated.len();

        for e in outdated {
            debug!("gc: removing outdated {}", e);
            self.entries.remove(&e);
        }

        if rm_cnt > 0 {
            info!("gc: removed {} obsolete entries", rm_cnt);
        }

        cnt
    }
}

#[derive(Debug)]
pub struct GcProperties {
    pub max_elements:        usize,
    pub max_lifetime:        Duration,
    pub sleep:                Duration,
}

async fn gc_runner(props: GcProperties, mut abort_ch: tokio::sync::watch::Receiver<()>) {
    const RETRY_DELAY: std::time::Duration = std::time::Duration::from_secs(1);

    loop {
        use std::sync::TryLockError;

        let sleep = {
            let cache = CACHE.try_write();

            match cache {
                Ok(mut cache) if !cache.is_empty()        => {
                    let cache_cnt = cache.gc_outdated(props.max_lifetime);

                    if cache_cnt > props.max_elements {
                        cache.gc_oldest(props.max_elements - cache_cnt)
                    }

                    props.sleep
                }
                Ok(_)                                => props.sleep,
                Err(TryLockError::WouldBlock)        => RETRY_DELAY,
                Err(e)                                => {
                    error!("cache gc failed with {:?}", e);
                    break;
                }
            }
        };

        if tokio::time::timeout(sleep, abort_ch.changed()).await.is_ok() {
            debug!("cache gc runner gracefully closed");
            break;
        }
    }
}

pub struct Cache();

impl Cache {
    #[instrument(level = "trace")]
    pub fn instanciate(tmpdir: &std::path::Path, props: GcProperties) {
        let mut cache = CACHE.write().unwrap();

        if cache.refcnt == 0 {
            let (tx, rx) = tokio::sync::watch::channel(());

            cache.tmpdir = tmpdir.into();
            cache.abort_ch = Some(tx);

            cache.gc = Some(tokio::task::spawn(gc_runner(props, rx)));
        }

        cache.refcnt += 1;
    }

    #[instrument(level = "trace")]
    // https://github.com/rust-lang/rust-clippy/issues/6446
    #[allow(clippy::await_holding_lock)]
    pub async fn close() {
        let mut cache = CACHE.write().unwrap();

        assert!(cache.refcnt > 0);

        cache.refcnt -= 1;

        if cache.refcnt == 0 {
            cache.entries.clear();

            let abort_ch = cache.abort_ch.take().unwrap();
            let gc = cache.gc.take().unwrap();

            drop(cache);

            abort_ch.send(()).unwrap();
            gc.await.unwrap();
        }
    }

    #[instrument(level = "trace", ret)]
    pub fn lookup_or_create(key: &url::Url) -> Entry {
        let mut cache = CACHE.write().unwrap();

        cache.lookup_or_create(key)
    }

    #[instrument(level = "trace", ret)]
    pub fn create(key: &url::Url) -> Entry {
        let mut cache = CACHE.write().unwrap();

        cache.create(key)
    }

    #[instrument(level = "trace")]
    pub fn replace(key: &url::Url, entry: &Entry) {
        let mut cache = CACHE.write().unwrap();

        cache.replace(key, entry)
    }

    #[instrument(level = "trace")]
    pub fn remove(key: &url::Url) {
        let mut cache = CACHE.write().unwrap();

        cache.remove(key)
    }

    pub fn get_client() -> Arc<reqwest::Client> {
        let cache = CACHE.read().unwrap();

        cache.get_client()
    }

    pub fn new_file() -> Result<std::fs::File> {
        let cache = CACHE.read().unwrap();

        Ok(tempfile::tempfile_in(&cache.tmpdir)?)
    }

    pub async fn dump() {
        let mut entries = Vec::new();

        {
            let cache = CACHE.read().unwrap();

            entries.reserve(cache.entries.len());
            entries.extend(cache.entries.values().cloned());
        }

        println!("Cache information ({} entries)", entries.len());

        for e in entries {
            println!("{}", e.read().await);
        }
    }

    pub async fn clear() {
        let mut cache = CACHE.write().unwrap();

        cache.clear();
    }
}

1	use std::{os::unix::prelude::AsRawFd, time::Instant};
2	use std::sync::Arc;
3	use tokio::sync::RwLock;
4	use std::collections::HashMap;
5	// use chrono::{ NaiveDateTime, Utc };
6	use std::time::Duration;
7
8	use super::http;
9	use http::Time;
10
11	use crate::{ Result, Error };
12	use crate::util::pretty_dump_wrap as pretty;
13
14	lazy_static::lazy_static!{
15	static ref CACHE: std::sync::RwLock<CacheImpl> = std::sync::RwLock::new(CacheImpl::new());
16	}
17
18	#[derive(Clone, Copy, Debug, Default)]	216✔
19	struct Stats {
20	pub tm: Duration,	108✔
21	}
22
23	impl Stats {
24	pub async fn chunk(&mut self, response: &mut reqwest::Response) -> reqwest::Result<Option<bytes::Bytes>>	464✔
25	{	464✔
26	let start = std::time::Instant::now();	232✔
27	let chunk = response.chunk().await;	232✔
28	self.tm += start.elapsed();	232✔
29
30	chunk
31	}	464✔
32	}
33
34	impl crate::util::PrettyDump for Stats {
35	fn pretty_dump(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {	×
36	f.write_fmt(format_args!("{:.3}s", self.tm.as_secs_f32()))	×
37	}	×
38	}
39
40	#[derive(Debug)]	×
41	enum State {
42	None,
43
44	Error(&'static str),	×
45
46	Init {
47	response: reqwest::Response,	×
48	},
49
50	HaveMeta {
51	response: reqwest::Response,	×
52	cache_info: http::CacheInfo,	×
53	file_size: Option<u64>,	×
54	stats: Stats,	×
55	},
56
57	Downloading {
58	response: reqwest::Response,	×
59	cache_info: http::CacheInfo,	×
60	file_size: Option<u64>,	×
61	file: std::fs::File,	×
62	file_pos: u64,	×
63	stats: Stats,	×
64	},
65
66	Complete {
67	cache_info: http::CacheInfo,	×
68	file: std::fs::File,	×
69	file_size: u64,	×
70	},
71
72	Refresh {
73	response: reqwest::Response,	×
74	cache_info: http::CacheInfo,	×
75	file: std::fs::File,	×
76	file_size: u64,	×
77	},
78	}
79
80	impl crate::util::PrettyDump for State {
81	fn pretty_dump(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {	400✔
82	match self {	400✔
83	State::None =>
84	f.write_str("no state"),	×
85	State::Error(e) =>	×
86	f.write_fmt(format_args!("error {:?}", e)),	×
87	State::Init { response } =>	×
88	f.write_fmt(format_args!("INIT({})", pretty(response))),	×
89	State::HaveMeta { response, cache_info, file_size, stats } =>	×
90	f.write_fmt(format_args!("META({}, {}, {}, {})",	×
91	pretty(response), pretty(cache_info),	×
92	pretty(file_size), pretty(stats))),	×
93	State::Downloading { response, cache_info, file_size, file, file_pos, stats } =>	×
94	f.write_fmt(format_args!("DOWNLOADING({}, {}, {}, {}@{}, {})",	×
95	pretty(response), pretty(cache_info),	×
96	pretty(file_size), pretty(file),	×
97	file_pos, pretty(stats))),	×
98	State::Complete { cache_info, file, file_size } =>	400✔
99	f.write_fmt(format_args!("COMPLETE({}, {}/{})",	800✔
100	pretty(cache_info), pretty(file), file_size)),	800✔
101	State::Refresh { response, cache_info, file, file_size } =>	×
102	f.write_fmt(format_args!("REFRESH({},.{}, [{}/{}])",	×
103	pretty(response), pretty(cache_info), pretty(file),	×
104	file_size)),
105	}
106	}	400✔
107	}
108
109	impl State {
110	pub fn take(&mut self, hint: &'static str) -> Self {	448✔
111	std::mem::replace(self, State::Error(hint))	448✔
112	}	448✔
113
114	pub fn is_none(&self) -> bool {	×
115	matches!(self, Self::None)	×
116	}	×
117
118	pub fn is_init(&self) -> bool {	1,398✔
119	matches!(self, Self::Init { .. })	1,398✔
120	}	1,398✔
121
122	pub fn is_error(&self) -> bool {	108✔
123	matches!(self, Self::Error(_))	108✔
124	}	108✔
125
126	pub fn is_refresh(&self) -> bool {	×
127	matches!(self, Self::Refresh { .. })	×
128	}	×
129
130	pub fn is_have_meta(&self) -> bool {	108✔
131	matches!(self, Self::HaveMeta { .. })	108✔
132	}	108✔
133
134	pub fn is_downloading(&self) -> bool {	108✔
135	matches!(self, Self::Downloading { .. })	108✔
136	}	108✔
137
138	pub fn is_complete(&self) -> bool {	×
139	matches!(self, Self::Complete { .. })	×
140	}	×
141
142	pub fn get_file_size(&self) -> Option<u64> {	108✔
143	match self {	108✔
144	Self::None \|
145	Self::Init { .. } => None,	×
146
147	Self::HaveMeta { file_size, .. } => *file_size,	108✔
148	Self::Downloading { file_size, .. } => *file_size,	×
149
150	Self::Complete { file_size, .. } => Some(*file_size),	×
151	Self::Refresh { file_size, .. } => Some(*file_size),	×
152
153	Self::Error(hint) => panic!("get_file_size called in error state ({hint})"),	×
154	}
155	}	108✔
156
157	pub fn get_cache_info(&self) -> Option<&http::CacheInfo> {	108✔
158	match self {	108✔
159	Self::None \|
160	Self::Error(_) \|
161	Self::Init { .. } => None,	92✔
162
163	Self::HaveMeta { cache_info, .. } \|	×
164	Self::Downloading { cache_info, .. } \|	×
165	Self::Complete { cache_info, .. } \|	16✔
166	Self::Refresh { cache_info, .. } => Some(cache_info),	16✔
167	}
168	}	108✔
169
170	fn read_file(file: &std::fs::File, ofs: u64, buf: &mut [u8], max: u64) -> Result<usize> {	1,058✔
171	use nix::libc;
172
173	assert!(max > ofs);	1,058✔
174
175	let len = (buf.len() as u64).min(max - ofs) as usize;	1,058✔
176	let buf_ptr = buf.as_mut_ptr() as *mut libc::c_void;	1,058✔
177
178	// TODO: this would be nice, but does not work because we can not get
179	// a mutable reference to 'file'
180	//file.flush()?;
181
182	let rc = unsafe { libc::pread(file.as_raw_fd(), buf_ptr, len, ofs as i64) };	1,058✔
183
184	if rc < 0 {	1,058✔
185	return Err(std::io::Error::last_os_error().into());	×
186	}
187
188	Ok(len)	1,058✔
189	}	1,058✔
190
191	pub fn read(&self, ofs: u64, buf: &mut [u8]) -> Result<Option<usize>> {	1,290✔
192	match &self {	2,348✔
193	State::Downloading { file, file_pos, .. } if ofs < *file_pos => {	1,182✔
194	Self::read_file(file, ofs, buf, *file_pos)	1,058✔
195	},
196
197	State::Complete { file, file_size, .. } if ofs < *file_size => {	×
198	Self::read_file(file, ofs, buf, *file_size)	×
199	}
200
201	State::Complete { file_size, .. } if ofs == *file_size => Ok(0),	×
202
203	State::Complete { file_size, .. } if ofs >= *file_size =>	×
204	Err(Error::Internal("file out-of-bound read")),	×
205
206	_ => return Ok(None)	232✔
207	}.map(Some)
208	}	1,290✔
209
210	pub fn is_outdated(&self, reftm: Instant, max_lt: Duration) -> bool {	×
211	match self.get_cache_info() {	×
212	None => true,	×
213	Some(info) => info.is_outdated(reftm, max_lt),	×
214	}
215	}	×
216	}
217
218	#[derive(Debug)]	×
219	pub struct EntryData {
220	pub key: url::Url,	×
221	state: State,
222	reftm: Time,	×
223	}
224
225	impl std::fmt::Display for EntryData {
226	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {	400✔
227	f.write_fmt(format_args!("{}: reftm={}, state={}", self.key,	1,200✔
228	pretty(&self.reftm.local), pretty(&self.state)))	800✔
229	}	400✔
230	}
231
232	impl EntryData {
233	pub fn new(url: &url::Url) -> Self {	92✔
234	Self {	92✔
235	key: url.clone(),	92✔
236	state: State::None,	92✔
237	reftm: Time::now(),	92✔
238	}
239	}	92✔
240
241	pub fn is_complete(&self) -> bool {	×
242	self.state.is_complete()	×
243	}	×
244
245	pub fn is_error(&self) -> bool {	108✔
246	self.state.is_error()	108✔
247	}	108✔
248
249	pub fn is_running(&self) -> bool {	108✔
250	self.state.is_have_meta() \|\| self.state.is_downloading()	108✔
251	}	108✔
252
253	pub fn update_localtm(&mut self) {	108✔
254	self.reftm = Time::now();	108✔
255	}	108✔
256
257	pub fn set_response(&mut self, response: reqwest::Response) {	108✔
258	self.state = match self.state.take("set_respone") {	216✔
259	State::None \|
260	State::Error(_) => State::Init { response },	108✔
261
262	State::Complete { cache_info, file, file_size } \|	×
263	State::Refresh { cache_info, file, file_size, .. } => State::Refresh {	×
264	cache_info: cache_info,	×
265	file: file,	×
266	file_size: file_size,	×
267	response: response,
268	},	×
269
270	s => panic!("unexpected state {s:?}"),	×
271	}
272	}	108✔
273
274	pub fn is_outdated(&self, reftm: Instant, max_lt: Duration) -> bool {	×
275	self.state.is_outdated(reftm, max_lt)	×
276	}	×
277
278	pub fn get_cache_info(&self) -> Option<&http::CacheInfo> {	×
279	self.state.get_cache_info()	×
280	}	×
281
282	pub async fn fill_meta(&mut self) -> Result<()> {	324✔
283	if !self.state.is_init() && !self.state.is_none() && !self.state.is_refresh() {	108✔
284	return Ok(());	×
285	}
286
287	self.state = match self.state.take("fill_meta") {	216✔
288	State::None => panic!("unexpected state"),	×
289
290	State::Init{ response } => {	108✔
291	let hdrs = response.headers();	108✔
292
293	State::HaveMeta {	108✔
294	cache_info: http::CacheInfo::new(self.reftm, hdrs)?,	108✔
295	file_size: response.content_length(),	108✔
296	response: response,	108✔
297	stats: Stats::default(),	108✔
298	}
299	},	108✔
300
301	State::Refresh { file, file_size, response, cache_info } => {	×
302	let hdrs = response.headers();	×
303
304	State::Complete {	×
305	cache_info: cache_info.update(self.reftm, hdrs)?,	×
306	file: file,	×
307	file_size: file_size,
308	}
309	},	×
310
311	_ => unreachable!(),	×
312	};	108✔
313
314	Ok(())	108✔
315	}	216✔
316
317	fn signal_complete(&self, stats: Stats) {	232✔
318	if let State::Complete { file_size, .. } = self.state {	232✔
319	info!("downloaded {} with {} bytes in {}ms", self.key, file_size, stats.tm.as_millis());	108✔
320	}
321	}	232✔
322
323	#[instrument(level = "trace")]	540✔
324	pub async fn get_filesize(&mut self) -> Result<u64> {	216✔
325	use std::io::Write;
326
327	if let Some(sz) = self.state.get_file_size() {	108✔
328	return Ok(sz);	108✔
329	}
330
331	match self.state.take("get_filesize") {	×
332	State::HaveMeta { mut response, file_size: None, mut stats, cache_info } => {	×
333	let mut file = Cache::new_file()?;	×
334	let mut pos = 0;	×
335
336
337	while let Some(chunk) = stats.chunk(&mut response).await? {	×
338	pos += chunk.len() as u64;	×
339	file.write_all(&chunk)?;	×
340	}	×
341
342	self.state = State::Complete {	×
343	file: file,	×
344	file_size: pos,	×
345	cache_info: cache_info,	×
346	};
347
348	self.signal_complete(stats);	×
349
350	Ok(pos)	×
351	},	×
352
353	State::Downloading { mut response, mut file, file_pos, file_size: None, mut stats, cache_info } => {	×
354	let mut pos = file_pos;	×
355
356	while let Some(chunk) = stats.chunk(&mut response).await? {	×
357	pos += chunk.len() as u64;	×
358	file.write_all(&chunk)?;	×
359	}	×
360
361	self.state = State::Complete {	×
362	file: file,	×
363	file_size: pos,	×
364	cache_info: cache_info,	×
365	};
366
367	self.signal_complete(stats);	×
368
369	Ok(pos)	×
370	}	×
371
372	s => panic!("unexpected state: {s:?}"),	×
373	}
374	}
375
376	pub fn fill_request(&self, req: reqwest::RequestBuilder) -> reqwest::RequestBuilder {	108✔
377	match self.state.get_cache_info() {	108✔
378	Some(info) => info.fill_request(self.reftm.mono, req),	16✔
379	None => req,	92✔
380	}
381	}	108✔
382
383	pub fn matches(&self, etag: Option<&str>) -> bool {	×
384	let cache_info = self.state.get_cache_info();	×
385
386	match cache_info.and_then(\|c\| c.not_after) {	×
387	Some(t) if t < Time::now().mono => return false,	×
388	_ => {},
389	}
390
391	let self_etag = match cache_info {	×
392	Some(c) => c.etag.as_ref(),	×
393	None => None,	×
394	};
395
396	match (self_etag, etag) {	×
397	(Some(a), Some(b)) if a == b => {},	×
398	(None, None) => {},
399	_ => return false,	×
400	}
401
402	true	×
403	}	×
404
405	pub fn invalidate(&mut self)	108✔
406	{
407	match &self.state {	108✔
408	State::Refresh { .. } => self.state = State::None,	×
409	State::Complete { .. } => self.state = State::None,	16✔
410	_ => {},
411	}
412	}	108✔
413
414	pub async fn read_some(&mut self, ofs: u64, buf: &mut [u8]) -> Result<usize>	2,580✔
415	{	2,580✔
416	use std::io::Write;
417
418	trace!("state={:?}, ofs={}, #buf={}", self.state, ofs, buf.len());	1,290✔
419
420	async fn fetch(response: &mut reqwest::Response, file: &mut std::fs::File,	464✔
421	buf: &mut [u8], stats: &mut Stats) -> Result<(usize, usize)> {	464✔
422	match stats.chunk(response).await? {	232✔
423	Some(data) => {	124✔
424	let len = buf.len().min(data.len());	124✔
425
426	buf[0..len].clone_from_slice(&data.as_ref()[0..len]);	124✔
427	file.write_all(&data)?;	124✔
428
429	// TODO: it would be better to do this in State::read_file()
430	file.flush()?;	124✔
431
432	Ok((len, data.len()))	124✔
433	},	124✔
434
435	None => Ok((0, 0))	108✔
436	}
437	}	464✔
438
439	if self.state.is_init() {	1,290✔
440	self.fill_meta().await?;	×
441	}
442
443	if let Some(sz) = self.state.read(ofs, buf)? {	1,290✔
444	return Ok(sz);	1,058✔
445	}
446
447	match self.state.take("read_some") {	232✔
448	State::HaveMeta { mut response, cache_info, file_size, mut stats } => {	108✔
449	let mut file = Cache::new_file()?;	108✔
450
451	let res = fetch(&mut response, &mut file, buf, &mut stats).await?;	108✔
452
453	self.state = match res {	216✔
454	(_, 0) => State::Complete {	4✔
455	cache_info: cache_info,	4✔
456	file: file,	4✔
457	file_size: 0,
458	},	4✔
459
460	(_, sz) => State::Downloading {	208✔
461	response: response,	104✔
462	cache_info: cache_info,	104✔
463	file_size: file_size,	104✔
464	file: file,	104✔
465	file_pos: sz as u64,
466	stats: stats,	104✔
467	}	104✔
468	};
469
470	self.signal_complete(stats);	108✔
471
472	Ok(res.0)	108✔
473	},	108✔
474
475	// catched by 'self.state.read()' above
476	State::Downloading { file_pos, .. } if ofs < file_pos => unreachable!(),	124✔
477
478	State::Downloading { mut response, cache_info, file_size, mut file, file_pos, mut stats } => {	124✔
479	let res = fetch(&mut response, &mut file, buf, &mut stats).await?;	124✔
480
481	self.state = match res {	248✔
482	(_, 0) => State::Complete {	104✔
483	cache_info: cache_info,	104✔
484	file: file,	104✔
485	file_size: file_pos,	104✔
486	},	104✔
487
488	(_, sz) => State::Downloading {	40✔
489	response: response,	20✔
490	cache_info: cache_info,	20✔
491	file_size: file_size,	20✔
492	file: file,	20✔
493	file_pos: file_pos + (sz as u64),	20✔
494	stats: stats,	20✔
495	}	20✔
496	};
497
498	self.signal_complete(stats);	124✔
499
500	Ok(res.0)	124✔
501	}	124✔
502
503	s => panic!("unexpected state: {s:?}"),	×
504	}
505
506	}	2,580✔
507	}
508
509	pub type Entry = Arc<RwLock<EntryData>>;
510
511	struct CacheImpl {
512	tmpdir: std::path::PathBuf,
513	entries: HashMap<url::Url, Entry>,
514	client: Arc<reqwest::Client>,
515	refcnt: u32,
516
517	abort_ch: Option<tokio::sync::watch::Sender<()>>,
518	gc: Option<tokio::task::JoinHandle<()>>,
519	}
520
521	pub enum LookupResult {
522	Found(Entry),
523	Missing,
524	}
525
526	impl CacheImpl {
527	fn new() -> Self {	2✔
528	Self {	2✔
529	tmpdir: std::env::temp_dir(),	2✔
530	entries: HashMap::new(),	2✔
531	client: Arc::new(reqwest::Client::new()),	2✔
532	abort_ch: None,	2✔
533	refcnt: 0,
534	gc: None,	2✔
535	}
536	}	2✔
537
538	pub fn is_empty(&self) -> bool {	2✔
539	self.entries.is_empty()	2✔
540	}	2✔
541
542	pub fn clear(&mut self) {	12✔
543	self.entries.clear();	12✔
544	}	12✔
545
546	pub fn get_client(&self) -> Arc<reqwest::Client> {	108✔
547	self.client.clone()	108✔
548	}	108✔
549
550	pub fn lookup_or_create(&mut self, key: &url::Url) -> Entry {	36✔
551	match self.entries.get(key) {	36✔
552	Some(v) => v.clone(),	16✔
553	None => self.create(key),	20✔
554	}
555	}	36✔
556
557	pub fn create(&mut self, key: &url::Url) -> Entry {	92✔
558	Entry::new(RwLock::new(EntryData::new(key)))	92✔
559	}	92✔
560
561	pub fn replace(&mut self, key: &url::Url, entry: &Entry) {	108✔
562	self.entries.insert(key.clone(), entry.clone());	108✔
563	}	108✔
564
565	pub fn remove(&mut self, key: &url::Url) {	×
566	self.entries.remove(key);	×
567	}	×
568
569	/// Removes the `num` oldest cache entries
570	pub fn gc_oldest(&mut self, mut num: usize) {	×
571	if num == 0 {	×
572	return;
573	}
574
575	let mut tmp = Vec::with_capacity(self.entries.len());	×
576
577	for (key, e) in &self.entries {	×
578	let entry = match e.try_read() {	×
579	Ok(e) => e,	×
580	_ => continue,
581	};	×
582
583	tmp.push((key.clone(), entry.get_cache_info().map(\|c\| c.local_time)));	×
584	}	×
585
586	tmp.sort_by(\|(_, tm_a), (_, tm_b)\| tm_a.cmp(tm_b));	×
587
588	let mut rm_cnt = 0;	×
589
590	for (key, _) in tmp {	×
591	if num == 0 {	×
592	break;
593	}
594
595	debug!("gc: removing old {}", key);	×
596	self.entries.remove(&key);	×
597	num -= 1;	×
598	rm_cnt += 1;	×
599	}	×
600
601	if rm_cnt > 0 {	×
602	info!("gc: removed {} old entries", rm_cnt);	×
603	}
604	}	×
605
606	/// Removes cache entries which are older than `max_lt`.
607	///
608	/// Returns the number of remaining cache entries.
609	pub fn gc_outdated(&mut self, max_lt: Duration) -> usize {	×
610	let now = Instant::now();	×
611	let mut outdated = Vec::new();	×
612	let mut cnt = 0;	×
613
614	for (key, e) in &self.entries {	×
615	match e.try_read().map(\|v\| v.is_outdated(now, max_lt)) {	×
616	Ok(true) => outdated.push(key.clone()),	×
617	_ => cnt += 1,	×
618	}
619	}
620
621	let rm_cnt = outdated.len();	×
622
623	for e in outdated {	×
624	debug!("gc: removing outdated {}", e);	×
625	self.entries.remove(&e);	×
626	}	×
627
628	if rm_cnt > 0 {	×
629	info!("gc: removed {} obsolete entries", rm_cnt);	×
630	}
631
632	cnt	×
633	}	×
634	}
635
636	#[derive(Debug)]	×
637	pub struct GcProperties {
638	pub max_elements: usize,
639	pub max_lifetime: Duration,	×
640	pub sleep: Duration,	×
641	}
642
643	async fn gc_runner(props: GcProperties, mut abort_ch: tokio::sync::watch::Receiver<()>) {	8✔
644	const RETRY_DELAY: std::time::Duration = std::time::Duration::from_secs(1);
645
646	loop {	2✔
647	use std::sync::TryLockError;
648
649	let sleep = {
650	let cache = CACHE.try_write();	2✔
651
652	match cache {	2✔
653	Ok(mut cache) if !cache.is_empty() => {	2✔
654	let cache_cnt = cache.gc_outdated(props.max_lifetime);	×
655
656	if cache_cnt > props.max_elements {	×
657	cache.gc_oldest(props.max_elements - cache_cnt)	×
658	}
659
660	props.sleep	×
661	}	×
662	Ok(_) => props.sleep,	2✔
663	Err(TryLockError::WouldBlock) => RETRY_DELAY,	×
664	Err(e) => {	×
665	error!("cache gc failed with {:?}", e);	×
666	break;
667	}	×
668	}
669	};	2✔
670
671	if tokio::time::timeout(sleep, abort_ch.changed()).await.is_ok() {	4✔
672	debug!("cache gc runner gracefully closed");	2✔
673	break;
674	}
675	}
676	}	4✔
677
678	pub struct Cache();
679
680	impl Cache {
681	#[instrument(level = "trace")]	14✔
682	pub fn instanciate(tmpdir: &std::path::Path, props: GcProperties) {	8✔
683	let mut cache = CACHE.write().unwrap();	2✔
684
685	if cache.refcnt == 0 {	2✔
686	let (tx, rx) = tokio::sync::watch::channel(());	2✔
687
688	cache.tmpdir = tmpdir.into();	2✔
689	cache.abort_ch = Some(tx);	2✔
690
691	cache.gc = Some(tokio::task::spawn(gc_runner(props, rx)));	2✔
692	}
693
694	cache.refcnt += 1;	2✔
695	}	2✔
696
697	#[instrument(level = "trace")]	14✔
698	// https://github.com/rust-lang/rust-clippy/issues/6446
699	#[allow(clippy::await_holding_lock)]
700	pub async fn close() {	2✔
701	let mut cache = CACHE.write().unwrap();	2✔
702
703	assert!(cache.refcnt > 0);	2✔
704
705	cache.refcnt -= 1;	2✔
706
707	if cache.refcnt == 0 {	2✔
708	cache.entries.clear();	2✔
709
710	let abort_ch = cache.abort_ch.take().unwrap();	2✔
711	let gc = cache.gc.take().unwrap();	2✔
712
713	drop(cache);	2✔
714
715	abort_ch.send(()).unwrap();	2✔
716	gc.await.unwrap();	4✔
717	}	2✔
718	}	2✔
719
720	#[instrument(level = "trace", ret)]	180✔
721	pub fn lookup_or_create(key: &url::Url) -> Entry {	36✔
722	let mut cache = CACHE.write().unwrap();	36✔
723
724	cache.lookup_or_create(key)	36✔
725	}	36✔
726
727	#[instrument(level = "trace", ret)]	360✔
728	pub fn create(key: &url::Url) -> Entry {	72✔
729	let mut cache = CACHE.write().unwrap();	72✔
730
731	cache.create(key)	72✔
732	}	72✔
733
734	#[instrument(level = "trace")]	756✔
735	pub fn replace(key: &url::Url, entry: &Entry) {	432✔
736	let mut cache = CACHE.write().unwrap();	108✔
737
738	cache.replace(key, entry)	108✔
739	}	108✔
740
741	#[instrument(level = "trace")]	×
742	pub fn remove(key: &url::Url) {	×
743	let mut cache = CACHE.write().unwrap();	×
744
745	cache.remove(key)	×
746	}	×
747
748	pub fn get_client() -> Arc<reqwest::Client> {	108✔
749	let cache = CACHE.read().unwrap();	108✔
750
751	cache.get_client()	108✔
752	}	108✔
753
754	pub fn new_file() -> Result<std::fs::File> {	108✔
755	let cache = CACHE.read().unwrap();	108✔
756
757	Ok(tempfile::tempfile_in(&cache.tmpdir)?)	108✔
758	}	108✔
759
760	pub async fn dump() {	336✔
761	let mut entries = Vec::new();	168✔
762
763	{
764	let cache = CACHE.read().unwrap();	168✔
765
766	entries.reserve(cache.entries.len());	168✔
767	entries.extend(cache.entries.values().cloned());	168✔
768	}	168✔
769
770	println!("Cache information ({} entries)", entries.len());	168✔
771
772	for e in entries {	568✔
773	println!("{}", e.read().await);	400✔
774	}	400✔
775	}	336✔
776
777	pub async fn clear() {	24✔
778	let mut cache = CACHE.write().unwrap();	12✔
779
780	cache.clear();	12✔
781	}	24✔
782	}

ensc / r-tftpd / 6956988158

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous