diff --git a/sqlx-core/src/any/connection/backend.rs b/sqlx-core/src/any/connection/backend.rs index 1ee743a37e..51a28f1f56 100644 --- a/sqlx-core/src/any/connection/backend.rs +++ b/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 { + None + } + #[doc(hidden)] fn flush(&mut self) -> BoxFuture<'_, crate::Result<()>>; diff --git a/sqlx-core/src/any/connection/mod.rs b/sqlx-core/src/any/connection/mod.rs index 894b109ccd..f5be8eff02 100644 --- a/sqlx-core/src/any/connection/mod.rs +++ b/sqlx-core/src/any/connection/mod.rs @@ -139,6 +139,10 @@ impl Connection for AnyConnection { self.backend.shrink_buffers() } + fn buffer_stats(&self) -> Option { + self.backend.buffer_stats() + } + #[doc(hidden)] fn flush(&mut self) -> impl Future> + Send + '_ { self.backend.flush() diff --git a/sqlx-core/src/connection.rs b/sqlx-core/src/connection.rs index b5f8138b2b..3539d36e3d 100644 --- a/sqlx-core/src/connection.rs +++ b/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 { + None + } + #[doc(hidden)] fn flush(&mut self) -> impl Future> + Send + '_; diff --git a/sqlx-core/src/net/mod.rs b/sqlx-core/src/net/mod.rs index f9c43668ab..fa01a08d57 100644 --- a/sqlx-core/src/net/mod.rs +++ b/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, }; diff --git a/sqlx-core/src/net/socket/buffered.rs b/sqlx-core/src/net/socket/buffered.rs index 6785e70879..a7ef361a2c 100644 --- a/sqlx-core/src/net/socket/buffered.rs +++ b/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 BufferedSocket { + /// 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(), + } + } } diff --git a/sqlx-core/src/net/socket/mod.rs b/sqlx-core/src/net/socket/mod.rs index 0f9aae61b4..7db217e86e 100644 --- a/sqlx-core/src/net/socket/mod.rs +++ b/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; diff --git a/sqlx-mysql/src/any.rs b/sqlx-mysql/src/any.rs index b0950e0b41..fd546ad408 100644 --- a/sqlx-mysql/src/any.rs +++ b/sqlx-mysql/src/any.rs @@ -61,6 +61,10 @@ impl AnyConnectionBackend for MySqlConnection { Connection::shrink_buffers(self); } + fn buffer_stats(&self) -> Option { + Connection::buffer_stats(self) + } + fn flush(&mut self) -> BoxFuture<'_, sqlx_core::Result<()>> { Connection::flush(self).boxed() } diff --git a/sqlx-mysql/src/connection/mod.rs b/sqlx-mysql/src/connection/mod.rs index 569ad32722..85c07c0910 100644 --- a/sqlx-mysql/src/connection/mod.rs +++ b/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 { + Some(self.inner.stream.buffer_stats()) + } } diff --git a/sqlx-postgres/src/any.rs b/sqlx-postgres/src/any.rs index 997f36db73..d2674670b8 100644 --- a/sqlx-postgres/src/any.rs +++ b/sqlx-postgres/src/any.rs @@ -63,6 +63,10 @@ impl AnyConnectionBackend for PgConnection { Connection::shrink_buffers(self); } + fn buffer_stats(&self) -> Option { + Connection::buffer_stats(self) + } + fn flush(&mut self) -> BoxFuture<'_, sqlx_core::Result<()>> { Connection::flush(self).boxed() } diff --git a/sqlx-postgres/src/connection/mod.rs b/sqlx-postgres/src/connection/mod.rs index d5db20ad05..041c5fec4c 100644 --- a/sqlx-postgres/src/connection/mod.rs +++ b/sqlx-postgres/src/connection/mod.rs @@ -234,6 +234,10 @@ impl Connection for PgConnection { self.inner.stream.shrink_buffers(); } + fn buffer_stats(&self) -> Option { + Some(self.inner.stream.buffer_stats()) + } + #[doc(hidden)] fn flush(&mut self) -> impl Future> + Send + '_ { self.wait_until_ready()