6391894226

Committed 03 Oct 2023 10:25AM UTC coverage: 70.425% (+0.05%) from 70.376%

Build # 6391894226

Build Type

push

github

Committed by

ensc

Commit Message

version 0.3.1

Run Details

1724 of 2448 relevant lines covered (70.42%)

383.71 hits per line

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

56.16

/mod-proxy/src/cache.rs

use std::os::unix::prelude::AsRawFd;
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 };

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
    }
}

#[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 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: Time, 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 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: Time, 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, 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()                        => 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>,
    is_dirty:        bool,
    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()),
            is_dirty:        false,
            abort_ch:        None,
            refcnt:        0,
            gc:                None,
        }
    }

    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)        => {
                self.is_dirty = true;
                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.is_dirty = true;
        self.entries.insert(key.clone(), entry.clone());
    }

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

    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.localtm.mono)));
        }

        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);
        }
    }

    pub fn gc_outdated(&mut self, max_lt: Duration) -> usize {
        let mut outdated = Vec::new();
        let now = Time::now();
        let mut cnt = 0;

        for (key, e) in &self.entries {
            cnt += 1;

            let entry = match e.try_read() {
                Ok(e)        => e,
                _        => continue,
            };

            if entry.is_outdated(now, max_lt) {
                outdated.push(key.clone());
            }
        }

        let mut rm_cnt = 0;

        for e in outdated {
            debug!("gc: removing outdated {}", e);
            self.entries.remove(&e);
            cnt -= 1;
            rm_cnt += 1;
        }

        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<()>) {
    loop {
        use std::sync::TryLockError;

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

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

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

                    cache.is_dirty = false;

                    props.sleep
                }
                Ok(_)                                => props.sleep,
                Err(TryLockError::WouldBlock)        => std::time::Duration::from_secs(1),
                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)?)
    }
}

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

ensc / r-tftpd / 6391894226

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