mirror of
https://github.com/neondatabase/neon.git
synced 2026-01-08 14:02:55 +00:00
d9cedb4a95
Follow up to #12701, which introduced a new regression. When profiling locally I noticed that writes have the tendency to always reallocate. On investigation I found that even if the `Connection`'s write buffer is empty, if it still shares the same data pointer as the `Client`'s write buffer then the client cannot reclaim it. The best way I found to fix this is to just drop the `Connection`'s write buffer each time we fully flush it. Additionally, I remembered that `BytesMut` has an `unsplit` method which is allows even better sharing over the previous optimisation I had when 'encoding'.
101 lines
2.4 KiB
Rust
101 lines
2.4 KiB
Rust
use std::io;
|
|
|
|
use bytes::BytesMut;
|
|
use fallible_iterator::FallibleIterator;
|
|
use postgres_protocol2::message::backend;
|
|
use tokio::sync::mpsc::UnboundedSender;
|
|
use tokio_util::codec::{Decoder, Encoder};
|
|
|
|
pub enum FrontendMessage {
|
|
Raw(BytesMut),
|
|
RecordNotices(RecordNotices),
|
|
}
|
|
|
|
pub struct RecordNotices {
|
|
pub sender: UnboundedSender<Box<str>>,
|
|
pub limit: usize,
|
|
}
|
|
|
|
pub enum BackendMessage {
|
|
Normal {
|
|
messages: BackendMessages,
|
|
ready: bool,
|
|
},
|
|
Async(backend::Message),
|
|
}
|
|
|
|
pub struct BackendMessages(BytesMut);
|
|
|
|
impl BackendMessages {
|
|
pub fn empty() -> BackendMessages {
|
|
BackendMessages(BytesMut::new())
|
|
}
|
|
}
|
|
|
|
impl FallibleIterator for BackendMessages {
|
|
type Item = backend::Message;
|
|
type Error = io::Error;
|
|
|
|
fn next(&mut self) -> io::Result<Option<backend::Message>> {
|
|
backend::Message::parse(&mut self.0)
|
|
}
|
|
}
|
|
|
|
pub struct PostgresCodec;
|
|
|
|
impl Encoder<BytesMut> for PostgresCodec {
|
|
type Error = io::Error;
|
|
|
|
fn encode(&mut self, item: BytesMut, dst: &mut BytesMut) -> io::Result<()> {
|
|
dst.unsplit(item);
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl Decoder for PostgresCodec {
|
|
type Item = BackendMessage;
|
|
type Error = io::Error;
|
|
|
|
fn decode(&mut self, src: &mut BytesMut) -> Result<Option<BackendMessage>, io::Error> {
|
|
let mut idx = 0;
|
|
|
|
let mut ready = false;
|
|
while let Some(header) = backend::Header::parse(&src[idx..])? {
|
|
let len = header.len() as usize + 1;
|
|
if src[idx..].len() < len {
|
|
break;
|
|
}
|
|
|
|
match header.tag() {
|
|
backend::NOTICE_RESPONSE_TAG
|
|
| backend::NOTIFICATION_RESPONSE_TAG
|
|
| backend::PARAMETER_STATUS_TAG => {
|
|
if idx == 0 {
|
|
let message = backend::Message::parse(src)?.unwrap();
|
|
return Ok(Some(BackendMessage::Async(message)));
|
|
} else {
|
|
break;
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
|
|
idx += len;
|
|
|
|
if header.tag() == backend::READY_FOR_QUERY_TAG {
|
|
ready = true;
|
|
break;
|
|
}
|
|
}
|
|
|
|
if idx == 0 {
|
|
Ok(None)
|
|
} else {
|
|
Ok(Some(BackendMessage::Normal {
|
|
messages: BackendMessages(src.split_to(idx)),
|
|
ready,
|
|
}))
|
|
}
|
|
}
|
|
}
|