From 6cfd3a57d1a4a7d3c9d8fa9d1947da3612488f26 Mon Sep 17 00:00:00 2001 From: Ryan Leckey Date: Thu, 11 Jul 2019 13:52:45 -0700 Subject: [PATCH] Optimize performance in several encode/decode programs --- sqlx-postgres-protocol/benches/decode.rs | 9 +- sqlx-postgres-protocol/benches/encode.rs | 36 +- .../src/backend_key_data.rs | 6 +- sqlx-postgres-protocol/src/data_row.rs | 8 +- sqlx-postgres-protocol/src/lib.rs | 2 +- sqlx-postgres-protocol/src/query.rs | 11 +- sqlx-postgres-protocol/src/ready_for_query.rs | 34 +- sqlx-postgres-protocol/src/response.rs | 321 +----------------- sqlx-postgres-protocol/src/startup_message.rs | 2 + sqlx-postgres-protocol/src/terminate.rs | 6 +- 10 files changed, 47 insertions(+), 388 deletions(-) diff --git a/sqlx-postgres-protocol/benches/decode.rs b/sqlx-postgres-protocol/benches/decode.rs index 1988ed2f..194f3aae 100644 --- a/sqlx-postgres-protocol/benches/decode.rs +++ b/sqlx-postgres-protocol/benches/decode.rs @@ -3,12 +3,13 @@ extern crate criterion; use bytes::Bytes; use criterion::{black_box, Criterion}; -use sqlx_postgres_protocol::{BackendKeyData, Decode, ParameterStatus, Response}; +use sqlx_postgres_protocol::{BackendKeyData, Decode, ParameterStatus, ReadyForQuery, Response}; fn criterion_benchmark(c: &mut Criterion) { const NOTICE_RESPONSE: &[u8] = b"SNOTICE\0VNOTICE\0C42710\0Mextension \"uuid-ossp\" already exists, skipping\0Fextension.c\0L1656\0RCreateExtension\0\0"; const PARAM_STATUS: &[u8] = b"session_authorization\0postgres\0"; const BACKEND_KEY_DATA: &[u8] = b"\0\0'\xc6\x89R\xc5+"; + const READY_FOR_QUERY: &[u8] = b"E"; c.bench_function("decode Response", |b| { b.iter(|| { @@ -28,6 +29,12 @@ fn criterion_benchmark(c: &mut Criterion) { let _ = ParameterStatus::decode(black_box(Bytes::from_static(PARAM_STATUS))).unwrap(); }) }); + + c.bench_function("decode ReadyForQuery", |b| { + b.iter(|| { + let _ = ReadyForQuery::decode(black_box(Bytes::from_static(READY_FOR_QUERY))).unwrap(); + }) + }); } criterion_group!(benches, criterion_benchmark); diff --git a/sqlx-postgres-protocol/benches/encode.rs b/sqlx-postgres-protocol/benches/encode.rs index 5df76640..f87a61ef 100644 --- a/sqlx-postgres-protocol/benches/encode.rs +++ b/sqlx-postgres-protocol/benches/encode.rs @@ -2,25 +2,9 @@ extern crate criterion; use criterion::Criterion; -use sqlx_postgres_protocol::{Encode, PasswordMessage, Response, Severity, StartupMessage}; +use sqlx_postgres_protocol::{Encode, PasswordMessage, Query, StartupMessage, Terminate}; fn criterion_benchmark(c: &mut Criterion) { - c.bench_function("encode Response::builder()", |b| { - let mut dst = Vec::with_capacity(1024); - b.iter(|| { - dst.clear(); - Response::builder() - .severity(Severity::Notice) - .code("42710") - .message("extension \"uuid-ossp\" already exists, skipping") - .file("extension.c") - .line(1656) - .routine("CreateExtension") - .encode(&mut dst) - .unwrap(); - }) - }); - c.bench_function("encode PasswordMessage::cleartext", |b| { let mut dst = Vec::with_capacity(1024); b.iter(|| { @@ -31,6 +15,22 @@ fn criterion_benchmark(c: &mut Criterion) { }) }); + c.bench_function("encode Query", |b| { + let mut dst = Vec::with_capacity(1024); + b.iter(|| { + dst.clear(); + Query::new("SELECT 1, 2, 3").encode(&mut dst).unwrap(); + }) + }); + + c.bench_function("encode Terminate", |b| { + let mut dst = Vec::with_capacity(1024); + b.iter(|| { + dst.clear(); + Terminate.encode(&mut dst).unwrap(); + }) + }); + c.bench_function("encode StartupMessage", |b| { let mut dst = Vec::with_capacity(1024); b.iter(|| { @@ -49,7 +49,7 @@ fn criterion_benchmark(c: &mut Criterion) { }) }); - c.bench_function("encode Password(MD5)", |b| { + c.bench_function("encode Password::md5", |b| { let mut dst = Vec::with_capacity(1024); b.iter(|| { dst.clear(); diff --git a/sqlx-postgres-protocol/src/backend_key_data.rs b/sqlx-postgres-protocol/src/backend_key_data.rs index 0388e7f5..5ee4de1f 100644 --- a/sqlx-postgres-protocol/src/backend_key_data.rs +++ b/sqlx-postgres-protocol/src/backend_key_data.rs @@ -1,7 +1,7 @@ use crate::Decode; -use byteorder::{BigEndian, ByteOrder}; use bytes::Bytes; use std::io; +use std::convert::TryInto; #[derive(Debug)] pub struct BackendKeyData { @@ -24,8 +24,8 @@ impl BackendKeyData { impl Decode for BackendKeyData { fn decode(src: Bytes) -> io::Result { - let process_id = BigEndian::read_u32(&src[..4]); - let secret_key = BigEndian::read_u32(&src[4..]); + let process_id = u32::from_be_bytes(src.as_ref()[0..4].try_into().unwrap()); + let secret_key = u32::from_be_bytes(src.as_ref()[4..8].try_into().unwrap()); Ok(Self { process_id, diff --git a/sqlx-postgres-protocol/src/data_row.rs b/sqlx-postgres-protocol/src/data_row.rs index 9470d72b..b75ec605 100644 --- a/sqlx-postgres-protocol/src/data_row.rs +++ b/sqlx-postgres-protocol/src/data_row.rs @@ -39,10 +39,6 @@ pub struct DataValues<'a> { impl<'a> Iterator for DataValues<'a> { type Item = Option<&'a [u8]>; - fn size_hint(&self) -> (usize, Option) { - (self.rem as usize, Some(self.rem as usize)) - } - fn next(&mut self) -> Option { if self.rem == 0 { return None; @@ -62,6 +58,10 @@ impl<'a> Iterator for DataValues<'a> { Some(value) } + + fn size_hint(&self) -> (usize, Option) { + (self.rem as usize, Some(self.rem as usize)) + } } impl<'a> ExactSizeIterator for DataValues<'a> {} diff --git a/sqlx-postgres-protocol/src/lib.rs b/sqlx-postgres-protocol/src/lib.rs index 0c771b9f..a6128015 100644 --- a/sqlx-postgres-protocol/src/lib.rs +++ b/sqlx-postgres-protocol/src/lib.rs @@ -28,7 +28,7 @@ pub use self::{ password_message::PasswordMessage, query::Query, ready_for_query::{ReadyForQuery, TransactionStatus}, - response::{Response, ResponseBuilder, Severity}, + response::{Response, Severity}, row_description::{FieldDescription, FieldDescriptions, RowDescription}, startup_message::StartupMessage, terminate::Terminate, diff --git a/sqlx-postgres-protocol/src/query.rs b/sqlx-postgres-protocol/src/query.rs index fec61516..960a2da5 100644 --- a/sqlx-postgres-protocol/src/query.rs +++ b/sqlx-postgres-protocol/src/query.rs @@ -1,11 +1,11 @@ use crate::Encode; -use bytes::BufMut; use std::io; #[derive(Debug)] pub struct Query<'a>(&'a str); impl<'a> Query<'a> { + #[inline] pub fn new(query: &'a str) -> Self { Self(query) } @@ -14,11 +14,10 @@ impl<'a> Query<'a> { impl Encode for Query<'_> { fn encode(&self, buf: &mut Vec) -> io::Result<()> { let len = self.0.len() + 4 + 1; - buf.reserve(len + 1); - buf.put_u8(b'Q'); - buf.put_u32_be(len as u32); - buf.put(self.0); - buf.put_u8(0); + buf.push(b'Q'); + buf.extend_from_slice(&(len as u32).to_be_bytes()); + buf.extend_from_slice(self.0.as_bytes()); + buf.push(0); Ok(()) } diff --git a/sqlx-postgres-protocol/src/ready_for_query.rs b/sqlx-postgres-protocol/src/ready_for_query.rs index 96f8e0c5..0e176042 100644 --- a/sqlx-postgres-protocol/src/ready_for_query.rs +++ b/sqlx-postgres-protocol/src/ready_for_query.rs @@ -1,5 +1,4 @@ -use crate::{Decode, Encode}; -use byteorder::{WriteBytesExt, BE}; +use crate::Decode; use bytes::Bytes; use std::io; @@ -22,21 +21,6 @@ pub struct ReadyForQuery { pub status: TransactionStatus, } -impl Encode for ReadyForQuery { - #[inline] - fn size_hint(&self) -> usize { - 6 - } - - fn encode(&self, buf: &mut Vec) -> io::Result<()> { - buf.write_u8(b'Z')?; - buf.write_u32::(5)?; - buf.write_u8(self.status as u8)?; - - Ok(()) - } -} - impl Decode for ReadyForQuery { fn decode(src: Bytes) -> io::Result { if src.len() != 1 { @@ -59,26 +43,12 @@ impl Decode for ReadyForQuery { #[cfg(test)] mod tests { use super::{ReadyForQuery, TransactionStatus}; - use crate::{Decode, Encode}; + use crate::Decode; use bytes::Bytes; use std::io; const READY_FOR_QUERY: &[u8] = b"E"; - #[test] - fn it_encodes_ready_for_query() -> io::Result<()> { - let message = ReadyForQuery { - status: TransactionStatus::Error, - }; - - let mut dst = Vec::with_capacity(message.size_hint()); - message.encode(&mut dst)?; - - assert_eq!(&dst[5..], READY_FOR_QUERY); - - Ok(()) - } - #[test] fn it_decodes_ready_for_query() -> io::Result<()> { let src = Bytes::from_static(READY_FOR_QUERY); diff --git a/sqlx-postgres-protocol/src/response.rs b/sqlx-postgres-protocol/src/response.rs index d9125802..e43635cd 100644 --- a/sqlx-postgres-protocol/src/response.rs +++ b/sqlx-postgres-protocol/src/response.rs @@ -1,9 +1,7 @@ -use crate::{decode::get_str, Decode, Encode}; -use byteorder::{BigEndian, WriteBytesExt}; +use crate::{decode::get_str, Decode}; use bytes::Bytes; use std::{ fmt, io, - ops::Range, pin::Pin, ptr::NonNull, str::{self, FromStr}, @@ -104,11 +102,6 @@ unsafe impl Send for Response {} unsafe impl Sync for Response {} impl Response { - #[inline] - pub fn builder() -> ResponseBuilder { - ResponseBuilder::new() - } - #[inline] pub fn severity(&self) -> Severity { self.severity @@ -232,26 +225,6 @@ impl fmt::Debug for Response { } } -impl Encode for Response { - #[inline] - fn size_hint(&self) -> usize { - self.storage.len() + 5 - } - - fn encode(&self, buf: &mut Vec) -> io::Result<()> { - if self.severity.is_error() { - buf.push(b'E'); - } else { - buf.push(b'N'); - } - - buf.write_u32::((4 + self.storage.len()) as u32)?; - buf.extend_from_slice(&self.storage); - - Ok(()) - } -} - impl Decode for Response { fn decode(src: Bytes) -> io::Result { let storage = Pin::new(src); @@ -423,306 +396,16 @@ impl Decode for Response { } } -pub struct ResponseBuilder { - storage: Vec, - severity: Option, - code: Option>, - message: Option>, - detail: Option>, - hint: Option>, - position: Option, - internal_position: Option, - internal_query: Option>, - where_: Option>, - schema: Option>, - table: Option>, - column: Option>, - data_type: Option>, - constraint: Option>, - file: Option>, - line: Option, - routine: Option>, -} - -impl Default for ResponseBuilder { - fn default() -> Self { - Self { - storage: Vec::with_capacity(256), - severity: None, - message: None, - code: None, - detail: None, - hint: None, - position: None, - internal_position: None, - internal_query: None, - where_: None, - schema: None, - table: None, - column: None, - data_type: None, - constraint: None, - file: None, - line: None, - routine: None, - } - } -} - -fn put_str(buf: &mut Vec, tag: u8, value: &str) -> Range { - buf.push(tag); - let beg = buf.len(); - buf.extend_from_slice(value.as_bytes()); - let end = buf.len(); - buf.push(0); - beg..end -} - -impl ResponseBuilder { - #[inline] - pub fn new() -> ResponseBuilder { - Self::default() - } - - #[inline] - pub fn severity(mut self, severity: Severity) -> Self { - let sev = severity.to_str(); - - let _ = put_str(&mut self.storage, b'S', sev); - let _ = put_str(&mut self.storage, b'V', sev); - - self.severity = Some(severity); - self - } - - #[inline] - pub fn message(mut self, message: &str) -> Self { - self.message = Some(put_str(&mut self.storage, b'M', message)); - self - } - - #[inline] - pub fn code(mut self, code: &str) -> Self { - self.code = Some(put_str(&mut self.storage, b'C', code)); - self - } - - #[inline] - pub fn detail(mut self, detail: &str) -> Self { - self.detail = Some(put_str(&mut self.storage, b'D', detail)); - self - } - - #[inline] - pub fn hint(mut self, hint: &str) -> Self { - self.hint = Some(put_str(&mut self.storage, b'H', hint)); - self - } - - #[inline] - pub fn position(mut self, position: usize) -> Self { - self.storage.push(b'P'); - // PANIC: Write to Vec is infallible - itoa::write(&mut self.storage, position).unwrap(); - self.storage.push(0); - - self.position = Some(position); - self - } - - #[inline] - pub fn internal_position(mut self, position: usize) -> Self { - self.storage.push(b'p'); - // PANIC: Write to Vec is infallible - itoa::write(&mut self.storage, position).unwrap(); - self.storage.push(0); - - self.internal_position = Some(position); - self - } - - #[inline] - pub fn internal_query(mut self, query: &str) -> Self { - self.internal_query = Some(put_str(&mut self.storage, b'q', query)); - self - } - - #[inline] - pub fn where_(mut self, where_: &str) -> Self { - self.where_ = Some(put_str(&mut self.storage, b'w', where_)); - self - } - - #[inline] - pub fn schema(mut self, schema: &str) -> Self { - self.schema = Some(put_str(&mut self.storage, b's', schema)); - self - } - - #[inline] - pub fn table(mut self, table: &str) -> Self { - self.table = Some(put_str(&mut self.storage, b't', table)); - self - } - - #[inline] - pub fn column(mut self, column: &str) -> Self { - self.column = Some(put_str(&mut self.storage, b'c', column)); - self - } - - #[inline] - pub fn data_type(mut self, data_type: &str) -> Self { - self.data_type = Some(put_str(&mut self.storage, b'd', data_type)); - self - } - - #[inline] - pub fn constraint(mut self, constraint: &str) -> Self { - self.constraint = Some(put_str(&mut self.storage, b'n', constraint)); - self - } - - #[inline] - pub fn file(mut self, file: &str) -> Self { - self.file = Some(put_str(&mut self.storage, b'F', file)); - self - } - - #[inline] - pub fn line(mut self, line: usize) -> Self { - self.storage.push(b'L'); - // PANIC: Write to Vec is infallible - itoa::write(&mut self.storage, line).unwrap(); - self.storage.push(0); - - self.line = Some(line); - self - } - - #[inline] - pub fn routine(mut self, routine: &str) -> Self { - self.routine = Some(put_str(&mut self.storage, b'R', routine)); - self - } - - pub fn build(mut self) -> Response { - // Add a \0 terminator - self.storage.push(0); - - // Freeze the storage and Pin so we can self-reference it - let storage = Pin::new(Bytes::from(self.storage)); - - let make_str_ref = |val: Option>| unsafe { - val.map(|r| NonNull::from(str::from_utf8_unchecked(&storage[r]))) - }; - - let code = make_str_ref(self.code); - let message = make_str_ref(self.message); - let detail = make_str_ref(self.detail); - let hint = make_str_ref(self.hint); - let internal_query = make_str_ref(self.internal_query); - let where_ = make_str_ref(self.where_); - let schema = make_str_ref(self.schema); - let table = make_str_ref(self.table); - let column = make_str_ref(self.column); - let data_type = make_str_ref(self.data_type); - let constraint = make_str_ref(self.constraint); - let file = make_str_ref(self.file); - let routine = make_str_ref(self.routine); - - Response { - storage, - // FIXME: Default and don't panic here - severity: self.severity.expect("`severity` required by protocol"), - code: code.expect("`code` required by protocol"), - message: message.expect("`message` required by protocol"), - detail, - hint, - internal_query, - where_, - schema, - table, - column, - data_type, - constraint, - file, - routine, - line: self.line, - position: self.position, - internal_position: self.internal_position, - } - } -} - -impl Encode for ResponseBuilder { - #[inline] - fn size_hint(&self) -> usize { - self.storage.len() + 6 - } - - fn encode(&self, buf: &mut Vec) -> io::Result<()> { - if self.severity.as_ref().map_or(false, |s| s.is_error()) { - buf.push(b'E'); - } else { - buf.push(b'N'); - } - - buf.write_u32::((5 + self.storage.len()) as u32)?; - buf.extend_from_slice(&self.storage); - buf.push(0); - - Ok(()) - } -} - #[cfg(test)] mod tests { use super::{Response, Severity}; - use crate::{Decode, Encode}; + use crate::Decode; use bytes::Bytes; use std::io; const RESPONSE: &[u8] = b"SNOTICE\0VNOTICE\0C42710\0Mextension \"uuid-ossp\" already exists, \ skipping\0Fextension.c\0L1656\0RCreateExtension\0\0"; - #[test] - fn it_encodes_response() -> io::Result<()> { - let message = Response::builder() - .severity(Severity::Notice) - .code("42710") - .message("extension \"uuid-ossp\" already exists, skipping") - .file("extension.c") - .line(1656) - .routine("CreateExtension") - .build(); - - let mut dst = Vec::with_capacity(message.size_hint()); - message.encode(&mut dst)?; - - assert_eq!(&dst[5..], RESPONSE); - - Ok(()) - } - - #[test] - fn it_encodes_response_builder() -> io::Result<()> { - let message = Response::builder() - .severity(Severity::Notice) - .code("42710") - .message("extension \"uuid-ossp\" already exists, skipping") - .file("extension.c") - .line(1656) - .routine("CreateExtension"); - - let mut dst = Vec::with_capacity(message.size_hint()); - message.encode(&mut dst)?; - - assert_eq!(&dst[5..], RESPONSE); - - Ok(()) - } - #[test] fn it_decodes_response() -> io::Result<()> { let src = Bytes::from_static(RESPONSE); diff --git a/sqlx-postgres-protocol/src/startup_message.rs b/sqlx-postgres-protocol/src/startup_message.rs index 20de04c0..70b15817 100644 --- a/sqlx-postgres-protocol/src/startup_message.rs +++ b/sqlx-postgres-protocol/src/startup_message.rs @@ -8,10 +8,12 @@ pub struct StartupMessage<'a> { } impl<'a> StartupMessage<'a> { + #[inline] pub fn new(params: &'a [(&'a str, &'a str)]) -> Self { Self { params } } + #[inline] pub fn params(&self) -> &'a [(&'a str, &'a str)] { self.params } diff --git a/sqlx-postgres-protocol/src/terminate.rs b/sqlx-postgres-protocol/src/terminate.rs index 6b624262..9d9763d9 100644 --- a/sqlx-postgres-protocol/src/terminate.rs +++ b/sqlx-postgres-protocol/src/terminate.rs @@ -1,5 +1,4 @@ use crate::Encode; -use bytes::BufMut; use std::io; #[derive(Debug)] @@ -7,9 +6,8 @@ pub struct Terminate; impl Encode for Terminate { fn encode(&self, buf: &mut Vec) -> io::Result<()> { - buf.reserve(5); - buf.put_u8(b'X'); - buf.put_u32_be(4); + buf.push(b'X'); + buf.extend_from_slice(&4_u32.to_be_bytes()); Ok(()) }