Add support for CopyBoth mode

[grim7reaper]

Does the groundwork to be able to talk replication with PostgreSQL, i.e. support "replication" connection and CopyBoth "protocol".

Fixes: #2923

[grim7reaper]

To test you can do something like that.

In a local DB, create a table and setup replication for it:

ALTER TABLE <your-table> REPLICA IDENTITY FULL;
CREATE PUBLICATION test_publication FOR TABLE <your-table>;

Then run
use a snippet like that:

use futures::prelude::*;
use sqlx::postgres::{PgConnectOptions, PgCopyBoth, PgPoolOptions, PgReplicationMode};

#[tokio::main]
async fn main() {
    // Connection must be configured with a replication mode!
    let options = PgConnectOptions::new()
        .host("0.0.0.0")
        .username("postgres")
        .password("postgres")
        .replication_mode(PgReplicationMode::Logical);

    let pool = PgPoolOptions::new()
        .connect_with(options)
        .await
        .expect("failed to connect to postgres");
    let starting_pos = "0/1573178"; // CHANGE THIS
    let query = format!(
        r#"START_REPLICATION SLOT "{}" LOGICAL {} ("proto_version" '1', "publication_names" '{}')"#,
        "test_slot", starting_pos, "test_publication",
    );
    let mut duplex = PgCopyBoth::new(&pool, query.as_str())
        .await
        .expect("start replication");
    // Read data from the server.
    while let Some(data) = duplex.next().await {
        println!("data: {:?}", data);
    }
}

Then do some insert in the table, and you should receive some bytes (new row inserted).

[grim7reaper]

Anything I can do to help/make the review easier for you? 🙂

[abonander]

I would prefer to see an API more like PgListener which holds a Pool internally. Maybe call it PgReplicationPool and use the existing functionality to create child pools so it can steal permits from the parent:

impl PgReplicationPool {
    pub async fn connect(url: &str, mode; PgReplicationMode) -> crate::Result<Self> {
        // ...
    }

    pub fn from_pool(pool: PgPool, mode: PgReplicationMode) -> Self {
        let pool_options = pool.options().clone();
        let connect_options = pool.connect_options()
            .clone()
            .replication_mode(mode);

        Self(pool_options.parent(pool).connect_lazy_with(connect_options))
    }

    pub async fn start_replication(&self, statement: &str) -> crate::Result<PgReplication> {
        // ...
    }
}

This is also missing a potential use-case where the user doesn't want to involve a Pool at all.

[abonander]

Splitting it into sender and receiver can be done at a higher level, IMO. You're asking the user to give up a lot of control in what is already a niche, low-level API.

[grim7reaper]

In what way do we give up a lot of control with this approach?
Having two handle seems more flexible (user can group them back into a structure if needed).
+
I thought that it was nice to have a "channel-like" API here as you have the upstream stream of command and the downstream stream of data.

---

If I go the single-object way I either need to have public fields of an into_sth to split it it two, right?

[grim7reaper]

For now I've made a structure with two public field (so that the caller can deconstruct if needed).
Does it sounds good to you or you had something else in mind?

[abonander]

Another reason not to do the sender/receiver split at this level, send_async() is not strictly cancel-safe: zesterer/flume#104 (comment)

[grim7reaper]

Hum good point!

How can I avoid/solve this?

[grim7reaper]

Any suggestion/pointer @abonander ?

[abonander]

How does the user get the CommandComplete responses with the number of rows changed?

[grim7reaper]

I'm not sure there is a CommandComplete involved in the case of CopyBoth: the command is never really complete, it's a stream of WAL events.
Seems like the graceful shutdown is an exchange of CopyDone message.

[Both the backend and the frontend may then send CopyData messages until either end sends a CopyDone message. After the client sends a CopyDone message, the connection goes from copy-both mode to copy-out mode, and the client may not send any more CopyData messages. Similarly, when the server sends a CopyDone message, the connection goes into copy-in mode, and the server may not send any more CopyData messages. After both sides have sent a CopyDone message, the copy mode is terminated, and the backend reverts to the command-processing mode.

[grim7reaper]

Thanks for the thorough review 🙏

I would prefer to see an API more like PgListener which holds a Pool internally. Maybe call it PgReplicationPool and use the existing functionality to create child pools so it can steal permits from the parent:

Ho interesting, didn't know about this "parent" thing. Will give it a try.
What does it mean to "steal permits from the parent"?

Your API definitely looks better, I'll try to rework the PR to go that way.
Not sure I'll have time to address the changes right now, but I'll try to get back to it next week if possible

[grim7reaper]

In what way do we give up a lot of control with this approach?
Having two handle seems more flexible (user can group them back into a structure if needed).
+
I thought that it was nice to have a "channel-like" API here as you have the upstream stream of command and the downstream stream of data.

---

If I go the single-object way I either need to have public fields of an into_sth to split it it two, right?

[grim7reaper]

Hum good point!

How can I avoid/solve this?

[abonander]

What does it mean to "steal permits from the parent"?

When one pool is a child of another, they effectively share a total max_connections.

To acquire a connection from the pool, the acquire() call must first get a permit from the pool semaphore, which represents the right to have a connection. It may then use this permit either to pop a connection from the idle queue, or open a new one if the pool size is below max_connections.

When a child pool's semaphore is exhausted, but it is below max_connections, it attempts to acquire a semaphore permit from the parent pool. When the checked-out connection using that permit is returned to the child pool, the permit is added to the child pool's semaphore, and further calls to acquire() behave normally. When a child pool is closed, any permits remaining in its semaphore are returned to the parent pool.

This does mean that the child pool can potentially exhaust the permits of the parent pool, so it's important that the sum of all child pools' max_connections are not larger than the parent's. This wasn't really a concern when designing the feature because it was only meant to be a part of #[sqlx::test]'s functionality, to allow multiple tests to run concurrently without exhausting the available connections on the server, so child pools were meant to be short-lived.

I thought about just having the child pool share the same semaphore as the parent pool, but that would have been a problem because the child pool couldn't keep connections in its own idle queue while returning permits to the parent, as other pools sharing the same parent wouldn't know about them, so the total number of open connections could easily exceed the maximum.

If exhaustion ends up becoming a problem, I think I would try addressing that by having the child pool return a permit to the parent when it closes a connection, instead of its own semaphore. And I would recommend having a short idle_timeout to make sure child pools can't hog permits for too long.

[grim7reaper]

Thanks for the detailed explanation

If exhaustion ends up becoming a problem, I think I would try addressing that by having the child pool return a permit to the parent when it closes a connection, instead of its own semaphore. And I would recommend having a short idle_timeout to make sure child pools can't hog permits for too long.

I don't think exhaustion would be an issue in this use case, since the replication only uses a single (long-lived) connection.