High-level API for cancellation

[sunng87]

We already have low-level API for cancellation ready. It has necessary handler to respond CancelRequest message. However, the application user will still need to wrap everything together in order to get it work. The idea is to provide a higher level API that has all necessary part for an automatic implementation of query cancellation.

It will carry following functionality or components:

• Backend Key generation, for both Protocol 3.0 and 3.2 secret key
• Session management: having a centralized registry for all active session so we can find the ClientInfo via backend key
• Capability to cancel an ongoing on_query execution in SimpleQueryHandler and ExtendedQueryHandler

Non-functional requirement:

1. We want to make this high-level API completely optional, it can be turned on via maybe handler wrappers or some toggles. Users are allowed to have their own session management, for instance it can be a distributed session manager, and we provide this high level API as reference implementation.
2. We want to try best not to make breaking changes to current APIs.

[Lilit0x]

Okay. I think this is pretty clear except for the non functional requirements. I'm struggling to see how to not make breaking chnages to the API because, let's say I intend to use a wrapper around SimpleQueryHandler:

pub struct CancelQueryHandler<H, R> {
    inner: H,
/// thsi is the registry you mentioned
    registry: Arc<R>,
}

impl<H, R> SimpleQueryHandler for CancelQueryHandler<H, R>
where
    H: SimpleQueryHandler,
    R: Registry,
{
    async fn on_query<C>(&self, client: &mut C, query: Query) -> PgWireResult<()> {
// create token, then, register it, then call do_query()
// the problem is that I still have to chnage the signature of do_query to accept the cancellation token
    }
}

Another thing is to create a new trait say, CancelSimpleQueryHandler, but then I have to aslo do it for ExtenedQueryHndler, so. i'll have to manage 4 traits for it.

This is what comes to me for now, I'll spend more time ruminating over it to see if I can get something better.
I hope you get what I'm saying.

[sunng87]

I have some rough idea about the design:

1. We will have a SessionManager trait acts as the registry in your example. It's responsible for:
   1. Generate process id and secret key in StartupHandler implementations. Especially for: https://github.com/sunng87/pgwire/blob/master/src/api/auth/mod.rs#L318 We will need a way to access SessionManager from various StrartupHandler implementations. Maybe we can add a function in StartupHandler trait like session_manager() -> Arc<dyn SessionManager>.
   2. Maintain mapping between the token and (pid, secret key) so we can get the token in CancelHandler implementation
   3. Clean up the mapping when connection closed
2. Create wrappers for SimpleQueryHandler and ExtendedQueryHandler implementations, that uses tokio::select for token and inner.do_query in its do_query. I'm not sure if we can store the token in ClientInfo. Note that in Postgres when we cancel a query, we actually cancel a connection, so we reuse the token if possible.

[Lilit0x]

This is straightforward, I'll get to work on it.

[sunng87]

@Lilit0x Thank you! Let me know if you have any questions.