A fork of attic a self-hostable Nix Binary Cache server
client: Refactor pushing to use a job queue
+322
-197
+21
Cargo.lock
+21
Cargo.lock
···
77
77
checksum = "8da52d66c7071e2e3fa2a1e5c6d088fec47b593032b254f5e980de8ea54454d6"
78
78
79
79
[[package]]
80
+
name = "async-channel"
81
+
version = "1.8.0"
82
+
source = "registry+https://github.com/rust-lang/crates.io-index"
83
+
checksum = "cf46fee83e5ccffc220104713af3292ff9bc7c64c7de289f66dae8e38d826833"
84
+
dependencies = [
85
+
"concurrent-queue",
86
+
"event-listener",
87
+
"futures-core",
88
+
]
89
+
90
+
[[package]]
80
91
name = "async-compression"
81
92
version = "0.3.15"
82
93
source = "registry+https://github.com/rust-lang/crates.io-index"
···
170
181
version = "0.1.0"
171
182
dependencies = [
172
183
"anyhow",
184
+
"async-channel",
173
185
"attic",
174
186
"bytes",
175
187
"clap 4.0.32",
···
1031
1043
dependencies = [
1032
1044
"termcolor",
1033
1045
"unicode-width",
1046
+
]
1047
+
1048
+
[[package]]
1049
+
name = "concurrent-queue"
1050
+
version = "2.1.0"
1051
+
source = "registry+https://github.com/rust-lang/crates.io-index"
1052
+
checksum = "c278839b831783b70278b14df4d45e1beb1aad306c07bb796637de9a0e323e8e"
1053
+
dependencies = [
1054
+
"crossbeam-utils",
1034
1055
]
1035
1056
1036
1057
[[package]]
+1
client/Cargo.toml
+1
client/Cargo.toml
+14
-197
client/src/command/push.rs
+14
-197
client/src/command/push.rs
···
1
1
use std::collections::{HashMap, HashSet};
2
-
use std::fmt::Write;
2
+
use std::cmp;
3
3
use std::path::PathBuf;
4
-
use std::pin::Pin;
5
4
use std::sync::Arc;
6
-
use std::task::{Context, Poll};
7
-
use std::time::{Duration, Instant};
8
5
9
6
use anyhow::{anyhow, Result};
10
-
use bytes::Bytes;
11
7
use clap::Parser;
12
8
use futures::future::join_all;
13
-
use futures::stream::{Stream, TryStreamExt};
14
-
use indicatif::{HumanBytes, MultiProgress, ProgressBar, ProgressState, ProgressStyle};
15
-
use tokio::sync::Semaphore;
9
+
use indicatif::MultiProgress;
16
10
17
11
use crate::api::ApiClient;
18
12
use crate::cache::{CacheName, CacheRef};
19
13
use crate::cli::Opts;
20
14
use crate::config::Config;
21
-
use attic::api::v1::upload_path::{UploadPathNarInfo, UploadPathResult, UploadPathResultKind};
22
-
use attic::error::AtticResult;
15
+
use crate::push::{Pusher, PushConfig};
23
16
use attic::nix_store::{NixStore, StorePath, StorePathHash, ValidPathInfo};
24
17
25
18
/// Push closures to a binary cache.
···
62
55
num_upstream: usize,
63
56
}
64
57
65
-
/// Wrapper to update a progress bar as a NAR is streamed.
66
-
struct NarStreamProgress<S> {
67
-
stream: S,
68
-
bar: ProgressBar,
69
-
}
70
-
71
-
/// Uploads a single path to a cache.
72
-
pub async fn upload_path(
73
-
store: Arc<NixStore>,
74
-
path_info: ValidPathInfo,
75
-
api: ApiClient,
76
-
cache: &CacheName,
77
-
mp: MultiProgress,
78
-
force_preamble: bool,
79
-
) -> Result<()> {
80
-
let path = &path_info.path;
81
-
let upload_info = {
82
-
let full_path = store
83
-
.get_full_path(path)
84
-
.to_str()
85
-
.ok_or_else(|| anyhow!("Path contains non-UTF-8"))?
86
-
.to_string();
87
-
88
-
let references = path_info
89
-
.references
90
-
.into_iter()
91
-
.map(|pb| {
92
-
pb.to_str()
93
-
.ok_or_else(|| anyhow!("Reference contains non-UTF-8"))
94
-
.map(|s| s.to_owned())
95
-
})
96
-
.collect::<Result<Vec<String>, anyhow::Error>>()?;
97
-
98
-
UploadPathNarInfo {
99
-
cache: cache.to_owned(),
100
-
store_path_hash: path.to_hash(),
101
-
store_path: full_path,
102
-
references,
103
-
system: None, // TODO
104
-
deriver: None, // TODO
105
-
sigs: path_info.sigs,
106
-
ca: path_info.ca,
107
-
nar_hash: path_info.nar_hash.to_owned(),
108
-
nar_size: path_info.nar_size as usize,
109
-
}
110
-
};
111
-
112
-
let template = format!(
113
-
"{{spinner}} {: <20.20} {{bar:40.green/blue}} {{human_bytes:10}} ({{average_speed}})",
114
-
path.name(),
115
-
);
116
-
let style = ProgressStyle::with_template(&template)
117
-
.unwrap()
118
-
.tick_chars("🕛🕐🕑🕒🕓🕔🕕🕖🕗🕘🕙🕚✅")
119
-
.progress_chars("██ ")
120
-
.with_key("human_bytes", |state: &ProgressState, w: &mut dyn Write| {
121
-
write!(w, "{}", HumanBytes(state.pos())).unwrap();
122
-
})
123
-
// Adapted from
124
-
// <https://github.com/console-rs/indicatif/issues/394#issuecomment-1309971049>
125
-
.with_key(
126
-
"average_speed",
127
-
|state: &ProgressState, w: &mut dyn Write| match (state.pos(), state.elapsed()) {
128
-
(pos, elapsed) if elapsed > Duration::ZERO => {
129
-
write!(w, "{}", average_speed(pos, elapsed)).unwrap();
130
-
}
131
-
_ => write!(w, "-").unwrap(),
132
-
},
133
-
);
134
-
let bar = mp.add(ProgressBar::new(path_info.nar_size));
135
-
bar.set_style(style);
136
-
let nar_stream = NarStreamProgress::new(store.nar_from_path(path.to_owned()), bar.clone())
137
-
.map_ok(Bytes::from);
138
-
139
-
let start = Instant::now();
140
-
match api
141
-
.upload_path(upload_info, nar_stream, force_preamble)
142
-
.await
143
-
{
144
-
Ok(r) => {
145
-
let r = r.unwrap_or(UploadPathResult {
146
-
kind: UploadPathResultKind::Uploaded,
147
-
file_size: None,
148
-
frac_deduplicated: None,
149
-
});
150
-
151
-
let info_string: String = match r.kind {
152
-
UploadPathResultKind::Deduplicated => "deduplicated".to_string(),
153
-
_ => {
154
-
let elapsed = start.elapsed();
155
-
let seconds = elapsed.as_secs_f64();
156
-
let speed = (path_info.nar_size as f64 / seconds) as u64;
157
-
158
-
let mut s = format!("{}/s", HumanBytes(speed));
159
-
160
-
if let Some(frac_deduplicated) = r.frac_deduplicated {
161
-
if frac_deduplicated > 0.01f64 {
162
-
s += &format!(", {:.1}% deduplicated", frac_deduplicated * 100.0);
163
-
}
164
-
}
165
-
166
-
s
167
-
}
168
-
};
169
-
170
-
mp.suspend(|| {
171
-
eprintln!(
172
-
"✅ {} ({})",
173
-
path.as_os_str().to_string_lossy(),
174
-
info_string
175
-
);
176
-
});
177
-
bar.finish_and_clear();
178
-
179
-
Ok(())
180
-
}
181
-
Err(e) => {
182
-
mp.suspend(|| {
183
-
eprintln!("❌ {}: {}", path.as_os_str().to_string_lossy(), e);
184
-
});
185
-
bar.finish_and_clear();
186
-
Err(e)
187
-
}
188
-
}
189
-
}
190
-
191
58
pub async fn run(opts: Opts) -> Result<()> {
192
59
let sub = opts.command.as_push().unwrap();
193
60
if sub.jobs == 0 {
···
239
106
);
240
107
}
241
108
109
+
let push_config = PushConfig {
110
+
num_workers: cmp::min(sub.jobs, plan.store_path_map.len()),
111
+
force_preamble: sub.force_preamble,
112
+
};
113
+
242
114
let mp = MultiProgress::new();
243
-
let upload_limit = Arc::new(Semaphore::new(sub.jobs));
244
-
let futures = plan
245
-
.store_path_map
246
-
.into_iter()
247
-
.map(|(_, path_info)| {
248
-
let store = store.clone();
249
-
let api = api.clone();
250
-
let mp = mp.clone();
251
-
let upload_limit = upload_limit.clone();
252
115
253
-
async move {
254
-
let permit = upload_limit.acquire().await?;
116
+
let pusher = Pusher::new(store, api, cache.to_owned(), mp, push_config);
117
+
for (_, path_info) in plan.store_path_map {
118
+
pusher.push(path_info).await?;
119
+
}
255
120
256
-
upload_path(
257
-
store.clone(),
258
-
path_info,
259
-
api,
260
-
cache,
261
-
mp.clone(),
262
-
sub.force_preamble,
263
-
)
264
-
.await?;
265
-
266
-
drop(permit);
267
-
Ok::<(), anyhow::Error>(())
268
-
}
269
-
})
270
-
.collect::<Vec<_>>();
271
-
272
-
futures::future::join_all(futures)
273
-
.await
274
-
.into_iter()
275
-
.collect::<Result<Vec<()>>>()?;
121
+
let results = pusher.wait().await;
122
+
results.into_iter().map(|(_, result)| result).collect::<Result<Vec<()>>>()?;
276
123
277
124
Ok(())
278
125
}
···
376
223
})
377
224
}
378
225
}
379
-
380
-
impl<S: Stream<Item = AtticResult<Vec<u8>>>> NarStreamProgress<S> {
381
-
fn new(stream: S, bar: ProgressBar) -> Self {
382
-
Self { stream, bar }
383
-
}
384
-
}
385
-
386
-
impl<S: Stream<Item = AtticResult<Vec<u8>>> + Unpin> Stream for NarStreamProgress<S> {
387
-
type Item = AtticResult<Vec<u8>>;
388
-
389
-
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
390
-
match Pin::new(&mut self.stream).as_mut().poll_next(cx) {
391
-
Poll::Ready(Some(data)) => {
392
-
if let Ok(data) = &data {
393
-
self.bar.inc(data.len() as u64);
394
-
}
395
-
396
-
Poll::Ready(Some(data))
397
-
}
398
-
other => other,
399
-
}
400
-
}
401
-
}
402
-
403
-
// Just the average, no fancy sliding windows that cause wild fluctuations
404
-
// <https://github.com/console-rs/indicatif/issues/394>
405
-
fn average_speed(bytes: u64, duration: Duration) -> String {
406
-
let speed = bytes as f64 * 1000_f64 / duration.as_millis() as f64;
407
-
format!("{}/s", HumanBytes(speed as u64))
408
-
}
+1
client/src/main.rs
+1
client/src/main.rs
+285
client/src/push.rs
+285
client/src/push.rs
···
1
+
//! Store path uploader.
2
+
//!
3
+
//! Multiple workers are spawned to upload store paths concurrently.
4
+
//!
5
+
//! TODO: Refactor out progress reporting and support a simple output style without progress bars
6
+
7
+
use std::collections::HashMap;
8
+
use std::fmt::Write;
9
+
use std::pin::Pin;
10
+
use std::sync::Arc;
11
+
use std::task::{Context, Poll};
12
+
use std::time::{Duration, Instant};
13
+
14
+
use anyhow::{anyhow, Result};
15
+
use async_channel as channel;
16
+
use bytes::Bytes;
17
+
use futures::stream::{Stream, TryStreamExt};
18
+
use futures::future::join_all;
19
+
use indicatif::{HumanBytes, MultiProgress, ProgressBar, ProgressState, ProgressStyle};
20
+
use tokio::task::{JoinHandle, spawn};
21
+
22
+
use attic::api::v1::upload_path::{UploadPathNarInfo, UploadPathResult, UploadPathResultKind};
23
+
use attic::cache::CacheName;
24
+
use attic::error::AtticResult;
25
+
use attic::nix_store::{NixStore, StorePath, ValidPathInfo};
26
+
use crate::api::ApiClient;
27
+
28
+
type JobSender = channel::Sender<ValidPathInfo>;
29
+
type JobReceiver = channel::Receiver<ValidPathInfo>;
30
+
31
+
/// Configuration for pushing store paths.
32
+
#[derive(Clone, Copy, Debug)]
33
+
pub struct PushConfig {
34
+
/// The number of workers to spawn.
35
+
pub num_workers: usize,
36
+
37
+
/// Whether to always include the upload info in the PUT payload.
38
+
pub force_preamble: bool,
39
+
}
40
+
41
+
/// A handle to push store paths to a cache.
42
+
///
43
+
/// The caller is responsible for computing closures and
44
+
/// checking for paths that already exist on the remote
45
+
/// cache.
46
+
pub struct Pusher {
47
+
workers: Vec<JoinHandle<HashMap<StorePath, Result<()>>>>,
48
+
sender: JobSender,
49
+
}
50
+
51
+
/// Wrapper to update a progress bar as a NAR is streamed.
52
+
struct NarStreamProgress<S> {
53
+
stream: S,
54
+
bar: ProgressBar,
55
+
}
56
+
57
+
impl Pusher {
58
+
pub fn new(store: Arc<NixStore>, api: ApiClient, cache: CacheName, mp: MultiProgress, config: PushConfig) -> Self {
59
+
let (sender, receiver) = channel::unbounded();
60
+
let mut workers = Vec::new();
61
+
62
+
for _ in 0..config.num_workers {
63
+
workers.push(spawn(worker(
64
+
receiver.clone(),
65
+
store.clone(),
66
+
api.clone(),
67
+
cache.clone(),
68
+
mp.clone(),
69
+
config.clone(),
70
+
)));
71
+
}
72
+
73
+
Self { workers, sender }
74
+
}
75
+
76
+
/// Sends a path to be pushed.
77
+
pub async fn push(&self, path_info: ValidPathInfo) -> Result<()> {
78
+
self.sender.send(path_info).await
79
+
.map_err(|e| anyhow!(e))
80
+
}
81
+
82
+
/// Waits for all workers to terminate, returning all results.
83
+
///
84
+
/// TODO: Stream the results with another channel
85
+
pub async fn wait(self) -> HashMap<StorePath, Result<()>> {
86
+
drop(self.sender);
87
+
88
+
let results = join_all(self.workers)
89
+
.await
90
+
.into_iter()
91
+
.map(|joinresult| joinresult.unwrap())
92
+
.fold(HashMap::new(), |mut acc, results| {
93
+
acc.extend(results);
94
+
acc
95
+
});
96
+
97
+
results
98
+
}
99
+
}
100
+
101
+
async fn worker(
102
+
receiver: JobReceiver,
103
+
store: Arc<NixStore>,
104
+
api: ApiClient,
105
+
cache: CacheName,
106
+
mp: MultiProgress,
107
+
config: PushConfig,
108
+
) -> HashMap<StorePath, Result<()>> {
109
+
let mut results = HashMap::new();
110
+
111
+
loop {
112
+
let path_info = match receiver.recv().await {
113
+
Ok(path_info) => path_info,
114
+
Err(_) => {
115
+
// channel is closed - we are done
116
+
break;
117
+
}
118
+
};
119
+
120
+
let store_path = path_info.path.clone();
121
+
122
+
let r = upload_path(
123
+
path_info,
124
+
store.clone(),
125
+
api.clone(),
126
+
&cache,
127
+
mp.clone(),
128
+
config.force_preamble,
129
+
).await;
130
+
131
+
results.insert(store_path, r);
132
+
}
133
+
134
+
results
135
+
}
136
+
137
+
/// Uploads a single path to a cache.
138
+
pub async fn upload_path(
139
+
path_info: ValidPathInfo,
140
+
store: Arc<NixStore>,
141
+
api: ApiClient,
142
+
cache: &CacheName,
143
+
mp: MultiProgress,
144
+
force_preamble: bool,
145
+
) -> Result<()> {
146
+
let path = &path_info.path;
147
+
let upload_info = {
148
+
let full_path = store
149
+
.get_full_path(path)
150
+
.to_str()
151
+
.ok_or_else(|| anyhow!("Path contains non-UTF-8"))?
152
+
.to_string();
153
+
154
+
let references = path_info
155
+
.references
156
+
.into_iter()
157
+
.map(|pb| {
158
+
pb.to_str()
159
+
.ok_or_else(|| anyhow!("Reference contains non-UTF-8"))
160
+
.map(|s| s.to_owned())
161
+
})
162
+
.collect::<Result<Vec<String>, anyhow::Error>>()?;
163
+
164
+
UploadPathNarInfo {
165
+
cache: cache.to_owned(),
166
+
store_path_hash: path.to_hash(),
167
+
store_path: full_path,
168
+
references,
169
+
system: None, // TODO
170
+
deriver: None, // TODO
171
+
sigs: path_info.sigs,
172
+
ca: path_info.ca,
173
+
nar_hash: path_info.nar_hash.to_owned(),
174
+
nar_size: path_info.nar_size as usize,
175
+
}
176
+
};
177
+
178
+
let template = format!(
179
+
"{{spinner}} {: <20.20} {{bar:40.green/blue}} {{human_bytes:10}} ({{average_speed}})",
180
+
path.name(),
181
+
);
182
+
let style = ProgressStyle::with_template(&template)
183
+
.unwrap()
184
+
.tick_chars("🕛🕐🕑🕒🕓🕔🕕🕖🕗🕘🕙🕚✅")
185
+
.progress_chars("██ ")
186
+
.with_key("human_bytes", |state: &ProgressState, w: &mut dyn Write| {
187
+
write!(w, "{}", HumanBytes(state.pos())).unwrap();
188
+
})
189
+
// Adapted from
190
+
// <https://github.com/console-rs/indicatif/issues/394#issuecomment-1309971049>
191
+
.with_key(
192
+
"average_speed",
193
+
|state: &ProgressState, w: &mut dyn Write| match (state.pos(), state.elapsed()) {
194
+
(pos, elapsed) if elapsed > Duration::ZERO => {
195
+
write!(w, "{}", average_speed(pos, elapsed)).unwrap();
196
+
}
197
+
_ => write!(w, "-").unwrap(),
198
+
},
199
+
);
200
+
let bar = mp.add(ProgressBar::new(path_info.nar_size));
201
+
bar.set_style(style);
202
+
let nar_stream = NarStreamProgress::new(store.nar_from_path(path.to_owned()), bar.clone())
203
+
.map_ok(Bytes::from);
204
+
205
+
let start = Instant::now();
206
+
match api
207
+
.upload_path(upload_info, nar_stream, force_preamble)
208
+
.await
209
+
{
210
+
Ok(r) => {
211
+
let r = r.unwrap_or(UploadPathResult {
212
+
kind: UploadPathResultKind::Uploaded,
213
+
file_size: None,
214
+
frac_deduplicated: None,
215
+
});
216
+
217
+
let info_string: String = match r.kind {
218
+
UploadPathResultKind::Deduplicated => "deduplicated".to_string(),
219
+
_ => {
220
+
let elapsed = start.elapsed();
221
+
let seconds = elapsed.as_secs_f64();
222
+
let speed = (path_info.nar_size as f64 / seconds) as u64;
223
+
224
+
let mut s = format!("{}/s", HumanBytes(speed));
225
+
226
+
if let Some(frac_deduplicated) = r.frac_deduplicated {
227
+
if frac_deduplicated > 0.01f64 {
228
+
s += &format!(", {:.1}% deduplicated", frac_deduplicated * 100.0);
229
+
}
230
+
}
231
+
232
+
s
233
+
}
234
+
};
235
+
236
+
mp.suspend(|| {
237
+
eprintln!(
238
+
"✅ {} ({})",
239
+
path.as_os_str().to_string_lossy(),
240
+
info_string
241
+
);
242
+
});
243
+
bar.finish_and_clear();
244
+
245
+
Ok(())
246
+
}
247
+
Err(e) => {
248
+
mp.suspend(|| {
249
+
eprintln!("❌ {}: {}", path.as_os_str().to_string_lossy(), e);
250
+
});
251
+
bar.finish_and_clear();
252
+
Err(e)
253
+
}
254
+
}
255
+
}
256
+
257
+
impl<S: Stream<Item = AtticResult<Vec<u8>>>> NarStreamProgress<S> {
258
+
fn new(stream: S, bar: ProgressBar) -> Self {
259
+
Self { stream, bar }
260
+
}
261
+
}
262
+
263
+
impl<S: Stream<Item = AtticResult<Vec<u8>>> + Unpin> Stream for NarStreamProgress<S> {
264
+
type Item = AtticResult<Vec<u8>>;
265
+
266
+
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
267
+
match Pin::new(&mut self.stream).as_mut().poll_next(cx) {
268
+
Poll::Ready(Some(data)) => {
269
+
if let Ok(data) = &data {
270
+
self.bar.inc(data.len() as u64);
271
+
}
272
+
273
+
Poll::Ready(Some(data))
274
+
}
275
+
other => other,
276
+
}
277
+
}
278
+
}
279
+
280
+
// Just the average, no fancy sliding windows that cause wild fluctuations
281
+
// <https://github.com/console-rs/indicatif/issues/394>
282
+
fn average_speed(bytes: u64, duration: Duration) -> String {
283
+
let speed = bytes as f64 * 1000_f64 / duration.as_millis() as f64;
284
+
format!("{}/s", HumanBytes(speed as u64))
285
+
}