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
|
use super::PathInfoService;
use crate::proto::PathInfo;
use async_stream::try_stream;
use data_encoding::BASE64;
use futures::stream::BoxStream;
use prost::Message;
use std::path::Path;
use tonic::async_trait;
use tracing::instrument;
use tracing::warn;
use tvix_castore::Error;
/// SledPathInfoService stores PathInfo in a [sled](https://github.com/spacejam/sled).
///
/// The PathInfo messages are stored as encoded protos, and keyed by their output hash,
/// as that's currently the only request type available.
pub struct SledPathInfoService {
db: sled::Db,
}
impl SledPathInfoService {
pub fn new<P: AsRef<Path>>(p: P) -> Result<Self, sled::Error> {
let config = sled::Config::default()
.use_compression(false) // is a required parameter
.path(p);
let db = config.open()?;
Ok(Self { db })
}
pub fn new_temporary() -> Result<Self, sled::Error> {
let config = sled::Config::default().temporary(true);
let db = config.open()?;
Ok(Self { db })
}
}
#[async_trait]
impl PathInfoService for SledPathInfoService {
#[instrument(level = "trace", skip_all, fields(path_info.digest = BASE64.encode(&digest)))]
async fn get(&self, digest: [u8; 20]) -> Result<Option<PathInfo>, Error> {
let resp = tokio::task::spawn_blocking({
let db = self.db.clone();
move || db.get(digest.as_slice())
})
.await?
.map_err(|e| {
warn!("failed to retrieve PathInfo: {}", e);
Error::StorageError(format!("failed to retrieve PathInfo: {}", e))
})?;
match resp {
None => Ok(None),
Some(data) => {
let path_info = PathInfo::decode(&*data).map_err(|e| {
warn!("failed to decode stored PathInfo: {}", e);
Error::StorageError(format!("failed to decode stored PathInfo: {}", e))
})?;
Ok(Some(path_info))
}
}
}
#[instrument(level = "trace", skip_all, fields(path_info.root_node = ?path_info.node))]
async fn put(&self, path_info: PathInfo) -> Result<PathInfo, Error> {
// Call validate on the received PathInfo message.
let store_path = path_info
.validate()
.map_err(|e| Error::InvalidRequest(format!("failed to validate PathInfo: {}", e)))?;
// In case the PathInfo is valid, we were able to parse a StorePath.
// Store it in the database, keyed by its digest.
// This overwrites existing PathInfo objects.
tokio::task::spawn_blocking({
let db = self.db.clone();
let k = *store_path.digest();
let data = path_info.encode_to_vec();
move || db.insert(k, data)
})
.await?
.map_err(|e| {
warn!("failed to insert PathInfo: {}", e);
Error::StorageError(format! {
"failed to insert PathInfo: {}", e
})
})?;
Ok(path_info)
}
fn list(&self) -> BoxStream<'static, Result<PathInfo, Error>> {
let db = self.db.clone();
let mut it = db.iter().values();
Box::pin(try_stream! {
// Don't block the executor while waiting for .next(), so wrap that
// in a spawn_blocking call.
// We need to pass around it to be able to reuse it.
while let (Some(elem), new_it) = tokio::task::spawn_blocking(move || {
(it.next(), it)
}).await? {
it = new_it;
let data = elem.map_err(|e| {
warn!("failed to retrieve PathInfo: {}", e);
Error::StorageError(format!("failed to retrieve PathInfo: {}", e))
})?;
let path_info = PathInfo::decode(&*data).map_err(|e| {
warn!("failed to decode stored PathInfo: {}", e);
Error::StorageError(format!("failed to decode stored PathInfo: {}", e))
})?;
yield path_info
}
})
}
}
|
