Skip to content

chore(spanner): implement dynamic channel pooling and transaction affinity #6728

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
olavloite wants to merge 4 commits into from
Closed

chore(spanner): implement dynamic channel pooling and transaction affinity #6728

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
d82af44
chore(spanner): implement dynamic channel pooling and transaction aff…
olavloite Sep 7, 2026
70b79d8
Merge branch 'main' into spanner-dynamic-channel-pool-and-affinity
olavloite Sep 8, 2026
deb0bbb
fix(spanner): fix build error after macro update
olavloite Sep 8, 2026
98920dc
Merge branch 'main' into spanner-dynamic-channel-pool-and-affinity
olavloite Sep 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 11 additions & 22 deletions src/spanner/src/batch_read_only_transaction.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -168,7 +168,7 @@ impl BatchReadOnlyTransaction {
.partition_query(
request,
crate::RequestOptions::default(),
self.inner.context.channel_hint,
self.inner.context.affinity(),
)
.await?;

Expand Down Expand Up @@ -231,7 +231,7 @@ impl BatchReadOnlyTransaction {
.partition_read(
request,
crate::RequestOptions::default(),
self.inner.context.channel_hint,
self.inner.context.affinity(),
)
.await?;

Expand Down Expand Up @@ -405,15 +405,11 @@ impl Partition {
req: &ExecuteSqlRequest,
gax_options: GaxRequestOptions,
) -> crate::Result<ResultSet> {
let channel_hint = client.next_channel_hint();
let gax_options = client.attach_request_id(gax_options, channel_hint);
let builder = client.execute_streaming_sql(req.clone(), gax_options, None);
let actual_gax_options = builder.options().clone();
let (stream, attempt_start_time) =
Self::execute_partition_stream(client, "ExecuteStreamingSql", || {
client
.execute_streaming_sql(req.clone(), gax_options.clone(), channel_hint)
.send()
})
.await?;
Self::execute_partition_stream(client, "ExecuteStreamingSql", move || builder.send())
Comment thread
olavloite marked this conversation as resolved.
Comment thread
olavloite marked this conversation as resolved.
.await?;
Comment thread
olavloite marked this conversation as resolved.
Comment thread
olavloite marked this conversation as resolved.
Comment on lines +411 to +412

@gemini-code-assist gemini-code-assist Bot Sep 8, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

The move || builder.send() closure is an FnOnce because builder.send() consumes builder. If execute_partition_stream is intended to perform retries (which is likely for transient errors or transaction aborts), it would need to call the closure multiple times, which is not possible with an FnOnce.

Since the stream builder (ExecuteStreamingSql) implements Clone, you can capture builder by reference and clone it inside the closure. This makes the closure an FnMut, allowing it to be called multiple times for retries.

Suggested change
Self::execute_partition_stream(client, "ExecuteStreamingSql", move || builder.send())
.await?;
Self::execute_partition_stream(client, "ExecuteStreamingSql", || builder.clone().send())
.await?;


ResultSet::create(ResultSetParams {
stream,
Expand All @@ -428,8 +424,7 @@ impl Partition {
session_name: req.session.clone(),
transaction_tag: None,
operation: StreamOperation::Query(req.clone()),
channel_hint,
gax_options,
gax_options: actual_gax_options,
Comment thread
olavloite marked this conversation as resolved.
method_name: "ExecuteStreamingSql",
attempt_start_time: Some(attempt_start_time),
operation_start_time: Some(attempt_start_time),
Expand All @@ -443,15 +438,10 @@ impl Partition {
req: &ReadRequest,
gax_options: GaxRequestOptions,
) -> crate::Result<ResultSet> {
let channel_hint = client.next_channel_hint();
let gax_options = client.attach_request_id(gax_options, channel_hint);
let builder = client.streaming_read(req.clone(), gax_options, None);
let actual_gax_options = builder.options().clone();
let (stream, attempt_start_time) =
Self::execute_partition_stream(client, "StreamingRead", || {
client
.streaming_read(req.clone(), gax_options.clone(), channel_hint)
.send()
})
.await?;
Self::execute_partition_stream(client, "StreamingRead", move || builder.send()).await?;
Comment thread
olavloite marked this conversation as resolved.
Comment thread
olavloite marked this conversation as resolved.
Comment thread
olavloite marked this conversation as resolved.
Comment thread
olavloite marked this conversation as resolved.

@gemini-code-assist gemini-code-assist Bot Sep 8, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

Similar to the execute_sql case, the move || builder.send() closure is an FnOnce, which prevents retries within execute_partition_stream. Since StreamingRead builder is also Clone, you should clone it inside the closure to make it an FnMut and support retries.

Suggested change
Self::execute_partition_stream(client, "StreamingRead", move || builder.send()).await?;
Self::execute_partition_stream(client, "StreamingRead", || builder.clone().send()).await?;


ResultSet::create(ResultSetParams {
stream,
Expand All @@ -466,8 +456,7 @@ impl Partition {
session_name: req.session.clone(),
transaction_tag: None,
operation: StreamOperation::Read(req.clone()),
channel_hint,
gax_options,
gax_options: actual_gax_options,
Comment thread
olavloite marked this conversation as resolved.
method_name: "StreamingRead",
attempt_start_time: Some(attempt_start_time),
operation_start_time: Some(attempt_start_time),
Expand Down
9 changes: 1 addition & 8 deletions src/spanner/src/batch_write_transaction.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -209,12 +209,10 @@ impl BatchWriteTransactionBuilder {
/// ```
pub fn build(self) -> BatchWriteTransaction {
let session_name = self.client.session_name();
let channel_hint = self.client.next_channel_hint();
let gax_options = apply_defaults(self.gax_options);
BatchWriteTransaction {
session_name,
client: self.client,
channel_hint,
transaction_tag: self.transaction_tag,
priority: self.priority,
exclude_txn_from_change_streams: self.exclude_txn_from_change_streams,
Expand All @@ -230,7 +228,6 @@ impl BatchWriteTransactionBuilder {
pub struct BatchWriteTransaction {
session_name: String,
client: DatabaseClient,
channel_hint: usize,
transaction_tag: Option<String>,
priority: Priority,
exclude_txn_from_change_streams: bool,
Expand Down Expand Up @@ -291,7 +288,6 @@ impl BatchWriteTransaction {
Ok(BatchWriteResponseStream {
client: self.client,
session_name: self.session_name,
channel_hint: self.channel_hint,
transaction_tag: self.transaction_tag,
priority: self.priority,
exclude_txn_from_change_streams: self.exclude_txn_from_change_streams,
Expand All @@ -317,7 +313,6 @@ impl BatchWriteTransaction {
pub struct BatchWriteResponseStream {
client: DatabaseClient,
session_name: String,
channel_hint: usize,
transaction_tag: Option<String>,
priority: Priority,
exclude_txn_from_change_streams: bool,
Expand Down Expand Up @@ -436,7 +431,7 @@ impl BatchWriteResponseStream {

let stream_result = self
.client
.batch_write(request, self.gax_options.clone(), self.channel_hint)
.batch_write(request, self.gax_options.clone(), None)
.send()
.await;

Expand Down Expand Up @@ -540,8 +535,6 @@ impl BatchWriteResponseStream {
match self.check_retry(error) {
Ok(()) => {
self.retry_count += 1;
// Rotate channel hint only when a retry is confirmed to distribute load across healthy connections.
self.channel_hint = self.client.next_channel_hint();
if let Some(policy) = self.gax_options.backoff_policy() {
let state = RetryState::new(true).set_attempt_count(self.retry_count as u32);
let delay = policy.on_failure(&state);
Expand Down
Loading
Loading