Postgres: Move io to background task.

[joeydewaal]

This PR moves all the io in the postgres driver to a background task. This should fix some long standing issues with cancellation safety and unblock/allow features like query pipelining, CopyBoth mode (#2924) and the ability to cancel queries. This also simplifies the dropping of PgTransaction, PgCopyIn and PgListener. There is still some work to do to increase performance but that's out of the scope of this PR.

Notable changes:

• SocketExt: An extension trait for Socket which has async methods for reading, writing, flushing and closing.
• BufferedSocket now has poll methods for reading messages.
• BufferedSocket has a Sink<&[u8]> impl to send bytes.
• Pipe is added; a temporary stream of responses from the background worker.
• Worker is added; the background worker that handles all the io.
  • The worker now handles ReadyForQuery messages because the wait_until_ready mechanism is removed.
• Shared is added; a shared structure between the connection and background worker.

This is a pretty big change that touches a lot in the postgres driver. I've added a lot of comments to explain my thinking. I'd recommend going commit by commit for reviewing. If there is anything else I can do to make this easier lmk.

Is this a breaking change?

No, this PR only has breaking changes in sqlx-core which is semver exempt. One unit test had to be changed for PgListener because the buffering mechanism is changed. I don't think that makes this a breaking change and I don't think Hyrum's law applies here.

[abonander]

I'll be honest, seeing this posted took the wind out of my sails a bit because I was really looking forward to experimenting with this, but there's so much other work that needs to be done anyway that I'm not going to get too hung up on it. I appreciate the initiative.

However, I don't have the time or energy to review this properly while trying to get 0.9.0 out the door so I think ideally we could land this in 0.9.1 or 0.9.2 since it should be backwards-compatible anyway.

[joeydewaal]

Right, I learned a lot while working on this. I mean there are many ways to solve this problem, if you'd like you can still experiment and see what solution you like best.

I can imagine that this takes some effort to review and don't mind when this is released, I'll rebase when needed.

[joeydewaal]

I know I'm getting way ahead of myself here but I just want to write this down somewhere: I'd be nice if the pool could benefit from query pipelining.

Instead of acquiring a conn -> send query -> wait for results -> release the conn. We could choose a conn and not acquire it so other tasks can still use it. A naive implementation would be to use an AtomicUsize and index into a Box<[DB::Connection]>. If the pool has a small connection count or a lot of tasks want to execute a query, the AtomicUsize would wrap around and send the query to a connection that is potentially still in use (and thus be query pipelining).

This would improve performance and also help with the infamous PoolTimedOut error.

I think the only breaking change this feature would need is to have the Executor trait take a &self.

My naive solution isn't a valid solution since we should have a way to acquire a connection to support Transaction. Maybe we could use an async RwLock or do something with ArcSwap.

I can't think of any other db or pooling crate that supports this behavior so it might be worth experimenting with this.

[abonander]

I have plans to completely re-architect Pool that would also explicitly allow pipelining, as well as significantly improve performance at high contention.

It's my next big project that I hope to land, if not in 0.9.0, then a subsequent release.

[joeydewaal]

Ah looks like you are one step ahead. Exciting stuff!

[joeydewaal]

So I've finally rebased this, here is a small todo list for after this is merged.

• PgCopyIn::send currently doesn't have any back-pressure, I think we can fix this by just waiting for the channel to closed because it's dropped in the bg task but I need to test this.
• Store the write buffer somewhere.
  • Currently a new buffer is allocated on every request which isn't great. We could move to a BytesMut and store it in the connection but that would require the ProtocolEncode trait to be changed from Vec<u8> to BytesMut. I was thinking of some kind of pool of buffers. If we use your custom channels (Austin), then we could potentially store them next to the write buffer so we can reuse those allocations as well.
• Have a custom Stream impl instead of the try_stream! macro. Then the Executor trait could move to an impl Stream instead of BoxStream.
• Instead of sending each response from the connection over the channel, read as many responses as we can and send the batch of responses in one message.
• The background task should be generic over the Socket it uses, which would allow for static dispatch instead of the current Box<dyn Socket>.