feat: add a buffer_stats() method on Connection to return allocated read/write buffer sizes

sqlx-core/src/any/connection/backend.rs

@@ -76,6 +76,13 @@ pub trait AnyConnectionBackend: std::any::Any + Debug + Send + 'static {
     /// [`Connection::shrink_buffers()`]: method@crate::connection::Connection::shrink_buffers
     fn shrink_buffers(&mut self);
 
+    /// Forward to [`Connection::buffer_stats()`].
+    ///
+    /// [`Connection::buffer_stats()`]: method@crate::connection::Connection::buffer_stats
+    fn buffer_stats(&self) -> Option<crate::net::BufferStats> {
+        None
+    }
+
     #[doc(hidden)]
     fn flush(&mut self) -> BoxFuture<'_, crate::Result<()>>;
 

sqlx-core/src/any/connection/mod.rs

@@ -139,6 +139,10 @@ impl Connection for AnyConnection {
         self.backend.shrink_buffers()
     }
 
+    fn buffer_stats(&self) -> Option<crate::net::BufferStats> {
+        self.backend.buffer_stats()
+    }
+
     #[doc(hidden)]
     fn flush(&mut self) -> impl Future<Output = Result<(), Error>> + Send + '_ {
         self.backend.flush()

sqlx-core/src/connection.rs

@@ -159,6 +159,17 @@ pub trait Connection: Send {
     /// allow the buffers to shrink.
     fn shrink_buffers(&mut self);
 
+    /// Returns statistics about the connection's internal buffer allocation.
+    ///
+    /// This can be used to monitor memory usage per connection. The default buffer
+    /// capacity is 8KB for both read and write buffers, but they may grow to
+    /// accommodate large queries or result sets.
+    ///
+    /// Returns `None` for databases that don't use buffered sockets (like SQLite).
+    fn buffer_stats(&self) -> Option<crate::net::BufferStats> {
+        None
+    }
+
     #[doc(hidden)]
     fn flush(&mut self) -> impl Future<Output = Result<(), Error>> + Send + '_;
 

sqlx-core/src/net/mod.rs

@@ -2,5 +2,6 @@ mod socket;
 pub mod tls;
 
 pub use socket::{
-    connect_tcp, connect_uds, BufferedSocket, Socket, SocketIntoBox, WithSocket, WriteBuffer,
+    connect_tcp, connect_uds, BufferStats, BufferedSocket, Socket, SocketIntoBox, WithSocket,
+    WriteBuffer,
 };

sqlx-core/src/net/socket/buffered.rs

@@ -326,4 +326,41 @@ impl ReadBuffer {
             self.available = BytesMut::with_capacity(DEFAULT_BUF_SIZE);
         }
     }
+
+    /// Returns the current allocated capacity of this buffer in bytes.
+    pub fn capacity(&self) -> usize {
+        self.read.capacity() + self.available.capacity()
+    }
+}
+
+/// Statistics about connection buffer allocation.
+///
+/// This can be used to monitor memory usage per connection for observability purposes.
+/// The default buffer capacity is 8KB for both read and write buffers, but they may grow
+/// to accommodate large queries or result sets.
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
+pub struct BufferStats {
+    /// Allocated capacity of the write buffer in bytes.
+    pub write_buffer_capacity: usize,
+    /// Allocated capacity of the read buffer in bytes.
+    pub read_buffer_capacity: usize,
+}
+
+impl WriteBuffer {
+    /// Returns the current allocated capacity of this buffer in bytes.
+    pub fn capacity(&self) -> usize {
+        self.buf.capacity()
+    }
+}
+
+impl<S: Socket> BufferedSocket<S> {
+    /// Returns statistics about the current buffer allocation.
+    ///
+    /// This can be useful for monitoring memory usage per connection.
+    pub fn buffer_stats(&self) -> BufferStats {
+        BufferStats {
+            write_buffer_capacity: self.write_buf.capacity(),
+            read_buffer_capacity: self.read_buf.capacity(),
+        }
+    }
 }

sqlx-core/src/net/socket/mod.rs

@@ -4,7 +4,7 @@ use std::path::Path;
 use std::pin::Pin;
 use std::task::{ready, Context, Poll};
 
-pub use buffered::{BufferedSocket, WriteBuffer};
+pub use buffered::{BufferStats, BufferedSocket, WriteBuffer};
 use bytes::BufMut;
 use cfg_if::cfg_if;
 

sqlx-mysql/src/any.rs

@@ -61,6 +61,10 @@ impl AnyConnectionBackend for MySqlConnection {
         Connection::shrink_buffers(self);
     }
 
+    fn buffer_stats(&self) -> Option<sqlx_core::net::BufferStats> {
+        Connection::buffer_stats(self)
+    }
+
     fn flush(&mut self) -> BoxFuture<'_, sqlx_core::Result<()>> {
         Connection::flush(self).boxed()
     }

sqlx-mysql/src/connection/mod.rs

@@ -136,4 +136,8 @@ impl Connection for MySqlConnection {
     fn shrink_buffers(&mut self) {
         self.inner.stream.shrink_buffers();
     }
+
+    fn buffer_stats(&self) -> Option<sqlx_core::net::BufferStats> {
+        Some(self.inner.stream.buffer_stats())
+    }
 }

sqlx-postgres/src/any.rs

@@ -63,6 +63,10 @@ impl AnyConnectionBackend for PgConnection {
         Connection::shrink_buffers(self);
     }
 
+    fn buffer_stats(&self) -> Option<sqlx_core::net::BufferStats> {
+        Connection::buffer_stats(self)
+    }
+
     fn flush(&mut self) -> BoxFuture<'_, sqlx_core::Result<()>> {
         Connection::flush(self).boxed()
     }

sqlx-postgres/src/connection/mod.rs

@@ -234,6 +234,10 @@ impl Connection for PgConnection {
         self.inner.stream.shrink_buffers();
     }
 
+    fn buffer_stats(&self) -> Option<sqlx_core::net::BufferStats> {
+        Some(self.inner.stream.buffer_stats())
+    }
+
     #[doc(hidden)]
     fn flush(&mut self) -> impl Future<Output = Result<(), Error>> + Send + '_ {
         self.wait_until_ready()