Line data Source code
1 : use std::f64;
2 : use std::num::NonZeroUsize;
3 : use std::sync::Arc;
4 : use std::sync::atomic::{AtomicU64, Ordering};
5 : use std::time::{Duration, Instant};
6 :
7 : use pageserver_api::models::HistoricLayerInfo;
8 : use pageserver_api::shard::TenantShardId;
9 : use pageserver_client::mgmt_api;
10 : use rand::seq::SliceRandom;
11 : use tokio::sync::{OwnedSemaphorePermit, mpsc};
12 : use tokio::task::JoinSet;
13 : use tokio_util::sync::CancellationToken;
14 : use tracing::{debug, info};
15 : use utils::id::{TenantTimelineId, TimelineId};
16 :
17 : /// Evict & on-demand download random layers.
18 : #[derive(clap::Parser)]
19 : pub(crate) struct Args {
20 : #[clap(long, default_value = "http://localhost:9898")]
21 0 : mgmt_api_endpoint: String,
22 : #[clap(long)]
23 : pageserver_jwt: Option<String>,
24 : #[clap(long)]
25 : runtime: Option<humantime::Duration>,
26 : #[clap(long, default_value = "1")]
27 0 : tasks_per_target: NonZeroUsize,
28 : #[clap(long, default_value = "1")]
29 0 : concurrency_per_target: NonZeroUsize,
30 : /// Probability for sending `latest=true` in the request (uniform distribution).
31 : #[clap(long)]
32 : limit_to_first_n_targets: Option<usize>,
33 : /// Before starting the benchmark, live-reconfigure the pageserver to use the given
34 : /// [`pageserver_api::models::virtual_file::IoEngineKind`].
35 : #[clap(long)]
36 : set_io_engine: Option<pageserver_api::models::virtual_file::IoEngineKind>,
37 0 : targets: Option<Vec<TenantTimelineId>>,
38 : }
39 :
40 0 : pub(crate) fn main(args: Args) -> anyhow::Result<()> {
41 0 : let rt = tokio::runtime::Builder::new_multi_thread()
42 0 : .enable_all()
43 0 : .build()?;
44 0 : let task = rt.spawn(main_impl(args));
45 0 : rt.block_on(task).unwrap().unwrap();
46 0 : Ok(())
47 0 : }
48 :
49 0 : #[derive(serde::Serialize)]
50 : struct Output {
51 : downloads_count: u64,
52 : downloads_bytes: u64,
53 : evictions_count: u64,
54 : timeline_restarts: u64,
55 : #[serde(with = "humantime_serde")]
56 : runtime: Duration,
57 : }
58 :
59 : #[derive(Debug, Default)]
60 : struct LiveStats {
61 : evictions_count: AtomicU64,
62 : downloads_count: AtomicU64,
63 : downloads_bytes: AtomicU64,
64 : timeline_restarts: AtomicU64,
65 : }
66 :
67 : impl LiveStats {
68 0 : fn eviction_done(&self) {
69 0 : self.evictions_count.fetch_add(1, Ordering::Relaxed);
70 0 : }
71 0 : fn download_done(&self, size: u64) {
72 0 : self.downloads_count.fetch_add(1, Ordering::Relaxed);
73 0 : self.downloads_bytes.fetch_add(size, Ordering::Relaxed);
74 0 : }
75 0 : fn timeline_restart_done(&self) {
76 0 : self.timeline_restarts.fetch_add(1, Ordering::Relaxed);
77 0 : }
78 : }
79 :
80 0 : async fn main_impl(args: Args) -> anyhow::Result<()> {
81 0 : let args: &'static Args = Box::leak(Box::new(args));
82 0 :
83 0 : let mgmt_api_client = Arc::new(pageserver_client::mgmt_api::Client::new(
84 0 : args.mgmt_api_endpoint.clone(),
85 0 : args.pageserver_jwt.as_deref(),
86 0 : ));
87 :
88 0 : if let Some(engine_str) = &args.set_io_engine {
89 0 : mgmt_api_client.put_io_engine(engine_str).await?;
90 0 : }
91 :
92 : // discover targets
93 0 : let timelines: Vec<TenantTimelineId> = crate::util::cli::targets::discover(
94 0 : &mgmt_api_client,
95 0 : crate::util::cli::targets::Spec {
96 0 : limit_to_first_n_targets: args.limit_to_first_n_targets,
97 0 : targets: args.targets.clone(),
98 0 : },
99 0 : )
100 0 : .await?;
101 :
102 0 : let token = CancellationToken::new();
103 0 : let mut tasks = JoinSet::new();
104 0 :
105 0 : let periodic_stats = Arc::new(LiveStats::default());
106 0 : let total_stats = Arc::new(LiveStats::default());
107 0 :
108 0 : let start = Instant::now();
109 0 : tasks.spawn({
110 0 : let periodic_stats = Arc::clone(&periodic_stats);
111 0 : let total_stats = Arc::clone(&total_stats);
112 0 : let cloned_token = token.clone();
113 0 : async move {
114 0 : let mut last_at = Instant::now();
115 : loop {
116 0 : if cloned_token.is_cancelled() {
117 0 : return;
118 0 : }
119 0 : tokio::time::sleep_until((last_at + Duration::from_secs(1)).into()).await;
120 0 : let now = Instant::now();
121 0 : let delta: Duration = now - last_at;
122 0 : last_at = now;
123 0 :
124 0 : let LiveStats {
125 0 : evictions_count,
126 0 : downloads_count,
127 0 : downloads_bytes,
128 0 : timeline_restarts,
129 0 : } = &*periodic_stats;
130 0 : let evictions_count = evictions_count.swap(0, Ordering::Relaxed);
131 0 : let downloads_count = downloads_count.swap(0, Ordering::Relaxed);
132 0 : let downloads_bytes = downloads_bytes.swap(0, Ordering::Relaxed);
133 0 : let timeline_restarts = timeline_restarts.swap(0, Ordering::Relaxed);
134 0 :
135 0 : total_stats.evictions_count.fetch_add(evictions_count, Ordering::Relaxed);
136 0 : total_stats.downloads_count.fetch_add(downloads_count, Ordering::Relaxed);
137 0 : total_stats.downloads_bytes.fetch_add(downloads_bytes, Ordering::Relaxed);
138 0 : total_stats.timeline_restarts.fetch_add(timeline_restarts, Ordering::Relaxed);
139 0 :
140 0 : let evictions_per_s = evictions_count as f64 / delta.as_secs_f64();
141 0 : let downloads_per_s = downloads_count as f64 / delta.as_secs_f64();
142 0 : let downloads_mibs_per_s = downloads_bytes as f64 / delta.as_secs_f64() / ((1 << 20) as f64);
143 0 :
144 0 : info!("evictions={evictions_per_s:.2}/s downloads={downloads_per_s:.2}/s download_bytes={downloads_mibs_per_s:.2}MiB/s timeline_restarts={timeline_restarts}");
145 : }
146 0 : }
147 0 : });
148 :
149 0 : for tl in timelines {
150 0 : for _ in 0..args.tasks_per_target.get() {
151 0 : tasks.spawn(timeline_actor(
152 0 : args,
153 0 : Arc::clone(&mgmt_api_client),
154 0 : tl,
155 0 : Arc::clone(&periodic_stats),
156 0 : token.clone(),
157 0 : ));
158 0 : }
159 : }
160 0 : if let Some(runtime) = args.runtime {
161 0 : tokio::spawn(async move {
162 0 : tokio::time::sleep(runtime.into()).await;
163 0 : token.cancel();
164 0 : });
165 0 : }
166 :
167 0 : while let Some(res) = tasks.join_next().await {
168 0 : res.unwrap();
169 0 : }
170 0 : let end = Instant::now();
171 0 : let duration: Duration = end - start;
172 0 :
173 0 : let output = {
174 0 : let LiveStats {
175 0 : evictions_count,
176 0 : downloads_count,
177 0 : downloads_bytes,
178 0 : timeline_restarts,
179 0 : } = &*total_stats;
180 0 : Output {
181 0 : downloads_count: downloads_count.load(Ordering::Relaxed),
182 0 : downloads_bytes: downloads_bytes.load(Ordering::Relaxed),
183 0 : evictions_count: evictions_count.load(Ordering::Relaxed),
184 0 : timeline_restarts: timeline_restarts.load(Ordering::Relaxed),
185 0 : runtime: duration,
186 0 : }
187 0 : };
188 0 : let output = serde_json::to_string_pretty(&output).unwrap();
189 0 : println!("{output}");
190 0 :
191 0 : Ok(())
192 0 : }
193 :
194 0 : async fn timeline_actor(
195 0 : args: &'static Args,
196 0 : mgmt_api_client: Arc<pageserver_client::mgmt_api::Client>,
197 0 : timeline: TenantTimelineId,
198 0 : live_stats: Arc<LiveStats>,
199 0 : token: CancellationToken,
200 0 : ) {
201 0 : // TODO: support sharding
202 0 : let tenant_shard_id = TenantShardId::unsharded(timeline.tenant_id);
203 :
204 : struct Timeline {
205 : joinset: JoinSet<()>,
206 : layers: Vec<mpsc::Sender<OwnedSemaphorePermit>>,
207 : concurrency: Arc<tokio::sync::Semaphore>,
208 : }
209 0 : while !token.is_cancelled() {
210 0 : debug!("restarting timeline");
211 0 : let layer_map_info = mgmt_api_client
212 0 : .layer_map_info(tenant_shard_id, timeline.timeline_id)
213 0 : .await
214 0 : .unwrap();
215 0 : let concurrency = Arc::new(tokio::sync::Semaphore::new(
216 0 : args.concurrency_per_target.get(),
217 0 : ));
218 0 :
219 0 : let mut joinset = JoinSet::new();
220 0 : let layers = layer_map_info
221 0 : .historic_layers
222 0 : .into_iter()
223 0 : .map(|historic_layer| {
224 0 : let (tx, rx) = mpsc::channel(1);
225 0 : joinset.spawn(layer_actor(
226 0 : tenant_shard_id,
227 0 : timeline.timeline_id,
228 0 : historic_layer,
229 0 : rx,
230 0 : Arc::clone(&mgmt_api_client),
231 0 : Arc::clone(&live_stats),
232 0 : ));
233 0 : tx
234 0 : })
235 0 : .collect::<Vec<_>>();
236 0 :
237 0 : let mut timeline = Timeline {
238 0 : joinset,
239 0 : layers,
240 0 : concurrency,
241 0 : };
242 0 :
243 0 : live_stats.timeline_restart_done();
244 :
245 0 : while !token.is_cancelled() {
246 0 : assert!(!timeline.joinset.is_empty());
247 0 : if let Some(res) = timeline.joinset.try_join_next() {
248 0 : debug!(?res, "a layer actor exited, should not happen");
249 0 : timeline.joinset.shutdown().await;
250 0 : break;
251 0 : }
252 :
253 0 : let mut permit = Some(
254 0 : Arc::clone(&timeline.concurrency)
255 0 : .acquire_owned()
256 0 : .await
257 0 : .unwrap(),
258 : );
259 :
260 : loop {
261 0 : let layer_tx = {
262 0 : let mut rng = rand::thread_rng();
263 0 : timeline.layers.choose_mut(&mut rng).expect("no layers")
264 0 : };
265 0 : match layer_tx.try_send(permit.take().unwrap()) {
266 0 : Ok(_) => break,
267 0 : Err(e) => match e {
268 0 : mpsc::error::TrySendError::Full(back) => {
269 0 : // TODO: retrying introduces bias away from slow downloaders
270 0 : permit.replace(back);
271 0 : }
272 0 : mpsc::error::TrySendError::Closed(_) => panic!(),
273 : },
274 : }
275 : }
276 : }
277 : }
278 0 : }
279 :
280 0 : async fn layer_actor(
281 0 : tenant_shard_id: TenantShardId,
282 0 : timeline_id: TimelineId,
283 0 : mut layer: HistoricLayerInfo,
284 0 : mut rx: mpsc::Receiver<tokio::sync::OwnedSemaphorePermit>,
285 0 : mgmt_api_client: Arc<mgmt_api::Client>,
286 0 : live_stats: Arc<LiveStats>,
287 0 : ) {
288 : #[derive(Clone, Copy)]
289 : enum Action {
290 : Evict,
291 : OnDemandDownload,
292 : }
293 :
294 0 : while let Some(_permit) = rx.recv().await {
295 0 : let action = if layer.is_remote() {
296 0 : Action::OnDemandDownload
297 : } else {
298 0 : Action::Evict
299 : };
300 :
301 0 : let did_it = match action {
302 : Action::Evict => {
303 0 : let did_it = mgmt_api_client
304 0 : .layer_evict(tenant_shard_id, timeline_id, layer.layer_file_name())
305 0 : .await
306 0 : .unwrap();
307 0 : live_stats.eviction_done();
308 0 : did_it
309 : }
310 : Action::OnDemandDownload => {
311 0 : let did_it = mgmt_api_client
312 0 : .layer_ondemand_download(tenant_shard_id, timeline_id, layer.layer_file_name())
313 0 : .await
314 0 : .unwrap();
315 0 : live_stats.download_done(layer.layer_file_size());
316 0 : did_it
317 : }
318 : };
319 0 : if !did_it {
320 0 : debug!("local copy of layer map appears out of sync, re-downloading");
321 0 : return;
322 0 : }
323 0 : debug!("did it");
324 0 : layer.set_remote(match action {
325 0 : Action::Evict => true,
326 0 : Action::OnDemandDownload => false,
327 : });
328 : }
329 0 : }
|
