Add Bound{Statement,Batch}

[nrxus]

This PR adds a version of BoundStatement and BoundBatch to hopefully help inspire your future design of these structs.

As discussed in: #941, you all wish to design these features yourselves (for totally understandable reasons) so this PR is here just to help inspire the future implementation, perhaps discuss decisions/API, or whatever else I may be able to help with. I will likely use this branch in my own work to solve my need for this API temporarily until an official version is up.

BoundStatement

pub struct BoundStatement {
    pub(crate) prepared: PreparedStatement,
    pub(crate) values: SerializedValues,
}

This is probably the easiest to understand. It is just a PreparedStatement with a bound SerializedValues. The only public way to create this struct is through PreparedStatement::bind(self, impl: SerializeRow), thus making it externally type safe at creation, while internally it erases the type and just keeps the byte buffer (SerializedValues).

This struct is useful to:

• Keep a prepared statement and its values together in a non-generic way by type erasing at the moment of creation.
• Allow users to get the Token related to the prepared statement + values combo without needing to double-serialize (once to get the token and then again when executing the statement).

All existing internal APIs where a PreparedStatement and its relevant values where passed in together are rewritten to use BoundStatement instead while all external APIs remain the same so as to not introduce breaking changes.

BoundBatch

pub struct BoundBatch {
    pub(crate) buffer: Vec<u8>,
    pub(crate) prepared: HashMap<Bytes, PreparedStatement>,
    first_prepared: Option<(PreparedStatement, Token)>,
    pub(crate) statements_len: u16,
    
    /* snip the less relevant fields */
}

This is where things begin to get interesting. While a naive implementation of BoundBatch that simply keeps a Vec<PreparedStatement> could work, that'd take away some of the current niceties of how Batch is serialized, namely when serializing a batch and its values it doesn't create a small buffer for each serialized row but instead serializes every statement + value into one joined buffer thus avoiding multiple small allocations. To let this current optimization live-on while still allowing each statement + value to be passed in sequentially the BoundBatch instead keeps a single buffer where it serializes each statement + value as they are passed in. When the bound batch is executed this large buffer is copied into the request buffer in one go. Since we are serializing as we go, we need to keep track a few extra details:

• A map of prepared statements, key'd by its ID: this allows batches to re-prepare its statements when the DB returns with an error about a statement not being prepared but expected it to be
• The first statement and its token: if that statement was prepared, so that the batch can use it for sharding purposes (batch goes to the shard of the first prepared statement)
• The length of its statements: needed to make sure we haven't reached the limit and so we can serialize the length when making the request.

This struct is useful to:

• Keep a batch and its values together in a non-generic way by type erasing at the moment of creation.
• More fool-proof in that the number of statements and values never go out-of-sync
• Allows adding a BoundStatement to avoid re-serializing values . This is useful because it is common for batches to be more efficient when they are all for the same token. So now a user can make a BoundStatement, use it to calculate its token, and based on that decide what BoundBatch instance to put the BoundStatement into. BoundBatch can easily copy the serialized bytes out of BoundStatement thus saving CPU resources.

Note that this implementation of BoundBatch allows for unprepared statements to be added if and only if there are no values associated for the unprepared statement. This is done to follow the existing logic of Batch where any unprepared statement that had values would be prepared prior to doing the request, now we are just making it explicit for the user to do have to prepared it if they want to use BoundBatch. The internal implementation of executing a Batch has been rewritten to first prepare the statements in the Batch (just as it used to do but earlier now) and then create a BoundBatch. This allows the logic of executing a batch to be all in one place so that all the existing tests of Batch end up testing BoundBatch execution as well.

Execute trait

This is a sealed trait (users of the crate can call its methods but not implement the trait) for types that can be executed on a Session without any additional values. For now this is directly implemented only on BoundBatch. Existing methods have been rewritten to use this trait but externally nothing has changed other than this trait being usable publicly. (so someone can call my_bound_batch.execute(&session).

ExecutePageable trait

Another sealed trait for types that be executed on a Session but are aware of pagiination. This is implemented directly on BoundStatement and (Statement, impl SerializeValeus). Any type that implements ExecutePageable will also auto-implement Execute with the implementation calling for the pageable methods with no page limit, and no initial paging state. This allows for BoundStatement and (Statement, impl SerializeValeus) to have the same three ways to be called now: no pagination from the start, no pagination from a saved point, and pagination from a saved point. This is done via a single method that has a const generic but could easily be refactored to two methods instead (that call for the generic method under the hood).

Other random notes

Because the scylla-cql crate is already in v1.x, and the Batch struct in that crate has its definition fully public, I couldn't change it at all so instead I had to create a new BatchV2 . I left the current implementation of Batch (de)serialization and instead made requests with the batch opcode (de)serialize using BatchV2. This allows code to compile but it is technically perhaps a breaking change in that if a foreign crate relies on scylla-cql deserializing to Batch instead of the new BatchV2 then their code will break at runtime. I can't think of a reason why someone would do that but I figured it was worth bringing up anyway.

Fixes: #941

Pre-review checklist

• I have split my patch into logically separate commits.
• All commit messages clearly explain what they change and why.
• I added relevant tests for new features and bug fixes.
• All commits compile, pass static checks and pass test.
• PR description sums up the changes and reasons why they should be introduced.
• I have provided docstrings for the public items that I want to introduce.
• I have adjusted the documentation in ./docs/source/.
• I added appropriate Fixes: annotations to PR description.

[github-actions]

cargo semver-checks found no API-breaking changes in this PR.
Checked commit: 2b5f6c6

[wprzytula]

Other random notes

Because the scylla-cql crate is already in v1.x, and the Batch struct in that crate has its definition fully public, I couldn't change it at all so instead I had to create a new BatchV2 . I left the current implementation of Batch (de)serialization and instead made requests with the batch opcode (de)serialize using BatchV2. This allows code to compile but it is technically perhaps a breaking change in that if a foreign crate relies on scylla-cql deserializing to Batch instead of the new BatchV2 then their code will break at runtime. I can't think of a reason why someone would do that but I figured it was worth bringing up anyway.

Your concern is justified. The scenario indeed may surprise the users with unexpected runtime behaviour.

There exists a general solution to this problem: all affected items from scylla-cql (by affected I mean having their logical behaviour altered) must be replaced with v2 versions. In this case, we could duplicate Request along with all functions that operate on it. Then, scylla would only use v2, whereas scylla-cql users would stay at the previous versions.

If you say that it's a lot of duplicated code, then I agree. If you say it's a huge effort to duplicate everything, I also agree.

Therefore, we likely prefer a non-general solution in this case, but rather one that is suited for this particular case. @Lorak-mmk has some idea IIUC.

[nrxus]

I've began using this branch at my job and we've noticed a very clear drop in how long it takes us to process our data and reduced our number of scaled instances as well.

For context, the specific workload is: we get a giant file with lots of different types of events that we want to save onto different tables, but there is often multiple events that go into the same table and partition. To do this we make batches that we keep in essentially a HashMap<Option<u64>, Batch>>. Using this branch, we've also been able to optimize some of our code to place values into the BoundBatch as we transform our data into our struct that implements SerializeRow instead of needing to keep a vector of all them to have them when the batch gets executed as we needed to do in the previous API.

Our average processing time for a file went from ~150ms to ~50ms, so a drop in 66%. Of course this is anecdotal and specific to my use case of batching so YMMV.

[wprzytula]

I stopped the review on the third commit for now.

[wprzytula]

💭 I think we'd prefer not to require cloning PreparedStatement (even it should be quite cheap) in execute_{unpaged,single_page}(). Could we instead introduce (for internal, but perhaps also external use) borrowed BoundStatements? Such that would borrow PreparedStatements, but own SerializedValues.

[wprzytula]

User might like to do something like this:

let prepared = session.prepare(...).await?;

let mut bound = Vec::new();
bound.push(prepared.bind_by_ref(...));
bound.push(prepared.bind_by_ref(...));
...

let results = futures::future::try_join_all(
    bound.into_iter()
        .map(|bound| session.execute_bound(bound))
)?;

[nrxus]

Yeah that should be very doable. I was assuming that cloning a PreparedStatement was cheap enough as to not matter but I can add a BoundStatementRef or something like that.

[nrxus]

I have modified the existing BoundStatement to hold a Cow<PreparedStatement> such that it allows a borrowed or an owned version of PreparedStatement. I then modified the existing PreparedStatement::bind to borrow the prepared statement by reference and added a new method that moves it instead.

[wprzytula]

❓ The commit message says that the BoundStatement is internal-only. At the same time, it's pub and it has pub methods.

[nrxus]

There is no public way to create it, the only method PreparedStatement::bind is pub(crate). But I am fine making the struct itself pub(crate) as well for this commit and then undoing it in the commit that makes it public.

[wprzytula]

🔧 The previous method names (e.g., extract_partition_key_and_calculate_token()), were verbose but informative. Their names warned that they required computations; the new names suggest getter-like cheap operation.

[nrxus]

Fair enough, and I am okay changing it back but if I could push back just slightly, the previous one was more expensive because it also serialized the values as part of that calculation. For this method the values have already been serialized so it is a cheaper calculation. Additionally based on the types it is already implied that some sort of calculation needs to happen at runtime since the operation can fail. Although... looking at the code it still a bit complex so I am fine adding calculate back in there or something like that.

[wprzytula]

❓ What does the name of this method mean?

[nrxus]

I couldn't think of a good name for it and I wanted to get the idea out there rather than bikeshed on it, my bad.

Basically this is the preparing of the batch that occurs at the moment a batch gets executed if any of its statements with values weren't already prepared. This logic already existed but it was done in a latter step and I had to do move it here since I wanted the execution of Batch to be built upon BoundBatch to make it easier to catch bugs instead of duplicating logic. BoundBatch does not allow for unprepared statements with values so this "last minute prepare" was moved here.

[wprzytula]

General thought:
You made PreparedStatement execution use BoundStatement internally. We probably prefer not to do that, because if we don't have to incur BoundStatement type-erasure-involved performance penalty, we don't wan to.
We prefer to use corresponding trait impls (SerializeRow, SerializeValue) to serialize those values straight to the networking buffer in scylla-cql, bypassing the middle man (buffer in BoundStatement).

[nrxus]

General thought: You made PreparedStatement execution use BoundStatement internally. We probably prefer not to do that, because if we don't have to incur BoundStatement type-erasure-involved performance penalty, we don't wan to. We prefer to use corresponding trait impls (SerializeRow, SerializeValue) to serialize those values straight to the networking buffer in scylla-cql, bypassing the middle man (buffer in BoundStatement).

That doesn't seem to match the existing implementation (in main) for executing a prepared statement. As far as I can tell, one of the first things the code does is serialize the statement.

Prepared statement unpaged:

scylla-rust-driver/scylla/src/client/session.rs

Line 1299 in 0085612

let serialized_values = prepared.serialize_values(&values)?;

Prepared statement single page:

scylla-rust-driver/scylla/src/client/session.rs

Line 1318 in 0085612

let serialized_values = prepared.serialize_values(&values)?;

Prepared statement iterating:

scylla-rust-driver/scylla/src/client/session.rs

Line 1438 in 0085612

let serialized_values = prepared.serialize_values(&values)?;

In fact, even for unprepared statements, as long as there are values, they eventually get prepared and serialized, instead of serializing straight into the networking buffer:

scylla-rust-driver/scylla/src/client/session.rs

Lines 1085 to 1086 in 0085612

	let prepared = connection.prepare(statement).await?;
	let serialized = prepared.serialize_values(values_ref)?;

The only case I see where we don't pre-serialize values is unprepared statements without values. But BoundStatement is explicitly for statements that were prepared and have a value so I don't think that matters here.

That all being said, you all do in fact do that for Batch, which is why I implemented the slightly complicated strategy of serializing as statements are pushed into one buffer to avoid the cost of many small buffers as you mentioned.