From 37ebdba835c2ae8e745f66d39931b9a6bf98578b Mon Sep 17 00:00:00 2001 From: Virgiel Date: Sat, 21 Oct 2023 00:25:46 +0200 Subject: [PATCH] Optimising codegen using an asynchronous client and pipelining --- Cargo.lock | 61 ++++++++++++++---------- benches/codegen.rs | 2 +- benches/execution/main.rs | 22 +++------ crates/cornucopia/Cargo.toml | 6 ++- crates/cornucopia/src/cli.rs | 4 +- crates/cornucopia/src/conn.rs | 35 ++++++++++---- crates/cornucopia/src/lib.rs | 10 ++-- crates/cornucopia/src/load_schema.rs | 6 +-- crates/cornucopia/src/prepare_queries.rs | 40 +++++++++++----- crates/cornucopia/src/type_registrar.rs | 2 +- crates/cornucopia/src/utils.rs | 4 +- crates/cornucopia/src/validation.rs | 2 +- test_integration/Cargo.toml | 2 + test_integration/src/codegen.rs | 2 +- test_integration/src/errors.rs | 2 +- test_integration/src/main.rs | 6 +-- test_integration/src/utils.rs | 6 ++- 17 files changed, 129 insertions(+), 83 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d0fe568e..293f26c1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -442,12 +442,14 @@ dependencies = [ "chumsky", "clap", "codegen_template", + "futures", "heck", "indexmap", "miette", - "postgres", "postgres-types", "thiserror", + "tokio", + "tokio-postgres", ] [[package]] @@ -481,9 +483,9 @@ dependencies = [ [[package]] name = "cpufeatures" -version = "0.2.9" +version = "0.2.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a17b76ff3a4162b0b27f354a0c87015ddad39d35f9c0c36607a3bdd175dde1f1" +checksum = "3fbc60abd742b35f2492f808e1abbb83d45f72db402e14c55057edc9c7b1e9e4" dependencies = [ "libc", ] @@ -854,9 +856,9 @@ dependencies = [ [[package]] name = "hashbrown" -version = "0.14.1" +version = "0.14.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7dfda62a12f55daeae5015f81b0baea145391cb4520f86c248fc615d72640d12" +checksum = "f93e7192158dbcda357bdec5fb5788eebf8bbac027f3f33e719d29135ae84156" [[package]] name = "heck" @@ -886,7 +888,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8adf3ddd720272c6ea8bf59463c04e0f93d0bbf7c5439b691bca2987e0270897" dependencies = [ "equivalent", - "hashbrown 0.14.1", + "hashbrown 0.14.2", ] [[package]] @@ -944,9 +946,9 @@ checksum = "da2479e8c062e40bf0066ffa0bc823de0a9368974af99c9f6df941d2c231e03f" [[package]] name = "lock_api" -version = "0.4.10" +version = "0.4.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c1cc9717a20b1bb222f333e6a92fd32f7d8a18ddc5a3191a11af45dcbf4dcd16" +checksum = "3c168f8615b12bc01f9c17e2eb0cc07dcae1940121185446edc3744920e8ef45" dependencies = [ "autocfg", "scopeguard", @@ -1093,13 +1095,13 @@ dependencies = [ [[package]] name = "parking_lot_core" -version = "0.9.8" +version = "0.9.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "93f00c865fe7cabf650081affecd3871070f26767e7b2070a3ffae14c654b447" +checksum = "4c42a9226546d68acdd9c0a280d17ce19bfe27a46bf68784e4066115788d008e" dependencies = [ "cfg-if", "libc", - "redox_syscall", + "redox_syscall 0.4.1", "smallvec", "windows-targets", ] @@ -1371,6 +1373,15 @@ dependencies = [ "bitflags 1.3.2", ] +[[package]] +name = "redox_syscall" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4722d768eff46b75989dd134e5c353f0d6296e5aaa3132e776cbdb56be7731aa" +dependencies = [ + "bitflags 1.3.2", +] + [[package]] name = "regex" version = "1.10.2" @@ -1468,9 +1479,9 @@ checksum = "dcf128d1287d2ea9d80910b5f1120d0b8eede3fbf1abe91c40d39ea7d51e6fda" [[package]] name = "rustix" -version = "0.38.19" +version = "0.38.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "745ecfa778e66b2b63c88a61cb36e0eea109e803b0b86bf9879fbc77c70e86ed" +checksum = "67ce50cb2e16c2903e30d1cbccfd8387a74b9d4c938b6a4c5ec6cc7556f7a8a0" dependencies = [ "bitflags 2.4.1", "errno", @@ -1601,9 +1612,9 @@ checksum = "b7c388c1b5e93756d0c740965c41e8822f866621d41acbdf6336a6a168f8840c" [[package]] name = "socket2" -version = "0.5.4" +version = "0.5.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4031e820eb552adee9295814c0ced9e5cf38ddf1e8b7d566d6de8e2538ea989e" +checksum = "7b5fac59a5cb5dd637972e5fca70daf0523c9067fcdc4842f053dae04a18f8e9" dependencies = [ "libc", "windows-sys", @@ -1709,7 +1720,7 @@ checksum = "cb94d2f3cc536af71caac6b6fcebf65860b347e7ce0cc9ebe8f70d3e521054ef" dependencies = [ "cfg-if", "fastrand", - "redox_syscall", + "redox_syscall 0.3.5", "rustix", "windows-sys", ] @@ -1748,10 +1759,12 @@ version = "0.1.0" dependencies = [ "clap", "cornucopia", + "futures", "owo-colors", "postgres", "serde", "tempfile", + "tokio-postgres", "toml 0.8.2", ] @@ -1768,18 +1781,18 @@ dependencies = [ [[package]] name = "thiserror" -version = "1.0.49" +version = "1.0.50" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1177e8c6d7ede7afde3585fd2513e611227efd6481bd78d2e82ba1ce16557ed4" +checksum = "f9a7210f5c9a7156bb50aa36aed4c95afb51df0df00713949448cf9e97d382d2" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "1.0.49" +version = "1.0.50" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "10712f02019e9288794769fba95cd6847df9874d49d871d062172f9dd41bc4cc" +checksum = "266b2e40bc00e5a6c09c3584011e08b06f123c00362c92b975ba9843aaaa14b8" dependencies = [ "proc-macro2", "quote", @@ -1954,9 +1967,9 @@ dependencies = [ [[package]] name = "tracing" -version = "0.1.39" +version = "0.1.40" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ee2ef2af84856a50c1d430afce2fdded0a4ec7eda868db86409b4543df0797f9" +checksum = "c3523ab5a71916ccf420eebdf5521fcef02141234bbc0b8a49f2fdc4544364ef" dependencies = [ "pin-project-lite", "tracing-attributes", @@ -2042,9 +2055,9 @@ checksum = "711b9620af191e0cdc7468a8d14e709c3dcdb115b36f838e601583af800a370a" [[package]] name = "uuid" -version = "1.4.1" +version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "79daa5ed5740825c40b389c5e50312b9c86df53fccd33f281df655642b43869d" +checksum = "88ad59a7560b41a70d191093a945f0b87bc1deeda46fb237479708a1d6b6cdfc" dependencies = [ "serde", ] diff --git a/benches/codegen.rs b/benches/codegen.rs index 4eabc30f..97b3c1e5 100644 --- a/benches/codegen.rs +++ b/benches/codegen.rs @@ -6,7 +6,7 @@ fn bench(c: &mut Criterion) { cornucopia::container::setup(false).unwrap(); let client = &mut cornucopia_conn().unwrap(); - cornucopia::load_schema(client, &["../codegen_test/schema.sql"]).unwrap(); + cornucopia::load_schema(client, &["../test_codegen/schema.sql"]).unwrap(); c.bench_function("codegen_sync", |b| { b.iter(|| { cornucopia::generate_live( diff --git a/benches/execution/main.rs b/benches/execution/main.rs index 64d07a86..441dc382 100644 --- a/benches/execution/main.rs +++ b/benches/execution/main.rs @@ -4,7 +4,6 @@ use cornucopia::conn::cornucopia_conn; use criterion::{BenchmarkId, Criterion}; use diesel::{Connection, PgConnection}; use postgres::{fallible_iterator::FallibleIterator, Client, NoTls}; -use tokio::runtime::Runtime; const QUERY_SIZE: &[usize] = &[1, 100, 10_000]; const INSERT_SIZE: &[usize] = &[1, 100, 1000]; @@ -128,23 +127,16 @@ fn prepare_full(client: &mut Client) { fn bench(c: &mut Criterion) { cornucopia::container::cleanup(false).ok(); cornucopia::container::setup(false).unwrap(); - let client = &mut cornucopia_conn().unwrap(); - let rt: &'static Runtime = Box::leak(Box::new(Runtime::new().unwrap())); - let async_client = &mut rt.block_on(async { - let (client, conn) = tokio_postgres::connect( - "postgresql://postgres:postgres@127.0.0.1:5435/postgres", - NoTls, - ) - .await - .unwrap(); - rt.spawn(conn); - client - }); + let client = &mut postgres::Client::connect( + "postgresql://postgres:postgres@127.0.0.1:5435/postgres", + NoTls, + ) + .unwrap(); + let async_client = &mut cornucopia_conn().unwrap(); let conn = &mut PgConnection::establish("postgresql://postgres:postgres@127.0.0.1:5435/postgres") .unwrap(); - - cornucopia::load_schema(client, &["usage/cornucopia_benches/schema.sql"]).unwrap(); + cornucopia::load_schema(async_client, &["execution/cornucopia_benches/schema.sql"]).unwrap(); { let mut group = c.benchmark_group("bench_trivial_query"); for size in QUERY_SIZE { diff --git a/crates/cornucopia/Cargo.toml b/crates/cornucopia/Cargo.toml index 04773127..bd111990 100644 --- a/crates/cornucopia/Cargo.toml +++ b/crates/cornucopia/Cargo.toml @@ -15,9 +15,13 @@ keywords = ["postgresql", "query", "generator", "sql", "tokio-postgres"] codegen_template = { path = "../codegen_template", version = "0.1.0" } # Postgres interaction -postgres = "0.19.4" +tokio-postgres = "0.7.8" postgres-types = "0.2.4" +# Async +tokio = { version = "1.28.1", features = ["rt-multi-thread"] } +futures = "0.3.28" + # Error handling and reporting thiserror = "1.0.38" miette = { version = "5.5.0", features = ["fancy"] } diff --git a/crates/cornucopia/src/cli.rs b/crates/cornucopia/src/cli.rs index 73692260..4860c758 100644 --- a/crates/cornucopia/src/cli.rs +++ b/crates/cornucopia/src/cli.rs @@ -64,8 +64,8 @@ pub fn run() -> Result<(), Error> { match action { Action::Live { url } => { - let mut client = conn::from_url(&url)?; - generate_live(&mut client, &queries_path, Some(&destination), settings)?; + let client = conn::from_url(&url)?; + generate_live(&client, &queries_path, Some(&destination), settings)?; } Action::Schema { schema_files } => { // Run the generate command. If the command is unsuccessful, cleanup Cornucopia's container diff --git a/crates/cornucopia/src/conn.rs b/crates/cornucopia/src/conn.rs index 83a54e0e..f9277f76 100644 --- a/crates/cornucopia/src/conn.rs +++ b/crates/cornucopia/src/conn.rs @@ -1,21 +1,36 @@ -use postgres::{Client, Config, NoTls}; +use tokio::runtime::Runtime; +use tokio_postgres::{Client, Config, NoTls}; use self::error::Error; /// Creates a non-TLS connection from a URL. pub(crate) fn from_url(url: &str) -> Result { - Ok(Client::connect(url, NoTls)?) + connect(url.parse()?) } /// Create a non-TLS connection to the container managed by Cornucopia. pub fn cornucopia_conn() -> Result { - Ok(Config::new() - .user("postgres") - .password("postgres") - .host("127.0.0.1") - .port(5435) - .dbname("postgres") - .connect(NoTls)?) + connect( + Config::new() + .user("postgres") + .password("postgres") + .host("127.0.0.1") + .port(5435) + .dbname("postgres") + .clone(), + ) +} + +fn connect(config: Config) -> Result { + let rt: &'static Runtime = Box::leak(Box::new( + Runtime::new().expect("Failed to start async Runtime"), + )); + let client = rt.block_on(async { + let (client, conn) = config.connect(NoTls).await.unwrap(); + rt.spawn(conn); + client + }); + Ok(client) } pub(crate) mod error { @@ -23,5 +38,5 @@ pub(crate) mod error { #[derive(Debug, thiserror::Error, Diagnostic)] #[error("Couldn't establish a connection with the database.")] - pub struct Error(#[from] pub postgres::Error); + pub struct Error(#[from] pub tokio_postgres::Error); } diff --git a/crates/cornucopia/src/lib.rs b/crates/cornucopia/src/lib.rs index b879ee28..ddf58e81 100644 --- a/crates/cornucopia/src/lib.rs +++ b/crates/cornucopia/src/lib.rs @@ -16,7 +16,7 @@ pub mod container; use std::path::Path; -use postgres::Client; +use tokio_postgres::Client; use codegen::generate as generate_internal; use error::WriteOutputError; @@ -43,7 +43,7 @@ pub struct CodegenSettings { /// the generated code will be written at that path. Code generation settings are /// set using the `settings` parameter. pub fn generate_live>( - client: &mut Client, + client: &Client, queries_path: P, destination: Option

, settings: CodegenSettings, @@ -84,9 +84,9 @@ pub fn generate_managed>( .map(parse_query_module) .collect::>()?; container::setup(podman)?; - let mut client = conn::cornucopia_conn()?; - load_schema(&mut client, schema_files)?; - let prepared_modules = prepare(&mut client, modules)?; + let client = conn::cornucopia_conn()?; + load_schema(&client, schema_files)?; + let prepared_modules = prepare(&client, modules)?; let generated_code = generate_internal(prepared_modules, settings); container::cleanup(podman)?; diff --git a/crates/cornucopia/src/load_schema.rs b/crates/cornucopia/src/load_schema.rs index f2616f90..0f7fe7b2 100644 --- a/crates/cornucopia/src/load_schema.rs +++ b/crates/cornucopia/src/load_schema.rs @@ -1,7 +1,7 @@ use std::path::Path; use miette::NamedSource; -use postgres::Client; +use tokio_postgres::Client; use crate::utils::db_err; @@ -10,14 +10,14 @@ use self::error::Error; /// Loads PostgreSQL schemas into a database. /// /// Takes a list of file paths as parameter and loads them in their given order. -pub fn load_schema>(client: &mut Client, paths: &[P]) -> Result<(), Error> { +pub fn load_schema>(client: &Client, paths: &[P]) -> Result<(), Error> { for path in paths { let path = path.as_ref(); let sql = std::fs::read_to_string(path).map_err(|err| Error::Io { path: path.to_string_lossy().to_string(), err, })?; - client.batch_execute(&sql).map_err(|err| { + futures::executor::block_on(client.batch_execute(&sql)).map_err(|err| { let msg = format!("{err:#}"); let src = NamedSource::new(path.to_string_lossy(), sql); if let Some((position, msg, help)) = db_err(&err) { diff --git a/crates/cornucopia/src/prepare_queries.rs b/crates/cornucopia/src/prepare_queries.rs index 9edf6eee..cda4c144 100644 --- a/crates/cornucopia/src/prepare_queries.rs +++ b/crates/cornucopia/src/prepare_queries.rs @@ -1,9 +1,10 @@ -use std::rc::Rc; +use std::{collections::HashMap, rc::Rc}; +use futures::{stream::FuturesUnordered, StreamExt}; use heck::ToUpperCamelCase; use indexmap::{map::Entry, IndexMap}; -use postgres::Client; use postgres_types::{Kind, Type}; +use tokio_postgres::{Client, Statement}; use crate::{ codegen::GenCtx, @@ -226,7 +227,8 @@ impl PreparedModule { } /// Prepares all modules -pub(crate) fn prepare(client: &mut Client, modules: Vec) -> Result { +pub(crate) fn prepare(client: &Client, modules: Vec) -> Result { + let stmts = prepare_sql(client, &modules); let mut registrar = TypeRegistrar::default(); let mut tmp = Preparation { modules: Vec::new(), @@ -240,7 +242,7 @@ pub(crate) fn prepare(client: &mut Client, modules: Vec) -> Result HashMap> { + let queries: FuturesUnordered<_> = modules + .iter() + .flat_map(|m| m.queries.iter().map(|q| q.sql_str.clone())) + .map(|query| async move { + let stmt = client.prepare(&query).await; + (query, stmt) + }) + .collect(); + let results: HashMap<_, _> = futures::executor::block_on(queries.collect()); + results +} + /// Prepares all queries in this module fn prepare_module( - client: &mut Client, + stmts: &HashMap>, module: Module, registrar: &mut TypeRegistrar, ) -> Result { @@ -332,7 +350,7 @@ fn prepare_module( for query in module.queries { prepare_query( - client, + stmts, &mut tmp_prepared_module, registrar, &module.types, @@ -348,7 +366,7 @@ fn prepare_module( /// Prepares a query fn prepare_query( - client: &mut Client, + stmts: &HashMap>, module: &mut PreparedModule, registrar: &mut TypeRegistrar, types: &[TypeAnnotation], @@ -363,9 +381,9 @@ fn prepare_query( module_info: &ModuleInfo, ) -> Result<(), Error> { // Prepare the statement - let stmt = client - .prepare(&sql_str) - .map_err(|e| Error::new_db_err(&e, module_info, &sql_span, &name))?; + let stmt = stmts[&sql_str] + .as_ref() + .map_err(|e| Error::new_db_err(e, module_info, &sql_span, &name))?; let (nullable_params_fields, params_name) = param.name_and_fields(types, &name, Some("Params")); let (nullable_row_fields, row_name) = row.name_and_fields(types, &name, None); @@ -477,7 +495,7 @@ pub(crate) mod error { impl Error { pub(crate) fn new_db_err( - err: &postgres::Error, + err: &tokio_postgres::Error, module_info: &ModuleInfo, query_span: &SourceSpan, query_name: &Span, diff --git a/crates/cornucopia/src/type_registrar.rs b/crates/cornucopia/src/type_registrar.rs index 1b7ca70f..6acf5c4d 100644 --- a/crates/cornucopia/src/type_registrar.rs +++ b/crates/cornucopia/src/type_registrar.rs @@ -455,7 +455,7 @@ pub(crate) mod error { #[derive(Debug, ThisError, Diagnostic)] #[error("Couldn't register SQL type.")] pub enum Error { - Db(#[from] postgres::Error), + Db(#[from] tokio_postgres::Error), UnsupportedPostgresType { #[source_code] src: NamedSource, diff --git a/crates/cornucopia/src/utils.rs b/crates/cornucopia/src/utils.rs index e9748da3..b21184d1 100644 --- a/crates/cornucopia/src/utils.rs +++ b/crates/cornucopia/src/utils.rs @@ -1,6 +1,6 @@ use indexmap::Equivalent; -use postgres::error::ErrorPosition; use postgres_types::Type; +use tokio_postgres::error::ErrorPosition; /// Allows us to query a map using type schema as key without having to own the key strings #[derive(PartialEq, Eq, Hash)] @@ -34,7 +34,7 @@ pub fn find_duplicate(slice: &[T], eq: fn(&T, &T) -> bool) -> Option<(&T, &T) } /// Extracts useful info from a `postgres`-generated error. -pub(crate) fn db_err(err: &postgres::Error) -> Option<(u32, String, Option)> { +pub(crate) fn db_err(err: &tokio_postgres::Error) -> Option<(u32, String, Option)> { if let Some(db_err) = err.as_db_error() { if let Some(ErrorPosition::Original(position)) = db_err.position() { Some(( diff --git a/crates/cornucopia/src/validation.rs b/crates/cornucopia/src/validation.rs index 9090e96a..6db45e60 100644 --- a/crates/cornucopia/src/validation.rs +++ b/crates/cornucopia/src/validation.rs @@ -9,8 +9,8 @@ use crate::{ use error::Error; use miette::SourceSpan; -use postgres::Column; use postgres_types::Type; +use tokio_postgres::Column; pub(crate) fn duplicate_nullable_ident( info: &ModuleInfo, diff --git a/test_integration/Cargo.toml b/test_integration/Cargo.toml index cd2e8fd4..0c904f27 100644 --- a/test_integration/Cargo.toml +++ b/test_integration/Cargo.toml @@ -20,6 +20,8 @@ clap = { version = "4.0.29", features = ["derive"] } # Postgres interaction postgres = { version = "0.19.4" } +tokio-postgres = "0.7.8" +futures = "0.3.28" # serde ## Test fixtures ser/de diff --git a/test_integration/src/codegen.rs b/test_integration/src/codegen.rs index 9d465407..cadf751f 100644 --- a/test_integration/src/codegen.rs +++ b/test_integration/src/codegen.rs @@ -9,7 +9,7 @@ use std::{env::set_current_dir, process::Command}; // Run codegen test, return true if all test are successful pub(crate) fn run_codegen_test( - client: &mut postgres::Client, + client: &tokio_postgres::Client, apply: bool, ) -> Result> { let mut successful = true; diff --git a/test_integration/src/errors.rs b/test_integration/src/errors.rs index d5531a1a..890592da 100644 --- a/test_integration/src/errors.rs +++ b/test_integration/src/errors.rs @@ -8,7 +8,7 @@ use crate::{ /// Run errors test, return true if all test are successful pub(crate) fn run_errors_test( - client: &mut postgres::Client, + client: &tokio_postgres::Client, apply: bool, ) -> Result> { let mut successful = true; diff --git a/test_integration/src/main.rs b/test_integration/src/main.rs index cf21be31..63aee305 100644 --- a/test_integration/src/main.rs +++ b/test_integration/src/main.rs @@ -44,9 +44,9 @@ fn test( container::cleanup(podman).ok(); container::setup(podman).unwrap(); let successful = std::panic::catch_unwind(|| { - let mut client = cornucopia::conn::cornucopia_conn().unwrap(); - display(run_errors_test(&mut client, apply_errors)).unwrap() - && display(run_codegen_test(&mut client, apply_codegen)).unwrap() + let client = cornucopia::conn::cornucopia_conn().unwrap(); + display(run_errors_test(&client, apply_errors)).unwrap() + && display(run_codegen_test(&client, apply_codegen)).unwrap() }); container::cleanup(podman).unwrap(); successful.unwrap() diff --git a/test_integration/src/utils.rs b/test_integration/src/utils.rs index 1ba234b1..48c0d480 100644 --- a/test_integration/src/utils.rs +++ b/test_integration/src/utils.rs @@ -5,8 +5,10 @@ use std::{ }; /// Reset the current database -pub(crate) fn reset_db(client: &mut postgres::Client) -> Result<(), postgres::Error> { - client.batch_execute("DROP SCHEMA public CASCADE;CREATE SCHEMA public;") +pub(crate) fn reset_db(client: &tokio_postgres::Client) -> Result<(), postgres::Error> { + futures::executor::block_on( + client.batch_execute("DROP SCHEMA public CASCADE;CREATE SCHEMA public;"), + ) } pub(crate) fn rustfmt_file(path: &Path) {