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