ffd7245e54
Benchmarks indicate that malloc accounts for a significant amount of the runtime of queries. The message buffer accounts for ~half of that (the other being channels), and this change should eliminate it.
57 lines
1.6 KiB
Rust
57 lines
1.6 KiB
Rust
use crate::client::{InnerClient, Responses};
|
|
use crate::codec::FrontendMessage;
|
|
use crate::connection::RequestMessages;
|
|
use crate::types::ToSql;
|
|
use crate::{query, Error, Statement};
|
|
use bytes::Bytes;
|
|
use futures::{ready, Stream};
|
|
use postgres_protocol::message::backend::Message;
|
|
use std::pin::Pin;
|
|
use std::task::{Context, Poll};
|
|
|
|
pub async fn copy_out<'a, I>(
|
|
client: &InnerClient,
|
|
statement: Statement,
|
|
params: I,
|
|
) -> Result<CopyStream, Error>
|
|
where
|
|
I: IntoIterator<Item = &'a dyn ToSql>,
|
|
I::IntoIter: ExactSizeIterator,
|
|
{
|
|
let buf = query::encode(client, &statement, params)?;
|
|
let responses = start(client, buf).await?;
|
|
Ok(CopyStream { responses })
|
|
}
|
|
|
|
async fn start(client: &InnerClient, buf: Bytes) -> Result<Responses, Error> {
|
|
let mut responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
|
|
|
|
match responses.next().await? {
|
|
Message::BindComplete => {}
|
|
_ => return Err(Error::unexpected_message()),
|
|
}
|
|
|
|
match responses.next().await? {
|
|
Message::CopyOutResponse(_) => {}
|
|
_ => return Err(Error::unexpected_message()),
|
|
}
|
|
|
|
Ok(responses)
|
|
}
|
|
|
|
pub struct CopyStream {
|
|
responses: Responses,
|
|
}
|
|
|
|
impl Stream for CopyStream {
|
|
type Item = Result<Bytes, Error>;
|
|
|
|
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
|
match ready!(self.responses.poll_next(cx)?) {
|
|
Message::CopyData(body) => Poll::Ready(Some(Ok(body.into_bytes()))),
|
|
Message::CopyDone => Poll::Ready(None),
|
|
_ => Poll::Ready(Some(Err(Error::unexpected_message()))),
|
|
}
|
|
}
|
|
}
|