Use the client buffer for cleanup

This commit is contained in:
Steven Fackler 2019-10-12 18:07:09 -07:00
parent 20adbaf36b
commit 9e42c67b7b
3 changed files with 17 additions and 14 deletions

View File

@ -2,7 +2,6 @@ use crate::client::InnerClient;
use crate::codec::FrontendMessage; use crate::codec::FrontendMessage;
use crate::connection::RequestMessages; use crate::connection::RequestMessages;
use crate::Statement; use crate::Statement;
use bytes::BytesMut;
use postgres_protocol::message::frontend; use postgres_protocol::message::frontend;
use std::sync::{Arc, Weak}; use std::sync::{Arc, Weak};
@ -15,10 +14,12 @@ struct Inner {
impl Drop for Inner { impl Drop for Inner {
fn drop(&mut self) { fn drop(&mut self) {
if let Some(client) = self.client.upgrade() { if let Some(client) = self.client.upgrade() {
let mut buf = BytesMut::new(); let buf = client.with_buf(|buf| {
frontend::close(b'P', &self.name, &mut buf).expect("portal name not valid"); frontend::close(b'P', &self.name, buf).unwrap();
frontend::sync(&mut buf); frontend::sync(buf);
let _ = client.send(RequestMessages::Single(FrontendMessage::Raw(buf.freeze()))); buf.take().freeze()
});
let _ = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)));
} }
} }
} }

View File

@ -2,7 +2,6 @@ use crate::client::InnerClient;
use crate::codec::FrontendMessage; use crate::codec::FrontendMessage;
use crate::connection::RequestMessages; use crate::connection::RequestMessages;
use crate::types::Type; use crate::types::Type;
use bytes::BytesMut;
use postgres_protocol::message::frontend; use postgres_protocol::message::frontend;
use std::sync::{Arc, Weak}; use std::sync::{Arc, Weak};
@ -16,10 +15,12 @@ struct StatementInner {
impl Drop for StatementInner { impl Drop for StatementInner {
fn drop(&mut self) { fn drop(&mut self) {
if let Some(client) = self.client.upgrade() { if let Some(client) = self.client.upgrade() {
let mut buf = BytesMut::new(); let buf = client.with_buf(|buf| {
frontend::close(b'S', &self.name, &mut buf).expect("statement name not valid"); frontend::close(b'S', &self.name, buf).unwrap();
frontend::sync(&mut buf); frontend::sync(buf);
let _ = client.send(RequestMessages::Single(FrontendMessage::Raw(buf.freeze()))); buf.take().freeze()
});
let _ = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)));
} }
} }
} }

View File

@ -14,7 +14,6 @@ use crate::{
use bytes::IntoBuf; use bytes::IntoBuf;
use futures::{TryStream, TryStreamExt}; use futures::{TryStream, TryStreamExt};
use postgres_protocol::message::frontend; use postgres_protocol::message::frontend;
use postgres_types::private::BytesMut;
use std::error; use std::error;
use tokio::io::{AsyncRead, AsyncWrite}; use tokio::io::{AsyncRead, AsyncWrite};
@ -34,17 +33,19 @@ impl<'a> Drop for Transaction<'a> {
return; return;
} }
let mut buf = BytesMut::new();
let query = if self.depth == 0 { let query = if self.depth == 0 {
"ROLLBACK".to_string() "ROLLBACK".to_string()
} else { } else {
format!("ROLLBACK TO sp{}", self.depth) format!("ROLLBACK TO sp{}", self.depth)
}; };
frontend::query(&query, &mut buf).unwrap(); let buf = self.client.inner().with_buf(|buf| {
frontend::query(&query, buf).unwrap();
buf.take().freeze()
});
let _ = self let _ = self
.client .client
.inner() .inner()
.send(RequestMessages::Single(FrontendMessage::Raw(buf.freeze()))); .send(RequestMessages::Single(FrontendMessage::Raw(buf)));
} }
} }