From 97a01bdcb47cf8c1f1aa444667a90c60d5567203 Mon Sep 17 00:00:00 2001 From: Christian Schwarz Date: Tue, 17 Dec 2024 21:01:03 +0100 Subject: [PATCH] pagebench: option for queue depth --- pageserver/client/src/page_service.rs | 34 +++++++---- .../pagebench/src/cmd/getpage_latest_lsn.rs | 57 +++++++++++-------- 2 files changed, 58 insertions(+), 33 deletions(-) diff --git a/pageserver/client/src/page_service.rs b/pageserver/client/src/page_service.rs index f9507fc47a..38f137852d 100644 --- a/pageserver/client/src/page_service.rs +++ b/pageserver/client/src/page_service.rs @@ -1,6 +1,9 @@ use std::pin::Pin; -use futures::SinkExt; +use futures::{ + stream::{SplitSink, SplitStream}, + SinkExt, StreamExt, +}; use pageserver_api::{ models::{ PagestreamBeMessage, PagestreamFeMessage, PagestreamGetPageRequest, @@ -10,7 +13,6 @@ use pageserver_api::{ }; use tokio::task::JoinHandle; use tokio_postgres::CopyOutStream; -use tokio_stream::StreamExt; use tokio_util::sync::CancellationToken; use utils::{ id::{TenantId, TimelineId}, @@ -62,13 +64,15 @@ impl Client { .client .copy_both_simple(&format!("pagestream_v2 {tenant_id} {timeline_id}")) .await?; + let (sink, stream) = copy_both.split(); // TODO: actually support splitting of the CopyBothDuplex so the lock inside this split adaptor goes away. let Client { cancel_on_client_drop, conn_task, client: _, } = self; Ok(PagestreamClient { - copy_both: Box::pin(copy_both), + sink, + stream, conn_task, cancel_on_client_drop, }) @@ -97,7 +101,8 @@ impl Client { /// Create using [`Client::pagestream`]. pub struct PagestreamClient { - copy_both: Pin>>, + sink: SplitSink, bytes::Bytes>, + stream: SplitStream>, cancel_on_client_drop: Option, conn_task: JoinHandle<()>, } @@ -110,10 +115,11 @@ pub struct RelTagBlockNo { impl PagestreamClient { pub async fn shutdown(self) { let Self { - copy_both, + sink, + stream, cancel_on_client_drop: cancel_conn_task, conn_task, - } = self; + } = { self }; // The `copy_both` contains internal channel sender, the receiver of which is polled by `conn_task`. // When `conn_task` observes the sender has been dropped, it sends a `FeMessage::CopyFail` into the connection. // (see https://github.com/neondatabase/rust-postgres/blob/2005bf79573b8add5cf205b52a2b208e356cc8b0/tokio-postgres/src/copy_both.rs#L56). @@ -133,21 +139,29 @@ impl PagestreamClient { // => https://github.com/neondatabase/neon/issues/6390 let _ = cancel_conn_task.unwrap(); conn_task.await.unwrap(); - drop(copy_both); + + // TODO: figure out what to do about sink & stream } pub async fn getpage( &mut self, req: PagestreamGetPageRequest, ) -> anyhow::Result { + self.getpage_send(req).await?; + self.getpage_recv().await + } + + pub async fn getpage_send(&mut self, req: PagestreamGetPageRequest) -> anyhow::Result<()> { let req = PagestreamFeMessage::GetPage(req); let req: bytes::Bytes = req.serialize(); // let mut req = tokio_util::io::ReaderStream::new(&req); let mut req = tokio_stream::once(Ok(req)); + self.sink.send_all(&mut req).await?; + Ok(()) + } - self.copy_both.send_all(&mut req).await?; - - let next: Option> = self.copy_both.next().await; + pub async fn getpage_recv(&mut self) -> anyhow::Result { + let next: Option> = self.stream.next().await; let next: bytes::Bytes = next.unwrap()?; let msg = PagestreamBeMessage::deserialize(next)?; diff --git a/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs b/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs index b2df01714d..b0a5810d01 100644 --- a/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs +++ b/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs @@ -13,7 +13,7 @@ use rand::prelude::*; use tokio::task::JoinSet; use tracing::info; -use std::collections::HashSet; +use std::collections::{HashSet, VecDeque}; use std::future::Future; use std::num::NonZeroUsize; use std::pin::Pin; @@ -63,6 +63,10 @@ pub(crate) struct Args { #[clap(long)] set_io_mode: Option, + /// Queue depth generated in each client. + #[clap(long, default_value = "1")] + queue_depth: NonZeroUsize, + targets: Option>, } @@ -298,6 +302,7 @@ async fn main_impl( start_work_barrier.wait().await; let client_start = Instant::now(); let mut ticks_processed = 0; + let mut inflight = VecDeque::new(); while !cancel.is_cancelled() { // Detect if a request took longer than the RPS rate if let Some(period) = &rps_period { @@ -311,28 +316,34 @@ async fn main_impl( ticks_processed = periods_passed_until_now; } - let start = Instant::now(); - let req = { - let mut rng = rand::thread_rng(); - let r = &ranges[weights.sample(&mut rng)]; - let key: i128 = rng.gen_range(r.start..r.end); - let key = Key::from_i128(key); - assert!(key.is_rel_block_key()); - let (rel_tag, block_no) = key - .to_rel_block() - .expect("we filter non-rel-block keys out above"); - PagestreamGetPageRequest { - request_lsn: if rng.gen_bool(args.req_latest_probability) { - Lsn::MAX - } else { - r.timeline_lsn - }, - not_modified_since: r.timeline_lsn, - rel: rel_tag, - blkno: block_no, - } - }; - client.getpage(req).await.unwrap(); + while inflight.len() < args.queue_depth.get() { + let start = Instant::now(); + let req = { + let mut rng = rand::thread_rng(); + let r = &ranges[weights.sample(&mut rng)]; + let key: i128 = rng.gen_range(r.start..r.end); + let key = Key::from_i128(key); + assert!(key.is_rel_block_key()); + let (rel_tag, block_no) = key + .to_rel_block() + .expect("we filter non-rel-block keys out above"); + PagestreamGetPageRequest { + request_lsn: if rng.gen_bool(args.req_latest_probability) { + Lsn::MAX + } else { + r.timeline_lsn + }, + not_modified_since: r.timeline_lsn, + rel: rel_tag, + blkno: block_no, + } + }; + client.getpage_send(req).await.unwrap(); + inflight.push_back(start); + } + + let start = inflight.pop_front().unwrap(); + client.getpage_recv().await.unwrap(); let end = Instant::now(); live_stats.request_done(); ticks_processed += 1;