Use the client buffer for cleanup

Author: sfackler

tokio-postgres/src/portal.rs

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

tokio-postgres/src/statement.rs

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

tokio-postgres/src/transaction.rs

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