Pick connections based on batch first statement's shard

[Ten0]

Fixes: #448
Helps #468

To achieve this, it was useful to rework the BatchValues trait.
With this rework, it's also possible to pretty easily provide cloneable Iterators over ValueList as BatchValues (through BatchValuesFromIterator) - so #499 may get closed if this gets merged (or this PR will require update if #499 gets merged before).

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 added appropriate Fixes: annotations to PR description.

[Ten0]

I'm not sure if this (plus the added Hash for Token) is enough to resolve #468, insight welcome.

[Ten0]

Hmm looks like it's not, as this specifically asks for shard, and not only token (which looks like it would improve batching abilities).
This is a good first step though.

[Ten0]

It's unclear for me how this should be done and likely would conflict with #491.

[psarna]

Don't worry about conflicts with #491, it's a larger effort and we'll make sure that it works correctly with all previously added features. As for #468, our driver is internally capable of computing which shard owns a particular token (if all the information is properly fetched and available), take a look at this test case of a Sharder:

scylla-rust-driver/scylla/src/routing.rs

Lines 144 to 166 in fd06928

	#[cfg(test)]
	mod tests {
	use super::Token;
	use super::{ShardCount, Sharder};
	use std::collections::HashSet;
	#[test]
	fn test_shard_of() {
	/* Test values taken from the gocql driver. */
	let sharder = Sharder::new(ShardCount::new(4).unwrap(), 12);
	assert_eq!(
	sharder.shard_of(Token {
	value: -9219783007514621794
	}),
	3
	);
	assert_eq!(
	sharder.shard_of(Token {
	value: 9222582454147032830
	}),
	3
	);
	}

So if you go deep enough, it's possible to expose this sharding information publicly.

[Ten0]

Sharder appears to be a property of a Node, so it looks like I'd have to go through building the query plan, getting the corresponding node (what if there are several nodes output by the plan? just take the first one?) then get the sharder for the node and eventually return something like a (node, shard) pair, but I don't know what we could use as node identifier, in order to allow users to group by that.

It looks like providing an interface for computing the token is a good first step still, as it will allow grouping at least by partition if not by shard, and exposing the shard interface can be done in a later PR.
As it's unlikely I'll have time for implementing the shard interface anytime soon, I'd suggest moving forward with this without this additional feature.

[cvybhu]

Alright I had a go at the patch. My initial reaction was that it's way too complicated, but after thinking this through I think I agree that this is the way. Here is a summary to gather thoughts and make reviewing easier for others.

We want a few things from the new BatchValues trait:

1. We want to be able to write all of them to the request BufMut
2. We want to be able to get SerializedValues for one of the statements inside the batch in order to calculate the token. Being able to get them only for the first ValueList would work for now, but maybe in the future we will improve the heuristic - it would be good to be able to get them for any chosen statement.
3. Writing to request and getting the values must be done at separate times, because calculating token happens in Session::batch, while request serialization is much deeper. That's why any iterators of ValueList must be Clone.
4. We want to get the number of BatchValues in Session::batch to validate that the user passed values for each statement inside the batch. For all common implementations (tuples, slices, vecs, etc) this is easy to get. For iterators they have to be an ExactSizeIterator or we could do .clone().count()
5. BatchValues should be creatable from an iterator of ValueLists to avoid collecting it and needlessly allocating

Seeing these requirements one could attempt to write:

trait BatchValues {
    type BatchValuesIter: BatchValuesIterator;

    fn iter(&self) -> Self::BatchValuesIter;
    fn len(&self) -> usize;
}

trait BatchValuesIterator: Clone {
    fn write_next_to_request (&mut self,buf:  &mut impl BufMut) -> Option<()>;
    fn next_serialized(&mut self) -> Option<SerializedResult<'_>>;
}

But there is a problem when the type BatchValuesIter has a refernce and requires a lifetime.
With GATs we could write:

trait BatchValues {
    type BatchValuesIter<'a>: BatchValuesIterator<'a>;

    fn iter<'a>(&'a self) -> Self::BatchValuesIter<'a>;
}

But because there are no GATs we end up with the workaround (which is really cool btw).

Looks very nice, I will read the rest and submit a review.

[Ten0]

Alright I had a go at the patch. My initial reaction was that it's way too complicated, but after thinking this through I think I agree that this is the way.

That was exactly my thought process, thanks a lot for the detailed write-up! 😃

which is really cool btw

Yeah at this point it has saved us more than once within Diesel 😊

[cvybhu]

Looks good, I only left some style comments.

Style is highly subjective, but let's try to keep it consistent with the rest of the code.

[cvybhu]

I'm worried that our tests might not test all code paths here.
Everywhere I look the first statement in a batch is shard aware.

In cases where the first PreparedStatement is not shard aware calculate_token will fail, This failure seems to be handled gracefully but still it would be nice to actually test it.

[Ten0]

In cases where the first PreparedStatement is not shard aware calculate_token will fail

Returns Ok(None), I'm not sure whether to even call that a fail ^^

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

Lines 1346 to 1348 in e484adc

	if !prepared.is_token_aware() {
	return Ok(None);
	}

Anyway, I'm just handling it the exact same way as is done in execute_paged:

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

Line 887 in dea8166

token: self.calculate_token(ps, first_serialized_value)?,

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

Lines 740 to 743 in e484adc

	let token = self.calculate_token(prepared, &serialized_values)?;
	let statement_info = Statement {
	token,

[cvybhu]

There is a merge conflict with scylla/src/transport/session.rs , please rebase.

[cvybhu]

Also could you squash these two commits together, both of them contain changes related to review comments but only one is named like it does.

[cvybhu]

I'm almost ready to merge this, but it would be good for someone else to take a look as well, @piodul @psarna .

[psarna]

@Ten0 could you rebase and get rid of the branch conflicts? @cvybhu looks good to me, I think I can follow the logic behind GAT workaround, and I see that you reviewed the code thoroughly, so feel free to merge once rebased (and the minor clippy issue is fixed)

[cvybhu]

Apart from rebasing, one more thing I would love to see before merging is #508 (comment)
The new interface gives us an iterator, but we don't test that after running out of items it returns None.
It isn't strictly required, but it would make the test complete and should be quick to add.

I will be on vacation next week, so I won't be able to merge it then.

[cvybhu]

The None case could be added to the other cases as well, but let's avoid spending extra time on the details.