From afa5502b2a94d4b74c81ea9c206d58270daeb12c Mon Sep 17 00:00:00 2001 From: Matthew Goroff Date: Fri, 13 Feb 2026 21:26:48 -0500 Subject: [PATCH 1/3] initial add buffer stats for connections --- sqlx-core/src/any/connection/backend.rs | 7 +++++ sqlx-core/src/any/connection/mod.rs | 4 +++ sqlx-core/src/connection.rs | 12 ++++++++ sqlx-core/src/net/mod.rs | 3 +- sqlx-core/src/net/socket/buffered.rs | 37 +++++++++++++++++++++++++ sqlx-core/src/net/socket/mod.rs | 2 +- sqlx-mysql/src/any.rs | 4 +++ sqlx-mysql/src/connection/mod.rs | 4 +++ sqlx-postgres/src/any.rs | 4 +++ sqlx-postgres/src/connection/mod.rs | 4 +++ 10 files changed, 79 insertions(+), 2 deletions(-) diff --git a/sqlx-core/src/any/connection/backend.rs b/sqlx-core/src/any/connection/backend.rs index 1ee743a37e..0a06d529fb 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) -> crate::net::BufferStats { + crate::net::BufferStats::default() + } + #[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..10b27588be 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) -> crate::net::BufferStats { + 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..fb823f8f7f 100644 --- a/sqlx-core/src/connection.rs +++ b/sqlx-core/src/connection.rs @@ -159,6 +159,18 @@ 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. + /// + /// For databases that don't use buffered sockets (like SQLite), this returns + /// a default `BufferStats` with zero values. + fn buffer_stats(&self) -> crate::net::BufferStats { + crate::net::BufferStats::default() + } + #[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..f6103da067 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, BufferedSocket, BufferStats, 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..a6bdb6c664 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::{BufferedSocket, BufferStats, 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..0d32a18803 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) -> sqlx_core::net::BufferStats { + 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..851097d01a 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) -> sqlx_core::net::BufferStats { + self.inner.stream.buffer_stats() + } } diff --git a/sqlx-postgres/src/any.rs b/sqlx-postgres/src/any.rs index 997f36db73..ce38349dba 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) -> sqlx_core::net::BufferStats { + 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..3eb4265d2f 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) -> sqlx_core::net::BufferStats { + self.inner.stream.buffer_stats() + } + #[doc(hidden)] fn flush(&mut self) -> impl Future> + Send + '_ { self.wait_until_ready() From d65a5d5efc65b67b70da664bd16acb433282d30a Mon Sep 17 00:00:00 2001 From: Matthew Goroff Date: Fri, 13 Feb 2026 22:01:52 -0500 Subject: [PATCH 2/3] use none instead of default of zeros --- sqlx-core/src/any/connection/backend.rs | 4 ++-- sqlx-core/src/any/connection/mod.rs | 2 +- sqlx-core/src/connection.rs | 7 +++---- sqlx-mysql/src/any.rs | 2 +- sqlx-mysql/src/connection/mod.rs | 4 ++-- sqlx-postgres/src/any.rs | 2 +- sqlx-postgres/src/connection/mod.rs | 4 ++-- 7 files changed, 12 insertions(+), 13 deletions(-) diff --git a/sqlx-core/src/any/connection/backend.rs b/sqlx-core/src/any/connection/backend.rs index 0a06d529fb..51a28f1f56 100644 --- a/sqlx-core/src/any/connection/backend.rs +++ b/sqlx-core/src/any/connection/backend.rs @@ -79,8 +79,8 @@ pub trait AnyConnectionBackend: std::any::Any + Debug + Send + 'static { /// Forward to [`Connection::buffer_stats()`]. /// /// [`Connection::buffer_stats()`]: method@crate::connection::Connection::buffer_stats - fn buffer_stats(&self) -> crate::net::BufferStats { - crate::net::BufferStats::default() + fn buffer_stats(&self) -> Option { + None } #[doc(hidden)] diff --git a/sqlx-core/src/any/connection/mod.rs b/sqlx-core/src/any/connection/mod.rs index 10b27588be..f5be8eff02 100644 --- a/sqlx-core/src/any/connection/mod.rs +++ b/sqlx-core/src/any/connection/mod.rs @@ -139,7 +139,7 @@ impl Connection for AnyConnection { self.backend.shrink_buffers() } - fn buffer_stats(&self) -> crate::net::BufferStats { + fn buffer_stats(&self) -> Option { self.backend.buffer_stats() } diff --git a/sqlx-core/src/connection.rs b/sqlx-core/src/connection.rs index fb823f8f7f..3539d36e3d 100644 --- a/sqlx-core/src/connection.rs +++ b/sqlx-core/src/connection.rs @@ -165,10 +165,9 @@ pub trait Connection: Send { /// capacity is 8KB for both read and write buffers, but they may grow to /// accommodate large queries or result sets. /// - /// For databases that don't use buffered sockets (like SQLite), this returns - /// a default `BufferStats` with zero values. - fn buffer_stats(&self) -> crate::net::BufferStats { - crate::net::BufferStats::default() + /// Returns `None` for databases that don't use buffered sockets (like SQLite). + fn buffer_stats(&self) -> Option { + None } #[doc(hidden)] diff --git a/sqlx-mysql/src/any.rs b/sqlx-mysql/src/any.rs index 0d32a18803..fd546ad408 100644 --- a/sqlx-mysql/src/any.rs +++ b/sqlx-mysql/src/any.rs @@ -61,7 +61,7 @@ impl AnyConnectionBackend for MySqlConnection { Connection::shrink_buffers(self); } - fn buffer_stats(&self) -> sqlx_core::net::BufferStats { + fn buffer_stats(&self) -> Option { Connection::buffer_stats(self) } diff --git a/sqlx-mysql/src/connection/mod.rs b/sqlx-mysql/src/connection/mod.rs index 851097d01a..85c07c0910 100644 --- a/sqlx-mysql/src/connection/mod.rs +++ b/sqlx-mysql/src/connection/mod.rs @@ -137,7 +137,7 @@ impl Connection for MySqlConnection { self.inner.stream.shrink_buffers(); } - fn buffer_stats(&self) -> sqlx_core::net::BufferStats { - self.inner.stream.buffer_stats() + 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 ce38349dba..d2674670b8 100644 --- a/sqlx-postgres/src/any.rs +++ b/sqlx-postgres/src/any.rs @@ -63,7 +63,7 @@ impl AnyConnectionBackend for PgConnection { Connection::shrink_buffers(self); } - fn buffer_stats(&self) -> sqlx_core::net::BufferStats { + fn buffer_stats(&self) -> Option { Connection::buffer_stats(self) } diff --git a/sqlx-postgres/src/connection/mod.rs b/sqlx-postgres/src/connection/mod.rs index 3eb4265d2f..041c5fec4c 100644 --- a/sqlx-postgres/src/connection/mod.rs +++ b/sqlx-postgres/src/connection/mod.rs @@ -234,8 +234,8 @@ impl Connection for PgConnection { self.inner.stream.shrink_buffers(); } - fn buffer_stats(&self) -> sqlx_core::net::BufferStats { - self.inner.stream.buffer_stats() + fn buffer_stats(&self) -> Option { + Some(self.inner.stream.buffer_stats()) } #[doc(hidden)] From 8457a783ec3e27290cb34b15be49f216af457686 Mon Sep 17 00:00:00 2001 From: Matthew Goroff Date: Fri, 13 Feb 2026 22:08:39 -0500 Subject: [PATCH 3/3] fmt --- sqlx-core/src/net/mod.rs | 2 +- sqlx-core/src/net/socket/mod.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/sqlx-core/src/net/mod.rs b/sqlx-core/src/net/mod.rs index f6103da067..fa01a08d57 100644 --- a/sqlx-core/src/net/mod.rs +++ b/sqlx-core/src/net/mod.rs @@ -2,6 +2,6 @@ mod socket; pub mod tls; pub use socket::{ - connect_tcp, connect_uds, BufferedSocket, BufferStats, Socket, SocketIntoBox, WithSocket, + connect_tcp, connect_uds, BufferStats, BufferedSocket, Socket, SocketIntoBox, WithSocket, WriteBuffer, }; diff --git a/sqlx-core/src/net/socket/mod.rs b/sqlx-core/src/net/socket/mod.rs index a6bdb6c664..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, BufferStats, WriteBuffer}; +pub use buffered::{BufferStats, BufferedSocket, WriteBuffer}; use bytes::BufMut; use cfg_if::cfg_if;