1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
|
use crate::error::Error;
use reqwest;
use rusqlite;
use rusqlite::types::ToSql;
use rusqlite::Connection;
use rusqlite::Error::SqliteFailure;
use rusqlite::ErrorCode::DatabaseLocked;
use serde;
use serde_json;
use std::path::Path;
use std::thread;
use std::time::Duration;
use thread_local::ThreadLocal;
#[derive(Debug)]
pub struct SimpleCache {
url: String,
conn: ThreadLocal<Result<Connection, rusqlite::Error>>,
pub cache_only: bool,
sem: tokio::sync::Semaphore,
}
impl SimpleCache {
pub fn new(db_path: impl AsRef<Path>) -> Result<Self, Error> {
Ok(Self {
url: format!("file:{}?cache=shared", db_path.as_ref().display()),
conn: ThreadLocal::new(),
cache_only: false,
sem: tokio::sync::Semaphore::new(32),
})
}
fn connect(&self) -> Result<Connection, rusqlite::Error> {
let conn = Connection::open(&self.url)?;
conn.execute_batch(
"
CREATE TABLE IF NOT EXISTS cache2 (key TEXT NOT NULL PRIMARY KEY, ver TEXT NOT NULL, data BLOB NOT NULL);
PRAGMA synchronous = 0;
PRAGMA JOURNAL_MODE = WAL;
PRAGMA read_uncommitted;",
)?;
Ok(conn)
}
#[inline]
fn with_connection<F, T>(&self, cb: F) -> Result<T, Error> where F: FnOnce(&Connection) -> Result<T, Error> {
let conn = self.conn.get_or(|| self.connect());
match conn {
Ok(conn) => cb(conn),
Err(err) => Err(Error::Other(err.to_string())),
}
}
pub async fn get_json<B>(&self, key: (&str, &str), url: impl AsRef<str>) -> Result<Option<B>, Error>
where B: for<'a> serde::Deserialize<'a> {
if let Some(data) = self.get_cached(key, url).await? {
match serde_json::from_slice(&data) {
Ok(res) => Ok(Some(res)),
Err(parse) => Err(Error::Parse(parse, data)),
}
} else {
Ok(None)
}
}
pub fn get(&self, key: (&str, &str)) -> Result<Option<Vec<u8>>, Error> {
Self::with_retries(|| self.get_inner(key))
}
fn with_retries<T>(mut cb: impl FnMut() -> Result<T, Error>) -> Result<T, Error> {
let mut retries = 5;
loop {
match cb() {
Err(Error::Db(SqliteFailure(ref e, _))) if retries > 0 && e.code == DatabaseLocked => {
eprintln!("Retrying: {}", e);
retries -= 1;
thread::sleep(Duration::from_secs(1));
},
err => return err,
}
}
}
fn get_inner(&self, key: (&str, &str)) -> Result<Option<Vec<u8>>, Error> {
self.with_connection(|conn| {
let mut q = conn.prepare_cached("SELECT data FROM cache2 WHERE key = ?1 AND ver = ?2")?;
let row: Result<Vec<u8>, _> = q.query_row(&[&key.0, &key.1], |r| r.get(0));
match row {
Ok(row) => Ok(Some(row)),
Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
Err(err) => Err(err)?,
}
})
}
pub async fn get_cached(&self, key: (&str, &str), url: impl AsRef<str>) -> Result<Option<Vec<u8>>, Error> {
Ok(if let Some(data) = self.get(key)? {
Some(data)
} else {
if self.cache_only {
None
} else {
let _s = self.sem.acquire().await;
let data = Self::fetch(url.as_ref()).await?;
self.set(key, &data)?;
Some(data)
}
})
}
pub fn delete(&self, key: (&str, &str)) -> Result<(), Error> {
self.with_connection(|conn| {
let mut q = conn.prepare_cached("DELETE FROM cache2 WHERE key = ?1")?;
q.execute(&[&key.0])?;
Ok(())
})
}
pub fn set(&self, key: (&str, &str), data: &[u8]) -> Result<(), Error> {
Self::with_retries(|| self.set_inner(key, data))
}
fn set_inner(&self, key: (&str, &str), data: &[u8]) -> Result<(), Error> {
self.with_connection(|conn| {
let mut q = conn.prepare_cached("INSERT OR REPLACE INTO cache2(key, ver, data) VALUES(?1, ?2, ?3)")?;
let arr: &[&dyn ToSql] = &[&key.0, &key.1, &data];
q.execute(arr)?;
Ok(())
})
}
pub(crate) async fn fetch(url: &str) -> Result<Vec<u8>, Error> {
println!("REQ {}", url);
let client = reqwest::Client::builder().build()?;
let res = client.get(url).header(reqwest::header::USER_AGENT, "crates.rs/1.0").send().await?;
if res.status() != reqwest::StatusCode::OK {
Err(res.status())?;
}
Ok(res.bytes().await?.to_vec())
}
}
|
