Memory leak / unlimited caching kills the server

[MartinKavik]

Bug Description

The connection (pool) "remembers" the query with the biggest memory size. It means when you insert some binary data (e.g. 100 MB) then it doesn't free the memory even though the data are already stored in the database.

I've tried to call persistent(false) or set various settings like max_lifetime but nothing really helped.

I've noticed the bug because my Rust app that uses sqlx through SeaORM was crashing on the server with the error "out of memory". The app uploads temporary binary data (~1-40MB) to Postgres and it allocates hundreds of MBs because of the sqlx bug until I set the maximum number of connections to a lower number with aggressive short lifetime and min_connections 0.

I've tried to debug it by myself by no luck unfortunately.

Minimal Reproduction

https://github.com/MartinKavik/sqlx-memory-debug
(instructions + a screenshot in the README)

Info

• SQLx version: 0.6
• SQLx features enabled: [ "runtime-tokio-rustls" , "postgres" ]
• Database server and version: Postgres 13.5 or newer (other DBs not tested)
• Operating system: desktop Kubuntu, server Ubuntu
• rustc --version: rustc 1.67.1 (d5a82bbd2 2023-02-07)

Thank you!

[DominicWrege]

Since postgres 15.2 my application grows in terms of memory consumption (2.5 GB)

[abonander]

It's not a leak but it is perhaps suboptimal behavior.

It has nothing to do with prepared statement caching; it's simply because we push whole messages to the stream's write buffer, growing it if necessary, but once those messages are flushed the excess capacity is retained because it's likely for that capacity to be needed again soon.

Ensuring we always push whole messages to the buffer helps prevent issues with cancellation, as we should never send a truncated or garbled message that might confuse the server.

However, if you're just using the pool directly it's unlikely that it will use the same connection twice in a row, so you have the issue of duplicating this excess allocation on potentially every query. That's likely why you're hitting OOM conditions.

As a temporary fix, you could .acquire() a connection from the pool, use it for this query, then .detach() and .close() it:

let mut conn = pool.acquire().await?;

sqlx::query!("INSERT INTO test_table VALUES ($1)", data)
    .execute(&mut *conn)
    .await?;

conn.detach().close().await?;

This may be error-prone if you're running with max_connections close to the Postgres server's connection limit (100) as the pool may open a new connection while this one is still being closed. We should probably add a direct .close() method to PoolConnection for situations like this.

You could also push smaller chunks with separate queries in a transaction.

You could use copy_in_raw() and feed the binary data directly to the database, although avoid doing one .send() call with the full thing as that would have the same problem. The COPY IN statement and binary format are documented here: https://www.postgresql.org/docs/current/sql-copy.html

It's a shame we don't have an API for large objects because I think that'd be the ideal thing to use here. You could use the server-side functions though.

As for a real solution, I'm not sure. We certainly could be smarter about how we manage the connection buffers, perhaps tracking the actual capacity we're using over time and shrinking it if it's not being used. That logic could probably live here, though we likely don't want to do that every flush as that could cause a lot of thrashing in the allocator.

We could add a .shrink_buffers() method to the Connection trait so the application can suggest when to free up memory, although again if you're calling that too often then the overhead of hitting allocator could become significant.

The buffering itself could be smarter such that we don't push whole messages to a single Vec<u8> but that would be a significant refactor.

[MartinKavik]

@abonander Thanks for detailed and interesting response!

---

You could also push smaller chunks with separate queries in a transaction.

Yes, it was my first workaround but I didn't like the changed app code.

---

However, if you're just using the pool directly

sqlx is hidden under the SeaORM layer in my case so basically no direct access is possible.

---

As for a real solution, I'm not sure. We certainly could be smarter about how we manage the connection buffers, perhaps tracking the actual capacity we're using over time and shrinking it if it's not being used. That logic could probably live here, though we likely don't want to do that every flush as that could cause a lot of thrashing in the allocator.

I was "playing" with the code in write_and_flush.rs but changing capacity of the buffer didn't come to my mind for some reasons 🤦

I've tested my app with slightly modified code in that file:

// ...
impl<'a, S> Drop for WriteAndFlush<'a, S> {
    fn drop(&mut self) {
        let buf_vec = self.buf.get_mut();
        // clear the buffer regardless of whether the flush succeeded or not
        buf_vec.clear();
        buf_vec.shrink_to(512); 
    }
}

and so far it seems to work as expected 👍
512 seems to be the default capacity and it's enough for almost all queries according my manual testing with my app.

If anybody want to test it, the change is in my fork, in the branch created from the tag 0.6.2:

[patch.crates-io]
sqlx-core = { git = "https://github.com/MartinKavik/sqlx", branch = "fix/buf_limit_512_from_v062" }

I can image I/we can just add that .shrink_to call with 512 constant replaced with a configurable (env?) variable to resolve this issue until proper solution is designed because it seems to be a bit edge-case but still a blocking edge-case for people like me. What do you think?

[abonander]

I'm adding a .shrink_buffers() method to Connection but it's going to have to land in 0.7.0 because we're in the middle of an alpha cycle for that release.

The default buffer size is 8192, that's pretty much what all of Tokio, async-std and std use for their buffered I/O.

[abonander]

I don't want to hardcode a shrink call to the buffer every time it's flushed because that could cause a lot of unnecessary thrashing in the global allocator, it could especially cause fragmentation issues in an allocator like glibc's.