feat(spanner): add support for static and dynamic channel pooling

src/spanner/README.md

@@ -69,6 +69,31 @@ export SPANNER_EMULATOR_HOST=localhost:9010
 The client builder automatically detects this variable, connects to the emulator
 endpoint, and configures anonymous credentials.
 
+### Configuring the Channel Pool
+
+By default, the client uses a static pool of 4 gRPC channels. You can configure
+the channel pool using `with_channel_pool`:
+
+```rust
+use google_cloud_spanner::client::{Spanner, SpannerBuilderExt};
+use google_cloud_spanner::channel_pool::{DynamicChannelPoolConfig, StaticChannelPoolConfig};
+
+# async fn sample() -> Result<(), google_cloud_spanner::Error> {
+// Custom static pool:
+let spanner = Spanner::builder()
+    .with_channel_pool(StaticChannelPoolConfig::new(8))
+    .build()
+    .await?;
+
+// Or dynamic load-based channel pool:
+let spanner = Spanner::builder()
+    .with_channel_pool(DynamicChannelPoolConfig::new())
+    .build()
+    .await?;
+# Ok(())
+# }
+```
+
 ## Session Management and Client Lifecycle
 
 The Spanner Rust client manages a long-lived multiplexed session under the hood.

src/spanner/src/batch_read_only_transaction.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use crate::channel_pool::ChannelTarget;
 use crate::database_client::DatabaseClient;
 use crate::model::{ExecuteSqlRequest, PartitionOptions, ReadRequest};
 use crate::precommit::PrecommitTokenTracker;
@@ -161,15 +162,15 @@ impl BatchReadOnlyTransaction {
             .set_transaction(selector.clone())
             .set_partition_options(options);
 
+        let target = ChannelTarget::from_affinity_or_hint(
+            self.inner.context.affinity(),
+            self.inner.context.channel_hint,
+        );
         let response = self
             .inner
             .context
             .client
-            .partition_query(
-                request,
-                crate::RequestOptions::default(),
-                self.inner.context.channel_hint,
-            )
+            .partition_query(request, crate::RequestOptions::default(), target)
             .await?;
 
         Ok(response
@@ -224,15 +225,15 @@ impl BatchReadOnlyTransaction {
             .set_transaction(selector.clone())
             .set_partition_options(options);
 
+        let target = ChannelTarget::from_affinity_or_hint(
+            self.inner.context.affinity(),
+            self.inner.context.channel_hint,
+        );
         let response = self
             .inner
             .context
             .client
-            .partition_read(
-                request,
-                crate::RequestOptions::default(),
-                self.inner.context.channel_hint,
-            )
+            .partition_read(request, crate::RequestOptions::default(), target)
             .await?;
 
         Ok(response

src/spanner/src/channel_pool/affinity.rs

@@ -55,16 +55,6 @@ impl TransactionAffinity {
         }
     }
 
-    /// Returns the provided affinity handle, or creates a new default `ReadOnly` affinity if `None`.
-    pub(crate) fn default_read_only(existing: Option<Arc<Self>>) -> Arc<Self> {
-        existing.unwrap_or_else(|| Arc::new(Self::new_read_only()))
-    }
-
-    /// Returns the provided affinity handle, or creates a new default `ReadWrite` affinity if `None`.
-    pub(crate) fn default_read_write(existing: Option<Arc<Self>>) -> Arc<Self> {
-        existing.unwrap_or_else(|| Arc::new(Self::new_read_write()))
-    }
-
     /// Returns `true` if this handle requires hard stickiness (Read/Write transactions).
     pub(crate) fn is_read_write(&self) -> bool {
         self.kind == AffinityKind::ReadWrite
@@ -102,6 +92,16 @@ impl TransactionAffinity {
         }
         *slot = Some(lease.rw_affinity_guard());
     }
+
+    /// Releases the active Read/Write transaction guard on the channel entry,
+    /// allowing draining channels to close once the transaction completes.
+    pub(crate) fn release_rw_guard(&self) {
+        let mut slot = self
+            .rw_guard
+            .lock()
+            .expect("affinity rw_guard lock poisoned");
+        *slot = None;
+    }
 }
 
 /// Stickiness kind for transaction channel affinity.
@@ -116,6 +116,64 @@ pub(crate) enum AffinityKind {
     ReadOnly,
 }
 
+/// Routing target for channel selection: either a round-robin hint or a transaction affinity handle.
+#[derive(Clone, Copy, Debug)]
+pub(crate) enum ChannelTarget<'a> {
+    Hint(usize),
+    Affinity(&'a TransactionAffinity),
+}
+
+impl<'a> ChannelTarget<'a> {
+    pub(crate) fn from_affinity_or_hint(
+        affinity: Option<&'a TransactionAffinity>,
+        hint: usize,
+    ) -> Self {
+        match affinity {
+            Some(affinity) => Self::Affinity(affinity),
+            None => Self::Hint(hint),
+        }
+    }
+}
+
+impl From<usize> for ChannelTarget<'_> {
+    fn from(hint: usize) -> Self {
+        Self::Hint(hint)
+    }
+}
+
+impl<'a> From<&'a TransactionAffinity> for ChannelTarget<'a> {
+    fn from(affinity: &'a TransactionAffinity) -> Self {
+        Self::Affinity(affinity)
+    }
+}
+
+impl<'a> From<&'a Arc<TransactionAffinity>> for ChannelTarget<'a> {
+    fn from(affinity: &'a Arc<TransactionAffinity>) -> Self {
+        Self::Affinity(affinity)
+    }
+}
+
+impl<'a> From<Option<&'a TransactionAffinity>> for ChannelTarget<'a> {
+    fn from(affinity: Option<&'a TransactionAffinity>) -> Self {
+        match affinity {
+            Some(affinity) => Self::Affinity(affinity),
+            None => Self::Hint(0),
+        }
+    }
+}
+
+impl<'a> From<(Option<&'a TransactionAffinity>, usize)> for ChannelTarget<'a> {
+    fn from((affinity, hint): (Option<&'a TransactionAffinity>, usize)) -> Self {
+        Self::from_affinity_or_hint(affinity, hint)
+    }
+}
+
+impl<'a> From<(usize, Option<&'a TransactionAffinity>)> for ChannelTarget<'a> {
+    fn from((hint, affinity): (usize, Option<&'a TransactionAffinity>)) -> Self {
+        Self::from_affinity_or_hint(affinity, hint)
+    }
+}
+
 #[cfg(test)]
 impl TransactionAffinity {
     pub(crate) fn set_entry_id(&self, entry_id: u64) {
@@ -155,6 +213,7 @@ mod tests {
     use crate::client::Channel;
     use crate::generated::gapic_dataplane::stub::Spanner as SpannerStub;
     use std::fmt::Debug;
+    use std::sync::Arc;
     use std::time::Duration;
 
     #[derive(Debug)]
@@ -255,37 +314,6 @@ mod tests {
         );
     }
 
-    #[test]
-    fn transaction_affinity_defaults() {
-        let default_read_only = TransactionAffinity::default_read_only(None);
-        assert!(
-            default_read_only.is_read_only(),
-            "default_read_only(None) must return ReadOnly affinity"
-        );
-
-        let custom_read_only = Arc::new(TransactionAffinity::new_read_only());
-        let passed_read_only =
-            TransactionAffinity::default_read_only(Some(Arc::clone(&custom_read_only)));
-        assert!(
-            Arc::ptr_eq(&custom_read_only, &passed_read_only),
-            "default_read_only(Some(handle)) must return existing handle without recreating"
-        );
-
-        let default_read_write = TransactionAffinity::default_read_write(None);
-        assert!(
-            default_read_write.is_read_write(),
-            "default_read_write(None) must return ReadWrite affinity"
-        );
-
-        let custom_read_write = Arc::new(TransactionAffinity::new_read_write());
-        let passed_read_write =
-            TransactionAffinity::default_read_write(Some(Arc::clone(&custom_read_write)));
-        assert!(
-            Arc::ptr_eq(&custom_read_write, &passed_read_write),
-            "default_read_write(Some(handle)) must return existing handle without recreating"
-        );
-    }
-
     #[test]
     fn attach_rw_guard_repinning_replaces_old_guard() {
         let channel1 = Channel::new_for_test(DummyStub);