diff --git a/googleapis/build.rs b/googleapis/build.rs index ddd03464..eeaa235f 100644 --- a/googleapis/build.rs +++ b/googleapis/build.rs @@ -21,6 +21,7 @@ fn generate(config: prost_build::Config, out_dir: impl AsRef) { .compile_with_config( config, &[ + "googleapis/google/rpc/error_details.proto", "googleapis/google/cloud/bigquery/storage/v1/storage.proto", "googleapis/google/storage/v2/storage.proto", "googleapis/google/pubsub/v1/pubsub.proto", diff --git a/googleapis/src/google.rpc.rs b/googleapis/src/google.rpc.rs index ae89e92c..841deb82 100644 --- a/googleapis/src/google.rpc.rs +++ b/googleapis/src/google.rpc.rs @@ -1,4 +1,374 @@ // This file is @generated by prost-build. +/// Describes the cause of the error with structured details. +/// +/// Example of an error when contacting the "pubsub.googleapis.com" API when it +/// is not enabled: +/// +/// { "reason": "API_DISABLED" +/// "domain": "googleapis.com" +/// "metadata": { +/// "resource": "projects/123", +/// "service": "pubsub.googleapis.com" +/// } +/// } +/// +/// This response indicates that the pubsub.googleapis.com API is not enabled. +/// +/// Example of an error that is returned when attempting to create a Spanner +/// instance in a region that is out of stock: +/// +/// { "reason": "STOCKOUT" +/// "domain": "spanner.googleapis.com", +/// "metadata": { +/// "availableRegions": "us-central1,us-east2" +/// } +/// } +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ErrorInfo { + /// The reason of the error. This is a constant value that identifies the + /// proximate cause of the error. Error reasons are unique within a particular + /// domain of errors. This should be at most 63 characters and match a + /// regular expression of `[A-Z][A-Z0-9_]+\[A-Z0-9\]`, which represents + /// UPPER_SNAKE_CASE. + #[prost(string, tag = "1")] + pub reason: ::prost::alloc::string::String, + /// The logical grouping to which the "reason" belongs. The error domain + /// is typically the registered service name of the tool or product that + /// generates the error. Example: "pubsub.googleapis.com". If the error is + /// generated by some common infrastructure, the error domain must be a + /// globally unique value that identifies the infrastructure. For Google API + /// infrastructure, the error domain is "googleapis.com". + #[prost(string, tag = "2")] + pub domain: ::prost::alloc::string::String, + /// Additional structured details about this error. + /// + /// Keys must match a regular expression of `[a-z][a-zA-Z0-9-_]+` but should + /// ideally be lowerCamelCase. Also, they must be limited to 64 characters in + /// length. When identifying the current value of an exceeded limit, the units + /// should be contained in the key, not the value. For example, rather than + /// `{"instanceLimit": "100/request"}`, should be returned as, + /// `{"instanceLimitPerRequest": "100"}`, if the client exceeds the number of + /// instances that can be created in a single (batch) request. + #[prost(map = "string, string", tag = "3")] + pub metadata: ::std::collections::HashMap< + ::prost::alloc::string::String, + ::prost::alloc::string::String, + >, +} +/// Describes when the clients can retry a failed request. Clients could ignore +/// the recommendation here or retry when this information is missing from error +/// responses. +/// +/// It's always recommended that clients should use exponential backoff when +/// retrying. +/// +/// Clients should wait until `retry_delay` amount of time has passed since +/// receiving the error response before retrying. If retrying requests also +/// fail, clients should use an exponential backoff scheme to gradually increase +/// the delay between retries based on `retry_delay`, until either a maximum +/// number of retries have been reached or a maximum retry delay cap has been +/// reached. +#[derive(Clone, Copy, PartialEq, ::prost::Message)] +pub struct RetryInfo { + /// Clients should wait at least this long between retrying the same request. + #[prost(message, optional, tag = "1")] + pub retry_delay: ::core::option::Option<::prost_types::Duration>, +} +/// Describes additional debugging info. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct DebugInfo { + /// The stack trace entries indicating where the error occurred. + #[prost(string, repeated, tag = "1")] + pub stack_entries: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, + /// Additional debugging information provided by the server. + #[prost(string, tag = "2")] + pub detail: ::prost::alloc::string::String, +} +/// Describes how a quota check failed. +/// +/// For example if a daily limit was exceeded for the calling project, +/// a service could respond with a QuotaFailure detail containing the project +/// id and the description of the quota limit that was exceeded. If the +/// calling project hasn't enabled the service in the developer console, then +/// a service could respond with the project id and set `service_disabled` +/// to true. +/// +/// Also see RetryInfo and Help types for other details about handling a +/// quota failure. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct QuotaFailure { + /// Describes all quota violations. + #[prost(message, repeated, tag = "1")] + pub violations: ::prost::alloc::vec::Vec, +} +/// Nested message and enum types in `QuotaFailure`. +pub mod quota_failure { + /// A message type used to describe a single quota violation. For example, a + /// daily quota or a custom quota that was exceeded. + #[derive(Clone, PartialEq, ::prost::Message)] + pub struct Violation { + /// The subject on which the quota check failed. + /// For example, "clientip:" or "project:". + #[prost(string, tag = "1")] + pub subject: ::prost::alloc::string::String, + /// A description of how the quota check failed. Clients can use this + /// description to find more about the quota configuration in the service's + /// public documentation, or find the relevant quota limit to adjust through + /// developer console. + /// + /// For example: "Service disabled" or "Daily Limit for read operations + /// exceeded". + #[prost(string, tag = "2")] + pub description: ::prost::alloc::string::String, + /// The API Service from which the `QuotaFailure.Violation` orginates. In + /// some cases, Quota issues originate from an API Service other than the one + /// that was called. In other words, a dependency of the called API Service + /// could be the cause of the `QuotaFailure`, and this field would have the + /// dependency API service name. + /// + /// For example, if the called API is Kubernetes Engine API + /// (container.googleapis.com), and a quota violation occurs in the + /// Kubernetes Engine API itself, this field would be + /// "container.googleapis.com". On the other hand, if the quota violation + /// occurs when the Kubernetes Engine API creates VMs in the Compute Engine + /// API (compute.googleapis.com), this field would be + /// "compute.googleapis.com". + #[prost(string, tag = "3")] + pub api_service: ::prost::alloc::string::String, + /// The metric of the violated quota. A quota metric is a named counter to + /// measure usage, such as API requests or CPUs. When an activity occurs in a + /// service, such as Virtual Machine allocation, one or more quota metrics + /// may be affected. + /// + /// For example, "compute.googleapis.com/cpus_per_vm_family", + /// "storage.googleapis.com/internet_egress_bandwidth". + #[prost(string, tag = "4")] + pub quota_metric: ::prost::alloc::string::String, + /// The id of the violated quota. Also know as "limit name", this is the + /// unique identifier of a quota in the context of an API service. + /// + /// For example, "CPUS-PER-VM-FAMILY-per-project-region". + #[prost(string, tag = "5")] + pub quota_id: ::prost::alloc::string::String, + /// The dimensions of the violated quota. Every non-global quota is enforced + /// on a set of dimensions. While quota metric defines what to count, the + /// dimensions specify for what aspects the counter should be increased. + /// + /// For example, the quota "CPUs per region per VM family" enforces a limit + /// on the metric "compute.googleapis.com/cpus_per_vm_family" on dimensions + /// "region" and "vm_family". And if the violation occurred in region + /// "us-central1" and for VM family "n1", the quota_dimensions would be, + /// + /// { + /// "region": "us-central1", + /// "vm_family": "n1", + /// } + /// + /// When a quota is enforced globally, the quota_dimensions would always be + /// empty. + #[prost(map = "string, string", tag = "6")] + pub quota_dimensions: ::std::collections::HashMap< + ::prost::alloc::string::String, + ::prost::alloc::string::String, + >, + /// The enforced quota value at the time of the `QuotaFailure`. + /// + /// For example, if the enforced quota value at the time of the + /// `QuotaFailure` on the number of CPUs is "10", then the value of this + /// field would reflect this quantity. + #[prost(int64, tag = "7")] + pub quota_value: i64, + /// The new quota value being rolled out at the time of the violation. At the + /// completion of the rollout, this value will be enforced in place of + /// quota_value. If no rollout is in progress at the time of the violation, + /// this field is not set. + /// + /// For example, if at the time of the violation a rollout is in progress + /// changing the number of CPUs quota from 10 to 20, 20 would be the value of + /// this field. + #[prost(int64, optional, tag = "8")] + pub future_quota_value: ::core::option::Option, + } +} +/// Describes what preconditions have failed. +/// +/// For example, if an RPC failed because it required the Terms of Service to be +/// acknowledged, it could list the terms of service violation in the +/// PreconditionFailure message. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct PreconditionFailure { + /// Describes all precondition violations. + #[prost(message, repeated, tag = "1")] + pub violations: ::prost::alloc::vec::Vec, +} +/// Nested message and enum types in `PreconditionFailure`. +pub mod precondition_failure { + /// A message type used to describe a single precondition failure. + #[derive(Clone, PartialEq, ::prost::Message)] + pub struct Violation { + /// The type of PreconditionFailure. We recommend using a service-specific + /// enum type to define the supported precondition violation subjects. For + /// example, "TOS" for "Terms of Service violation". + #[prost(string, tag = "1")] + pub r#type: ::prost::alloc::string::String, + /// The subject, relative to the type, that failed. + /// For example, "google.com/cloud" relative to the "TOS" type would indicate + /// which terms of service is being referenced. + #[prost(string, tag = "2")] + pub subject: ::prost::alloc::string::String, + /// A description of how the precondition failed. Developers can use this + /// description to understand how to fix the failure. + /// + /// For example: "Terms of service not accepted". + #[prost(string, tag = "3")] + pub description: ::prost::alloc::string::String, + } +} +/// Describes violations in a client request. This error type focuses on the +/// syntactic aspects of the request. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct BadRequest { + /// Describes all violations in a client request. + #[prost(message, repeated, tag = "1")] + pub field_violations: ::prost::alloc::vec::Vec, +} +/// Nested message and enum types in `BadRequest`. +pub mod bad_request { + /// A message type used to describe a single bad request field. + #[derive(Clone, PartialEq, ::prost::Message)] + pub struct FieldViolation { + /// A path that leads to a field in the request body. The value will be a + /// sequence of dot-separated identifiers that identify a protocol buffer + /// field. + /// + /// Consider the following: + /// + /// message CreateContactRequest { + /// message EmailAddress { + /// enum Type { + /// TYPE_UNSPECIFIED = 0; + /// HOME = 1; + /// WORK = 2; + /// } + /// + /// optional string email = 1; + /// repeated EmailType type = 2; + /// } + /// + /// string full_name = 1; + /// repeated EmailAddress email_addresses = 2; + /// } + /// + /// In this example, in proto `field` could take one of the following values: + /// + /// * `full_name` for a violation in the `full_name` value + /// * `email_addresses\[1\].email` for a violation in the `email` field of the + /// first `email_addresses` message + /// * `email_addresses\[3\].type\[2\]` for a violation in the second `type` + /// value in the third `email_addresses` message. + /// + /// In JSON, the same values are represented as: + /// + /// * `fullName` for a violation in the `fullName` value + /// * `emailAddresses\[1\].email` for a violation in the `email` field of the + /// first `emailAddresses` message + /// * `emailAddresses\[3\].type\[2\]` for a violation in the second `type` + /// value in the third `emailAddresses` message. + #[prost(string, tag = "1")] + pub field: ::prost::alloc::string::String, + /// A description of why the request element is bad. + #[prost(string, tag = "2")] + pub description: ::prost::alloc::string::String, + /// The reason of the field-level error. This is a constant value that + /// identifies the proximate cause of the field-level error. It should + /// uniquely identify the type of the FieldViolation within the scope of the + /// google.rpc.ErrorInfo.domain. This should be at most 63 + /// characters and match a regular expression of `[A-Z][A-Z0-9_]+\[A-Z0-9\]`, + /// which represents UPPER_SNAKE_CASE. + #[prost(string, tag = "3")] + pub reason: ::prost::alloc::string::String, + /// Provides a localized error message for field-level errors that is safe to + /// return to the API consumer. + #[prost(message, optional, tag = "4")] + pub localized_message: ::core::option::Option, + } +} +/// Contains metadata about the request that clients can attach when filing a bug +/// or providing other forms of feedback. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct RequestInfo { + /// An opaque string that should only be interpreted by the service generating + /// it. For example, it can be used to identify requests in the service's logs. + #[prost(string, tag = "1")] + pub request_id: ::prost::alloc::string::String, + /// Any data that was used to serve this request. For example, an encrypted + /// stack trace that can be sent back to the service provider for debugging. + #[prost(string, tag = "2")] + pub serving_data: ::prost::alloc::string::String, +} +/// Describes the resource that is being accessed. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ResourceInfo { + /// A name for the type of resource being accessed, e.g. "sql table", + /// "cloud storage bucket", "file", "Google calendar"; or the type URL + /// of the resource: e.g. "type.googleapis.com/google.pubsub.v1.Topic". + #[prost(string, tag = "1")] + pub resource_type: ::prost::alloc::string::String, + /// The name of the resource being accessed. For example, a shared calendar + /// name: "example.com_4fghdhgsrgh@group.calendar.google.com", if the current + /// error is + /// [google.rpc.Code.PERMISSION_DENIED][google.rpc.Code.PERMISSION_DENIED]. + #[prost(string, tag = "2")] + pub resource_name: ::prost::alloc::string::String, + /// The owner of the resource (optional). + /// For example, "user:" or "project:". + #[prost(string, tag = "3")] + pub owner: ::prost::alloc::string::String, + /// Describes what error is encountered when accessing this resource. + /// For example, updating a cloud project may require the `writer` permission + /// on the developer console project. + #[prost(string, tag = "4")] + pub description: ::prost::alloc::string::String, +} +/// Provides links to documentation or for performing an out of band action. +/// +/// For example, if a quota check failed with an error indicating the calling +/// project hasn't enabled the accessed service, this can contain a URL pointing +/// directly to the right place in the developer console to flip the bit. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct Help { + /// URL(s) pointing to additional information on handling the current error. + #[prost(message, repeated, tag = "1")] + pub links: ::prost::alloc::vec::Vec, +} +/// Nested message and enum types in `Help`. +pub mod help { + /// Describes a URL link. + #[derive(Clone, PartialEq, ::prost::Message)] + pub struct Link { + /// Describes what the link offers. + #[prost(string, tag = "1")] + pub description: ::prost::alloc::string::String, + /// The URL of the link. + #[prost(string, tag = "2")] + pub url: ::prost::alloc::string::String, + } +} +/// Provides a localized error message that is safe to return to the user +/// which can be attached to an RPC error. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct LocalizedMessage { + /// The locale used following the specification defined at + /// + /// Examples are: "en-US", "fr-CH", "es-MX" + #[prost(string, tag = "1")] + pub locale: ::prost::alloc::string::String, + /// The localized error message in the above locale. + #[prost(string, tag = "2")] + pub message: ::prost::alloc::string::String, +} /// The `Status` type defines a logical error model that is suitable for /// different programming environments, including REST APIs and RPC APIs. It is /// used by [gRPC](). Each `Status` message contains diff --git a/spanner/Cargo.toml b/spanner/Cargo.toml index e19012d5..3f266deb 100644 --- a/spanner/Cargo.toml +++ b/spanner/Cargo.toml @@ -16,12 +16,14 @@ doctest = false [dependencies] tracing = "0.1" prost-types = "0.14" +prost = "0.14" tokio = "1.32" time = { version="0.3", features = ["std", "macros", "formatting", "parsing"] } thiserror = "2.0" parking_lot = "0.12" base64 = "0.22" -tokio-util = "0.7" +serde = { version = "1.0", optional = true, features = ["derive"] } +tokio-util = { version = "0.7", features = ["rt"] } bigdecimal = { version="0.4", features=["serde"] } token-source = "1.0" @@ -36,6 +38,7 @@ tokio = { version="1.32", features=["rt-multi-thread"] } tracing-subscriber = { version="0.3", features=["env-filter"] } serial_test = "3.1" ctor = "0.5" +tonic = { version = "0.13", default-features = false, features = ["transport"] } google-cloud-auth = { package = "gcloud-auth", path="../foundation/auth", default-features=false, features=["rustls-tls"]} [features] diff --git a/spanner/src/client.rs b/spanner/src/client.rs index 61f2b262..8d9eba7f 100644 --- a/spanner/src/client.rs +++ b/spanner/src/client.rs @@ -7,13 +7,13 @@ use std::time::Duration; use google_cloud_gax::conn::{ConnectionOptions, Environment}; use google_cloud_gax::grpc::{Code, Status}; -use google_cloud_gax::retry::{invoke_fn, TryAs}; +use google_cloud_gax::retry::TryAs; use google_cloud_googleapis::spanner::v1::{commit_request, transaction_options, Mutation, TransactionOptions}; use token_source::NoopTokenSourceProvider; use crate::apiv1::conn_pool::{ConnectionManager, SPANNER}; -use crate::retry::TransactionRetrySetting; -use crate::session::{ManagedSession, SessionConfig, SessionError, SessionManager}; +use crate::retry::{invoke_with_session_retry, SessionRetryAction, TransactionRetrySetting}; +use crate::session::{is_session_not_found_status, ManagedSession, SessionConfig, SessionError, SessionManager}; use crate::statement::Statement; use crate::transaction::{CallOptions, QueryOptions}; use crate::transaction_ro::{BatchReadOnlyTransaction, ReadOnlyTransaction}; @@ -275,9 +275,20 @@ impl Client { &self, options: ReadOnlyTransactionOption, ) -> Result { - let session = self.get_session().await?; - let result = ReadOnlyTransaction::begin(session, options.timestamp_bound, options.call_options).await?; - Ok(result) + let timestamp_bound = options.timestamp_bound; + let call_options = options.call_options; + loop { + let session = self.get_session().await?; + match ReadOnlyTransaction::begin(session, timestamp_bound.clone(), call_options.clone()).await { + Ok(tx) => return Ok(tx), + Err(status) => { + if is_session_not_found_status(&status) { + continue; + } + return Err(Error::GRPC(status)); + } + } + } } /// batch_read_only_transaction returns a BatchReadOnlyTransaction that can be used @@ -297,9 +308,20 @@ impl Client { &self, options: ReadOnlyTransactionOption, ) -> Result { - let session = self.get_session().await?; - let result = BatchReadOnlyTransaction::begin(session, options.timestamp_bound, options.call_options).await?; - Ok(result) + let timestamp_bound = options.timestamp_bound; + let call_options = options.call_options; + loop { + let session = self.get_session().await?; + match BatchReadOnlyTransaction::begin(session, timestamp_bound.clone(), call_options.clone()).await { + Ok(tx) => return Ok(tx), + Err(status) => { + if is_session_not_found_status(&status) { + continue; + } + return Err(Error::GRPC(status)); + } + } + } } /// partitioned_update executes a DML statement in parallel across the database, @@ -328,30 +350,35 @@ impl Client { stmt: Statement, options: PartitionedUpdateOption, ) -> Result { - let ro = TransactionRetrySetting::new(vec![Code::Aborted, Code::Internal]); - let session = Some(self.get_session().await?); - - // reuse session - invoke_fn( - Some(ro), - |session| async { - let mut tx = match ReadWriteTransaction::begin_partitioned_dml( - session.unwrap(), - options.begin_options.clone(), - options.transaction_tag.clone(), - ) - .await - { - Ok(tx) => tx, - Err(e) => return Err((Error::GRPC(e.status), Some(e.session))), - }; - tx.update_with_option(stmt.clone(), options.query_options.clone().unwrap_or_default()) + let helper = self.session_retry_helper(TransactionRetrySetting::new(vec![Code::Aborted, Code::Internal])); + helper + .run(|session| { + let stmt = stmt.clone(); + let options = options.clone(); + async move { + let mut tx = match ReadWriteTransaction::begin_partitioned_dml( + session, + options.begin_options.clone(), + options.transaction_tag.clone(), + ) .await - .map_err(|e| (Error::GRPC(e), tx.take_session())) - }, - session, - ) - .await + { + Ok(tx) => tx, + Err(e) => return Err(SessionRetryAction::retry(Error::GRPC(e.status), e.session)), + }; + match tx + .update_with_option(stmt, options.query_options.clone().unwrap_or_default()) + .await + { + Ok(rows) => Ok(rows), + Err(status) => match tx.take_session() { + Some(session) => Err(SessionRetryAction::retry(Error::GRPC(status), session)), + None => Err(SessionRetryAction::fail(Error::GRPC(status))), + }, + } + } + }) + .await } /// apply_at_least_once may attempt to apply mutations more than once; if @@ -381,25 +408,24 @@ impl Client { ms: Vec, options: CommitOptions, ) -> Result, Error> { - let ro = TransactionRetrySetting::default(); - let mut session = self.get_session().await?; - - invoke_fn( - Some(ro), - |session| async { - let tx = commit_request::Transaction::SingleUseTransaction(TransactionOptions { - exclude_txn_from_change_streams: false, - mode: Some(transaction_options::Mode::ReadWrite(transaction_options::ReadWrite::default())), - isolation_level: IsolationLevel::Unspecified as i32, - }); - match commit(session, ms.clone(), tx, options.clone()).await { - Ok(s) => Ok(Some(s.into())), - Err(e) => Err((Error::GRPC(e), session)), + let helper = self.session_retry_helper(TransactionRetrySetting::new(vec![Code::Aborted, Code::Internal])); + helper + .run(|mut session| { + let ms = ms.clone(); + let options = options.clone(); + async move { + let tx = commit_request::Transaction::SingleUseTransaction(TransactionOptions { + exclude_txn_from_change_streams: false, + mode: Some(transaction_options::Mode::ReadWrite(transaction_options::ReadWrite::default())), + isolation_level: IsolationLevel::Unspecified as i32, + }); + match commit(&mut session, ms, tx, options.clone()).await { + Ok(s) => Ok(Some(s.into())), + Err(status) => Err(SessionRetryAction::retry(Error::GRPC(status), session)), + } } - }, - &mut session, - ) - .await + }) + .await } /// Apply applies a list of mutations atomically to the database. @@ -526,22 +552,21 @@ impl Client { F: for<'tx> Fn(&'tx mut ReadWriteTransaction) -> Pin> + Send + 'tx>>, { let (bo, co, tag) = Client::split_read_write_transaction_option(options); - - let ro = TransactionRetrySetting::default(); - let session = Some(self.get_session().await?); - // must reuse session - invoke_fn( - Some(ro), - |session| async { - let mut tx = self - .create_read_write_transaction::(session, bo.clone(), tag.clone()) - .await?; - let result = f(&mut tx).await; - tx.finish(result, Some(co.clone())).await - }, - session, - ) - .await + let client = self; + let helper = self.session_retry_helper(TransactionRetrySetting::default()); + helper + .run(|session| { + let bo = bo.clone(); + let co = co.clone(); + let tag = tag.clone(); + let func = &f; + async move { + let mut tx = client.create_read_write_transaction::(session, bo, tag).await?; + let result = func(&mut tx).await; + tx.finish(result, Some(co)).await.map_err(SessionRetryAction::from) + } + }) + .await } /// begin_read_write_transaction creates new ReadWriteTransaction. @@ -610,37 +635,35 @@ impl Client { E: TryAs + From + From, { let (bo, co, tag) = Client::split_read_write_transaction_option(options); - - let ro = TransactionRetrySetting::default(); - let session = Some(self.get_session().await?); - - // reuse session - invoke_fn( - Some(ro), - |session| async { - let mut tx = self - .create_read_write_transaction::(session, bo.clone(), tag.clone()) - .await?; - let result = f(&mut tx); - tx.finish(result, Some(co.clone())).await - }, - session, - ) - .await + let client = self; + let helper = self.session_retry_helper(TransactionRetrySetting::default()); + helper + .run(|session| { + let bo = bo.clone(); + let co = co.clone(); + let tag = tag.clone(); + let func = &f; + async move { + let mut tx = client.create_read_write_transaction::(session, bo, tag).await?; + let result = func(&mut tx); + tx.finish(result, Some(co)).await.map_err(SessionRetryAction::from) + } + }) + .await } async fn create_read_write_transaction( &self, - session: Option, + session: ManagedSession, bo: CallOptions, transaction_tag: Option, - ) -> Result)> + ) -> Result> where E: TryAs + From + From, { - ReadWriteTransaction::begin(session.unwrap(), bo, transaction_tag) + ReadWriteTransaction::begin(session, bo, transaction_tag) .await - .map_err(|e| (E::from(e.status), Some(e.session))) + .map_err(|e| SessionRetryAction::retry(E::from(e.status), e.session)) } async fn get_session(&self) -> Result { @@ -652,4 +675,38 @@ impl Client { ) -> (CallOptions, CommitOptions, Option) { (options.begin_options, options.commit_options, options.transaction_tag) } + + fn session_retry_helper(&self, retry: TransactionRetrySetting) -> SessionRetryHelper<'_> { + SessionRetryHelper { + sessions: &self.sessions, + retry, + } + } +} + +struct SessionRetryHelper<'a> { + sessions: &'a Arc, + retry: TransactionRetrySetting, +} + +impl<'a> SessionRetryHelper<'a> { + async fn run(&self, mut f: F) -> Result + where + E: TryAs + From + From, + Fut: Future>>, + F: FnMut(ManagedSession) -> Fut, + { + let session = Some(self.sessions.get().await?); + invoke_with_session_retry(Some(self.retry.clone()), session, |session| { + let session = session.expect("session missing"); + let fut = f(session); + async move { + match fut.await { + Ok(value) => Ok(value), + Err(action) => Err(action.into()), + } + } + }) + .await + } } diff --git a/spanner/src/reader.rs b/spanner/src/reader.rs index ed798366..ca56428d 100644 --- a/spanner/src/reader.rs +++ b/spanner/src/reader.rs @@ -1,5 +1,7 @@ use std::collections::{HashMap, VecDeque}; +use std::ops::{Deref, DerefMut}; use std::sync::Arc; +use std::time::Duration; use prost_types::{value::Kind, Value}; @@ -8,7 +10,7 @@ use google_cloud_googleapis::spanner::v1::struct_type::Field; use google_cloud_googleapis::spanner::v1::{ExecuteSqlRequest, PartialResultSet, ReadRequest, ResultSetMetadata}; use crate::row::Row; -use crate::session::SessionHandle; +use crate::session::{is_session_not_found_status, ManagedSession, SessionHandle}; use crate::transaction::CallOptions; pub trait Reader: Send + Sync { @@ -21,6 +23,8 @@ pub trait Reader: Send + Sync { fn update_token(&mut self, resume_token: Vec); fn can_resume(&self) -> bool; + + fn update_session(&mut self, session: &SessionHandle); } pub struct StatementReader { @@ -47,6 +51,10 @@ impl Reader for StatementReader { fn can_resume(&self) -> bool { self.enable_resume && !self.request.resume_token.is_empty() } + + fn update_session(&mut self, session: &SessionHandle) { + self.request.session = session.session.name.clone(); + } } pub struct TableReader { @@ -72,6 +80,10 @@ impl Reader for TableReader { fn can_resume(&self) -> bool { !self.request.resume_token.is_empty() } + + fn update_session(&mut self, session: &SessionHandle) { + self.request.session = session.session.name.clone(); + } } pub struct ResultSet { @@ -182,10 +194,12 @@ where T: Reader, { streaming: Streaming, - session: &'a mut SessionHandle, + session_owner: &'a mut ManagedSession, reader: T, rs: ResultSet, reader_option: Option, + allow_session_refresh: bool, + session_wait_timeout: Option, } impl<'a, T> RowIterator<'a, T> @@ -193,11 +207,28 @@ where T: Reader, { pub(crate) async fn new( - session: &'a mut SessionHandle, - reader: T, + session: &'a mut ManagedSession, + mut reader: T, option: Option, + allow_session_refresh: bool, ) -> Result, Status> { - let streaming = reader.read(session, option).await?.into_inner(); + let session_wait_timeout = option.as_ref().and_then(|opts| opts.session_renew_timeout()); + let streaming = loop { + let call_option = option.clone(); + let result = reader.read(session.deref_mut(), call_option).await; + match session.deref_mut().invalidate_if_needed(result).await { + Ok(response) => break response.into_inner(), + Err(status) => { + tracing::warn!("session read attempt failed: {:?}", status); + if allow_session_refresh && is_session_not_found_status(&status) { + session.renew(session_wait_timeout).await.map_err(Status::from)?; + reader.update_session(session.deref()); + continue; + } + return Err(status); + } + } + }; let rs = ResultSet { fields: Arc::new(vec![]), index: Arc::new(HashMap::new()), @@ -206,10 +237,12 @@ where }; Ok(Self { streaming, - session, + session_owner: session, reader, rs, reader_option: None, + allow_session_refresh, + session_wait_timeout, }) } @@ -218,17 +251,31 @@ where } async fn try_recv(&mut self, option: Option) -> Result { - // try getting records from server - let maybe_result_set = match self.streaming.message().await { - Ok(s) => s, - Err(e) => { - if !self.reader.can_resume() { - return Err(e); + let maybe_result_set = loop { + match self.streaming.message().await { + Ok(s) => break s, + Err(status) => { + if self.allow_session_refresh && !self.reader.can_resume() && is_session_not_found_status(&status) { + tracing::warn!("streaming saw Session not found; attempting to renew session"); + self.session_owner + .renew(self.session_wait_timeout) + .await + .map_err(Status::from)?; + self.reader.update_session(self.session_owner.deref()); + let call_option = option.clone(); + let result = self.reader.read(self.session_owner.deref_mut(), call_option).await?; + self.streaming = result.into_inner(); + continue; + } + + if !self.reader.can_resume() { + return Err(status); + } + tracing::debug!("streaming error: {}. resume reading by resume_token", status); + let call_option = option.clone(); + let result = self.reader.read(self.session_owner.deref_mut(), call_option).await?; + self.streaming = result.into_inner(); } - tracing::debug!("streaming error: {}. resume reading by resume_token", e); - let result = self.reader.read(self.session, option).await?; - self.streaming = result.into_inner(); - self.streaming.message().await? } }; diff --git a/spanner/src/retry.rs b/spanner/src/retry.rs index 1d6271fc..d985686d 100644 --- a/spanner/src/retry.rs +++ b/spanner/src/retry.rs @@ -1,9 +1,12 @@ +use std::future::Future; use std::iter::Take; use std::marker::PhantomData; use google_cloud_gax::grpc::{Code, Status}; use google_cloud_gax::retry::{CodeCondition, Condition, ExponentialBackoff, Retry, RetrySetting, TryAs}; +use crate::session::{is_session_not_found_status, ManagedSession, SessionError}; + pub struct TransactionCondition where E: TryAs, @@ -29,6 +32,9 @@ where { return false; } + if code == Code::NotFound { + return is_session_not_found_status(status); + } return self.inner.should_retry(error); } false @@ -125,6 +131,102 @@ impl Default for TransactionRetrySetting { } } +/// Result wrapper used by higher-level helpers to describe whether a failed +/// attempt should be retried with a refreshed session. +pub enum SessionRetryAction { + Retry { error: E, session: Box }, + Fail(E), +} + +impl SessionRetryAction { + pub fn retry(error: E, session: ManagedSession) -> Self { + SessionRetryAction::Retry { + error, + session: Box::new(session), + } + } + + pub fn fail(error: E) -> Self { + SessionRetryAction::Fail(error) + } +} + +impl From> for (E, Option) { + fn from(action: SessionRetryAction) -> Self { + match action { + SessionRetryAction::Retry { error, session } => (error, Some(*session)), + SessionRetryAction::Fail(error) => (error, None), + } + } +} + +impl From<(E, Option)> for SessionRetryAction { + fn from(value: (E, Option)) -> Self { + match value { + (error, Some(session)) => SessionRetryAction::Retry { + error, + session: Box::new(session), + }, + (error, None) => SessionRetryAction::Fail(error), + } + } +} + +/// Retry helper that understands how to renew sessions when Cloud Spanner reports +/// `Session not found`. The closure receives ownership of the current session +/// (wrapped in `Option`) and must return it on failure so it can be retried. +pub async fn invoke_with_session_retry( + retry: Option, + mut session: Option, + mut f: F, +) -> Result +where + E: TryAs + From + From, + Fut: Future)>>, + F: FnMut(Option) -> Fut, +{ + let retry = retry.unwrap_or_default(); + let mut strategy = >>::strategy(&retry); + let mut condition = >>::condition(&retry); + + loop { + let current_session = session.take(); + match f(current_session).await { + Ok(value) => return Ok(value), + Err((err, returned_session)) => { + session = returned_session; + if !condition.should_retry(&err) { + return Err(err); + } + + let mut skip_delay = false; + if let Some(status) = err.try_as() { + if is_session_not_found_status(status) { + if let Some(ref mut managed_session) = session { + if let Err(renew_err) = managed_session.renew(None).await { + return Err(E::from(renew_err)); + } + skip_delay = true; + } else { + return Err(err); + } + } + } + + if skip_delay { + continue; + } + + if let Some(duration) = strategy.next() { + tokio::time::sleep(duration).await; + } else { + return Err(err); + } + } + } + } +} + #[cfg(test)] mod tests { use google_cloud_gax::grpc::{Code, Status}; @@ -142,4 +244,17 @@ mod tests { let err = &Error::GRPC(Status::new(Code::Aborted, "")); assert!(default.condition().should_retry(err)); } + + #[test] + fn test_session_not_found_should_retry() { + let err = &Error::GRPC(Status::new( + Code::NotFound, + "Session not found: projects/local/instances/test/databases/db/sessions/session-id", + )); + let default = TransactionRetrySetting::default(); + assert!( + default.condition().should_retry(err), + "Session not found should be treated as retryable" + ); + } } diff --git a/spanner/src/session.rs b/spanner/src/session.rs index 7e1598b2..a8504189 100644 --- a/spanner/src/session.rs +++ b/spanner/src/session.rs @@ -1,4 +1,5 @@ use std::collections::VecDeque; +use std::io::Cursor; use std::mem; use std::ops::{Deref, DerefMut}; use std::sync::Arc; @@ -12,19 +13,104 @@ use tokio::sync::{mpsc, oneshot}; use tokio::task::{JoinHandle, JoinSet}; use tokio::time::{sleep, timeout}; use tokio_util::sync::CancellationToken; +use tokio_util::task::TaskTracker; use google_cloud_gax::grpc::metadata::MetadataMap; use google_cloud_gax::grpc::{Code, Status}; use google_cloud_gax::retry::TryAs; +use google_cloud_googleapis::rpc::{self, ResourceInfo}; use google_cloud_googleapis::spanner::v1::{BatchCreateSessionsRequest, DeleteSessionRequest, Session}; +use prost::Message; use crate::apiv1::conn_pool::ConnectionManager; use crate::apiv1::spanner_client::{ping_query_request, Client}; +const RESOURCE_INFO_TYPE_URL: &str = "type.googleapis.com/google.rpc.ResourceInfo"; +const SESSION_RESOURCE_TYPE: &str = "type.googleapis.com/google.spanner.v1.Session"; + +pub(crate) fn is_session_not_found_status(status: &Status) -> bool { + if status.code() != Code::NotFound { + return false; + } + if matches_resource_type(status, SESSION_RESOURCE_TYPE) { + return true; + } + status.message().contains("Session not found:") +} + +fn matches_resource_type(status: &Status, expected: &str) -> bool { + match extract_resource_type(status) { + Some(resource_type) => resource_type == expected, + None => false, + } +} + +fn extract_resource_type(status: &Status) -> Option { + let details = status.details(); + if details.is_empty() { + return None; + } + let mut cursor = Cursor::new(details); + let parsed = rpc::Status::decode(&mut cursor).ok()?; + parsed + .details + .into_iter() + .find_map(|any| decode_resource_info(any).map(|info| info.resource_type)) +} + +fn decode_resource_info(any: prost_types::Any) -> Option { + if any.type_url != RESOURCE_INFO_TYPE_URL { + return None; + } + let mut cursor = Cursor::new(any.value); + ResourceInfo::decode(&mut cursor).ok() +} + +#[cfg(test)] +mod status_tests { + use super::*; + + #[test] + fn detects_session_not_found_from_resource_info() { + let status = build_status_with_resource_info(SESSION_RESOURCE_TYPE); + assert!(is_session_not_found_status(&status)); + } + + #[test] + fn falls_back_to_message_when_details_missing() { + let status = Status::new(Code::NotFound, "Session not found: foo"); + assert!(is_session_not_found_status(&status)); + } + + fn build_status_with_resource_info(resource_type: &str) -> Status { + let info = ResourceInfo { + resource_type: resource_type.to_string(), + resource_name: "projects/p/instances/i/databases/d/sessions/s".to_string(), + owner: String::new(), + description: String::new(), + }; + let mut info_buf = Vec::new(); + info.encode(&mut info_buf).unwrap(); + let detail = prost_types::Any { + type_url: RESOURCE_INFO_TYPE_URL.to_string(), + value: info_buf, + }; + let grpc_status = rpc::Status { + code: Code::NotFound as i32, + message: String::from("Session not found"), + details: vec![detail], + }; + let mut details_buf = Vec::new(); + grpc_status.encode(&mut details_buf).unwrap(); + Status::with_details(Code::NotFound, "Session not found", details_buf.into()) + } +} + /// Session pub struct SessionHandle { pub session: Session, pub spanner_client: Client, + task_tracker: Arc, valid: bool, deleted: bool, last_used_at: Instant, @@ -34,10 +120,16 @@ pub struct SessionHandle { } impl SessionHandle { - pub(crate) fn new(session: Session, spanner_client: Client, now: Instant) -> SessionHandle { + pub(crate) fn new( + session: Session, + spanner_client: Client, + task_tracker: Arc, + now: Instant, + ) -> SessionHandle { SessionHandle { session, spanner_client, + task_tracker, valid: true, deleted: false, last_used_at: now, @@ -51,7 +143,7 @@ impl SessionHandle { match arg { Ok(s) => Ok(s), Err(e) => { - if e.code() == Code::NotFound && e.message().contains("Session not found:") { + if is_session_not_found_status(&e) { tracing::debug!("session invalidate {}", self.session.name); self.delete().await; } @@ -61,15 +153,21 @@ impl SessionHandle { } async fn delete(&mut self) { + if self.deleted { + self.valid = false; + return; + } self.valid = false; - let session_name = &self.session.name; - let request = DeleteSessionRequest { - name: session_name.to_string(), - }; - match self.spanner_client.delete_session(request, None).await { - Ok(_) => self.deleted = true, - Err(e) => tracing::error!("failed to delete session {}, {:?}", session_name, e), - }; + self.deleted = true; + let session_name = self.session.name.to_string(); + let mut client = self.spanner_client.clone(); + let tracker = self.task_tracker.clone(); + tracker.spawn(async move { + let request = DeleteSessionRequest { name: session_name }; + if let Err(e) = client.delete_session(request, None).await { + tracing::error!("failed to delete session, {:?}", e); + } + }); } } @@ -86,12 +184,36 @@ impl ManagedSession { session: Some(session), } } + + pub(crate) async fn renew(&mut self, wait_timeout: Option) -> Result<(), SessionError> { + if let Some(mut session) = self.session.take() { + if session.valid { + // Ensure the orphaned session is deleted before returning it to the pool so + // it will not be handed out again. + session.delete().await; + } + self.session_pool.recycle(session); + } + let wait_timeout = wait_timeout.unwrap_or(self.session_pool.config.session_get_timeout); + match self.session_pool.acquire_with_timeout(wait_timeout).await { + Ok(mut replacement) => { + self.session = replacement.session.take(); + Ok(()) + } + Err(SessionError::SessionGetTimeout) => { + tracing::warn!("session renewal timed out while waiting for a new session"); + Err(SessionError::SessionGetTimeout) + } + Err(e) => Err(e), + } + } } impl Drop for ManagedSession { fn drop(&mut self) { - let session = self.session.take().unwrap(); - self.session_pool.recycle(session); + if let Some(session) = self.session.take() { + self.session_pool.recycle(session); + } } } @@ -152,6 +274,7 @@ impl Sessions { } fn release(&mut self, session: SessionHandle) { + assert!(self.num_inuse > 0, "release called without an in-use session"); self.num_inuse -= 1; if session.valid { self.available_sessions.push_back(session); @@ -204,6 +327,7 @@ struct SessionPool { inner: Arc>, session_creation_sender: UnboundedSender, config: Arc, + task_tracker: Arc, } impl SessionPool { @@ -212,8 +336,9 @@ impl SessionPool { conn_pool: &ConnectionManager, session_creation_sender: UnboundedSender, config: Arc, + task_tracker: Arc, ) -> Result { - let available_sessions = Self::init_pool(database, conn_pool, config.min_opened).await?; + let available_sessions = Self::init_pool(database, conn_pool, config.min_opened, task_tracker.clone()).await?; Ok(SessionPool { inner: Arc::new(RwLock::new(Sessions { available_sessions, @@ -224,6 +349,7 @@ impl SessionPool { })), session_creation_sender, config, + task_tracker, }) } @@ -231,6 +357,7 @@ impl SessionPool { database: String, conn_pool: &ConnectionManager, min_opened: usize, + task_tracker: Arc, ) -> Result, Status> { let channel_num = conn_pool.num(); let creation_count_per_channel = min_opened / channel_num; @@ -247,7 +374,8 @@ impl SessionPool { }; let next_client = conn_pool.conn().with_metadata(client_metadata(&database)); let database = database.clone(); - tasks.spawn(async move { batch_create_sessions(next_client, &database, creation_count).await }); + let tracker = task_tracker.clone(); + tasks.spawn(async move { batch_create_sessions(next_client, &database, creation_count, tracker).await }); } while let Some(r) = tasks.join_next().await { let new_sessions = r.map_err(|e| Status::from_error(e.into()))??; @@ -268,6 +396,10 @@ impl SessionPool { /// The client on the waiting list will be notified when another client's session has finished and /// when the process of replenishing the available sessions is complete. async fn acquire(&self) -> Result { + self.acquire_with_timeout(self.config.session_get_timeout).await + } + + async fn acquire_with_timeout(&self, wait_timeout: Duration) -> Result { loop { let (on_session_acquired, session_count) = { let mut sessions = self.inner.write(); @@ -291,7 +423,7 @@ impl SessionPool { } // Wait for the session available notification. - match timeout(self.config.session_get_timeout, on_session_acquired).await { + match timeout(wait_timeout, on_session_acquired).await { Ok(Ok(())) => { let mut sessions = self.inner.write(); if let Some(mut s) = sessions.take() { @@ -447,10 +579,21 @@ impl TryAs for SessionError { } } +impl From for Status { + fn from(value: SessionError) -> Self { + match value { + SessionError::GRPC(status) => status, + SessionError::SessionGetTimeout => Status::new(Code::DeadlineExceeded, "session get time out"), + SessionError::FailedToCreateSession => Status::new(Code::Internal, "failed to create session"), + } + } +} + pub(crate) struct SessionManager { session_pool: SessionPool, cancel: CancellationToken, tasks: Mutex>>, + task_tracker: Arc, } impl SessionManager { @@ -461,7 +604,15 @@ impl SessionManager { ) -> Result, Status> { let database = database.into(); let (sender, receiver) = mpsc::unbounded_channel(); - let session_pool = SessionPool::new(database.clone(), &conn_pool, sender, Arc::new(config.clone())).await?; + let task_tracker = Arc::new(TaskTracker::new()); + let session_pool = SessionPool::new( + database.clone(), + &conn_pool, + sender, + Arc::new(config.clone()), + task_tracker.clone(), + ) + .await?; let cancel = CancellationToken::new(); let task_session_cleaner = Self::spawn_health_check_task(config, session_pool.clone(), cancel.clone()); @@ -472,6 +623,7 @@ impl SessionManager { session_pool, cancel, tasks: Mutex::new(vec![task_session_cleaner, task_session_creator]), + task_tracker, }; Ok(Arc::new(sm)) } @@ -494,6 +646,8 @@ impl SessionManager { let _ = task.await; } self.session_pool.close().await; + self.task_tracker.close(); + self.task_tracker.wait().await; } fn spawn_session_creation_task( @@ -505,6 +659,7 @@ impl SessionManager { ) -> JoinHandle<()> { tokio::spawn(async move { let mut tasks = JoinSet::default(); + let tracker = session_pool.task_tracker.clone(); loop { select! { biased; @@ -516,7 +671,13 @@ impl SessionManager { Some(session_count) => { let client = conn_pool.conn().with_metadata(client_metadata(&database)); let database = database.clone(); - tasks.spawn(async move { (session_count, batch_create_sessions(client, &database, session_count).await) }); + let tracker_clone = tracker.clone(); + tasks.spawn(async move { + ( + session_count, + batch_create_sessions(client, &database, session_count, tracker_clone).await + ) + }); }, None => continue }, @@ -614,10 +775,13 @@ async fn batch_create_sessions( spanner_client: Client, database: &str, mut remaining_create_count: usize, + task_tracker: Arc, ) -> Result, Status> { let mut created = Vec::with_capacity(remaining_create_count); while remaining_create_count > 0 { - let sessions = batch_create_session(spanner_client.clone(), database, remaining_create_count).await?; + let sessions = + batch_create_session(spanner_client.clone(), database, remaining_create_count, task_tracker.clone()) + .await?; // Spanner could return less sessions than requested. // In that case, we should do another call using the same gRPC channel. let actually_created = sessions.len(); @@ -631,6 +795,7 @@ async fn batch_create_session( mut spanner_client: Client, database: &str, session_count: usize, + task_tracker: Arc, ) -> Result, Status> { let request = BatchCreateSessionsRequest { database: database.to_string(), @@ -645,7 +810,7 @@ async fn batch_create_session( Ok(response .session .into_iter() - .map(|s| SessionHandle::new(s, spanner_client.clone(), now)) + .map(|s| SessionHandle::new(s, spanner_client.clone(), task_tracker.clone(), now)) .collect::>()) } @@ -664,7 +829,7 @@ mod tests { use parking_lot::RwLock; use serial_test::serial; use tokio::time::sleep; - use tokio_util::sync::CancellationToken; + use tokio_util::{sync::CancellationToken, task::TaskTracker}; use google_cloud_gax::conn::{ConnectionOptions, Environment}; use google_cloud_googleapis::spanner::v1::ExecuteSqlRequest; @@ -1106,7 +1271,8 @@ mod tests { .unwrap(); let client = cm.conn().with_metadata(client_metadata(DATABASE)); let session_count = 125; - let result = batch_create_sessions(client.clone(), DATABASE, session_count).await; + let tracker = Arc::new(TaskTracker::new()); + let result = batch_create_sessions(client.clone(), DATABASE, session_count, tracker).await; match result { Ok(created) => { assert_eq!(session_count, created.len()); diff --git a/spanner/src/transaction.rs b/spanner/src/transaction.rs index 44441803..23d0e478 100644 --- a/spanner/src/transaction.rs +++ b/spanner/src/transaction.rs @@ -1,5 +1,5 @@ -use std::ops::DerefMut; use std::sync::atomic::AtomicI64; +use std::time::Duration; use prost_types::Struct; @@ -7,8 +7,8 @@ use google_cloud_gax::grpc::Status; use google_cloud_gax::retry::RetrySetting; use google_cloud_googleapis::spanner::v1::request_options::Priority; use google_cloud_googleapis::spanner::v1::{ - execute_sql_request::QueryMode, execute_sql_request::QueryOptions as ExecuteQueryOptions, ExecuteSqlRequest, - ReadRequest, RequestOptions, TransactionSelector, + execute_sql_request::QueryMode, execute_sql_request::QueryOptions as ExecuteQueryOptions, transaction_selector, + ExecuteSqlRequest, ReadRequest, RequestOptions, TransactionSelector, }; use crate::key::{Key, KeySet}; @@ -24,6 +24,12 @@ pub struct CallOptions { pub retry: Option, } +impl CallOptions { + pub fn session_renew_timeout(&self) -> Option { + self.retry.as_ref().and_then(|setting| setting.max_delay) + } +} + #[derive(Clone)] pub struct ReadOptions { /// The index to use for reading. If non-empty, you can only read columns @@ -160,12 +166,13 @@ impl Transaction { directed_read_options: None, last_statement: false, }; - let session = self.session.as_mut().unwrap().deref_mut(); let reader = StatementReader { enable_resume: options.enable_resume, request, }; - RowIterator::new(session, reader, Some(options.call_options)).await + let allow_refresh = self.can_retry_on_session_not_found(); + let session = self.as_mut_session(); + RowIterator::new(session, reader, Some(options.call_options), allow_refresh).await } /// read returns a RowIterator for reading multiple rows from the database. @@ -226,9 +233,10 @@ impl Transaction { lock_hint: 0, }; - let session = self.as_mut_session(); let reader = TableReader { request }; - RowIterator::new(session, reader, Some(options.call_options)).await + let allow_refresh = self.can_retry_on_session_not_found(); + let session = self.as_mut_session(); + RowIterator::new(session, reader, Some(options.call_options), allow_refresh).await } /// read returns a RowIterator for reading multiple rows from the database. @@ -264,6 +272,13 @@ impl Transaction { reader.next().await } + pub(crate) fn can_retry_on_session_not_found(&self) -> bool { + matches!( + self.transaction_selector.selector, + Some(transaction_selector::Selector::SingleUse(_)) + ) + } + pub(crate) fn get_session_name(&self) -> String { self.session.as_ref().unwrap().session.name.to_string() } diff --git a/spanner/src/transaction_ro.rs b/spanner/src/transaction_ro.rs index 28203561..ce26c4fb 100644 --- a/spanner/src/transaction_ro.rs +++ b/spanner/src/transaction_ro.rs @@ -6,11 +6,11 @@ use time::OffsetDateTime; use crate::key::KeySet; use crate::reader::{Reader, RowIterator, StatementReader, TableReader}; -use crate::session::ManagedSession; +use crate::session::{is_session_not_found_status, ManagedSession}; use crate::statement::Statement; use crate::transaction::{CallOptions, QueryOptions, ReadOptions, Transaction}; use crate::value::TimestampBound; -use google_cloud_gax::grpc::Status; +use google_cloud_gax::grpc::{Code, Status}; use google_cloud_googleapis::spanner::v1::transaction_options::IsolationLevel; use google_cloud_googleapis::spanner::v1::{ transaction_options, transaction_selector, BeginTransactionRequest, DirectedReadOptions, ExecuteSqlRequest, @@ -37,6 +37,11 @@ pub struct ReadOnlyTransaction { pub rts: Option, } +pub(crate) struct ReadOnlyBeginError { + pub(crate) status: Status, + pub(crate) session: ManagedSession, +} + impl Deref for ReadOnlyTransaction { type Target = Transaction; @@ -72,10 +77,21 @@ impl ReadOnlyTransaction { /// begin starts a snapshot read-only Transaction on Cloud Spanner. pub async fn begin( - mut session: ManagedSession, + session: ManagedSession, tb: TimestampBound, options: CallOptions, ) -> Result { + match ReadOnlyTransaction::begin_internal(session, tb, options).await { + Ok(tx) => Ok(tx), + Err(err) => Err(err.status), + } + } + + pub(crate) async fn begin_internal( + mut session: ManagedSession, + tb: TimestampBound, + options: CallOptions, + ) -> Result { let request = BeginTransactionRequest { session: session.session.name.to_string(), options: Some(TransactionOptions { @@ -105,7 +121,7 @@ impl ReadOnlyTransaction { rts: Some(OffsetDateTime::from(st)), }) } - Err(e) => Err(e), + Err(e) => Err(ReadOnlyBeginError { status: e, session }), } } } @@ -121,6 +137,8 @@ pub struct Partition { /// same snapshot of the database. pub struct BatchReadOnlyTransaction { base_tx: ReadOnlyTransaction, + timestamp_bound: TimestampBound, + call_options: CallOptions, } impl Deref for BatchReadOnlyTransaction { @@ -143,8 +161,12 @@ impl BatchReadOnlyTransaction { tb: TimestampBound, options: CallOptions, ) -> Result { - let tx = ReadOnlyTransaction::begin(session, tb, options).await?; - Ok(BatchReadOnlyTransaction { base_tx: tx }) + let tx = ReadOnlyTransaction::begin(session, tb.clone(), options.clone()).await?; + Ok(BatchReadOnlyTransaction { + base_tx: tx, + timestamp_bound: tb, + call_options: options, + }) } /// partition_read returns a list of Partitions that can be used to read rows from @@ -176,54 +198,80 @@ impl BatchReadOnlyTransaction { data_boost_enabled: bool, directed_read_options: Option, ) -> Result>, Status> { + let table_name = table.to_string(); let columns: Vec = columns.iter().map(|x| x.to_string()).collect(); - let inner_keyset = keys.into().inner; - let request = PartitionReadRequest { - session: self.get_session_name(), - transaction: Some(self.transaction_selector.clone()), - table: table.to_string(), - index: ro.index.clone(), - columns: columns.clone(), - key_set: Some(inner_keyset.clone()), - partition_options: po, - }; - let result = match self - .as_mut_session() - .spanner_client - .partition_read(request, ro.call_options.retry) - .await - { - Ok(r) => Ok(r - .into_inner() - .partitions - .into_iter() - .map(|x| Partition { - reader: TableReader { - request: ReadRequest { - session: self.get_session_name(), - transaction: Some(self.transaction_selector.clone()), - table: table.to_string(), - index: ro.index.clone(), - columns: columns.clone(), - key_set: Some(inner_keyset.clone()), - limit: ro.limit, - resume_token: vec![], - partition_token: x.partition_token, - request_options: Transaction::create_request_options( - ro.call_options.priority, - self.base_tx.transaction_tag.clone(), - ), - directed_read_options: directed_read_options.clone(), - data_boost_enabled, - order_by: 0, - lock_hint: 0, - }, - }, - }) - .collect()), - Err(e) => Err(e), - }; - self.as_mut_session().invalidate_if_needed(result).await + let keyset = keys.into(); + let inner_keyset = keyset.inner; + let partition_options = po; + let directed = directed_read_options; + let ReadOptions { + index, + limit, + call_options, + } = ro; + loop { + let request = PartitionReadRequest { + session: self.get_session_name(), + transaction: Some(self.transaction_selector.clone()), + table: table_name.clone(), + index: index.clone(), + columns: columns.clone(), + key_set: Some(inner_keyset.clone()), + partition_options, + }; + match self + .as_mut_session() + .spanner_client + .partition_read(request, call_options.retry.clone()) + .await + { + Ok(response) => { + let partitions = response + .into_inner() + .partitions + .into_iter() + .map(|x| Partition { + reader: TableReader { + request: ReadRequest { + session: self.get_session_name(), + transaction: Some(self.transaction_selector.clone()), + table: table_name.clone(), + index: index.clone(), + columns: columns.clone(), + key_set: Some(inner_keyset.clone()), + limit, + resume_token: vec![], + partition_token: x.partition_token, + request_options: Transaction::create_request_options( + call_options.priority, + self.base_tx.transaction_tag.clone(), + ), + directed_read_options: directed.clone(), + data_boost_enabled, + order_by: 0, + lock_hint: 0, + }, + }, + }) + .collect::>>(); + return self.as_mut_session().invalidate_if_needed(Ok(partitions)).await; + } + Err(status) => { + if is_session_not_found_status(&status) { + let _ = self + .as_mut_session() + .invalidate_if_needed::<()>(Err(status.clone())) + .await; + self.restart_snapshot().await?; + continue; + } + return self + .as_mut_session() + .invalidate_if_needed::>>(Err(status)) + .await; + } + } + } } /// partition_query returns a list of Partitions that can be used to execute a query against the database. @@ -241,56 +289,81 @@ impl BatchReadOnlyTransaction { data_boost_enabled: bool, directed_read_options: Option, ) -> Result>, Status> { - let request = PartitionQueryRequest { - session: self.get_session_name(), - transaction: Some(self.transaction_selector.clone()), - sql: stmt.sql.clone(), - params: Some(prost_types::Struct { - fields: stmt.params.clone(), - }), - param_types: stmt.param_types.clone(), - partition_options: po, - }; - let result = match self - .as_mut_session() - .spanner_client - .partition_query(request.clone(), qo.call_options.retry.clone()) - .await - { - Ok(r) => Ok(r - .into_inner() - .partitions - .into_iter() - .map(|x| Partition { - reader: StatementReader { - enable_resume: qo.enable_resume, - request: ExecuteSqlRequest { - session: self.get_session_name(), - transaction: Some(self.transaction_selector.clone()), - sql: stmt.sql.clone(), - params: Some(prost_types::Struct { - fields: stmt.params.clone(), - }), - param_types: stmt.param_types.clone(), - resume_token: vec![], - query_mode: 0, - partition_token: x.partition_token, - seqno: 0, - query_options: qo.optimizer_options.clone(), - request_options: Transaction::create_request_options( - qo.call_options.priority, - self.base_tx.transaction_tag.clone(), - ), - data_boost_enabled, - directed_read_options: directed_read_options.clone(), - last_statement: false, - }, - }, - }) - .collect()), - Err(e) => Err(e), - }; - self.as_mut_session().invalidate_if_needed(result).await + let partition_options = po; + let directed = directed_read_options; + let QueryOptions { + mode: _, + optimizer_options, + call_options, + enable_resume, + } = qo; + loop { + let request = PartitionQueryRequest { + session: self.get_session_name(), + transaction: Some(self.transaction_selector.clone()), + sql: stmt.sql.clone(), + params: Some(prost_types::Struct { + fields: stmt.params.clone(), + }), + param_types: stmt.param_types.clone(), + partition_options, + }; + match self + .as_mut_session() + .spanner_client + .partition_query(request, call_options.retry.clone()) + .await + { + Ok(response) => { + let partitions = response + .into_inner() + .partitions + .into_iter() + .map(|x| Partition { + reader: StatementReader { + enable_resume, + request: ExecuteSqlRequest { + session: self.get_session_name(), + transaction: Some(self.transaction_selector.clone()), + sql: stmt.sql.clone(), + params: Some(prost_types::Struct { + fields: stmt.params.clone(), + }), + param_types: stmt.param_types.clone(), + resume_token: vec![], + query_mode: 0, + partition_token: x.partition_token, + seqno: 0, + query_options: optimizer_options.clone(), + request_options: Transaction::create_request_options( + call_options.priority, + self.base_tx.transaction_tag.clone(), + ), + data_boost_enabled, + directed_read_options: directed.clone(), + last_statement: false, + }, + }, + }) + .collect::>>(); + return self.as_mut_session().invalidate_if_needed(Ok(partitions)).await; + } + Err(status) => { + if is_session_not_found_status(&status) { + let _ = self + .as_mut_session() + .invalidate_if_needed::<()>(Err(status.clone())) + .await; + self.restart_snapshot().await?; + continue; + } + return self + .as_mut_session() + .invalidate_if_needed::>>(Err(status)) + .await; + } + } + } } /// execute runs a single Partition obtained from partition_read or partition_query. @@ -299,7 +372,36 @@ impl BatchReadOnlyTransaction { partition: Partition, option: Option, ) -> Result, Status> { + let allow_refresh = self.base_tx.can_retry_on_session_not_found(); let session = self.as_mut_session(); - RowIterator::new(session, partition.reader, option).await + RowIterator::new(session, partition.reader, option, allow_refresh).await + } + + async fn restart_snapshot(&mut self) -> Result<(), Status> { + let mut session = self + .base_tx + .take_session() + .ok_or_else(|| Status::new(Code::Internal, "missing session while restarting batch transaction"))?; + loop { + match ReadOnlyTransaction::begin_internal(session, self.timestamp_bound.clone(), self.call_options.clone()) + .await + { + Ok(tx) => { + self.base_tx = tx; + return Ok(()); + } + Err(err) => { + if is_session_not_found_status(&err.status) { + session = err.session; + session + .renew(self.call_options.session_renew_timeout()) + .await + .map_err(Status::from)?; + continue; + } + return Err(err.status); + } + } + } } } diff --git a/spanner/tests/session_refresh_test.rs b/spanner/tests/session_refresh_test.rs new file mode 100644 index 00000000..9a2fe0bb --- /dev/null +++ b/spanner/tests/session_refresh_test.rs @@ -0,0 +1,313 @@ +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; + +use serial_test::serial; +use time::OffsetDateTime; +use tokio::time::{timeout, Duration}; + +use common::*; +use google_cloud_gax::conn::{ConnectionOptions, Environment}; +use google_cloud_googleapis::spanner::v1::{DeleteSessionRequest, ListSessionsRequest}; +use google_cloud_spanner::apiv1::conn_pool::{ConnectionManager, SPANNER}; +use google_cloud_spanner::apiv1::spanner_client::Client as LowLevelClient; +use google_cloud_spanner::client::Error as ClientError; +use google_cloud_spanner::key::Key; +use google_cloud_spanner::statement::Statement; + +mod common; + +const DATABASE: &str = "projects/local-project/instances/test-instance/databases/local-database"; + +#[ctor::ctor] +fn init() { + let filter = tracing_subscriber::filter::EnvFilter::from_default_env() + .add_directive("google_cloud_spanner=trace".parse().unwrap()); + let _ = tracing_subscriber::fmt().with_env_filter(filter).try_init(); + std::env::set_var("SPANNER_EMULATOR_HOST", "localhost:9010"); +} + +#[tokio::test] +#[serial] +async fn single_read_recovers_from_deleted_session() { + purge_sessions().await; + let client = create_data_client().await; + + let now = OffsetDateTime::now_utc(); + let user_id = format!("sess_user_{}", now.unix_timestamp()); + let ms = vec![create_user_mutation(&user_id, &now)]; + client.apply(ms).await.unwrap(); + + let mut single = client.single().await.unwrap(); + + delete_only_session().await; + + let read_result = timeout(Duration::from_secs(30), single.read("User", &["UserId"], Key::new(&user_id))) + .await + .expect("read timed out"); + assert!( + read_result.is_ok(), + "expected client to transparently replace a deleted session, but saw: {:?}", + read_result.err() + ); +} + +#[tokio::test] +#[serial] +async fn single_query_recovers_from_deleted_session() { + purge_sessions().await; + let client = create_data_client().await; + + let now = OffsetDateTime::now_utc(); + let user_id = format!("sess_query_user_{}", now.unix_timestamp()); + let ms = vec![create_user_mutation(&user_id, &now)]; + client.apply(ms).await.unwrap(); + + let mut single = client.single().await.unwrap(); + + delete_only_session().await; + + let mut stmt = Statement::new("SELECT UserId FROM User WHERE UserId = @UserId"); + stmt.add_param("UserId", &user_id); + + let mut iter = timeout(Duration::from_secs(30), single.query(stmt)) + .await + .expect("query timed out") + .expect("query failed after session refresh"); + + let row = timeout(Duration::from_secs(30), iter.next()) + .await + .expect("row fetch timed out") + .expect("row fetch failed") + .expect("expected at least one row"); + + let fetched_user_id = row.column_by_name::("UserId").unwrap(); + assert_eq!( + fetched_user_id, user_id, + "expected iterator to surface the inserted row after session refresh" + ); +} + +#[tokio::test] +#[serial] +async fn read_write_tx_recovers_from_deleted_session() { + purge_sessions().await; + let client = create_data_client().await; + + let now = OffsetDateTime::now_utc(); + let user_id = format!("rw_sess_user_{}", now.unix_timestamp()); + let ms = vec![create_user_mutation(&user_id, &now)]; + client.apply(ms).await.unwrap(); + + let delete_once = Arc::new(AtomicBool::new(false)); + let tx_result = timeout( + Duration::from_secs(30), + client.read_write_transaction({ + let delete_once = delete_once.clone(); + move |tx| { + let user_id = user_id.clone(); + let delete_once = delete_once.clone(); + Box::pin(async move { + if !delete_once.swap(true, Ordering::SeqCst) { + delete_only_session().await; + } + let mut reader = tx.read("User", &["UserId"], vec![Key::new(&user_id)]).await?; + while reader.next().await?.is_some() {} + Ok::<(), ClientError>(()) + }) + } + }), + ) + .await + .expect("transaction timed out"); + + assert!( + tx_result.is_ok(), + "expected read_write_transaction to replace a deleted session automatically, saw: {:?}", + tx_result.err() + ); +} + +#[tokio::test] +#[serial] +async fn read_only_tx_recovers_from_deleted_session() { + purge_sessions().await; + let client = create_data_client().await; + + let now = OffsetDateTime::now_utc(); + let user_id = format!("ro_sess_user_{}", now.unix_timestamp()); + let ms = vec![create_user_mutation(&user_id, &now)]; + client.apply(ms).await.unwrap(); + + delete_only_session().await; + + let mut tx = client.read_only_transaction().await.unwrap(); + + let mut reader = timeout(Duration::from_secs(30), tx.read("User", &["UserId"], vec![Key::new(&user_id)])) + .await + .expect("read-only read timed out") + .expect("read-only read failed after session refresh"); + + let row = timeout(Duration::from_secs(30), reader.next()) + .await + .expect("read-only iterator timed out") + .expect("read-only iterator failed") + .expect("expected at least one row"); + + let fetched_user_id = row.column_by_name::("UserId").unwrap(); + assert_eq!( + fetched_user_id, user_id, + "expected read_only_transaction to surface the inserted row after session refresh" + ); +} + +#[tokio::test] +#[serial] +async fn partitioned_update_recovers_from_deleted_session() { + purge_sessions().await; + let client = create_data_client().await; + + let now = OffsetDateTime::now_utc(); + let user_id = format!("pdml_sess_user_{}", now.unix_timestamp()); + let ms = vec![create_user_mutation(&user_id, &now)]; + client.apply(ms).await.unwrap(); + + delete_only_session().await; + + let mut stmt = Statement::new("UPDATE User SET NullableString = @Value WHERE UserId = @UserId"); + let updated_value = format!("updated {}", user_id); + stmt.add_param("UserId", &user_id); + stmt.add_param("Value", &updated_value); + + let updated_rows = timeout(Duration::from_secs(30), client.partitioned_update(stmt)) + .await + .expect("partitioned_update timed out") + .expect("partitioned_update failed after session refresh"); + + assert_eq!( + updated_rows, 1, + "partitioned_update should modify the inserted row even after refreshing the session" + ); +} + +#[tokio::test] +#[serial] +async fn batch_partition_read_recovers_from_deleted_session() { + purge_sessions().await; + let client = create_data_client().await; + + let now = OffsetDateTime::now_utc(); + let user_id = format!("batch_sess_user_{}", now.unix_timestamp()); + let ms = vec![create_user_mutation(&user_id, &now)]; + client.apply(ms).await.unwrap(); + + let mut batch_tx = client.batch_read_only_transaction().await.unwrap(); + + delete_only_session().await; + + let partitions = timeout( + Duration::from_secs(30), + batch_tx.partition_read("User", &["UserId"], vec![Key::new(&user_id)]), + ) + .await + .expect("partition_read timed out") + .expect("partition_read failed after session refresh"); + + assert!(!partitions.is_empty(), "expected at least one partition to be returned"); + + let mut found_row = false; + for partition in partitions { + let mut rows = timeout(Duration::from_secs(30), batch_tx.execute(partition, None)) + .await + .expect("batch execute timed out") + .expect("batch execute failed after session refresh"); + loop { + let maybe_row = timeout(Duration::from_secs(30), rows.next()) + .await + .expect("partition row fetch timed out") + .expect("partition row fetch failed"); + match maybe_row { + Some(row) => { + let fetched_user_id = row.column_by_name::("UserId").unwrap(); + if fetched_user_id == user_id { + found_row = true; + break; + } + } + None => break, + } + } + if found_row { + break; + } + } + assert!( + found_row, + "expected batch partition read to surface the inserted row after session refresh" + ); +} + +async fn delete_only_session() { + let mut raw = new_spanner_client().await; + let response = raw + .list_sessions( + ListSessionsRequest { + database: DATABASE.to_string(), + page_size: 0, + page_token: String::new(), + filter: String::new(), + }, + None, + ) + .await + .expect("list sessions"); + let sessions = response.into_inner().sessions; + assert!( + !sessions.is_empty(), + "expected at least one session to be present before deletion" + ); + raw.delete_session( + DeleteSessionRequest { + name: sessions[0].name.clone(), + }, + None, + ) + .await + .expect("delete session"); +} + +async fn purge_sessions() { + let mut raw = new_spanner_client().await; + let mut page_token = String::new(); + loop { + let response = raw + .list_sessions( + ListSessionsRequest { + database: DATABASE.to_string(), + page_size: 0, + page_token: page_token.clone(), + filter: String::new(), + }, + None, + ) + .await + .expect("list sessions"); + let inner = response.into_inner(); + for session in inner.sessions { + let _ = raw + .delete_session(DeleteSessionRequest { name: session.name }, None) + .await; + } + if inner.next_page_token.is_empty() { + break; + } + page_token = inner.next_page_token; + } +} + +async fn new_spanner_client() -> LowLevelClient { + let host = std::env::var("SPANNER_EMULATOR_HOST").unwrap_or_else(|_| "localhost:9010".to_string()); + let cm = ConnectionManager::new(1, &Environment::Emulator(host), SPANNER, &ConnectionOptions::default()) + .await + .expect("create spanner connection manager"); + cm.conn() +}