test(spanner): retry internal errors from the Emulator (#5840)

The latest version of the Spanner Emulator can return an INTERNAL error
if DDL statements and read/write transactions are executed concurrently.
This change adds a specific error to the set of errors that are retried
by the Spanner client when it is running agains the Emulator. This
additional check is only done if the client is connected to the
Emulator, and not if the client is connected to real Spanner.

Note: The test that failed and triggered the creation of this ticket, is
not the actual culprit, and a similar failure could in theory happen for
any of the other tests for the Spanner client that execute a read/write
transaction.

Fixes #5826

Author: olavloite

src/spanner/src/client.rs

@@ -20,7 +20,9 @@ use crate::model::{
 };
 use crate::server_streaming::builder;
 use gaxi::options::{ClientConfig, Credentials};
+use google_cloud_auth::credentials::anonymous;
 use google_cloud_gax::client_builder::ClientBuilder as GaxClientBuilder;
+use google_cloud_gax::client_builder::internal::new_builder;
 use google_cloud_gax::options::{
     RequestOptions as GaxRequestOptions, internal::RequestOptionsExt as _,
 };
@@ -49,6 +51,7 @@ pub struct Spanner {
     pub(crate) channels: Vec<Channel>,
     pub(crate) counter: std::sync::Arc<AtomicUsize>,
     pub(crate) config: ClientConfig,
+    pub(crate) is_emulator: bool,
 }
 
 /// A factory for constructing `Spanner` clients.
@@ -58,7 +61,21 @@ impl google_cloud_gax::client_builder::internal::ClientFactory for Factory {
     type Client = Spanner;
     type Credentials = Credentials;
 
-    async fn build(self, config: ClientConfig) -> crate::ClientBuilderResult<Self::Client> {
+    async fn build(self, mut config: ClientConfig) -> crate::ClientBuilderResult<Self::Client> {
+        let mut is_emulator = false;
+        if let Some(endpoint) = std::env::var("SPANNER_EMULATOR_HOST")
+            .ok()
+            .filter(|s| !s.is_empty())
+        {
+            is_emulator = true;
+            if config.endpoint.is_none() {
+                config.endpoint = Some(parse_emulator_endpoint(&endpoint));
+            }
+            if config.cred.is_none() {
+                config.cred = Some(anonymous::Builder::new().build());
+            }
+        }
+
         let num_channels = std::env::var("SPANNER_NUM_CHANNELS")
             .ok()
             .and_then(|s| s.parse::<usize>().ok())
@@ -73,6 +90,7 @@ impl google_cloud_gax::client_builder::internal::ClientFactory for Factory {
             channels,
             counter: std::sync::Arc::new(AtomicUsize::new(0)),
             config,
+            is_emulator,
         })
     }
 }
@@ -175,22 +193,7 @@ impl Spanner {
     /// detects and connects to the Spanner emulator if the `SPANNER_EMULATOR_HOST`
     /// environment variable is set.
     pub fn builder() -> ClientBuilder {
-        let builder = google_cloud_gax::client_builder::internal::new_builder(Factory);
-        // The Spanner client should automatically use the Spanner emulator if the
-        // SPANNER_EMULATOR_HOST environment variable is set.
-        let Some(endpoint) = std::env::var("SPANNER_EMULATOR_HOST")
-            .ok()
-            .filter(|s| !s.is_empty())
-        else {
-            return builder;
-        };
-
-        // Determine if we need to prefix the endpoint with a scheme
-        let full_endpoint = parse_emulator_endpoint(&endpoint);
-
-        builder
-            .with_endpoint(full_endpoint)
-            .with_credentials(google_cloud_auth::credentials::anonymous::Builder::new().build())
+        new_builder(Factory)
     }
 
     /// Returns a builder for the [DatabaseAdmin] client.
@@ -221,11 +224,7 @@ impl Spanner {
         C: Clone + From<Credentials>,
     {
         if let Some(ref endpoint) = self.config.endpoint {
-            let is_emulator = std::env::var("SPANNER_EMULATOR_HOST")
-                .ok()
-                .filter(|s| !s.is_empty())
-                .is_some();
-            let ep = map_emulator_admin_endpoint(endpoint, is_emulator);
+            let ep = map_emulator_admin_endpoint(endpoint, self.is_emulator);
             builder = builder.with_endpoint(ep);
         }
         if let Some(ref cred) = self.config.cred {
@@ -279,9 +278,14 @@ impl Spanner {
             }],
             counter: std::sync::Arc::new(AtomicUsize::new(0)),
             config: ClientConfig::default(),
+            is_emulator: false,
         }
     }
 
+    pub(crate) fn is_emulator(&self) -> bool {
+        self.is_emulator
+    }
+
     pub(crate) fn get_channel(&self, hint: usize) -> &Channel {
         let idx = hint % self.channels.len();
         &self.channels[idx]

src/spanner/src/database_client.rs

@@ -54,6 +54,10 @@ pub struct DatabaseClient {
 }
 
 impl DatabaseClient {
+    pub(crate) fn is_emulator(&self) -> bool {
+        self.spanner.is_emulator()
+    }
+
     /// Returns a builder for a single-use read-only transaction.
     ///
     /// # Example

src/spanner/src/partitioned_dml_transaction.rs

@@ -180,9 +180,9 @@ impl PartitionedDmlTransaction {
         let base_request = statement.into_request();
         let channel_hint = self.client.spanner.next_channel_hint();
         let client = self.client;
+        let is_emulator = client.is_emulator();
 
-        // Execute the statement and retry if the transaction is aborted by Spanner.
-        retry_aborted(&*self.retry_policy, || {
+        let action = || {
             let begin_request = begin_request.clone();
             let base_request = base_request.clone();
             let session_name = session_name.clone();
@@ -214,8 +214,9 @@ impl PartitionedDmlTransaction {
 
                 extract_lower_bound_update_count_from_stream(stream).await
             }
-        })
-        .await
+        };
+
+        retry_aborted(&*self.retry_policy, action, is_emulator).await
     }
 
     fn amend_gax_options(&self, options: &mut GaxRequestOptions) {

src/spanner/src/read_write_transaction.rs

@@ -63,7 +63,7 @@ use wkt::Duration;
 /// A builder for [ReadWriteTransaction].
 #[derive(Clone, Debug)]
 pub(crate) struct ReadWriteTransactionBuilder {
-    client: DatabaseClient,
+    pub(crate) client: DatabaseClient,
     options: TransactionOptions,
     transaction_tag: Option<String>,
     max_commit_delay: Option<Duration>,