feat: add admin FORCE_RELOAD command

[jkaczman]

Adds an admin FORCE_RELOAD command, which acts as a normal RELOAD plus guarantees all in-flight transactions are terminated prior to the command's conclusion (and subsequent transactions are sent to the new Pools).

TODO:

• Benchmark

PR for documentation of this command: pgdogdev/docs#111
Fixes #1472.

[codecov]

Codecov Report

❌ Patch coverage is 97.63780% with 3 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
pgdog/src/backend/pool/connection/mod.rs	88.23%	2 Missing ⚠️
pgdog/src/frontend/client/query_engine/query.rs	91.66%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[jkaczman]

I wanted to get some input on this before implementing a solution for handling errors. I think that if we're guaranteeing all transactions are terminated, then we probably want to exit early with an error to the client and not reload.

However, what if there were some circumstance (e.g. no free connection slots on the server), which prevents us from terminating forever? To handle cases like that, should we have two different variants of this command (or perhaps a parameter) to "override" a potential failure in terminating a transaction?

[levkk]

I'm curious what you think about an alternative approach.

Instead of making Postgres cancel queries, do it from our side. You can add a select! here, for example:

pgdog/pgdog/src/frontend/client/query_engine/query.rs

Lines 69 to 73 in 7594279

	let result = safe_timeout(
	query_timeout,
	self.client_server_exchange(context, query_planner),
	)
	.await;

The other branch can be a CancellationToken you get from Cluster (should be in-scope already) that FORCE_RELOAD cancels; that way, you can make this command pool-specific if you wanted to. This will stop waiting on a in-flight query. Once cancelled, execute a function on the server connection to force-close it.

For clients that are idle-in-transaction, you could do the same here:

pgdog/pgdog/src/frontend/client/mod.rs

Line 658 in 7594279

match safe_timeout(idle_timeout, self.stream_buffer.read(&mut self.stream)).await {

except you just send the client an error and exit, no need to do anything to the server connections I don't think; the Guard will clean up the abandoned transaction.

Do all of this after setting the pool to offline, so no new queries are accepted.

I'm thinking this should be faster than opening a new conn to Postgres for each connection pool, and also guarantees that this command is processed by us. pg_terminate_backend may not terminate the connection immediately (although it usually does).

[jkaczman]

I definitely agree; I wasn't a fan of it either. Thanks!

I pushed a commit with some revisions representing a draft of your suggestions. Also, I added another test for cancelling during a query using pg_sleep (in addition to an idle transaction).

that way, you can make this command pool-specific if you wanted to

I added internal support for per-pool cancellation, but it definitely adds some extra complexity (the Box::pin, futures::select_all, a Binding method, ..). Is this something we want to support in the command now, or is that structure we want for potentially doing something in the future?

[levkk]

I think you should make these Cluster-scoped. Each client will be connected to at most one cluster, so you won't need a Vec anymore (1 less allocation in the hot path).

Also, you're cancelling all pools every time, so having it live on the pools isn't necessary, just the cluster will do.

[levkk]

Good stuff!!