scylladb / scylla-rust-driver

Async CQL driver for Rust, optimized for ScyllaDB!
Apache License 2.0
564 stars 99 forks source link

Avoid allocations in `CachingSession::execute` #475

Open colin-grapl opened 2 years ago

colin-grapl commented 2 years ago

Currently CachingSession::execute takes impl Into<Query>, but internally it only ever uses functions/methods that take &Query.

I ran into this because I was trying to pool my service's allocations for Query but realized it was a waste since the clone happens anyways.

So I started by tracking the lifetime of a Query passed into CachingSession::execute

All the way down to (and including) the SerializedRequest::make, it seems that ownership of the Query contents is unnecessary.

A bunch of allocations can be avoided by refactoring things a bit:

  1. CachingSession::execute does not require ownership since, as far as I can tell, there's no unconditional point in the call graph where ownership of the Query string is required
  2. CachingSession::add_prepared_statement could avoid the clone when preparing the session if Session::prepare didn't require impl Into<Query>
  3. Session::query_paged doesn't need Query, which has implications for a whole bunch of apis triggering unnecessary allocations

There are a few options to fix this.

  1. Take impl AsRef<Query>
  2. Take &Query
  3. Refactor Query to use Cow<str> instead of String
  4. Create a new type, QueryRef

Both (1) and (2) share the same problems. They're breaking API changes and they break code like:

session.execute("myquery", &());

I imagine that the ability to provide static strings like that is pretty desirable, so I think it's fair to discount those approaches.

Approach (3) would allow Query to be generic over an owned vs unowned string. Query would become:

#[derive(Clone)]
pub struct Query<'a> {
    pub(crate) config: StatementConfig,

    pub contents: std::cow::Cow<'a, String>,
    page_size: Option<i32>,
}

A helper method to create a Query<'static> would probably be useful.

The main problem I see with this approach is the API breakage would permeate through every API where a Query is taken, and users who are working with Query structs directly will have to add a lifetime (although they could do so with 'static pretty easily).

I actually think this is still the best approach despite that as it lets the caller manage the buffer.

(4) Is probably the best in terms of not having any API breakage

A new type would be created, QueryRef,

#[derive(Clone)]
pub struct QueryRef<'a> {
    pub(crate) config: StatementConfig,

    pub contents: &'a str,
    page_size: Option<i32>,
}

There'd be impls for AsRef<QueryRef> for Query and From<QueryRef> for Query.

Initially I thought this would be the ideal approach but I'm not sure.

Or something else entirely, idk.

colin-grapl commented 2 years ago

https://github.com/colin-grapl/scylla-rust-driver/tree/remove-allocations

Implemented approach (3) here. This lets me manage the buffer in my own code and removes some clones.

psarna commented 2 years ago

@cvybhu please take a look

Lorak-mmk commented 3 months ago

We'll be refactoring whole Session API soon (most likely by introducing Executable trait and BoundStatement) so I don't think it makes sens to work on it now. I'll attach it to umbrella issue about Session API changes so that we take this issue into consideration when designing new API