From 44e7301cf48cebb9b6e8d5777bf1fd0db3ac74fa Mon Sep 17 00:00:00 2001 From: Caio Date: Sat, 10 Feb 2024 23:09:03 -0300 Subject: [PATCH] Update the `wtx` crate --- diesel_bench/Cargo.toml | 2 +- diesel_bench/benches/wtx.rs | 115 ++++++++++++++---------------------- 2 files changed, 46 insertions(+), 71 deletions(-) diff --git a/diesel_bench/Cargo.toml b/diesel_bench/Cargo.toml index f0593b304390..b9663883dd79 100644 --- a/diesel_bench/Cargo.toml +++ b/diesel_bench/Cargo.toml @@ -25,7 +25,7 @@ futures = {version = "0.3", optional = true} diesel-async = {version = "0.4", optional = true, default-features = false} criterion-perf-events = { version = "0.4", optional = true} perfcnt = {version = "0.8", optional = true} -wtx = { default-features = false, features = ["atoi", "postgres", "simdutf8", "std", "tokio"], optional = true, version = "0.12" } +wtx = { default-features = false, features = ["atoi", "memchr", "postgres", "simdutf8", "std", "tokio"], optional = true, version = "0.14" } [dependencies.diesel] path = "../diesel" diff --git a/diesel_bench/benches/wtx.rs b/diesel_bench/benches/wtx.rs index 6f04efbaf8e8..565bffebb3bf 100644 --- a/diesel_bench/benches/wtx.rs +++ b/diesel_bench/benches/wtx.rs @@ -1,13 +1,12 @@ use crate::Bencher; -use std::collections::HashMap; -use std::fmt::Write; +use std::{collections::HashMap, fmt::Write}; use tokio::{net::TcpStream, runtime::Runtime}; use wtx::{ database::{ client::postgres::{Config, Executor, ExecutorBuffer}, - Encode, Executor as _, Record as _, + Executor as _, Record as _, }, - misc::UriPartsRef, + misc::{Either, UriRef}, rng::StdRng, }; @@ -85,21 +84,15 @@ pub fn bench_loading_associations_sequentially(b: &mut Bencher) { let mut posts_query = String::from("SELECT id, title, user_id, body FROM posts WHERE user_id IN("); - let mut users_ids = Vec::with_capacity(LEN); - concat( - users.iter().enumerate(), - &mut posts_query, - |local_str, (idx, &User { id, .. })| { - local_str.write_fmt(format_args!("${}", idx + 1)).unwrap(); - users_ids.push(id); - }, - ); - posts_query += ")"; + concat(0..users.len(), &mut posts_query, |local_str, idx| { + local_str.write_fmt(format_args!("${}", idx + 1)).unwrap(); + }); + posts_query.push(')'); let mut posts = Vec::with_capacity(LEN); - conn.fetch_many_with_stmt::( + conn.fetch_many_with_stmt( posts_query.as_str(), - users_ids.as_slice(), + &mut users.iter().map(|user| user.id), |record| { posts.push(Post { body: record.decode_opt(3).unwrap(), @@ -115,21 +108,15 @@ pub fn bench_loading_associations_sequentially(b: &mut Bencher) { let mut comments_query = String::from("SELECT id, post_id, text FROM comments WHERE post_id IN("); - let mut posts_ids = Vec::with_capacity(LEN); - concat( - posts.iter().enumerate(), - &mut comments_query, - |local_str, (idx, &Post { id, .. })| { - local_str.write_fmt(format_args!("${}", idx + 1)).unwrap(); - posts_ids.push(id); - }, - ); - comments_query += ")"; + concat(0..posts.len(), &mut comments_query, |local_str, idx| { + local_str.write_fmt(format_args!("${}", idx + 1)).unwrap(); + }); + comments_query.push(')'); let mut comments = Vec::with_capacity(LEN); - conn.fetch_many_with_stmt::( + conn.fetch_many_with_stmt( comments_query.as_str(), - posts_ids.as_slice(), + &mut posts.iter().map(|post| post.id), |record| { comments.push(Comment { id: record.decode(0).unwrap(), @@ -197,7 +184,7 @@ pub fn bench_medium_complex_query(b: &mut Bencher, size: usize) { b.iter(|| { runtime.block_on(async { let mut _rslt = Vec::with_capacity(size); - conn.fetch_many_with_stmt::(stmt_hash, ("black",), |record| { + conn.fetch_many_with_stmt(stmt_hash, ("black",), |record| { let user = User { id: record.decode(0).unwrap(), name: record.decode(1).unwrap(), @@ -237,8 +224,8 @@ pub fn bench_trivial_query(b: &mut Bencher, size: usize) { (conn, stmt_hash) }); b.iter(|| { - let mut users = Vec::with_capacity(size); runtime.block_on(async { + let mut users = Vec::with_capacity(size); conn.fetch_many_with_stmt(stmt_hash, (), |record| { users.push(User { id: record.decode(0).unwrap(), @@ -266,34 +253,31 @@ where } } -async fn connection() -> Executor { +async fn connection() -> Executor { dotenvy::dotenv().ok(); let url = dotenvy::var("POSTGRES_DATABASE_URL") .or_else(|_| dotenvy::var("DATABASE_URL")) .expect("DATABASE_URL must be set in order to run tests"); - let up = UriPartsRef::new(url.as_str()); + let uri = UriRef::new(url.as_str()); let mut rng = StdRng::default(); let mut conn = Executor::connect( - &Config::from_uri_parts(&up).unwrap(), + &Config::from_uri(&uri).unwrap(), ExecutorBuffer::with_default_params(&mut rng), &mut rng, - TcpStream::connect(up.host()).await.unwrap(), + TcpStream::connect(uri.host()).await.unwrap(), + ) + .await + .unwrap(); + conn.execute( + "TRUNCATE TABLE comments CASCADE;TRUNCATE TABLE posts CASCADE;TRUNCATE TABLE users CASCADE", + |_| {}, ) .await .unwrap(); - conn.execute("TRUNCATE TABLE comments CASCADE", |_| {}) - .await - .unwrap(); - conn.execute("TRUNCATE TABLE posts CASCADE", |_| {}) - .await - .unwrap(); - conn.execute("TRUNCATE TABLE users CASCADE", |_| {}) - .await - .unwrap(); conn } -async fn insert_posts(conn: &mut Executor) { +async fn insert_posts(conn: &mut Executor) { let mut users_ids = Vec::with_capacity(N); conn.fetch_many_with_stmt("SELECT id FROM users", (), |record| { users_ids.push(record.decode(0).unwrap()); @@ -305,11 +289,17 @@ async fn insert_posts(conn: &mut Executor)), + ] + }) }) .collect::>(); - let mut insert_stmt = String::from("INSERT INTO posts(title, user_id, body) VALUES"); + let mut insert_stmt = String::from("INSERT INTO posts(title, user_id, body) VALUES "); concat( 0..params.len(), &mut insert_stmt, @@ -325,45 +315,30 @@ async fn insert_posts(conn: &mut Executor)| { - let a: &dyn Encode<_, _> = &el.0; - let b: &dyn Encode<_, _> = &el.1; - let c: &dyn Encode<_, _> = &el.2; - [a, b, c] - }) - .collect::>>(); - - conn.execute_with_stmt::(insert_stmt.as_str(), params_ref.as_slice()) + conn.execute_with_stmt(insert_stmt.as_str(), &mut params.into_iter().flatten()) .await .unwrap(); } async fn insert_users( - conn: &mut Executor, + conn: &mut Executor, hair_color_init: impl Fn(usize) -> Option<&'static str>, ) { let mut query = String::from("INSERT INTO users (name, hair_color) VALUES"); - let mut params = Vec::with_capacity(2 * N); - concat(0..N, &mut query, |local_query, idx| { local_query .write_fmt(format_args!("(${}, ${})", 2 * idx + 1, 2 * idx + 2)) .unwrap(); - params.push((format!("User {idx}"), hair_color_init(idx))); }); - let params_ref = params - .iter() - .flat_map(|el: &(String, Option<&'static str>)| { - let a: &dyn Encode<_, _> = &el.0; - let b: &dyn Encode<_, _> = &el.1; - [a, b] - }) - .collect::>(); + let mut params = (0..N).into_iter().flat_map(|idx| { + [ + Either::Left(format!("User {idx}")), + Either::Right(hair_color_init(idx)), + ] + }); - conn.execute_with_stmt::(query.as_str(), params_ref.as_slice()) + conn.execute_with_stmt(query.as_str(), &mut params) .await .unwrap(); }