WIP: libpq-based client

depends on https://github.com/neondatabase/rust-postgres/pull/25

pageserver/src/bin/getpage_bench_libpq.rs

@@ -0,0 +1,306 @@
+use clap::Parser;
+use futures::{SinkExt, TryStreamExt};
+use hyper::client::conn::Parts;
+use hyper::client::HttpConnector;
+use hyper::{Client, Uri};
+use pageserver::page_cache::PAGE_SZ;
+use pageserver::pgdatadir_mapping::{is_rel_block_key, key_to_rel_block};
+use pageserver::{repository, tenant};
+use pageserver_api::models::{
+    PagestreamBeMessage, PagestreamFeMessage, PagestreamGetPageRequest, PagestreamGetPageResponse,
+};
+use pageserver_api::reltag::RelTag;
+use rand::prelude::*;
+use scopeguard::defer;
+use std::env::args;
+use std::future::Future;
+use std::str::FromStr;
+use std::sync::{Arc, Mutex};
+use std::thread;
+use tokio::sync::mpsc::{channel, Sender};
+use tokio::sync::Mutex as AsyncMutex;
+use tokio::task::JoinHandle;
+use tokio_stream::{Stream, StreamExt};
+use utils::completion;
+use utils::lsn::Lsn;
+
+struct Key(repository::Key);
+
+impl std::str::FromStr for Key {
+    type Err = anyhow::Error;
+
+    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
+        repository::Key::from_hex(s).map(Key)
+    }
+}
+
+struct KeyRange {
+    start: Key,
+    end: Key,
+}
+
+impl KeyRange {
+    fn len(&self) -> i128 {
+        self.end.0.to_i128() - self.start.0.to_i128()
+    }
+}
+
+struct RelTagBlockNo {
+    rel_tag: RelTag,
+    block_no: u32,
+}
+
+#[derive(clap::Parser)]
+struct Args {
+    #[clap(long, default_value = "http://localhost:9898")]
+    ps_endpoint: String,
+    // tenant_id: String,
+    // timeline_id: String,
+    num_tasks: usize,
+    num_requests: usize,
+    tenants: Option<Vec<String>>,
+    #[clap(long)]
+    pick_n_tenants: Option<usize>,
+}
+
+#[tokio::main]
+async fn main() {
+    let args: &'static Args = Box::leak(Box::new(Args::parse()));
+
+    // std::env::set_var("RUST_LOG", "info,tokio_postgres=trace");
+    // tracing_subscriber::fmt::init();
+
+    let client = Client::new();
+
+    let tenants = if let Some(tenants) = &args.tenants {
+        tenants.clone()
+    } else {
+        // let tenant_id = "b97965931096047b2d54958756baee7b";
+        // let timeline_id = "2868f84a8d166779e4c651b116c45059";
+
+        let resp = client
+            .get(Uri::try_from(&format!("{}/v1/tenant", args.ps_endpoint)).unwrap())
+            .await
+            .unwrap();
+
+        let body = hyper::body::to_bytes(resp).await.unwrap();
+        let tenants: serde_json::Value = serde_json::from_slice(&body).unwrap();
+        let mut out = Vec::new();
+        for t in tenants.as_array().unwrap() {
+            if let Some(limit) = args.pick_n_tenants {
+                if out.len() >= limit {
+                    break;
+                }
+            }
+            out.push(t.get("id").unwrap().as_str().unwrap().to_owned());
+        }
+        if let Some(limit) = args.pick_n_tenants {
+            assert_eq!(out.len(), limit);
+        }
+        out
+    };
+
+    let mut tenant_timelines = Vec::new();
+    for tenant_id in tenants {
+        let resp = client
+            .get(
+                Uri::try_from(&format!(
+                    "{}/v1/tenant/{}/timeline",
+                    args.ps_endpoint, tenant_id
+                ))
+                .unwrap(),
+            )
+            .await
+            .unwrap();
+
+        let body = hyper::body::to_bytes(resp).await.unwrap();
+        let timelines: serde_json::Value = serde_json::from_slice(&body).unwrap();
+        for t in timelines.as_array().unwrap() {
+            let timeline_id = t.get("timeline_id").unwrap().as_str().unwrap().to_owned();
+            tenant_timelines.push((tenant_id.clone(), timeline_id));
+        }
+    }
+    println!("tenant_timelines:\n{:?}", tenant_timelines);
+
+    let mut tasks = Vec::new();
+    for (tenant_id, timeline_id) in tenant_timelines {
+        let t = tokio::spawn(timeline(args, client.clone(), tenant_id, timeline_id));
+        tasks.push(t);
+    }
+
+    for t in tasks {
+        t.await.unwrap();
+    }
+}
+
+fn timeline(
+    args: &'static Args,
+    client: Client<HttpConnector, hyper::Body>,
+    tenant_id: String,
+    timeline_id: String,
+) -> impl Future<Output = ()> + Send + Sync {
+    async move {
+        let mut resp = client
+            .get(
+                Uri::try_from(&format!(
+                    "{}/v1/tenant/{}/timeline/{}/keyspace",
+                    args.ps_endpoint, tenant_id, timeline_id
+                ))
+                .unwrap(),
+            )
+            .await
+            .unwrap();
+        if !resp.status().is_success() {
+            panic!("Failed to get keyspace: {resp:?}");
+        }
+        let body = hyper::body::to_bytes(resp).await.unwrap();
+        let keyspace: serde_json::Value = serde_json::from_slice(&body).unwrap();
+        let lsn: Lsn = keyspace["at_lsn"].as_str().unwrap().parse().unwrap();
+
+        let ranges = keyspace["keys"]
+            .as_array()
+            .unwrap()
+            .iter()
+            .filter_map(|r| {
+                let r = r.as_array().unwrap();
+                assert_eq!(r.len(), 2);
+                let start = Key::from_str(r[0].as_str().unwrap()).unwrap();
+                let end = Key::from_str(r[1].as_str().unwrap()).unwrap();
+                // filter out non-relblock keys
+                match (is_rel_block_key(start.0), is_rel_block_key(end.0)) {
+                    (true, true) => Some(KeyRange { start, end }),
+                    (true, false) | (false, true) => {
+                        unimplemented!("split up range")
+                    }
+                    (false, false) => None,
+                }
+            })
+            .collect::<Vec<_>>();
+
+        // weighted ranges
+        let weights = ranges.iter().map(|r| r.len()).collect::<Vec<_>>();
+
+        let ranges = Arc::new(ranges);
+        let weights = Arc::new(weights);
+
+        let mut tasks = Vec::<JoinHandle<()>>::new();
+
+        let start = std::time::Instant::now();
+
+        for i in 0..args.num_tasks {
+            let ranges = ranges.clone();
+            let weights = weights.clone();
+            let client = client.clone();
+            let tenant_id = tenant_id.clone();
+            let timeline_id = timeline_id.clone();
+            let task = tokio::spawn(async move {
+                let mut client = getpage_client::Client::new(tenant_id, timeline_id)
+                    .await
+                    .unwrap();
+                for i in 0..args.num_requests {
+                    let key = {
+                        let mut rng = rand::thread_rng();
+                        let r = ranges.choose_weighted(&mut rng, |r| r.len()).unwrap();
+                        let key: i128 = rng.gen_range((r.start.0.to_i128()..r.end.0.to_i128()));
+                        let key = repository::Key::from_i128(key);
+                        // XXX filter these out when we iterate the keyspace
+                        assert!(
+                            is_rel_block_key(key),
+                            "we filter non-relblock keys out above"
+                        );
+                        let (rel_tag, block_no) = key_to_rel_block(key).expect("we just checked");
+                        RelTagBlockNo { rel_tag, block_no }
+                    };
+                    client.getpage(key, lsn).await.unwrap();
+                }
+            });
+            tasks.push(task);
+        }
+
+        for task in tasks {
+            task.await.unwrap();
+        }
+
+        let elapsed = start.elapsed();
+        println!(
+            "RPS: {:.0}",
+            (args.num_requests * args.num_tasks) as f64 / elapsed.as_secs_f64()
+        );
+    }
+}
+
+mod getpage_client {
+    use std::pin::Pin;
+
+    use futures::SinkExt;
+    use pageserver_api::models::{
+        PagestreamBeMessage, PagestreamFeMessage, PagestreamGetPageRequest,
+        PagestreamGetPageResponse,
+    };
+    use tokio::task::JoinHandle;
+    use tokio_stream::StreamExt;
+    use utils::lsn::Lsn;
+
+    use crate::RelTagBlockNo;
+
+    pub(crate) struct Client {
+        copy_both: Pin<Box<tokio_postgres::CopyBothDuplex<bytes::Bytes>>>,
+        conn_task: JoinHandle<()>,
+    }
+
+    impl Client {
+        pub fn new(
+            tenant_id: String,
+            timeline_id: String,
+        ) -> impl std::future::Future<Output = anyhow::Result<Self>> + Send {
+            async move {
+                let (client, connection) =
+                    tokio_postgres::connect("host=localhost port=64000", postgres::NoTls).await?;
+
+                let conn_task = tokio::spawn(async move {
+                    let res = connection.await;
+                    res.unwrap();
+                });
+
+                let copy_both: tokio_postgres::CopyBothDuplex<bytes::Bytes> = client
+                    .copy_both_simple(&format!("pagestream {tenant_id} {timeline_id}"))
+                    .await?;
+
+                Ok(Self {
+                    copy_both: Box::pin(copy_both),
+                    conn_task,
+                })
+            }
+        }
+
+        pub async fn getpage(
+            &mut self,
+            key: RelTagBlockNo,
+            lsn: Lsn,
+        ) -> anyhow::Result<PagestreamGetPageResponse> {
+            let req = PagestreamGetPageRequest {
+                latest: false,
+                rel: key.rel_tag,
+                blkno: key.block_no,
+                lsn,
+            };
+            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.copy_both.send_all(&mut req).await?;
+
+            let next: Option<Result<bytes::Bytes, _>> = self.copy_both.next().await;
+            let next = next.unwrap().unwrap();
+
+            match PagestreamBeMessage::deserialize(next)? {
+                PagestreamBeMessage::Exists(_) => todo!(),
+                PagestreamBeMessage::Nblocks(_) => todo!(),
+                PagestreamBeMessage::GetPage(p) => Ok(p),
+                PagestreamBeMessage::Error(_) => todo!(),
+                PagestreamBeMessage::DbSize(_) => todo!(),
+            }
+        }
+    }
+}

pageserver/src/pgdatadir_mapping.rs

@@ -1694,6 +1694,7 @@ const AUX_FILES_KEY: Key = Key {
 // Reverse mappings for a few Keys.
 // These are needed by WAL redo manager.
 
+/// Guaranteed to return `Ok()` if [[is_rel_block_key]] returns `true` for `key`.
 pub fn key_to_rel_block(key: Key) -> anyhow::Result<(RelTag, BlockNumber)> {
     Ok(match key.field1 {
         0x00 => (
@@ -1709,7 +1710,8 @@ pub fn key_to_rel_block(key: Key) -> anyhow::Result<(RelTag, BlockNumber)> {
     })
 }
 
-fn is_rel_block_key(key: Key) -> bool {
+/// See [[key_to_rel_block]].
+pub fn is_rel_block_key(key: Key) -> bool {
     key.field1 == 0x00 && key.field4 != 0
 }