diff --git a/generator/generator_config.proto b/generator/generator_config.proto index 84582f76c40a6..a30dab96e678b 100644 --- a/generator/generator_config.proto +++ b/generator/generator_config.proto @@ -184,8 +184,9 @@ message ServiceConfiguration { repeated BespokeMethod bespoke_methods = 30; // If set to `true`, the stub class methods will include an additional - // parameter of type `google::cloud::bigtable_internal::OperationContext&`. - bool experimental_bigtable_operation_context = 31; + // parameter of type + // `google::cloud::_internal::OperationContext&`. + bool experimental_operation_context = 31; } message DiscoveryDocumentDefinedProduct { diff --git a/generator/generator_config.textproto b/generator/generator_config.textproto index ead53ac05c05b..1d368e62c7786 100644 --- a/generator/generator_config.textproto +++ b/generator/generator_config.textproto @@ -584,7 +584,7 @@ service { omit_connection: true omit_stub_factory: true generate_round_robin_decorator: true - experimental_bigtable_operation_context: true + experimental_operation_context: true omitted_rpcs: [ "GenerateInitialChangeStreamPartitions", "ReadChangeStream" @@ -3907,6 +3907,7 @@ service { omit_client: true omit_connection: true omit_stub_factory: true + experimental_operation_context: true omitted_rpcs: [ "GetSession", "ListSessions", diff --git a/generator/internal/codegen_utils_test.cc b/generator/internal/codegen_utils_test.cc index 7b52663d35382..63790fadb73b7 100644 --- a/generator/internal/codegen_utils_test.cc +++ b/generator/internal/codegen_utils_test.cc @@ -323,13 +323,13 @@ TEST(ProcessCommandLineArgs, ProcessExperimental) { EXPECT_THAT(*result, Contains(Pair("experimental", "true"))); } -TEST(ProcessCommandLineArgs, ProcessExperimentalBigtableOperationContext) { +TEST(ProcessCommandLineArgs, ProcessExperimentalOperationContext) { auto result = ProcessCommandLineArgs( "product_path=google/cloud/bigtable/" - ",experimental_bigtable_operation_context=true"); + ",experimental_operation_context=true"); ASSERT_THAT(result, IsOk()); - EXPECT_THAT(*result, Contains(Pair("experimental_bigtable_operation_context", - "true"))); + EXPECT_THAT(*result, + Contains(Pair("experimental_operation_context", "true"))); } TEST(ProcessCommandLineArgs, ProcessServiceNameMapping) { diff --git a/generator/internal/descriptor_utils.cc b/generator/internal/descriptor_utils.cc index dbf667d62474c..a9fd8615e3763 100644 --- a/generator/internal/descriptor_utils.cc +++ b/generator/internal/descriptor_utils.cc @@ -850,27 +850,25 @@ VarsDictionary CreateServiceVars( SetRetryStatusCodeExpression(vars); vars["transient_errors_comment"] = TransientErrorsComment(vars); SetLongrunningOperationServiceVars(descriptor, vars); - auto const experimental_bigtable_operation_context = - vars.find("experimental_bigtable_operation_context"); - if (experimental_bigtable_operation_context != vars.end() && - experimental_bigtable_operation_context->second == "true") { + auto const experimental_operation_context = + vars.find("experimental_operation_context"); + if (experimental_operation_context != vars.end() && + experimental_operation_context->second == "true") { + auto const& ns = vars.find("product_internal_namespace")->second; vars["op_ctx_decl"] = - ",\n google::cloud::bigtable_internal::OperationContext& " - "operation_context"; + absl::StrCat(",\n ", ns, "::OperationContext& operation_context"); vars["op_ctx_arg"] = ", operation_context"; vars["op_ctx_cap"] = ", &operation_context"; vars["op_ctx_stub_decl"] = - ",\n google::cloud::bigtable_internal::OperationContext&"; + absl::StrCat(",\n ", ns, "::OperationContext&"); vars["op_ctx_shared_decl"] = - ",\n " - "std::shared_ptr " - "operation_context"; + absl::StrCat(",\n std::shared_ptr<", ns, + "::OperationContext> operation_context"); vars["op_ctx_shared_arg"] = ", std::move(operation_context)"; vars["op_ctx_shared_cap"] = ", operation_context = std::move(operation_context)"; vars["op_ctx_shared_stub_decl"] = - ",\n " - "std::shared_ptr"; + absl::StrCat(",\n std::shared_ptr<", ns, "::OperationContext>"); } else { vars["op_ctx_decl"] = ""; vars["op_ctx_arg"] = ""; diff --git a/generator/internal/service_code_generator.cc b/generator/internal/service_code_generator.cc index b7b1dad62bf38..8f99fe2251cd6 100644 --- a/generator/internal/service_code_generator.cc +++ b/generator/internal/service_code_generator.cc @@ -105,8 +105,8 @@ bool ServiceCodeGenerator::IsExperimental() const { return iter != vars().end() && iter->second == "true"; } -bool ServiceCodeGenerator::HasExperimentalBigtableOperationContext() const { - auto iter = vars().find("experimental_bigtable_operation_context"); +bool ServiceCodeGenerator::HasExperimentalOperationContext() const { + auto iter = vars().find("experimental_operation_context"); return iter != vars().end() && iter->second == "true"; } diff --git a/generator/internal/service_code_generator.h b/generator/internal/service_code_generator.h index f85fdfeabd5dd..2635706f07aab 100644 --- a/generator/internal/service_code_generator.h +++ b/generator/internal/service_code_generator.h @@ -128,9 +128,9 @@ class ServiceCodeGenerator : public GeneratorInterface { bool IsExperimental() const; /** - * Determines if the service enables experimental Bigtable OperationContext. + * Determines if the service enables experimental OperationContext. */ - bool HasExperimentalBigtableOperationContext() const; + bool HasExperimentalOperationContext() const; /** * Determines if the service contains at least one method that requires diff --git a/generator/internal/stub_generator.cc b/generator/internal/stub_generator.cc index 205599a8ce9fa..7e6247f5b1d19 100644 --- a/generator/internal/stub_generator.cc +++ b/generator/internal/stub_generator.cc @@ -85,15 +85,15 @@ Status StubGenerator::GenerateHeader() { HeaderProtobufGenCodeIncludes( {vars("proto_grpc_header_path"), include_lro_header ? "google/longrunning/operations.grpc.pb.h" : ""}); + HeaderLocalIncludes( + {HasExperimentalOperationContext() + ? absl::StrCat(vars("product_path"), "internal/operation_context.h") + : ""}); HeaderSystemIncludes({"memory", "utility"}); HeaderGrpcPortsDefInclude(); auto result = HeaderOpenNamespaces(NamespaceType::kInternal); if (!result.ok()) return result; - if (HasExperimentalBigtableOperationContext()) { - HeaderPrint("\nclass OperationContext;\n"); - } - // Abstract interface Stub base class HeaderPrint( // clang-format off "\n" diff --git a/generator/standalone_main.cc b/generator/standalone_main.cc index fc4ee09150668..3ef12319db596 100644 --- a/generator/standalone_main.cc +++ b/generator/standalone_main.cc @@ -303,9 +303,9 @@ std::vector> GenerateCodeFromProtos( args.emplace_back( "--cpp_codegen_opt=generate_round_robin_decorator=true"); } - if (service.experimental_bigtable_operation_context()) { + if (service.experimental_operation_context()) { args.emplace_back( - "--cpp_codegen_opt=experimental_bigtable_operation_context=true"); + "--cpp_codegen_opt=experimental_operation_context=true"); } args.emplace_back("--cpp_codegen_opt=service_endpoint_env_var=" + service.service_endpoint_env_var()); diff --git a/google/cloud/bigtable/internal/bigtable_auth_decorator.cc b/google/cloud/bigtable/internal/bigtable_auth_decorator.cc index d1987006b0e68..5e97fe06bf7be 100644 --- a/google/cloud/bigtable/internal/bigtable_auth_decorator.cc +++ b/google/cloud/bigtable/internal/bigtable_auth_decorator.cc @@ -41,8 +41,7 @@ std::unique_ptr context, Options const& options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { using ErrorStream = ::google::cloud::internal::StreamingReadRpcError< google::bigtable::v2::ReadRowsResponse>; auto status = auth_->ConfigureContext(*context); @@ -56,8 +55,7 @@ std::unique_ptr context, Options const& options, google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { using ErrorStream = ::google::cloud::internal::StreamingReadRpcError< google::bigtable::v2::SampleRowKeysResponse>; auto status = auth_->ConfigureContext(*context); @@ -69,7 +67,7 @@ BigtableAuth::SampleRowKeys( StatusOr BigtableAuth::MutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::MutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { auto status = auth_->ConfigureContext(context); if (!status.ok()) return status; return child_->MutateRow(context, options, request, operation_context); @@ -80,8 +78,7 @@ std::unique_ptr context, Options const& options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { using ErrorStream = ::google::cloud::internal::StreamingReadRpcError< google::bigtable::v2::MutateRowsResponse>; auto status = auth_->ConfigureContext(*context); @@ -94,7 +91,7 @@ StatusOr BigtableAuth::CheckAndMutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { auto status = auth_->ConfigureContext(context); if (!status.ok()) return status; return child_->CheckAndMutateRow(context, options, request, @@ -104,7 +101,7 @@ BigtableAuth::CheckAndMutateRow( StatusOr BigtableAuth::PingAndWarm( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PingAndWarmRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { auto status = auth_->ConfigureContext(context); if (!status.ok()) return status; return child_->PingAndWarm(context, options, request, operation_context); @@ -114,7 +111,7 @@ StatusOr BigtableAuth::ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { auto status = auth_->ConfigureContext(context); if (!status.ok()) return status; return child_->ReadModifyWriteRow(context, options, request, @@ -124,7 +121,7 @@ BigtableAuth::ReadModifyWriteRow( StatusOr BigtableAuth::PrepareQuery( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PrepareQueryRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { auto status = auth_->ConfigureContext(context); if (!status.ok()) return status; return child_->PrepareQuery(context, options, request, operation_context); @@ -135,8 +132,7 @@ std::unique_ptr context, Options const& options, google::bigtable::v2::ExecuteQueryRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { using ErrorStream = ::google::cloud::internal::StreamingReadRpcError< google::bigtable::v2::ExecuteQueryResponse>; auto status = auth_->ConfigureContext(*context); @@ -149,7 +145,7 @@ StatusOr BigtableAuth::GetClientConfiguration( grpc::ClientContext& context, Options const& options, google::bigtable::v2::GetClientConfigurationRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { auto status = auth_->ConfigureContext(context); if (!status.ok()) return status; return child_->GetClientConfiguration(context, options, request, @@ -163,8 +159,7 @@ BigtableAuth::AsyncOpenTable( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { using StreamAuth = google::cloud::internal::AsyncStreamingReadWriteRpcAuth< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>; @@ -186,8 +181,7 @@ BigtableAuth::AsyncOpenAuthorizedView( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { using StreamAuth = google::cloud::internal::AsyncStreamingReadWriteRpcAuth< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>; @@ -209,8 +203,7 @@ BigtableAuth::AsyncOpenMaterializedView( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { using StreamAuth = google::cloud::internal::AsyncStreamingReadWriteRpcAuth< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>; @@ -232,8 +225,7 @@ BigtableAuth::AsyncReadRows( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { using StreamAuth = google::cloud::internal::AsyncStreamingReadRpcAuth< google::bigtable::v2::ReadRowsResponse>; @@ -255,8 +247,7 @@ BigtableAuth::AsyncSampleRowKeys( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { using StreamAuth = google::cloud::internal::AsyncStreamingReadRpcAuth< google::bigtable::v2::SampleRowKeysResponse>; @@ -277,8 +268,7 @@ BigtableAuth::AsyncMutateRow( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::MutateRowRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return auth_->AsyncConfigureContext(std::move(context)) .then([cq, child = child_, options = std::move(options), request, operation_context = std::move(operation_context)]( @@ -303,8 +293,7 @@ BigtableAuth::AsyncMutateRows( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { using StreamAuth = google::cloud::internal::AsyncStreamingReadRpcAuth< google::bigtable::v2::MutateRowsResponse>; @@ -325,8 +314,7 @@ BigtableAuth::AsyncCheckAndMutateRow( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return auth_->AsyncConfigureContext(std::move(context)) .then([cq, child = child_, options = std::move(options), request, operation_context = std::move(operation_context)]( @@ -350,8 +338,7 @@ BigtableAuth::AsyncPingAndWarm( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::PingAndWarmRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return auth_->AsyncConfigureContext(std::move(context)) .then([cq, child = child_, options = std::move(options), request, operation_context = std::move(operation_context)]( @@ -375,8 +362,7 @@ BigtableAuth::AsyncReadModifyWriteRow( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return auth_->AsyncConfigureContext(std::move(context)) .then([cq, child = child_, options = std::move(options), request, operation_context = std::move(operation_context)]( @@ -400,8 +386,7 @@ BigtableAuth::AsyncPrepareQuery( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::PrepareQueryRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return auth_->AsyncConfigureContext(std::move(context)) .then([cq, child = child_, options = std::move(options), request, operation_context = std::move(operation_context)]( diff --git a/google/cloud/bigtable/internal/bigtable_auth_decorator.h b/google/cloud/bigtable/internal/bigtable_auth_decorator.h index 955ed16f7a349..2c38966b791c2 100644 --- a/google/cloud/bigtable/internal/bigtable_auth_decorator.h +++ b/google/cloud/bigtable/internal/bigtable_auth_decorator.h @@ -45,136 +45,124 @@ class BigtableAuth : public BigtableStub { google::bigtable::v2::ReadRowsResponse>> ReadRows(std::shared_ptr context, Options const& options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr + std::shared_ptr operation_context) override; std::unique_ptr> - SampleRowKeys( - std::shared_ptr context, Options const& options, - google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr - operation_context) override; + SampleRowKeys(std::shared_ptr context, + Options const& options, + google::bigtable::v2::SampleRowKeysRequest const& request, + std::shared_ptr + operation_context) override; StatusOr MutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::MutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; std::unique_ptr> MutateRows(std::shared_ptr context, Options const& options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr + std::shared_ptr operation_context) override; StatusOr CheckAndMutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; StatusOr PingAndWarm( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PingAndWarmRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; StatusOr ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; StatusOr PrepareQuery( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PrepareQueryRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; std::unique_ptr> - ExecuteQuery( - std::shared_ptr context, Options const& options, - google::bigtable::v2::ExecuteQueryRequest const& request, - std::shared_ptr - operation_context) override; + ExecuteQuery(std::shared_ptr context, + Options const& options, + google::bigtable::v2::ExecuteQueryRequest const& request, + std::shared_ptr + operation_context) override; StatusOr GetClientConfiguration( grpc::ClientContext& context, Options const& options, google::bigtable::v2::GetClientConfigurationRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>> - AsyncOpenTable( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) override; + AsyncOpenTable(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>> - AsyncOpenAuthorizedView( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) override; + AsyncOpenAuthorizedView(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>> - AsyncOpenMaterializedView( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) override; + AsyncOpenMaterializedView(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::ReadRowsResponse>> - AsyncReadRows( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr - operation_context) override; + AsyncReadRows(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::ReadRowsRequest const& request, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::SampleRowKeysResponse>> - AsyncSampleRowKeys( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr - operation_context) override; + AsyncSampleRowKeys(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::SampleRowKeysRequest const& request, + std::shared_ptr + operation_context) override; future> AsyncMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::MutateRowRequest const& request, - std::shared_ptr - operation_context) override; + std::shared_ptr operation_context) + override; std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::MutateRowsResponse>> - AsyncMutateRows( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr - operation_context) override; + AsyncMutateRows(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::MutateRowsRequest const& request, + std::shared_ptr + operation_context) override; future> AsyncCheckAndMutateRow( @@ -182,16 +170,16 @@ class BigtableAuth : public BigtableStub { std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - std::shared_ptr - operation_context) override; + std::shared_ptr operation_context) + override; future> AsyncPingAndWarm( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::PingAndWarmRequest const& request, - std::shared_ptr - operation_context) override; + std::shared_ptr operation_context) + override; future> AsyncReadModifyWriteRow( @@ -199,17 +187,16 @@ class BigtableAuth : public BigtableStub { std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - std::shared_ptr - operation_context) override; + std::shared_ptr operation_context) + override; future> - AsyncPrepareQuery( - google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request, - std::shared_ptr - operation_context) override; + AsyncPrepareQuery(google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::PrepareQueryRequest const& request, + std::shared_ptr + operation_context) override; private: std::shared_ptr auth_; diff --git a/google/cloud/bigtable/internal/bigtable_channel_refresh.cc b/google/cloud/bigtable/internal/bigtable_channel_refresh.cc index a484dde4c483d..76df90e866d4f 100644 --- a/google/cloud/bigtable/internal/bigtable_channel_refresh.cc +++ b/google/cloud/bigtable/internal/bigtable_channel_refresh.cc @@ -24,7 +24,7 @@ std::unique_ptr client_context, Options const& options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return child_->ReadRows(std::move(client_context), options, request, std::move(operation_context)); } @@ -34,7 +34,7 @@ std::unique_ptr client_context, Options const& options, google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return child_->SampleRowKeys(std::move(client_context), options, request, std::move(operation_context)); } @@ -43,7 +43,7 @@ StatusOr BigtableChannelRefresh::MutateRow( grpc::ClientContext& client_context, Options const& options, google::bigtable::v2::MutateRowRequest const& request, - OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { return child_->MutateRow(client_context, options, request, operation_context); } @@ -52,7 +52,7 @@ std::unique_ptr client_context, Options const& options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return child_->MutateRows(std::move(client_context), options, request, std::move(operation_context)); } @@ -61,7 +61,7 @@ StatusOr BigtableChannelRefresh::CheckAndMutateRow( grpc::ClientContext& client_context, Options const& options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { return child_->CheckAndMutateRow(client_context, options, request, operation_context); } @@ -70,7 +70,7 @@ StatusOr BigtableChannelRefresh::PingAndWarm( grpc::ClientContext& client_context, Options const& options, google::bigtable::v2::PingAndWarmRequest const& request, - OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { return child_->PingAndWarm(client_context, options, request, operation_context); } @@ -79,7 +79,7 @@ StatusOr BigtableChannelRefresh::ReadModifyWriteRow( grpc::ClientContext& client_context, Options const& options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { return child_->ReadModifyWriteRow(client_context, options, request, operation_context); } @@ -88,7 +88,7 @@ StatusOr BigtableChannelRefresh::PrepareQuery( grpc::ClientContext& client_context, Options const& options, google::bigtable::v2::PrepareQueryRequest const& request, - OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { return child_->PrepareQuery(client_context, options, request, operation_context); } @@ -98,7 +98,7 @@ std::unique_ptr client_context, Options const& options, google::bigtable::v2::ExecuteQueryRequest const& request, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return child_->ExecuteQuery(std::move(client_context), options, request, std::move(operation_context)); } @@ -110,7 +110,7 @@ BigtableChannelRefresh::AsyncReadRows( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return child_->AsyncReadRows(cq, std::move(context), std::move(options), request, std::move(operation_context)); } @@ -122,7 +122,7 @@ BigtableChannelRefresh::AsyncSampleRowKeys( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return child_->AsyncSampleRowKeys(cq, std::move(context), std::move(options), request, std::move(operation_context)); } @@ -133,7 +133,7 @@ BigtableChannelRefresh::AsyncMutateRow( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::MutateRowRequest const& request, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return child_->AsyncMutateRow(cq, std::move(context), std::move(options), request, std::move(operation_context)); } @@ -145,7 +145,7 @@ BigtableChannelRefresh::AsyncMutateRows( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return child_->AsyncMutateRows(cq, std::move(context), std::move(options), request, std::move(operation_context)); } @@ -156,7 +156,7 @@ BigtableChannelRefresh::AsyncCheckAndMutateRow( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return child_->AsyncCheckAndMutateRow(cq, std::move(context), std::move(options), request, std::move(operation_context)); @@ -168,7 +168,7 @@ BigtableChannelRefresh::AsyncPingAndWarm( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::PingAndWarmRequest const& request, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return child_->AsyncPingAndWarm(cq, std::move(context), std::move(options), request, std::move(operation_context)); } @@ -179,7 +179,7 @@ BigtableChannelRefresh::AsyncReadModifyWriteRow( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return child_->AsyncReadModifyWriteRow(cq, std::move(context), std::move(options), request, std::move(operation_context)); @@ -191,7 +191,7 @@ BigtableChannelRefresh::AsyncPrepareQuery( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::PrepareQueryRequest const& request, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return child_->AsyncPrepareQuery(cq, std::move(context), std::move(options), request, std::move(operation_context)); } @@ -200,7 +200,7 @@ StatusOr BigtableChannelRefresh::GetClientConfiguration( grpc::ClientContext& context, Options const& options, google::bigtable::v2::GetClientConfigurationRequest const& request, - OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { return child_->GetClientConfiguration(context, options, request, operation_context); } @@ -212,7 +212,7 @@ BigtableChannelRefresh::AsyncOpenTable( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return child_->AsyncOpenTable(cq, std::move(context), std::move(options), std::move(operation_context)); } @@ -224,7 +224,7 @@ BigtableChannelRefresh::AsyncOpenAuthorizedView( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return child_->AsyncOpenAuthorizedView( cq, std::move(context), std::move(options), std::move(operation_context)); } @@ -236,7 +236,7 @@ BigtableChannelRefresh::AsyncOpenMaterializedView( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return child_->AsyncOpenMaterializedView( cq, std::move(context), std::move(options), std::move(operation_context)); } diff --git a/google/cloud/bigtable/internal/bigtable_channel_refresh.h b/google/cloud/bigtable/internal/bigtable_channel_refresh.h index f9ce9fe9fa3a0..9719e3caed7a3 100644 --- a/google/cloud/bigtable/internal/bigtable_channel_refresh.h +++ b/google/cloud/bigtable/internal/bigtable_channel_refresh.h @@ -47,53 +47,57 @@ class BigtableChannelRefresh : public BigtableStub { ReadRows(std::shared_ptr client_context, Options const& options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr operation_context) override; + std::shared_ptr + operation_context) override; std::unique_ptr> SampleRowKeys(std::shared_ptr client_context, Options const& options, google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr operation_context) override; + std::shared_ptr + operation_context) override; StatusOr MutateRow( grpc::ClientContext& client_context, Options const& options, google::bigtable::v2::MutateRowRequest const& request, - OperationContext& operation_context) override; + bigtable_internal::OperationContext& operation_context) override; std::unique_ptr> MutateRows(std::shared_ptr client_context, Options const& options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr operation_context) override; + std::shared_ptr + operation_context) override; StatusOr CheckAndMutateRow( grpc::ClientContext& client_context, Options const& options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - OperationContext& operation_context) override; + bigtable_internal::OperationContext& operation_context) override; StatusOr PingAndWarm( grpc::ClientContext& client_context, Options const& options, google::bigtable::v2::PingAndWarmRequest const& request, - OperationContext& operation_context) override; + bigtable_internal::OperationContext& operation_context) override; StatusOr ReadModifyWriteRow( grpc::ClientContext& client_context, Options const& options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - OperationContext& operation_context) override; + bigtable_internal::OperationContext& operation_context) override; StatusOr PrepareQuery( grpc::ClientContext& client, Options const& options, google::bigtable::v2::PrepareQueryRequest const& request, - OperationContext& operation_context) override; + bigtable_internal::OperationContext& operation_context) override; std::unique_ptr> ExecuteQuery(std::shared_ptr client_context, Options const& options, google::bigtable::v2::ExecuteQueryRequest const& request, - std::shared_ptr operation_context) override; + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::ReadRowsResponse>> @@ -101,23 +105,25 @@ class BigtableChannelRefresh : public BigtableStub { std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr operation_context) override; + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::SampleRowKeysResponse>> - AsyncSampleRowKeys( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr operation_context) override; + AsyncSampleRowKeys(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::SampleRowKeysRequest const& request, + std::shared_ptr + operation_context) override; future> AsyncMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions, google::bigtable::v2::MutateRowRequest const& request, - std::shared_ptr operation_context) override; + std::shared_ptr operation_context) + override; std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::MutateRowsResponse>> @@ -125,7 +131,8 @@ class BigtableChannelRefresh : public BigtableStub { std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr operation_context) override; + std::shared_ptr + operation_context) override; future> AsyncCheckAndMutateRow( @@ -133,14 +140,16 @@ class BigtableChannelRefresh : public BigtableStub { std::shared_ptr context, google::cloud::internal::ImmutableOptions, google::bigtable::v2::CheckAndMutateRowRequest const& request, - std::shared_ptr operation_context) override; + std::shared_ptr operation_context) + override; future> AsyncPingAndWarm( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::PingAndWarmRequest const& request, - std::shared_ptr operation_context) override; + std::shared_ptr operation_context) + override; future> AsyncReadModifyWriteRow( @@ -148,20 +157,21 @@ class BigtableChannelRefresh : public BigtableStub { std::shared_ptr context, google::cloud::internal::ImmutableOptions, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - std::shared_ptr operation_context) override; + std::shared_ptr operation_context) + override; future> - AsyncPrepareQuery( - google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request, - std::shared_ptr operation_context) override; + AsyncPrepareQuery(google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::PrepareQueryRequest const& request, + std::shared_ptr + operation_context) override; StatusOr GetClientConfiguration( grpc::ClientContext& context, Options const& options, google::bigtable::v2::GetClientConfigurationRequest const& request, - OperationContext& operation_context) override; + bigtable_internal::OperationContext& operation_context) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, @@ -169,25 +179,26 @@ class BigtableChannelRefresh : public BigtableStub { AsyncOpenTable(google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - std::shared_ptr operation_context) override; + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>> - AsyncOpenAuthorizedView( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - std::shared_ptr operation_context) override; + AsyncOpenAuthorizedView(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>> - AsyncOpenMaterializedView( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - std::shared_ptr operation_context) override; + AsyncOpenMaterializedView(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + std::shared_ptr + operation_context) override; private: std::shared_ptr child_; diff --git a/google/cloud/bigtable/internal/bigtable_logging_decorator.cc b/google/cloud/bigtable/internal/bigtable_logging_decorator.cc index 3b7d75b658c68..ba27cd94f044f 100644 --- a/google/cloud/bigtable/internal/bigtable_logging_decorator.cc +++ b/google/cloud/bigtable/internal/bigtable_logging_decorator.cc @@ -48,8 +48,7 @@ std::unique_ptr context, Options const& options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return google::cloud::internal::LogWrapper( [this, operation_context = std::move(operation_context)]( std::shared_ptr context, Options const& options, @@ -75,8 +74,7 @@ std::unique_ptr context, Options const& options, google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return google::cloud::internal::LogWrapper( [this, operation_context = std::move(operation_context)]( std::shared_ptr context, Options const& options, @@ -100,7 +98,7 @@ BigtableLogging::SampleRowKeys( StatusOr BigtableLogging::MutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::MutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { return google::cloud::internal::LogWrapper( [this, &operation_context]( grpc::ClientContext& context, Options const& options, @@ -115,8 +113,7 @@ std::unique_ptr context, Options const& options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return google::cloud::internal::LogWrapper( [this, operation_context = std::move(operation_context)]( std::shared_ptr context, Options const& options, @@ -141,7 +138,7 @@ StatusOr BigtableLogging::CheckAndMutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { return google::cloud::internal::LogWrapper( [this, &operation_context]( grpc::ClientContext& context, Options const& options, @@ -156,7 +153,7 @@ StatusOr BigtableLogging::PingAndWarm( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PingAndWarmRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { return google::cloud::internal::LogWrapper( [this, &operation_context]( grpc::ClientContext& context, Options const& options, @@ -171,7 +168,7 @@ StatusOr BigtableLogging::ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { return google::cloud::internal::LogWrapper( [this, &operation_context]( grpc::ClientContext& context, Options const& options, @@ -186,7 +183,7 @@ StatusOr BigtableLogging::PrepareQuery( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PrepareQueryRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { return google::cloud::internal::LogWrapper( [this, &operation_context]( grpc::ClientContext& context, Options const& options, @@ -202,8 +199,7 @@ std::unique_ptr context, Options const& options, google::bigtable::v2::ExecuteQueryRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return google::cloud::internal::LogWrapper( [this, operation_context = std::move(operation_context)]( std::shared_ptr context, Options const& options, @@ -228,7 +224,7 @@ StatusOr BigtableLogging::GetClientConfiguration( grpc::ClientContext& context, Options const& options, google::bigtable::v2::GetClientConfigurationRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { return google::cloud::internal::LogWrapper( [this, &operation_context]( grpc::ClientContext& context, Options const& options, @@ -246,8 +242,7 @@ BigtableLogging::AsyncOpenTable( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { using LoggingStream = ::google::cloud::internal::AsyncStreamingReadWriteRpcLogging< google::bigtable::v2::SessionRequest, @@ -271,8 +266,7 @@ BigtableLogging::AsyncOpenAuthorizedView( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { using LoggingStream = ::google::cloud::internal::AsyncStreamingReadWriteRpcLogging< google::bigtable::v2::SessionRequest, @@ -296,8 +290,7 @@ BigtableLogging::AsyncOpenMaterializedView( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { using LoggingStream = ::google::cloud::internal::AsyncStreamingReadWriteRpcLogging< google::bigtable::v2::SessionRequest, @@ -321,8 +314,7 @@ BigtableLogging::AsyncReadRows( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { using LoggingStream = ::google::cloud::internal::AsyncStreamingReadRpcLogging< google::bigtable::v2::ReadRowsResponse>; @@ -347,8 +339,7 @@ BigtableLogging::AsyncSampleRowKeys( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { using LoggingStream = ::google::cloud::internal::AsyncStreamingReadRpcLogging< google::bigtable::v2::SampleRowKeysResponse>; @@ -372,8 +363,7 @@ BigtableLogging::AsyncMutateRow( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::MutateRowRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return google::cloud::internal::LogWrapper( [this, operation_context = std::move(operation_context)]( google::cloud::CompletionQueue& cq, @@ -395,8 +385,7 @@ BigtableLogging::AsyncMutateRows( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { using LoggingStream = ::google::cloud::internal::AsyncStreamingReadRpcLogging< google::bigtable::v2::MutateRowsResponse>; @@ -420,8 +409,7 @@ BigtableLogging::AsyncCheckAndMutateRow( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return google::cloud::internal::LogWrapper( [this, operation_context = std::move(operation_context)]( google::cloud::CompletionQueue& cq, @@ -442,8 +430,7 @@ BigtableLogging::AsyncPingAndWarm( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::PingAndWarmRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return google::cloud::internal::LogWrapper( [this, operation_context = std::move(operation_context)]( google::cloud::CompletionQueue& cq, @@ -464,8 +451,7 @@ BigtableLogging::AsyncReadModifyWriteRow( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return google::cloud::internal::LogWrapper( [this, operation_context = std::move(operation_context)]( google::cloud::CompletionQueue& cq, @@ -486,8 +472,7 @@ BigtableLogging::AsyncPrepareQuery( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::PrepareQueryRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return google::cloud::internal::LogWrapper( [this, operation_context = std::move(operation_context)]( google::cloud::CompletionQueue& cq, diff --git a/google/cloud/bigtable/internal/bigtable_logging_decorator.h b/google/cloud/bigtable/internal/bigtable_logging_decorator.h index 96fa28fb19026..9e04ed81d0c49 100644 --- a/google/cloud/bigtable/internal/bigtable_logging_decorator.h +++ b/google/cloud/bigtable/internal/bigtable_logging_decorator.h @@ -45,136 +45,124 @@ class BigtableLogging : public BigtableStub { google::bigtable::v2::ReadRowsResponse>> ReadRows(std::shared_ptr context, Options const& options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr + std::shared_ptr operation_context) override; std::unique_ptr> - SampleRowKeys( - std::shared_ptr context, Options const& options, - google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr - operation_context) override; + SampleRowKeys(std::shared_ptr context, + Options const& options, + google::bigtable::v2::SampleRowKeysRequest const& request, + std::shared_ptr + operation_context) override; StatusOr MutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::MutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; std::unique_ptr> MutateRows(std::shared_ptr context, Options const& options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr + std::shared_ptr operation_context) override; StatusOr CheckAndMutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; StatusOr PingAndWarm( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PingAndWarmRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; StatusOr ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; StatusOr PrepareQuery( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PrepareQueryRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; std::unique_ptr> - ExecuteQuery( - std::shared_ptr context, Options const& options, - google::bigtable::v2::ExecuteQueryRequest const& request, - std::shared_ptr - operation_context) override; + ExecuteQuery(std::shared_ptr context, + Options const& options, + google::bigtable::v2::ExecuteQueryRequest const& request, + std::shared_ptr + operation_context) override; StatusOr GetClientConfiguration( grpc::ClientContext& context, Options const& options, google::bigtable::v2::GetClientConfigurationRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>> - AsyncOpenTable( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) override; + AsyncOpenTable(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>> - AsyncOpenAuthorizedView( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) override; + AsyncOpenAuthorizedView(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>> - AsyncOpenMaterializedView( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) override; + AsyncOpenMaterializedView(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::ReadRowsResponse>> - AsyncReadRows( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr - operation_context) override; + AsyncReadRows(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::ReadRowsRequest const& request, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::SampleRowKeysResponse>> - AsyncSampleRowKeys( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr - operation_context) override; + AsyncSampleRowKeys(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::SampleRowKeysRequest const& request, + std::shared_ptr + operation_context) override; future> AsyncMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::MutateRowRequest const& request, - std::shared_ptr - operation_context) override; + std::shared_ptr operation_context) + override; std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::MutateRowsResponse>> - AsyncMutateRows( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr - operation_context) override; + AsyncMutateRows(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::MutateRowsRequest const& request, + std::shared_ptr + operation_context) override; future> AsyncCheckAndMutateRow( @@ -182,16 +170,16 @@ class BigtableLogging : public BigtableStub { std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - std::shared_ptr - operation_context) override; + std::shared_ptr operation_context) + override; future> AsyncPingAndWarm( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::PingAndWarmRequest const& request, - std::shared_ptr - operation_context) override; + std::shared_ptr operation_context) + override; future> AsyncReadModifyWriteRow( @@ -199,17 +187,16 @@ class BigtableLogging : public BigtableStub { std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - std::shared_ptr - operation_context) override; + std::shared_ptr operation_context) + override; future> - AsyncPrepareQuery( - google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request, - std::shared_ptr - operation_context) override; + AsyncPrepareQuery(google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::PrepareQueryRequest const& request, + std::shared_ptr + operation_context) override; private: std::shared_ptr child_; diff --git a/google/cloud/bigtable/internal/bigtable_metadata_decorator.cc b/google/cloud/bigtable/internal/bigtable_metadata_decorator.cc index 5ca7f51cec268..5fef22c35cfb9 100644 --- a/google/cloud/bigtable/internal/bigtable_metadata_decorator.cc +++ b/google/cloud/bigtable/internal/bigtable_metadata_decorator.cc @@ -54,8 +54,7 @@ std::unique_ptr context, Options const& options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { std::vector params; params.reserve(3); @@ -112,8 +111,7 @@ std::unique_ptr context, Options const& options, google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { std::vector params; params.reserve(3); @@ -168,7 +166,7 @@ BigtableMetadata::SampleRowKeys( StatusOr BigtableMetadata::MutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::MutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { std::vector params; params.reserve(2); @@ -209,8 +207,7 @@ std::unique_ptr context, Options const& options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { std::vector params; params.reserve(2); @@ -251,7 +248,7 @@ StatusOr BigtableMetadata::CheckAndMutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { std::vector params; params.reserve(2); @@ -292,7 +289,7 @@ StatusOr BigtableMetadata::PingAndWarm( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PingAndWarmRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { std::vector params; params.reserve(2); @@ -326,7 +323,7 @@ StatusOr BigtableMetadata::ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { std::vector params; params.reserve(2); @@ -367,7 +364,7 @@ StatusOr BigtableMetadata::PrepareQuery( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PrepareQueryRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { std::vector params; params.reserve(2); @@ -402,8 +399,7 @@ std::unique_ptr context, Options const& options, google::bigtable::v2::ExecuteQueryRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { std::vector params; params.reserve(2); @@ -438,7 +434,7 @@ StatusOr BigtableMetadata::GetClientConfiguration( grpc::ClientContext& context, Options const& options, google::bigtable::v2::GetClientConfigurationRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { SetMetadata(context, options); return child_->GetClientConfiguration(context, options, request, operation_context); @@ -451,8 +447,7 @@ BigtableMetadata::AsyncOpenTable( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { SetMetadata(*context, *options); return child_->AsyncOpenTable(cq, std::move(context), std::move(options), std::move(operation_context)); @@ -465,8 +460,7 @@ BigtableMetadata::AsyncOpenAuthorizedView( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { SetMetadata(*context, *options); return child_->AsyncOpenAuthorizedView( cq, std::move(context), std::move(options), std::move(operation_context)); @@ -479,8 +473,7 @@ BigtableMetadata::AsyncOpenMaterializedView( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { SetMetadata(*context, *options); return child_->AsyncOpenMaterializedView( cq, std::move(context), std::move(options), std::move(operation_context)); @@ -493,8 +486,7 @@ BigtableMetadata::AsyncReadRows( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { std::vector params; params.reserve(3); @@ -553,8 +545,7 @@ BigtableMetadata::AsyncSampleRowKeys( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { std::vector params; params.reserve(3); @@ -612,8 +603,7 @@ BigtableMetadata::AsyncMutateRow( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::MutateRowRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { std::vector params; params.reserve(2); @@ -657,8 +647,7 @@ BigtableMetadata::AsyncMutateRows( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { std::vector params; params.reserve(2); @@ -701,8 +690,7 @@ BigtableMetadata::AsyncCheckAndMutateRow( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { std::vector params; params.reserve(2); @@ -746,8 +734,7 @@ BigtableMetadata::AsyncPingAndWarm( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::PingAndWarmRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { std::vector params; params.reserve(2); @@ -784,8 +771,7 @@ BigtableMetadata::AsyncReadModifyWriteRow( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { std::vector params; params.reserve(2); @@ -829,8 +815,7 @@ BigtableMetadata::AsyncPrepareQuery( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::PrepareQueryRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { std::vector params; params.reserve(2); diff --git a/google/cloud/bigtable/internal/bigtable_metadata_decorator.h b/google/cloud/bigtable/internal/bigtable_metadata_decorator.h index 93546ccafaee4..080087684fe2d 100644 --- a/google/cloud/bigtable/internal/bigtable_metadata_decorator.h +++ b/google/cloud/bigtable/internal/bigtable_metadata_decorator.h @@ -45,136 +45,124 @@ class BigtableMetadata : public BigtableStub { google::bigtable::v2::ReadRowsResponse>> ReadRows(std::shared_ptr context, Options const& options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr + std::shared_ptr operation_context) override; std::unique_ptr> - SampleRowKeys( - std::shared_ptr context, Options const& options, - google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr - operation_context) override; + SampleRowKeys(std::shared_ptr context, + Options const& options, + google::bigtable::v2::SampleRowKeysRequest const& request, + std::shared_ptr + operation_context) override; StatusOr MutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::MutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; std::unique_ptr> MutateRows(std::shared_ptr context, Options const& options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr + std::shared_ptr operation_context) override; StatusOr CheckAndMutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; StatusOr PingAndWarm( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PingAndWarmRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; StatusOr ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; StatusOr PrepareQuery( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PrepareQueryRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; std::unique_ptr> - ExecuteQuery( - std::shared_ptr context, Options const& options, - google::bigtable::v2::ExecuteQueryRequest const& request, - std::shared_ptr - operation_context) override; + ExecuteQuery(std::shared_ptr context, + Options const& options, + google::bigtable::v2::ExecuteQueryRequest const& request, + std::shared_ptr + operation_context) override; StatusOr GetClientConfiguration( grpc::ClientContext& context, Options const& options, google::bigtable::v2::GetClientConfigurationRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>> - AsyncOpenTable( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) override; + AsyncOpenTable(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>> - AsyncOpenAuthorizedView( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) override; + AsyncOpenAuthorizedView(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>> - AsyncOpenMaterializedView( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) override; + AsyncOpenMaterializedView(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::ReadRowsResponse>> - AsyncReadRows( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr - operation_context) override; + AsyncReadRows(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::ReadRowsRequest const& request, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::SampleRowKeysResponse>> - AsyncSampleRowKeys( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr - operation_context) override; + AsyncSampleRowKeys(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::SampleRowKeysRequest const& request, + std::shared_ptr + operation_context) override; future> AsyncMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::MutateRowRequest const& request, - std::shared_ptr - operation_context) override; + std::shared_ptr operation_context) + override; std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::MutateRowsResponse>> - AsyncMutateRows( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr - operation_context) override; + AsyncMutateRows(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::MutateRowsRequest const& request, + std::shared_ptr + operation_context) override; future> AsyncCheckAndMutateRow( @@ -182,16 +170,16 @@ class BigtableMetadata : public BigtableStub { std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - std::shared_ptr - operation_context) override; + std::shared_ptr operation_context) + override; future> AsyncPingAndWarm( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::PingAndWarmRequest const& request, - std::shared_ptr - operation_context) override; + std::shared_ptr operation_context) + override; future> AsyncReadModifyWriteRow( @@ -199,17 +187,16 @@ class BigtableMetadata : public BigtableStub { std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - std::shared_ptr - operation_context) override; + std::shared_ptr operation_context) + override; future> - AsyncPrepareQuery( - google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request, - std::shared_ptr - operation_context) override; + AsyncPrepareQuery(google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::PrepareQueryRequest const& request, + std::shared_ptr + operation_context) override; private: void SetMetadata(grpc::ClientContext& context, Options const& options, diff --git a/google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.cc b/google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.cc index 2cb501b017754..632be385db5b4 100644 --- a/google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.cc +++ b/google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.cc @@ -13,6 +13,7 @@ // limitations under the License. #include "google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.h" +#include "google/cloud/bigtable/internal/operation_context.h" #include "google/cloud/async_streaming_read_write_rpc.h" #include "google/cloud/internal/async_streaming_read_rpc.h" #include "google/cloud/internal/streaming_read_rpc.h" @@ -103,7 +104,7 @@ class AsyncStreamingReadWriteRpcTracking template Response UnaryHelper(std::shared_ptr>& pool, - OperationContext& oc, + bigtable_internal::OperationContext& oc, std::function const& fn) { SelectedChannel selection = pool->GetChannelRandomTwoLeastUsed(); @@ -118,7 +119,8 @@ Response UnaryHelper(std::shared_ptr>& pool, template Response AsyncHelper(std::shared_ptr>& pool, - std::shared_ptr const& operation_context, + std::shared_ptr const& + operation_context, std::function const& fn) { SelectedChannel selection = pool->GetChannelRandomTwoLeastUsed(); @@ -136,7 +138,8 @@ Response AsyncHelper(std::shared_ptr>& pool, template std::unique_ptr> StreamingHelper( std::shared_ptr>& pool, - std::shared_ptr const& operation_context, + std::shared_ptr const& + operation_context, std::function>( BigtableStub&)> const& fn) { SelectedChannel selection = @@ -159,7 +162,8 @@ std::unique_ptr> StreamingHelper( template std::unique_ptr> AsyncStreamingHelper( std::shared_ptr>& pool, - std::shared_ptr const& operation_context, + std::shared_ptr const& + operation_context, std::function>( BigtableStub&)> const& fn) { SelectedChannel selection = @@ -182,7 +186,8 @@ std::unique_ptr> AsyncStreamingHelper( template std::unique_ptr> AsyncStreamingHelper(std::shared_ptr>& pool, - std::shared_ptr const& operation_context, + std::shared_ptr const& + operation_context, std::function>(BigtableStub&)> const& fn) { SelectedChannel selection = @@ -211,7 +216,7 @@ std::unique_ptr< BigtableRandomTwoLeastUsed::ReadRows( std::shared_ptr context, Options const& options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return StreamingHelper( pool_, operation_context, [&, context = std::move(context), @@ -226,7 +231,7 @@ std::unique_ptr< BigtableRandomTwoLeastUsed::SampleRowKeys( std::shared_ptr context, Options const& options, google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return StreamingHelper( pool_, operation_context, [&, context = std::move(context), @@ -240,7 +245,7 @@ StatusOr BigtableRandomTwoLeastUsed::MutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::MutateRowRequest const& request, - OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { return UnaryHelper>( pool_, operation_context, [&](BigtableStub& stub) { return stub.MutateRow(context, options, request, operation_context); @@ -252,7 +257,7 @@ std::unique_ptr< BigtableRandomTwoLeastUsed::MutateRows( std::shared_ptr context, Options const& options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return StreamingHelper( pool_, operation_context, [&, context = std::move(context), @@ -266,7 +271,7 @@ StatusOr BigtableRandomTwoLeastUsed::CheckAndMutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { return UnaryHelper>( pool_, operation_context, [&](BigtableStub& stub) { return stub.CheckAndMutateRow(context, options, request, @@ -278,7 +283,7 @@ StatusOr BigtableRandomTwoLeastUsed::PingAndWarm( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PingAndWarmRequest const& request, - OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { return UnaryHelper>( pool_, operation_context, [&](BigtableStub& stub) { return stub.PingAndWarm(context, options, request, operation_context); @@ -289,7 +294,7 @@ StatusOr BigtableRandomTwoLeastUsed::ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { return UnaryHelper< StatusOr>( pool_, operation_context, [&](BigtableStub& stub) { @@ -302,7 +307,7 @@ StatusOr BigtableRandomTwoLeastUsed::PrepareQuery( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PrepareQueryRequest const& request, - OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { return UnaryHelper>( pool_, operation_context, [&](BigtableStub& stub) { return stub.PrepareQuery(context, options, request, operation_context); @@ -314,7 +319,7 @@ std::unique_ptr< BigtableRandomTwoLeastUsed::ExecuteQuery( std::shared_ptr context, Options const& options, google::bigtable::v2::ExecuteQueryRequest const& request, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return StreamingHelper( pool_, operation_context, [&, context = std::move(context), @@ -330,7 +335,7 @@ BigtableRandomTwoLeastUsed::AsyncReadRows( CompletionQueue const& cq, std::shared_ptr context, internal::ImmutableOptions options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return AsyncStreamingHelper( pool_, operation_context, [&, context = std::move(context), options = std::move(options), @@ -346,7 +351,7 @@ BigtableRandomTwoLeastUsed::AsyncSampleRowKeys( CompletionQueue const& cq, std::shared_ptr context, internal::ImmutableOptions options, google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return AsyncStreamingHelper( pool_, operation_context, [&, context = std::move(context), options = std::move(options), @@ -362,7 +367,7 @@ BigtableRandomTwoLeastUsed::AsyncMutateRow( CompletionQueue& cq, std::shared_ptr context, internal::ImmutableOptions options, google::bigtable::v2::MutateRowRequest const& request, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return AsyncHelper>>( pool_, operation_context, [&, context = std::move(context), options = std::move(options), @@ -378,7 +383,7 @@ BigtableRandomTwoLeastUsed::AsyncMutateRows( CompletionQueue const& cq, std::shared_ptr context, internal::ImmutableOptions options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return AsyncStreamingHelper( pool_, operation_context, [&, context = std::move(context), options = std::move(options), @@ -393,7 +398,7 @@ BigtableRandomTwoLeastUsed::AsyncCheckAndMutateRow( CompletionQueue& cq, std::shared_ptr context, internal::ImmutableOptions options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return AsyncHelper< future>>( pool_, operation_context, @@ -410,7 +415,7 @@ BigtableRandomTwoLeastUsed::AsyncPingAndWarm( CompletionQueue& cq, std::shared_ptr context, internal::ImmutableOptions options, google::bigtable::v2::PingAndWarmRequest const& request, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return AsyncHelper< future>>( pool_, operation_context, @@ -426,7 +431,7 @@ BigtableRandomTwoLeastUsed::AsyncReadModifyWriteRow( CompletionQueue& cq, std::shared_ptr context, internal::ImmutableOptions options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return AsyncHelper< future>>( pool_, operation_context, @@ -443,7 +448,7 @@ BigtableRandomTwoLeastUsed::AsyncPrepareQuery( CompletionQueue& cq, std::shared_ptr context, internal::ImmutableOptions options, google::bigtable::v2::PrepareQueryRequest const& request, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return AsyncHelper< future>>( pool_, operation_context, @@ -459,7 +464,7 @@ StatusOr BigtableRandomTwoLeastUsed::GetClientConfiguration( grpc::ClientContext& context, Options const& options, google::bigtable::v2::GetClientConfigurationRequest const& request, - OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { return UnaryHelper>( pool_, operation_context, [&](BigtableStub& stub) { return stub.GetClientConfiguration(context, options, request, @@ -473,7 +478,7 @@ std::unique_ptr< BigtableRandomTwoLeastUsed::AsyncOpenTable( CompletionQueue const& cq, std::shared_ptr context, internal::ImmutableOptions options, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return AsyncStreamingHelper( pool_, operation_context, @@ -490,7 +495,7 @@ std::unique_ptr< BigtableRandomTwoLeastUsed::AsyncOpenAuthorizedView( CompletionQueue const& cq, std::shared_ptr context, internal::ImmutableOptions options, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return AsyncStreamingHelper( pool_, operation_context, @@ -508,7 +513,7 @@ std::unique_ptr< BigtableRandomTwoLeastUsed::AsyncOpenMaterializedView( CompletionQueue const& cq, std::shared_ptr context, internal::ImmutableOptions options, - std::shared_ptr operation_context) { + std::shared_ptr operation_context) { return AsyncStreamingHelper( pool_, operation_context, diff --git a/google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.h b/google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.h index 48ddd67cf0d77..10ffb9915b544 100644 --- a/google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.h +++ b/google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.h @@ -43,53 +43,57 @@ class BigtableRandomTwoLeastUsed : public BigtableStub { internal::StreamingReadRpc> ReadRows(std::shared_ptr context, Options const& options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr operation_context) override; + std::shared_ptr + operation_context) override; std::unique_ptr< internal::StreamingReadRpc> SampleRowKeys(std::shared_ptr context, Options const& options, google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr operation_context) override; + std::shared_ptr + operation_context) override; StatusOr MutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::MutateRowRequest const& request, - OperationContext& operation_context) override; + bigtable_internal::OperationContext& operation_context) override; std::unique_ptr< internal::StreamingReadRpc> MutateRows(std::shared_ptr context, Options const& options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr operation_context) override; + std::shared_ptr + operation_context) override; StatusOr CheckAndMutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - OperationContext& operation_context) override; + bigtable_internal::OperationContext& operation_context) override; StatusOr PingAndWarm( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PingAndWarmRequest const& request, - OperationContext& operation_context) override; + bigtable_internal::OperationContext& operation_context) override; StatusOr ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - OperationContext& operation_context) override; + bigtable_internal::OperationContext& operation_context) override; StatusOr PrepareQuery( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PrepareQueryRequest const& request, - OperationContext& operation_context) override; + bigtable_internal::OperationContext& operation_context) override; std::unique_ptr< internal::StreamingReadRpc> ExecuteQuery(std::shared_ptr context, Options const& options, google::bigtable::v2::ExecuteQueryRequest const& request, - std::shared_ptr operation_context) override; + std::shared_ptr + operation_context) override; std::unique_ptr< internal::AsyncStreamingReadRpc> @@ -97,21 +101,24 @@ class BigtableRandomTwoLeastUsed : public BigtableStub { std::shared_ptr context, internal::ImmutableOptions options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr operation_context) override; + std::shared_ptr + operation_context) override; std::unique_ptr> - AsyncSampleRowKeys( - CompletionQueue const& cq, std::shared_ptr context, - internal::ImmutableOptions options, - google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr operation_context) override; + AsyncSampleRowKeys(CompletionQueue const& cq, + std::shared_ptr context, + internal::ImmutableOptions options, + google::bigtable::v2::SampleRowKeysRequest const& request, + std::shared_ptr + operation_context) override; future> AsyncMutateRow( CompletionQueue& cq, std::shared_ptr context, internal::ImmutableOptions options, google::bigtable::v2::MutateRowRequest const& request, - std::shared_ptr operation_context) override; + std::shared_ptr operation_context) + override; std::unique_ptr< internal::AsyncStreamingReadRpc> @@ -119,39 +126,44 @@ class BigtableRandomTwoLeastUsed : public BigtableStub { std::shared_ptr context, internal::ImmutableOptions options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr operation_context) override; + std::shared_ptr + operation_context) override; future> AsyncCheckAndMutateRow( CompletionQueue& cq, std::shared_ptr context, internal::ImmutableOptions options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - std::shared_ptr operation_context) override; + std::shared_ptr operation_context) + override; future> AsyncPingAndWarm( CompletionQueue& cq, std::shared_ptr context, internal::ImmutableOptions options, google::bigtable::v2::PingAndWarmRequest const& request, - std::shared_ptr operation_context) override; + std::shared_ptr operation_context) + override; future> AsyncReadModifyWriteRow( CompletionQueue& cq, std::shared_ptr context, internal::ImmutableOptions options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - std::shared_ptr operation_context) override; + std::shared_ptr operation_context) + override; future> - AsyncPrepareQuery( - CompletionQueue& cq, std::shared_ptr context, - internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request, - std::shared_ptr operation_context) override; + AsyncPrepareQuery(CompletionQueue& cq, + std::shared_ptr context, + internal::ImmutableOptions options, + google::bigtable::v2::PrepareQueryRequest const& request, + std::shared_ptr + operation_context) override; StatusOr GetClientConfiguration( grpc::ClientContext& context, Options const& options, google::bigtable::v2::GetClientConfigurationRequest const& request, - OperationContext& operation_context) override; + bigtable_internal::OperationContext& operation_context) override; std::unique_ptr< AsyncStreamingReadWriteRpc context, internal::ImmutableOptions options, - std::shared_ptr operation_context) override; + std::shared_ptr + operation_context) override; std::unique_ptr< AsyncStreamingReadWriteRpc> - AsyncOpenAuthorizedView( - CompletionQueue const& cq, std::shared_ptr context, - internal::ImmutableOptions options, - std::shared_ptr operation_context) override; + AsyncOpenAuthorizedView(CompletionQueue const& cq, + std::shared_ptr context, + internal::ImmutableOptions options, + std::shared_ptr + operation_context) override; std::unique_ptr< AsyncStreamingReadWriteRpc> - AsyncOpenMaterializedView( - CompletionQueue const& cq, std::shared_ptr context, - internal::ImmutableOptions options, - std::shared_ptr operation_context) override; + AsyncOpenMaterializedView(CompletionQueue const& cq, + std::shared_ptr context, + internal::ImmutableOptions options, + std::shared_ptr + operation_context) override; private: std::shared_ptr> pool_; diff --git a/google/cloud/bigtable/internal/bigtable_round_robin_decorator.cc b/google/cloud/bigtable/internal/bigtable_round_robin_decorator.cc index 5df35deb3f3fd..4fccd463fb9cb 100644 --- a/google/cloud/bigtable/internal/bigtable_round_robin_decorator.cc +++ b/google/cloud/bigtable/internal/bigtable_round_robin_decorator.cc @@ -38,8 +38,7 @@ std::unique_ptr context, Options const& options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return Child()->ReadRows(std::move(context), options, request, std::move(operation_context)); } @@ -49,8 +48,7 @@ std::unique_ptr context, Options const& options, google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return Child()->SampleRowKeys(std::move(context), options, request, std::move(operation_context)); } @@ -58,7 +56,7 @@ BigtableRoundRobin::SampleRowKeys( StatusOr BigtableRoundRobin::MutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::MutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { return Child()->MutateRow(context, options, request, operation_context); } @@ -67,8 +65,7 @@ std::unique_ptr context, Options const& options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return Child()->MutateRows(std::move(context), options, request, std::move(operation_context)); } @@ -77,7 +74,7 @@ StatusOr BigtableRoundRobin::CheckAndMutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { return Child()->CheckAndMutateRow(context, options, request, operation_context); } @@ -86,7 +83,7 @@ StatusOr BigtableRoundRobin::PingAndWarm( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PingAndWarmRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { return Child()->PingAndWarm(context, options, request, operation_context); } @@ -94,7 +91,7 @@ StatusOr BigtableRoundRobin::ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { return Child()->ReadModifyWriteRow(context, options, request, operation_context); } @@ -103,7 +100,7 @@ StatusOr BigtableRoundRobin::PrepareQuery( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PrepareQueryRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { return Child()->PrepareQuery(context, options, request, operation_context); } @@ -112,8 +109,7 @@ std::unique_ptr context, Options const& options, google::bigtable::v2::ExecuteQueryRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return Child()->ExecuteQuery(std::move(context), options, request, std::move(operation_context)); } @@ -122,7 +118,7 @@ StatusOr BigtableRoundRobin::GetClientConfiguration( grpc::ClientContext& context, Options const& options, google::bigtable::v2::GetClientConfigurationRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { return Child()->GetClientConfiguration(context, options, request, operation_context); } @@ -134,8 +130,7 @@ BigtableRoundRobin::AsyncOpenTable( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return Child()->AsyncOpenTable(cq, std::move(context), std::move(options), std::move(operation_context)); } @@ -147,8 +142,7 @@ BigtableRoundRobin::AsyncOpenAuthorizedView( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return Child()->AsyncOpenAuthorizedView( cq, std::move(context), std::move(options), std::move(operation_context)); } @@ -160,8 +154,7 @@ BigtableRoundRobin::AsyncOpenMaterializedView( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return Child()->AsyncOpenMaterializedView( cq, std::move(context), std::move(options), std::move(operation_context)); } @@ -173,8 +166,7 @@ BigtableRoundRobin::AsyncReadRows( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return Child()->AsyncReadRows(cq, std::move(context), std::move(options), request, std::move(operation_context)); } @@ -186,8 +178,7 @@ BigtableRoundRobin::AsyncSampleRowKeys( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return Child()->AsyncSampleRowKeys(cq, std::move(context), std::move(options), request, std::move(operation_context)); } @@ -198,8 +189,7 @@ BigtableRoundRobin::AsyncMutateRow( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::MutateRowRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return Child()->AsyncMutateRow(cq, std::move(context), std::move(options), request, std::move(operation_context)); } @@ -211,8 +201,7 @@ BigtableRoundRobin::AsyncMutateRows( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return Child()->AsyncMutateRows(cq, std::move(context), std::move(options), request, std::move(operation_context)); } @@ -223,8 +212,7 @@ BigtableRoundRobin::AsyncCheckAndMutateRow( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return Child()->AsyncCheckAndMutateRow(cq, std::move(context), std::move(options), request, std::move(operation_context)); @@ -236,8 +224,7 @@ BigtableRoundRobin::AsyncPingAndWarm( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::PingAndWarmRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return Child()->AsyncPingAndWarm(cq, std::move(context), std::move(options), request, std::move(operation_context)); } @@ -248,8 +235,7 @@ BigtableRoundRobin::AsyncReadModifyWriteRow( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return Child()->AsyncReadModifyWriteRow(cq, std::move(context), std::move(options), request, std::move(operation_context)); @@ -261,8 +247,7 @@ BigtableRoundRobin::AsyncPrepareQuery( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::PrepareQueryRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { return Child()->AsyncPrepareQuery(cq, std::move(context), std::move(options), request, std::move(operation_context)); } diff --git a/google/cloud/bigtable/internal/bigtable_round_robin_decorator.h b/google/cloud/bigtable/internal/bigtable_round_robin_decorator.h index 95bf02e7c06d5..01403a14a38f4 100644 --- a/google/cloud/bigtable/internal/bigtable_round_robin_decorator.h +++ b/google/cloud/bigtable/internal/bigtable_round_robin_decorator.h @@ -43,136 +43,124 @@ class BigtableRoundRobin : public BigtableStub { google::bigtable::v2::ReadRowsResponse>> ReadRows(std::shared_ptr context, Options const& options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr + std::shared_ptr operation_context) override; std::unique_ptr> - SampleRowKeys( - std::shared_ptr context, Options const& options, - google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr - operation_context) override; + SampleRowKeys(std::shared_ptr context, + Options const& options, + google::bigtable::v2::SampleRowKeysRequest const& request, + std::shared_ptr + operation_context) override; StatusOr MutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::MutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; std::unique_ptr> MutateRows(std::shared_ptr context, Options const& options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr + std::shared_ptr operation_context) override; StatusOr CheckAndMutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; StatusOr PingAndWarm( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PingAndWarmRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; StatusOr ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; StatusOr PrepareQuery( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PrepareQueryRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; std::unique_ptr> - ExecuteQuery( - std::shared_ptr context, Options const& options, - google::bigtable::v2::ExecuteQueryRequest const& request, - std::shared_ptr - operation_context) override; + ExecuteQuery(std::shared_ptr context, + Options const& options, + google::bigtable::v2::ExecuteQueryRequest const& request, + std::shared_ptr + operation_context) override; StatusOr GetClientConfiguration( grpc::ClientContext& context, Options const& options, google::bigtable::v2::GetClientConfigurationRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>> - AsyncOpenTable( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) override; + AsyncOpenTable(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>> - AsyncOpenAuthorizedView( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) override; + AsyncOpenAuthorizedView(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>> - AsyncOpenMaterializedView( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) override; + AsyncOpenMaterializedView(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::ReadRowsResponse>> - AsyncReadRows( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr - operation_context) override; + AsyncReadRows(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::ReadRowsRequest const& request, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::SampleRowKeysResponse>> - AsyncSampleRowKeys( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr - operation_context) override; + AsyncSampleRowKeys(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::SampleRowKeysRequest const& request, + std::shared_ptr + operation_context) override; future> AsyncMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::MutateRowRequest const& request, - std::shared_ptr - operation_context) override; + std::shared_ptr operation_context) + override; std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::MutateRowsResponse>> - AsyncMutateRows( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr - operation_context) override; + AsyncMutateRows(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::MutateRowsRequest const& request, + std::shared_ptr + operation_context) override; future> AsyncCheckAndMutateRow( @@ -180,16 +168,16 @@ class BigtableRoundRobin : public BigtableStub { std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - std::shared_ptr - operation_context) override; + std::shared_ptr operation_context) + override; future> AsyncPingAndWarm( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::PingAndWarmRequest const& request, - std::shared_ptr - operation_context) override; + std::shared_ptr operation_context) + override; future> AsyncReadModifyWriteRow( @@ -197,17 +185,16 @@ class BigtableRoundRobin : public BigtableStub { std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - std::shared_ptr - operation_context) override; + std::shared_ptr operation_context) + override; future> - AsyncPrepareQuery( - google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request, - std::shared_ptr - operation_context) override; + AsyncPrepareQuery(google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::PrepareQueryRequest const& request, + std::shared_ptr + operation_context) override; private: std::shared_ptr Child(); diff --git a/google/cloud/bigtable/internal/bigtable_stub.cc b/google/cloud/bigtable/internal/bigtable_stub.cc index ccad440812389..026d523c21add 100644 --- a/google/cloud/bigtable/internal/bigtable_stub.cc +++ b/google/cloud/bigtable/internal/bigtable_stub.cc @@ -40,7 +40,7 @@ std::unique_ptr context, Options const&, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr) { + std::shared_ptr) { auto stream = grpc_stub_->ReadRows(context.get(), request); return std::make_unique>(std::move(context), @@ -52,7 +52,7 @@ std::unique_ptr context, Options const&, google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr) { + std::shared_ptr) { auto stream = grpc_stub_->SampleRowKeys(context.get(), request); return std::make_unique>(std::move(context), @@ -63,7 +63,7 @@ StatusOr DefaultBigtableStub::MutateRow( grpc::ClientContext& context, Options const&, google::bigtable::v2::MutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext&) { + bigtable_internal::OperationContext&) { google::bigtable::v2::MutateRowResponse response; auto status = grpc_stub_->MutateRow(&context, request, &response); if (!status.ok()) { @@ -77,7 +77,7 @@ std::unique_ptr context, Options const&, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr) { + std::shared_ptr) { auto stream = grpc_stub_->MutateRows(context.get(), request); return std::make_unique>(std::move(context), @@ -88,7 +88,7 @@ StatusOr DefaultBigtableStub::CheckAndMutateRow( grpc::ClientContext& context, Options const&, google::bigtable::v2::CheckAndMutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext&) { + bigtable_internal::OperationContext&) { google::bigtable::v2::CheckAndMutateRowResponse response; auto status = grpc_stub_->CheckAndMutateRow(&context, request, &response); if (!status.ok()) { @@ -101,7 +101,7 @@ StatusOr DefaultBigtableStub::PingAndWarm( grpc::ClientContext& context, Options const&, google::bigtable::v2::PingAndWarmRequest const& request, - google::cloud::bigtable_internal::OperationContext&) { + bigtable_internal::OperationContext&) { google::bigtable::v2::PingAndWarmResponse response; auto status = grpc_stub_->PingAndWarm(&context, request, &response); if (!status.ok()) { @@ -114,7 +114,7 @@ StatusOr DefaultBigtableStub::ReadModifyWriteRow( grpc::ClientContext& context, Options const&, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - google::cloud::bigtable_internal::OperationContext&) { + bigtable_internal::OperationContext&) { google::bigtable::v2::ReadModifyWriteRowResponse response; auto status = grpc_stub_->ReadModifyWriteRow(&context, request, &response); if (!status.ok()) { @@ -127,7 +127,7 @@ StatusOr DefaultBigtableStub::PrepareQuery( grpc::ClientContext& context, Options const&, google::bigtable::v2::PrepareQueryRequest const& request, - google::cloud::bigtable_internal::OperationContext&) { + bigtable_internal::OperationContext&) { google::bigtable::v2::PrepareQueryResponse response; auto status = grpc_stub_->PrepareQuery(&context, request, &response); if (!status.ok()) { @@ -141,7 +141,7 @@ std::unique_ptr context, Options const&, google::bigtable::v2::ExecuteQueryRequest const& request, - std::shared_ptr) { + std::shared_ptr) { auto stream = grpc_stub_->ExecuteQuery(context.get(), request); return std::make_unique>(std::move(context), @@ -152,7 +152,7 @@ StatusOr DefaultBigtableStub::GetClientConfiguration( grpc::ClientContext& context, Options const&, google::bigtable::v2::GetClientConfigurationRequest const& request, - google::cloud::bigtable_internal::OperationContext&) { + bigtable_internal::OperationContext&) { google::bigtable::v2::ClientConfiguration response; auto status = grpc_stub_->GetClientConfiguration(&context, request, &response); @@ -169,7 +169,7 @@ DefaultBigtableStub::AsyncOpenTable( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - std::shared_ptr) { + std::shared_ptr) { return google::cloud::internal::MakeStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>( @@ -186,7 +186,7 @@ DefaultBigtableStub::AsyncOpenAuthorizedView( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - std::shared_ptr) { + std::shared_ptr) { return google::cloud::internal::MakeStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>( @@ -203,7 +203,7 @@ DefaultBigtableStub::AsyncOpenMaterializedView( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - std::shared_ptr) { + std::shared_ptr) { return google::cloud::internal::MakeStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>( @@ -220,7 +220,7 @@ DefaultBigtableStub::AsyncReadRows( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr) { + std::shared_ptr) { return google::cloud::internal::MakeStreamingReadRpc< google::bigtable::v2::ReadRowsRequest, google::bigtable::v2::ReadRowsResponse>( @@ -239,7 +239,7 @@ DefaultBigtableStub::AsyncSampleRowKeys( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr) { + std::shared_ptr) { return google::cloud::internal::MakeStreamingReadRpc< google::bigtable::v2::SampleRowKeysRequest, google::bigtable::v2::SampleRowKeysResponse>( @@ -258,7 +258,7 @@ DefaultBigtableStub::AsyncMutateRow( // NOLINTNEXTLINE(performance-unnecessary-value-param) google::cloud::internal::ImmutableOptions, google::bigtable::v2::MutateRowRequest const& request, - std::shared_ptr) { + std::shared_ptr) { return internal::MakeUnaryRpcImpl( cq, @@ -277,7 +277,7 @@ DefaultBigtableStub::AsyncMutateRows( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr) { + std::shared_ptr) { return google::cloud::internal::MakeStreamingReadRpc< google::bigtable::v2::MutateRowsRequest, google::bigtable::v2::MutateRowsResponse>( @@ -296,7 +296,7 @@ DefaultBigtableStub::AsyncCheckAndMutateRow( // NOLINTNEXTLINE(performance-unnecessary-value-param) google::cloud::internal::ImmutableOptions, google::bigtable::v2::CheckAndMutateRowRequest const& request, - std::shared_ptr) { + std::shared_ptr) { return internal::MakeUnaryRpcImpl< google::bigtable::v2::CheckAndMutateRowRequest, google::bigtable::v2::CheckAndMutateRowResponse>( @@ -316,7 +316,7 @@ DefaultBigtableStub::AsyncPingAndWarm( // NOLINTNEXTLINE(performance-unnecessary-value-param) google::cloud::internal::ImmutableOptions, google::bigtable::v2::PingAndWarmRequest const& request, - std::shared_ptr) { + std::shared_ptr) { return internal::MakeUnaryRpcImpl( cq, @@ -335,7 +335,7 @@ DefaultBigtableStub::AsyncReadModifyWriteRow( // NOLINTNEXTLINE(performance-unnecessary-value-param) google::cloud::internal::ImmutableOptions, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - std::shared_ptr) { + std::shared_ptr) { return internal::MakeUnaryRpcImpl< google::bigtable::v2::ReadModifyWriteRowRequest, google::bigtable::v2::ReadModifyWriteRowResponse>( @@ -355,7 +355,7 @@ DefaultBigtableStub::AsyncPrepareQuery( // NOLINTNEXTLINE(performance-unnecessary-value-param) google::cloud::internal::ImmutableOptions, google::bigtable::v2::PrepareQueryRequest const& request, - std::shared_ptr) { + std::shared_ptr) { return internal::MakeUnaryRpcImpl( cq, diff --git a/google/cloud/bigtable/internal/bigtable_stub.h b/google/cloud/bigtable/internal/bigtable_stub.h index cd3c8a6785e07..d619724fda26c 100644 --- a/google/cloud/bigtable/internal/bigtable_stub.h +++ b/google/cloud/bigtable/internal/bigtable_stub.h @@ -19,6 +19,7 @@ #ifndef GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_INTERNAL_BIGTABLE_STUB_H #define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_INTERNAL_BIGTABLE_STUB_H +#include "google/cloud/bigtable/internal/operation_context.h" #include "google/cloud/async_streaming_read_write_rpc.h" #include "google/cloud/completion_queue.h" #include "google/cloud/future.h" @@ -39,8 +40,6 @@ namespace cloud { namespace bigtable_internal { GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN -class OperationContext; - class BigtableStub { public: virtual ~BigtableStub() = 0; @@ -49,140 +48,127 @@ class BigtableStub { google::bigtable::v2::ReadRowsResponse>> ReadRows(std::shared_ptr context, Options const& options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr + std::shared_ptr operation_context) = 0; virtual std::unique_ptr> - SampleRowKeys( - std::shared_ptr context, Options const& options, - google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr - operation_context) = 0; + SampleRowKeys(std::shared_ptr context, + Options const& options, + google::bigtable::v2::SampleRowKeysRequest const& request, + std::shared_ptr + operation_context) = 0; virtual StatusOr MutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::MutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& - operation_context) = 0; + bigtable_internal::OperationContext& operation_context) = 0; virtual std::unique_ptr> MutateRows(std::shared_ptr context, Options const& options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr + std::shared_ptr operation_context) = 0; virtual StatusOr CheckAndMutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& - operation_context) = 0; + bigtable_internal::OperationContext& operation_context) = 0; virtual StatusOr PingAndWarm( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PingAndWarmRequest const& request, - google::cloud::bigtable_internal::OperationContext& - operation_context) = 0; + bigtable_internal::OperationContext& operation_context) = 0; virtual StatusOr ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& - operation_context) = 0; + bigtable_internal::OperationContext& operation_context) = 0; virtual StatusOr PrepareQuery( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PrepareQueryRequest const& request, - google::cloud::bigtable_internal::OperationContext& - operation_context) = 0; + bigtable_internal::OperationContext& operation_context) = 0; virtual std::unique_ptr> - ExecuteQuery( - std::shared_ptr context, Options const& options, - google::bigtable::v2::ExecuteQueryRequest const& request, - std::shared_ptr - operation_context) = 0; + ExecuteQuery(std::shared_ptr context, + Options const& options, + google::bigtable::v2::ExecuteQueryRequest const& request, + std::shared_ptr + operation_context) = 0; virtual StatusOr GetClientConfiguration( grpc::ClientContext& context, Options const& options, google::bigtable::v2::GetClientConfigurationRequest const& request, - google::cloud::bigtable_internal::OperationContext& - operation_context) = 0; + bigtable_internal::OperationContext& operation_context) = 0; virtual std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>> - AsyncOpenTable( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) = 0; + AsyncOpenTable(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + std::shared_ptr + operation_context) = 0; virtual std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>> - AsyncOpenAuthorizedView( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) = 0; + AsyncOpenAuthorizedView(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + std::shared_ptr + operation_context) = 0; virtual std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>> - AsyncOpenMaterializedView( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) = 0; + AsyncOpenMaterializedView(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + std::shared_ptr + operation_context) = 0; virtual std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::ReadRowsResponse>> - AsyncReadRows( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr - operation_context) = 0; + AsyncReadRows(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::ReadRowsRequest const& request, + std::shared_ptr + operation_context) = 0; virtual std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::SampleRowKeysResponse>> - AsyncSampleRowKeys( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr - operation_context) = 0; + AsyncSampleRowKeys(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::SampleRowKeysRequest const& request, + std::shared_ptr + operation_context) = 0; virtual future> - AsyncMutateRow( - google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowRequest const& request, - std::shared_ptr - operation_context) = 0; + AsyncMutateRow(google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::MutateRowRequest const& request, + std::shared_ptr + operation_context) = 0; virtual std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::MutateRowsResponse>> - AsyncMutateRows( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr - operation_context) = 0; + AsyncMutateRows(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::MutateRowsRequest const& request, + std::shared_ptr + operation_context) = 0; virtual future> AsyncCheckAndMutateRow( @@ -190,17 +176,16 @@ class BigtableStub { std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - std::shared_ptr + std::shared_ptr operation_context) = 0; virtual future> - AsyncPingAndWarm( - google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PingAndWarmRequest const& request, - std::shared_ptr - operation_context) = 0; + AsyncPingAndWarm(google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::PingAndWarmRequest const& request, + std::shared_ptr + operation_context) = 0; virtual future> AsyncReadModifyWriteRow( @@ -208,17 +193,16 @@ class BigtableStub { std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - std::shared_ptr + std::shared_ptr operation_context) = 0; virtual future> - AsyncPrepareQuery( - google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request, - std::shared_ptr - operation_context) = 0; + AsyncPrepareQuery(google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::PrepareQueryRequest const& request, + std::shared_ptr + operation_context) = 0; }; class DefaultBigtableStub : public BigtableStub { @@ -231,136 +215,124 @@ class DefaultBigtableStub : public BigtableStub { google::bigtable::v2::ReadRowsResponse>> ReadRows(std::shared_ptr context, Options const& options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr + std::shared_ptr operation_context) override; std::unique_ptr> - SampleRowKeys( - std::shared_ptr context, Options const& options, - google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr - operation_context) override; + SampleRowKeys(std::shared_ptr context, + Options const& options, + google::bigtable::v2::SampleRowKeysRequest const& request, + std::shared_ptr + operation_context) override; StatusOr MutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::MutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; std::unique_ptr> MutateRows(std::shared_ptr context, Options const& options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr + std::shared_ptr operation_context) override; StatusOr CheckAndMutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; StatusOr PingAndWarm( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PingAndWarmRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; StatusOr ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; StatusOr PrepareQuery( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PrepareQueryRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; std::unique_ptr> - ExecuteQuery( - std::shared_ptr context, Options const& options, - google::bigtable::v2::ExecuteQueryRequest const& request, - std::shared_ptr - operation_context) override; + ExecuteQuery(std::shared_ptr context, + Options const& options, + google::bigtable::v2::ExecuteQueryRequest const& request, + std::shared_ptr + operation_context) override; StatusOr GetClientConfiguration( grpc::ClientContext& context, Options const& options, google::bigtable::v2::GetClientConfigurationRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>> - AsyncOpenTable( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) override; + AsyncOpenTable(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>> - AsyncOpenAuthorizedView( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) override; + AsyncOpenAuthorizedView(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>> - AsyncOpenMaterializedView( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) override; + AsyncOpenMaterializedView(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::ReadRowsResponse>> - AsyncReadRows( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr - operation_context) override; + AsyncReadRows(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::ReadRowsRequest const& request, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::SampleRowKeysResponse>> - AsyncSampleRowKeys( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr - operation_context) override; + AsyncSampleRowKeys(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::SampleRowKeysRequest const& request, + std::shared_ptr + operation_context) override; future> AsyncMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::MutateRowRequest const& request, - std::shared_ptr - operation_context) override; + std::shared_ptr operation_context) + override; std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::MutateRowsResponse>> - AsyncMutateRows( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr - operation_context) override; + AsyncMutateRows(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::MutateRowsRequest const& request, + std::shared_ptr + operation_context) override; future> AsyncCheckAndMutateRow( @@ -368,16 +340,16 @@ class DefaultBigtableStub : public BigtableStub { std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - std::shared_ptr - operation_context) override; + std::shared_ptr operation_context) + override; future> AsyncPingAndWarm( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::PingAndWarmRequest const& request, - std::shared_ptr - operation_context) override; + std::shared_ptr operation_context) + override; future> AsyncReadModifyWriteRow( @@ -385,17 +357,16 @@ class DefaultBigtableStub : public BigtableStub { std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - std::shared_ptr - operation_context) override; + std::shared_ptr operation_context) + override; future> - AsyncPrepareQuery( - google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request, - std::shared_ptr - operation_context) override; + AsyncPrepareQuery(google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::PrepareQueryRequest const& request, + std::shared_ptr + operation_context) override; private: std::unique_ptr grpc_stub_; diff --git a/google/cloud/bigtable/internal/bigtable_tracing_stub.cc b/google/cloud/bigtable/internal/bigtable_tracing_stub.cc index 8463db58e7039..45603795384a5 100644 --- a/google/cloud/bigtable/internal/bigtable_tracing_stub.cc +++ b/google/cloud/bigtable/internal/bigtable_tracing_stub.cc @@ -40,8 +40,7 @@ std::unique_ptr context, Options const& options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "ReadRows"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(*context, *propagator_); @@ -57,8 +56,7 @@ std::unique_ptr context, Options const& options, google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "SampleRowKeys"); auto scope = opentelemetry::trace::Scope(span); @@ -74,7 +72,7 @@ StatusOr BigtableTracingStub::MutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::MutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "MutateRow"); auto scope = opentelemetry::trace::Scope(span); @@ -89,8 +87,7 @@ std::unique_ptr context, Options const& options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "MutateRows"); auto scope = opentelemetry::trace::Scope(span); @@ -106,7 +103,7 @@ StatusOr BigtableTracingStub::CheckAndMutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "CheckAndMutateRow"); auto scope = opentelemetry::trace::Scope(span); @@ -120,7 +117,7 @@ StatusOr BigtableTracingStub::PingAndWarm( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PingAndWarmRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "PingAndWarm"); auto scope = opentelemetry::trace::Scope(span); @@ -134,7 +131,7 @@ StatusOr BigtableTracingStub::ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "ReadModifyWriteRow"); auto scope = opentelemetry::trace::Scope(span); @@ -148,7 +145,7 @@ StatusOr BigtableTracingStub::PrepareQuery( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PrepareQueryRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "PrepareQuery"); auto scope = opentelemetry::trace::Scope(span); @@ -163,8 +160,7 @@ std::unique_ptr context, Options const& options, google::bigtable::v2::ExecuteQueryRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "ExecuteQuery"); auto scope = opentelemetry::trace::Scope(span); @@ -180,7 +176,7 @@ StatusOr BigtableTracingStub::GetClientConfiguration( grpc::ClientContext& context, Options const& options, google::bigtable::v2::GetClientConfigurationRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) { + bigtable_internal::OperationContext& operation_context) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "GetClientConfiguration"); auto scope = opentelemetry::trace::Scope(span); @@ -196,8 +192,7 @@ std::unique_ptr< BigtableTracingStub::AsyncOpenTable( CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "OpenTable"); internal::OTelScope scope(span); @@ -216,8 +211,7 @@ std::unique_ptr< BigtableTracingStub::AsyncOpenAuthorizedView( CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "OpenAuthorizedView"); internal::OTelScope scope(span); @@ -236,8 +230,7 @@ std::unique_ptr< BigtableTracingStub::AsyncOpenMaterializedView( CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "OpenMaterializedView"); internal::OTelScope scope(span); @@ -257,8 +250,7 @@ BigtableTracingStub::AsyncReadRows( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "ReadRows"); internal::OTelScope scope(span); internal::InjectTraceContext(*context, *propagator_); @@ -276,8 +268,7 @@ BigtableTracingStub::AsyncSampleRowKeys( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "SampleRowKeys"); internal::OTelScope scope(span); @@ -295,8 +286,7 @@ BigtableTracingStub::AsyncMutateRow( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::MutateRowRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "MutateRow"); internal::OTelScope scope(span); @@ -313,8 +303,7 @@ BigtableTracingStub::AsyncMutateRows( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "MutateRows"); internal::OTelScope scope(span); @@ -332,8 +321,7 @@ BigtableTracingStub::AsyncCheckAndMutateRow( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "CheckAndMutateRow"); internal::OTelScope scope(span); @@ -349,8 +337,7 @@ BigtableTracingStub::AsyncPingAndWarm( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::PingAndWarmRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "PingAndWarm"); internal::OTelScope scope(span); @@ -366,8 +353,7 @@ BigtableTracingStub::AsyncReadModifyWriteRow( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "ReadModifyWriteRow"); internal::OTelScope scope(span); @@ -383,8 +369,7 @@ BigtableTracingStub::AsyncPrepareQuery( std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::PrepareQueryRequest const& request, - std::shared_ptr - operation_context) { + std::shared_ptr operation_context) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "PrepareQuery"); internal::OTelScope scope(span); diff --git a/google/cloud/bigtable/internal/bigtable_tracing_stub.h b/google/cloud/bigtable/internal/bigtable_tracing_stub.h index 43a244dbc80bd..e393add68c3ca 100644 --- a/google/cloud/bigtable/internal/bigtable_tracing_stub.h +++ b/google/cloud/bigtable/internal/bigtable_tracing_stub.h @@ -43,136 +43,124 @@ class BigtableTracingStub : public BigtableStub { google::bigtable::v2::ReadRowsResponse>> ReadRows(std::shared_ptr context, Options const& options, google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr + std::shared_ptr operation_context) override; std::unique_ptr> - SampleRowKeys( - std::shared_ptr context, Options const& options, - google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr - operation_context) override; + SampleRowKeys(std::shared_ptr context, + Options const& options, + google::bigtable::v2::SampleRowKeysRequest const& request, + std::shared_ptr + operation_context) override; StatusOr MutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::MutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; std::unique_ptr> MutateRows(std::shared_ptr context, Options const& options, google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr + std::shared_ptr operation_context) override; StatusOr CheckAndMutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; StatusOr PingAndWarm( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PingAndWarmRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; StatusOr ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; StatusOr PrepareQuery( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PrepareQueryRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; std::unique_ptr> - ExecuteQuery( - std::shared_ptr context, Options const& options, - google::bigtable::v2::ExecuteQueryRequest const& request, - std::shared_ptr - operation_context) override; + ExecuteQuery(std::shared_ptr context, + Options const& options, + google::bigtable::v2::ExecuteQueryRequest const& request, + std::shared_ptr + operation_context) override; StatusOr GetClientConfiguration( grpc::ClientContext& context, Options const& options, google::bigtable::v2::GetClientConfigurationRequest const& request, - google::cloud::bigtable_internal::OperationContext& operation_context) - override; + bigtable_internal::OperationContext& operation_context) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>> - AsyncOpenTable( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) override; + AsyncOpenTable(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>> - AsyncOpenAuthorizedView( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) override; + AsyncOpenAuthorizedView(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, google::bigtable::v2::SessionResponse>> - AsyncOpenMaterializedView( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - std::shared_ptr - operation_context) override; + AsyncOpenMaterializedView(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::ReadRowsResponse>> - AsyncReadRows( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadRowsRequest const& request, - std::shared_ptr - operation_context) override; + AsyncReadRows(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::ReadRowsRequest const& request, + std::shared_ptr + operation_context) override; std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::SampleRowKeysResponse>> - AsyncSampleRowKeys( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::SampleRowKeysRequest const& request, - std::shared_ptr - operation_context) override; + AsyncSampleRowKeys(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::SampleRowKeysRequest const& request, + std::shared_ptr + operation_context) override; future> AsyncMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::MutateRowRequest const& request, - std::shared_ptr - operation_context) override; + std::shared_ptr operation_context) + override; std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::MutateRowsResponse>> - AsyncMutateRows( - google::cloud::CompletionQueue const& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowsRequest const& request, - std::shared_ptr - operation_context) override; + AsyncMutateRows(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::MutateRowsRequest const& request, + std::shared_ptr + operation_context) override; future> AsyncCheckAndMutateRow( @@ -180,16 +168,16 @@ class BigtableTracingStub : public BigtableStub { std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::CheckAndMutateRowRequest const& request, - std::shared_ptr - operation_context) override; + std::shared_ptr operation_context) + override; future> AsyncPingAndWarm( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::PingAndWarmRequest const& request, - std::shared_ptr - operation_context) override; + std::shared_ptr operation_context) + override; future> AsyncReadModifyWriteRow( @@ -197,17 +185,16 @@ class BigtableTracingStub : public BigtableStub { std::shared_ptr context, google::cloud::internal::ImmutableOptions options, google::bigtable::v2::ReadModifyWriteRowRequest const& request, - std::shared_ptr - operation_context) override; + std::shared_ptr operation_context) + override; future> - AsyncPrepareQuery( - google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request, - std::shared_ptr - operation_context) override; + AsyncPrepareQuery(google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::PrepareQueryRequest const& request, + std::shared_ptr + operation_context) override; private: std::shared_ptr child_; diff --git a/google/cloud/bigtable/internal/dynamic_channel_pool_test.cc b/google/cloud/bigtable/internal/dynamic_channel_pool_test.cc index c9e9d771ca2b2..95718fdc49c12 100644 --- a/google/cloud/bigtable/internal/dynamic_channel_pool_test.cc +++ b/google/cloud/bigtable/internal/dynamic_channel_pool_test.cc @@ -153,7 +153,7 @@ TEST_F(DynamicChannelPoolTest, SelectLeastUsedFromTwoChannels) { EXPECT_CALL(*mock_stub_0, CheckAndMutateRow) .WillOnce([](grpc::ClientContext&, Options const&, google::bigtable::v2::CheckAndMutateRowRequest const&, - OperationContext&) { + bigtable_internal::OperationContext&) { google::bigtable::v2::CheckAndMutateRowResponse response; response.set_predicate_matched(true); return response; @@ -222,7 +222,7 @@ TEST_F(DynamicChannelPoolTest, OneInitialChannel) { EXPECT_CALL(*mock_stub_0, CheckAndMutateRow) .WillOnce([](grpc::ClientContext&, Options const&, google::bigtable::v2::CheckAndMutateRowRequest const&, - OperationContext&) { + bigtable_internal::OperationContext&) { google::bigtable::v2::CheckAndMutateRowResponse response; response.set_predicate_matched(true); return response; @@ -281,7 +281,7 @@ TEST_F(DynamicChannelPoolTest, EmptyInitialPool) { EXPECT_CALL(*mock_stub, CheckAndMutateRow) .WillOnce([](grpc::ClientContext&, Options const&, google::bigtable::v2::CheckAndMutateRowRequest const&, - OperationContext&) { + bigtable_internal::OperationContext&) { google::bigtable::v2::CheckAndMutateRowResponse response; response.set_predicate_matched(true); return response; @@ -766,7 +766,7 @@ TEST_F(DynamicChannelPoolTest, HandleBadChannelsTwoChannelsOneBad) { EXPECT_CALL(*mock_stub, CheckAndMutateRow) .WillOnce([](grpc::ClientContext&, Options const&, google::bigtable::v2::CheckAndMutateRowRequest const&, - OperationContext&) { + bigtable_internal::OperationContext&) { google::bigtable::v2::CheckAndMutateRowResponse response; response.set_predicate_matched(true); return response; @@ -848,7 +848,7 @@ TEST_F(DynamicChannelPoolTest, HandleBadChannelsTwoChannelsOtherOneBad) { EXPECT_CALL(*mock_stub, CheckAndMutateRow) .WillOnce([](grpc::ClientContext&, Options const&, google::bigtable::v2::CheckAndMutateRowRequest const&, - OperationContext&) { + bigtable_internal::OperationContext&) { google::bigtable::v2::CheckAndMutateRowResponse response; response.set_predicate_matched(true); return response; @@ -932,7 +932,7 @@ TEST_F(DynamicChannelPoolTest, HandleBadChannelsThreeChannelsOneBad) { EXPECT_CALL(*mock_stub_0, CheckAndMutateRow) .WillOnce([](grpc::ClientContext&, Options const&, google::bigtable::v2::CheckAndMutateRowRequest const&, - OperationContext&) { + bigtable_internal::OperationContext&) { google::bigtable::v2::CheckAndMutateRowResponse response; response.set_predicate_matched(true); return response; @@ -1022,7 +1022,7 @@ TEST_F(DynamicChannelPoolTest, HandleBadChannelsAllChannelsBad) { EXPECT_CALL(*mock_stub, CheckAndMutateRow) .WillOnce([](grpc::ClientContext&, Options const&, google::bigtable::v2::CheckAndMutateRowRequest const&, - OperationContext&) { + bigtable_internal::OperationContext&) { google::bigtable::v2::CheckAndMutateRowResponse response; response.set_predicate_matched(true); return response; diff --git a/google/cloud/bigtable/internal/operation_context.h b/google/cloud/bigtable/internal/operation_context.h index 739f34242671e..f58291409c78e 100644 --- a/google/cloud/bigtable/internal/operation_context.h +++ b/google/cloud/bigtable/internal/operation_context.h @@ -17,6 +17,7 @@ #include "google/cloud/bigtable/version.h" #include "google/cloud/internal/clock.h" +#include "google/cloud/internal/operation_context.h" #include "google/cloud/status.h" #include #include @@ -56,7 +57,7 @@ class Metric; * } * @endcode */ -class OperationContext { +class OperationContext : public ::google::cloud::internal::OperationContext { public: using Clock = ::google::cloud::internal::SteadyClock; @@ -69,12 +70,12 @@ class OperationContext { // Called when a stub is selected from a channel pool. void StubSelection(StubSelectionParams const& params); // Called before each RPC attempt. - void PreCall(grpc::ClientContext& client_context); + void PreCall(grpc::ClientContext& client_context) override; // Called after receiving RPC attempt response. void PostCall(grpc::ClientContext const& client_context, - google::cloud::Status const& status); + google::cloud::Status const& status) override; // A hook that executes at the end of a client operation. - void OnDone(Status const& status); + void OnDone(Status const& status) override; // Called during operations that allow the user to iterate over data // synchronously or asynchronously. void ElementRequest(grpc::ClientContext const& client_context); diff --git a/google/cloud/bigtable/internal/stub_manager_test.cc b/google/cloud/bigtable/internal/stub_manager_test.cc index 8e217e9f87c48..ed194a7023b1a 100644 --- a/google/cloud/bigtable/internal/stub_manager_test.cc +++ b/google/cloud/bigtable/internal/stub_manager_test.cc @@ -43,7 +43,7 @@ TEST(StubManagerTest, NoAffinity) { EXPECT_CALL(*mock, MutateRow) .WillOnce([&](grpc::ClientContext&, Options const&, google::bigtable::v2::MutateRowRequest const& request, - OperationContext&) { + bigtable_internal::OperationContext&) { EXPECT_THAT(request.table_name(), Eq(expected_table_name)); return google::bigtable::v2::MutateRowResponse{}; }); @@ -69,7 +69,7 @@ TEST(StubManagerTest, AffinityToExistingInstance) { .WillOnce([instance_name = instance.FullName()]( grpc::ClientContext&, Options const&, google::bigtable::v2::MutateRowRequest const& request, - OperationContext&) { + bigtable_internal::OperationContext&) { EXPECT_THAT(request.table_name(), StartsWith(instance_name)); return google::bigtable::v2::MutateRowResponse{}; }); @@ -111,7 +111,7 @@ TEST(StubManagerTest, AffinityToMissingInstance) { .WillOnce([instance_name = std::string{instance_name}]( grpc::ClientContext&, Options const&, google::bigtable::v2::MutateRowRequest const& request, - OperationContext&) { + bigtable_internal::OperationContext&) { EXPECT_THAT(request.table_name(), StartsWith(instance_name)); return google::bigtable::v2::MutateRowResponse{}; }); diff --git a/google/cloud/bigtable/testing/mock_bigtable_stub.h b/google/cloud/bigtable/testing/mock_bigtable_stub.h index 8bd14c9beb11a..3a3b39dddbd2c 100644 --- a/google/cloud/bigtable/testing/mock_bigtable_stub.h +++ b/google/cloud/bigtable/testing/mock_bigtable_stub.h @@ -16,6 +16,7 @@ #define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_TESTING_MOCK_BIGTABLE_STUB_H #include "google/cloud/bigtable/internal/bigtable_stub.h" +#include "google/cloud/bigtable/internal/operation_context.h" #include "google/cloud/testing_util/mock_async_streaming_read_rpc.h" #include diff --git a/google/cloud/google_cloud_cpp_grpc_utils.bzl b/google/cloud/google_cloud_cpp_grpc_utils.bzl index d565700024f2c..193036f6510a9 100644 --- a/google/cloud/google_cloud_cpp_grpc_utils.bzl +++ b/google/cloud/google_cloud_cpp_grpc_utils.bzl @@ -67,6 +67,7 @@ google_cloud_cpp_grpc_utils_hdrs = [ "internal/grpc_service_account_authentication.h", "internal/log_wrapper.h", "internal/minimal_iam_credentials_stub.h", + "internal/operation_context.h", "internal/populate_grpc_options.h", "internal/resumable_streaming_read_rpc.h", "internal/retry_loop.h", diff --git a/google/cloud/google_cloud_cpp_grpc_utils.cmake b/google/cloud/google_cloud_cpp_grpc_utils.cmake index f8b099bec2f0c..bf7b75f0a31aa 100644 --- a/google/cloud/google_cloud_cpp_grpc_utils.cmake +++ b/google/cloud/google_cloud_cpp_grpc_utils.cmake @@ -88,6 +88,7 @@ add_library( internal/log_wrapper.h internal/minimal_iam_credentials_stub.cc internal/minimal_iam_credentials_stub.h + internal/operation_context.h internal/populate_grpc_options.cc internal/populate_grpc_options.h internal/resumable_streaming_read_rpc.h diff --git a/google/cloud/internal/operation_context.h b/google/cloud/internal/operation_context.h new file mode 100644 index 0000000000000..ed179a99926d6 --- /dev/null +++ b/google/cloud/internal/operation_context.h @@ -0,0 +1,62 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_INTERNAL_OPERATION_CONTEXT_H +#define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_INTERNAL_OPERATION_CONTEXT_H + +#include "google/cloud/status.h" +#include "google/cloud/version.h" +#include + +namespace google { +namespace cloud { +GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN +namespace internal { + +/** + * An abstract base class for service-specific operation contexts across + * retries. + * + * This provides lifecycle hooks (`PreCall`, `PostCall`, `OnDone`) for generated + * stubs and decorators. + */ +class OperationContext { + public: + virtual ~OperationContext() = default; + + // Called before each RPC attempt to inject headers and update state. + virtual void PreCall(grpc::ClientContext& context) = 0; + + // Called after receiving an RPC attempt response. + virtual void PostCall(grpc::ClientContext const& context, + Status const& status) = 0; + + // Called when the overall logical operation completes across all attempts. + virtual void OnDone(Status const& status) = 0; +}; + +class NoopOperationContext : public OperationContext { + public: + ~NoopOperationContext() override = default; + void PreCall(grpc::ClientContext&) override {} + void PostCall(grpc::ClientContext const&, Status const&) override {} + void OnDone(Status const&) override {} +}; + +} // namespace internal +GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END +} // namespace cloud +} // namespace google + +#endif // GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_INTERNAL_OPERATION_CONTEXT_H diff --git a/google/cloud/spanner/CMakeLists.txt b/google/cloud/spanner/CMakeLists.txt index a4b54cf773fd7..da4ad27aeaaa0 100644 --- a/google/cloud/spanner/CMakeLists.txt +++ b/google/cloud/spanner/CMakeLists.txt @@ -133,6 +133,8 @@ add_library( internal/logging_result_set_reader.h internal/merge_chunk.cc internal/merge_chunk.h + internal/operation_context.cc + internal/operation_context.h internal/partial_result_set_reader.h internal/partial_result_set_resume.cc internal/partial_result_set_resume.h @@ -150,6 +152,10 @@ add_library( internal/spanner_logging_decorator.h internal/spanner_metadata_decorator.cc internal/spanner_metadata_decorator.h + internal/spanner_operation_context_factory.cc + internal/spanner_operation_context_factory.h + internal/spanner_request_id.cc + internal/spanner_request_id.h internal/spanner_stub.cc internal/spanner_stub.h internal/spanner_stub_factory.cc @@ -458,10 +464,13 @@ function (spanner_client_define_tests) internal/instance_admin_metadata_test.cc internal/logging_result_set_reader_test.cc internal/merge_chunk_test.cc + internal/operation_context_test.cc internal/partial_result_set_resume_test.cc internal/partial_result_set_source_test.cc internal/route_to_leader_test.cc internal/session_pool_test.cc + internal/spanner_operation_context_factory_test.cc + internal/spanner_request_id_test.cc internal/spanner_stub_factory_test.cc internal/status_utils_test.cc internal/transaction_impl_test.cc diff --git a/google/cloud/spanner/benchmarks/multiple_rows_cpu_benchmark.cc b/google/cloud/spanner/benchmarks/multiple_rows_cpu_benchmark.cc index 92c881035e736..09106c0c5927b 100644 --- a/google/cloud/spanner/benchmarks/multiple_rows_cpu_benchmark.cc +++ b/google/cloud/spanner/benchmarks/multiple_rows_cpu_benchmark.cc @@ -16,6 +16,7 @@ #include "google/cloud/spanner/benchmarks/benchmarks_config.h" #include "google/cloud/spanner/client.h" #include "google/cloud/spanner/internal/defaults.h" +#include "google/cloud/spanner/internal/operation_context.h" #include "google/cloud/spanner/internal/route_to_leader.h" #include "google/cloud/spanner/internal/session_pool.h" #include "google/cloud/spanner/internal/spanner_stub_factory.h" @@ -545,8 +546,9 @@ class ReadExperiment : public BasicExperiment { spanner_internal::RouteToLeader(context); // always for CreateSession google::spanner::v1::CreateSessionRequest request{}; request.set_database(database.FullName()); - auto response = - stub->CreateSession(context, google::cloud::Options{}, request); + spanner_internal::OperationContext op_context(nullptr, 0, ""); + auto response = stub->CreateSession(context, google::cloud::Options{}, + request, op_context); if (response) return response->name(); last_status = response.status(); } @@ -590,8 +592,10 @@ class ReadExperiment : public BasicExperiment { int row_count = 0; std::vector row; row.resize(columns.size()); - auto stream = stub->StreamingRead(std::make_shared(), - google::cloud::Options{}, request); + auto stream = stub->StreamingRead( + std::make_shared(), google::cloud::Options{}, + request, + std::make_shared(nullptr, 0, "")); for (;;) { google::spanner::v1::PartialResultSet result; auto status = stream->Read(&result); @@ -688,8 +692,9 @@ class SelectExperiment : public BasicExperiment { spanner_internal::RouteToLeader(context); // always for CreateSession google::spanner::v1::CreateSessionRequest request{}; request.set_database(database.FullName()); - auto response = - stub->CreateSession(context, google::cloud::Options{}, request); + spanner_internal::OperationContext op_context(nullptr, 0, ""); + auto response = stub->CreateSession(context, google::cloud::Options{}, + request, op_context); if (response) return response->name(); last_status = response.status(); } @@ -737,9 +742,10 @@ class SelectExperiment : public BasicExperiment { int row_count = 0; std::vector row; row.resize(ExperimentImpl::kColumnCount); - auto stream = - stub->ExecuteStreamingSql(std::make_shared(), - google::cloud::Options{}, request); + auto stream = stub->ExecuteStreamingSql( + std::make_shared(), google::cloud::Options{}, + request, + std::make_shared(nullptr, 0, "")); for (;;) { google::spanner::v1::PartialResultSet result; auto status = stream->Read(&result); @@ -852,8 +858,9 @@ class UpdateExperiment : public BasicExperiment { spanner_internal::RouteToLeader(context); // always for CreateSession google::spanner::v1::CreateSessionRequest request{}; request.set_database(database.FullName()); - auto response = - stub->CreateSession(context, google::cloud::Options{}, request); + spanner_internal::OperationContext op_context(nullptr, 0, ""); + auto response = stub->CreateSession(context, google::cloud::Options{}, + request, op_context); if (response) return response->name(); last_status = response.status(); } @@ -913,8 +920,9 @@ class UpdateExperiment : public BasicExperiment { google::cloud::Status status; { grpc::ClientContext context; - auto response = - stub->ExecuteSql(context, google::cloud::Options{}, request); + spanner_internal::OperationContext op_context(nullptr, 0, ""); + auto response = stub->ExecuteSql(context, google::cloud::Options{}, + request, op_context); if (response) { row_count = static_cast(response->stats().row_count_lower_bound()); @@ -929,8 +937,9 @@ class UpdateExperiment : public BasicExperiment { google::spanner::v1::CommitRequest commit_request; commit_request.set_session(*session); commit_request.set_transaction_id(transaction_id); - auto response = - stub->Commit(context, google::cloud::Options{}, commit_request); + spanner_internal::OperationContext op_context(nullptr, 0, ""); + auto response = stub->Commit(context, google::cloud::Options{}, + commit_request, op_context); if (!response) status = std::move(response).status(); } @@ -1041,8 +1050,9 @@ class MutationExperiment : public BasicExperiment { spanner_internal::RouteToLeader(context); // always for CreateSession google::spanner::v1::CreateSessionRequest request{}; request.set_database(database.FullName()); - auto response = - stub->CreateSession(context, google::cloud::Options{}, request); + spanner_internal::OperationContext op_context(nullptr, 0, ""); + auto response = stub->CreateSession(context, google::cloud::Options{}, + request, op_context); if (response) return response->name(); last_status = response.status(); } @@ -1103,8 +1113,9 @@ class MutationExperiment : public BasicExperiment { *row.add_values() = spanner_internal::ToProto(spanner::Value(std::move(v))).second; } - auto response = - stub->Commit(context, google::cloud::Options{}, commit_request); + spanner_internal::OperationContext op_context(nullptr, 0, ""); + auto response = stub->Commit(context, google::cloud::Options{}, + commit_request, op_context); auto const usage = timer.Sample(); samples.push_back(RowCpuSample{ diff --git a/google/cloud/spanner/google_cloud_cpp_spanner.bzl b/google/cloud/spanner/google_cloud_cpp_spanner.bzl index 4ee5cfe878095..f6d0bc6d62f95 100644 --- a/google/cloud/spanner/google_cloud_cpp_spanner.bzl +++ b/google/cloud/spanner/google_cloud_cpp_spanner.bzl @@ -72,6 +72,7 @@ google_cloud_cpp_spanner_hdrs = [ "internal/instance_admin_stub.h", "internal/logging_result_set_reader.h", "internal/merge_chunk.h", + "internal/operation_context.h", "internal/partial_result_set_reader.h", "internal/partial_result_set_resume.h", "internal/partial_result_set_source.h", @@ -81,6 +82,8 @@ google_cloud_cpp_spanner_hdrs = [ "internal/spanner_auth_decorator.h", "internal/spanner_logging_decorator.h", "internal/spanner_metadata_decorator.h", + "internal/spanner_operation_context_factory.h", + "internal/spanner_request_id.h", "internal/spanner_stub.h", "internal/spanner_stub_factory.h", "internal/spanner_tracing_stub.h", @@ -164,6 +167,7 @@ google_cloud_cpp_spanner_srcs = [ "internal/instance_admin_stub.cc", "internal/logging_result_set_reader.cc", "internal/merge_chunk.cc", + "internal/operation_context.cc", "internal/partial_result_set_resume.cc", "internal/partial_result_set_source.cc", "internal/route_to_leader.cc", @@ -172,6 +176,8 @@ google_cloud_cpp_spanner_srcs = [ "internal/spanner_auth_decorator.cc", "internal/spanner_logging_decorator.cc", "internal/spanner_metadata_decorator.cc", + "internal/spanner_operation_context_factory.cc", + "internal/spanner_request_id.cc", "internal/spanner_stub.cc", "internal/spanner_stub_factory.cc", "internal/spanner_tracing_stub.cc", diff --git a/google/cloud/spanner/integration_tests/session_pool_integration_test.cc b/google/cloud/spanner/integration_tests/session_pool_integration_test.cc index ef3036e083e77..8d785d6d3e5d1 100644 --- a/google/cloud/spanner/integration_tests/session_pool_integration_test.cc +++ b/google/cloud/spanner/integration_tests/session_pool_integration_test.cc @@ -14,6 +14,8 @@ #include "google/cloud/internal/disable_deprecation_warnings.inc" #include "google/cloud/spanner/internal/defaults.h" #include "google/cloud/spanner/internal/session_pool.h" +#include "google/cloud/spanner/internal/spanner_operation_context_factory.h" +#include "google/cloud/spanner/internal/spanner_request_id.h" #include "google/cloud/spanner/internal/spanner_stub_factory.h" #include "google/cloud/spanner/options.h" #include "google/cloud/spanner/testing/database_integration_test.h" @@ -34,22 +36,24 @@ struct SessionPoolFriendForTest { std::map const& labels, std::string const& role, int num_sessions) { internal::OptionsSpan span(session_pool->opts_); - return session_pool->AsyncBatchCreateSessions(cq, stub, labels, role, - num_sessions); + return session_pool->AsyncBatchCreateSessions(cq, stub, /*channel_id=*/0, + labels, role, num_sessions); } static future AsyncDeleteSession( std::shared_ptr const& session_pool, CompletionQueue& cq, std::shared_ptr const& stub, std::string session_name) { internal::OptionsSpan span(session_pool->opts_); - return session_pool->AsyncDeleteSession(cq, stub, std::move(session_name)); + return session_pool->AsyncDeleteSession(cq, stub, /*channel_id=*/0, + std::move(session_name)); } static future> AsyncRefreshSession( std::shared_ptr const& session_pool, CompletionQueue& cq, std::shared_ptr const& stub, std::string session_name) { internal::OptionsSpan span(session_pool->opts_); - return session_pool->AsyncRefreshSession(cq, stub, std::move(session_name)); + return session_pool->AsyncRefreshSession(cq, stub, /*channel_id=*/0, + std::move(session_name)); } }; @@ -73,7 +77,12 @@ TEST_F(SessionPoolIntegrationTest, SessionAsyncCRUD) { auto stub = CreateDefaultSpannerStub( db, internal::CreateAuthenticationStrategy(cq, opts), opts, /*channel_id=*/0); - auto session_pool = MakeSessionPool(db, {stub}, cq, opts); + auto context_factory = + std::make_shared( + /*client_id=*/1, std::make_shared( + spanner_internal::ProcessRandomId())); + auto session_pool = + MakeSessionPool(db, {stub}, cq, std::move(context_factory), opts); // Make an asynchronous request, but immediately block until the response // arrives diff --git a/google/cloud/spanner/internal/channel.h b/google/cloud/spanner/internal/channel.h index b0ccd0ea661c0..b903b3a913253 100644 --- a/google/cloud/spanner/internal/channel.h +++ b/google/cloud/spanner/internal/channel.h @@ -17,6 +17,7 @@ #include "google/cloud/spanner/internal/spanner_stub.h" #include "google/cloud/spanner/version.h" +#include #include namespace google { @@ -29,14 +30,15 @@ GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN */ struct Channel { /// @p stub_param must not be nullptr - explicit Channel(std::shared_ptr stub_param) - : stub(std::move(stub_param)) {} + Channel(std::shared_ptr stub_param, std::uint32_t channel_id) + : stub(std::move(stub_param)), channel_id(channel_id) {} // This class is not copyable or movable. Channel(Channel const&) = delete; Channel& operator=(Channel const&) = delete; std::shared_ptr const stub; + std::uint32_t const channel_id; int session_count = 0; }; diff --git a/google/cloud/spanner/internal/connection_impl.cc b/google/cloud/spanner/internal/connection_impl.cc index d77fe853ab8cb..79a83b93c4933 100644 --- a/google/cloud/spanner/internal/connection_impl.cc +++ b/google/cloud/spanner/internal/connection_impl.cc @@ -18,6 +18,8 @@ #include "google/cloud/spanner/internal/partial_result_set_resume.h" #include "google/cloud/spanner/internal/partial_result_set_source.h" #include "google/cloud/spanner/internal/route_to_leader.h" +#include "google/cloud/spanner/internal/spanner_operation_context_factory.h" +#include "google/cloud/spanner/internal/spanner_request_id.h" #include "google/cloud/spanner/internal/status_utils.h" #include "google/cloud/spanner/options.h" #include "google/cloud/spanner/query_partition.h" @@ -129,10 +131,13 @@ class DefaultPartialResultSetReader : public PartialResultSetReader { public: DefaultPartialResultSetReader( std::shared_ptr context, + std::shared_ptr operation_context, std::unique_ptr< internal::StreamingReadRpc> reader) - : context_(std::move(context)), reader_(std::move(reader)) {} + : context_(std::move(context)), + operation_context_(std::move(operation_context)), + reader_(std::move(reader)) {} ~DefaultPartialResultSetReader() override = default; @@ -145,6 +150,9 @@ class DefaultPartialResultSetReader : public PartialResultSetReader { if (opt_status.has_value()) { final_status_ = *std::move(opt_status); + if (operation_context_) { + operation_context_->PostCall(*context_, final_status_); + } return false; } return true; @@ -154,6 +162,7 @@ class DefaultPartialResultSetReader : public PartialResultSetReader { private: std::shared_ptr context_; + std::shared_ptr operation_context_; std::unique_ptr< internal::StreamingReadRpc> reader_; @@ -377,8 +386,12 @@ ConnectionImpl::ConnectionImpl( : db_(std::move(db)), background_threads_(std::move(background_threads)), opts_(internal::MergeOptions(std::move(opts), Connection::options())), + context_factory_(std::make_shared( + NextClientId(), + std::make_shared(ProcessRandomId()))), session_pool_(MakeSessionPool(db_, std::move(stubs), - background_threads_->cq(), opts_)) {} + background_threads_->cq(), context_factory_, + opts_)) {} spanner::RowStream ConnectionImpl::Read(ReadParams params) { return Visit( @@ -517,6 +530,10 @@ spanner::BatchedCommitResultStream ConnectionImpl::BatchWrite( * Helper function that ensures `session` holds a valid `Session`, or returns * an error if `session` is empty and no `Session` can be allocated. */ +Status ConnectionImpl::PrepareSession(SessionHolder& session) { + return PrepareSession(session, Session::Mode::kPooled); +} + Status ConnectionImpl::PrepareSession(SessionHolder& session, Session::Mode mode) { if (!session) { @@ -535,7 +552,7 @@ Status ConnectionImpl::PrepareSession(SessionHolder& session, return Status(); } -std::shared_ptr ConnectionImpl::GetStubBasedOnSessionMode( +StubAndChannel ConnectionImpl::GetStubBasedOnSessionMode( Session& session, TransactionContext& ctx) { if (session.is_multiplexed()) { return session_pool_->GetStub(session, ctx); @@ -570,18 +587,26 @@ StatusOr ConnectionImpl::BeginTransaction( *begin.mutable_mutation_key() = *mutation; } - auto stub = GetStubBasedOnSessionMode(*session, ctx); + auto stub_and_channel = GetStubBasedOnSessionMode(*session, ctx); + auto op_context = context_factory_->BeginTransaction(); + op_context.BindChannel(stub_and_channel.channel_id); + auto const& current = internal::CurrentOptions(); auto response = RetryLoop( RetryPolicyPrototype(current)->clone(), BackoffPolicyPrototype(current)->clone(), Idempotency::kIdempotent, - [&stub, route_to_leader = ctx.route_to_leader]( + [&stub = stub_and_channel.stub, &op_context, + route_to_leader = ctx.route_to_leader]( grpc::ClientContext& context, Options const& options, google::spanner::v1::BeginTransactionRequest const& request) { if (route_to_leader) RouteToLeader(context); - return stub->BeginTransaction(context, options, request); + op_context.PreCall(context); + auto s = stub->BeginTransaction(context, options, request, op_context); + op_context.PostCall(context, s.status()); + return s; }, current, begin, func); + op_context.OnDone(response.status()); if (!response) { auto status = std::move(response).status(); if (IsSessionNotFound(status)) session->set_bad(); @@ -644,23 +669,26 @@ spanner::RowStream ConnectionImpl::ReadImpl( }), params.directed_read_option); - // Capture a copy of `stub` to ensure the `shared_ptr<>` remains valid through - // the lifetime of the lambda. - auto stub = GetStubBasedOnSessionMode(*session, ctx); + auto stub_and_channel = GetStubBasedOnSessionMode(*session, ctx); auto const tracing_enabled = RpcStreamTracingEnabled(); auto const& tracing_options = RpcTracingOptions(); - auto factory = [stub, request, route_to_leader = ctx.route_to_leader, - tracing_enabled, + auto op_context = + std::make_shared(context_factory_->StreamingRead()); + op_context->BindChannel(stub_and_channel.channel_id); + + auto factory = [stub = stub_and_channel.stub, request, op_context, + route_to_leader = ctx.route_to_leader, tracing_enabled, tracing_options](std::string const& resume_token) mutable { if (!resume_token.empty()) request->set_resume_token(resume_token); auto context = std::make_shared(); auto const& options = internal::CurrentOptions(); internal::ConfigureContext(*context, options); if (route_to_leader) RouteToLeader(*context); - auto stream = stub->StreamingRead(context, options, *request); + op_context->PreCall(*context); + auto stream = stub->StreamingRead(context, options, *request, op_context); std::unique_ptr reader = - std::make_unique(std::move(context), - std::move(stream)); + std::make_unique( + std::move(context), op_context, std::move(stream)); if (tracing_enabled) { reader = std::make_unique(std::move(reader), tracing_options); @@ -671,7 +699,7 @@ spanner::RowStream ConnectionImpl::ReadImpl( auto rpc = std::make_unique( factory, Idempotency::kIdempotent, RetryPolicyPrototype()->clone(), BackoffPolicyPrototype()->clone()); - auto reader = PartialResultSetSource::Create(std::move(rpc)); + auto reader = PartialResultSetSource::Create(std::move(rpc), op_context); if (reader.ok()) { ctx.precommit_token = (*reader)->PrecommitToken(); } @@ -730,22 +758,30 @@ StatusOr> ConnectionImpl::PartitionReadImpl( *request.mutable_key_set() = ToProto(params.keys); *request.mutable_partition_options() = ToProto(partition_options); - auto stub = GetStubBasedOnSessionMode(*session, ctx); + auto stub_and_channel = GetStubBasedOnSessionMode(*session, ctx); + auto op_context = context_factory_->PartitionRead(); + op_context.BindChannel(stub_and_channel.channel_id); + auto const& current = internal::CurrentOptions(); for (;;) { auto response = RetryLoop( RetryPolicyPrototype()->clone(), BackoffPolicyPrototype()->clone(), Idempotency::kIdempotent, - [&stub](grpc::ClientContext& context, Options const& options, - google::spanner::v1::PartitionReadRequest const& request) { + [&stub = stub_and_channel.stub, &op_context]( + grpc::ClientContext& context, Options const& options, + google::spanner::v1::PartitionReadRequest const& request) { RouteToLeader(context); // always for PartitionRead() - return stub->PartitionRead(context, options, request); + op_context.PreCall(context); + auto s = stub->PartitionRead(context, options, request, op_context); + op_context.PostCall(context, s.status()); + return s; }, current, request, __func__); if (selector->has_begin()) { if (response.ok()) { if (!response->has_transaction()) { selector = MissingTransactionStatus(__func__); + op_context.OnDone(selector.status()); return selector.status(); } selector->set_id(response->transaction().id()); @@ -763,6 +799,7 @@ StatusOr> ConnectionImpl::PartitionReadImpl( if (!response.ok()) { auto status = std::move(response).status(); + op_context.OnDone(status); if (IsSessionNotFound(status)) session->set_bad(); return status; } @@ -780,6 +817,7 @@ StatusOr> ConnectionImpl::PartitionReadImpl( params.keys, params.columns, data_boost, params.read_options)); } + op_context.OnDone(Status{}); return read_partitions; } } @@ -877,30 +915,34 @@ ResultType ConnectionImpl::CommonQueryImpl( if (!prepare_status.ok()) { return MakeStatusOnlyResult(std::move(prepare_status)); } - // Capture a copy of of these to ensure the `shared_ptr<>` remains valid - // through the lifetime of the lambda. Note that the local variables are a - // reference to avoid increasing refcounts twice, but the capture is by value. - auto stub = GetStubBasedOnSessionMode(*session, ctx); + auto stub_and_channel = GetStubBasedOnSessionMode(*session, ctx); auto const& retry_policy_prototype = RetryPolicyPrototype(); auto const& backoff_policy_prototype = BackoffPolicyPrototype(); auto const tracing_enabled = RpcStreamTracingEnabled(); auto const& tracing_options = RpcTracingOptions(); + auto op_context = std::make_shared( + context_factory_->ExecuteStreamingSql()); + op_context->BindChannel(stub_and_channel.channel_id); + auto retry_resume_fn = - [stub, retry_policy_prototype, backoff_policy_prototype, - route_to_leader = ctx.route_to_leader, tracing_enabled, + [stub = stub_and_channel.stub, op_context, retry_policy_prototype, + backoff_policy_prototype, route_to_leader = ctx.route_to_leader, + tracing_enabled, tracing_options](google::spanner::v1::ExecuteSqlRequest& request) mutable -> StatusOr> { - auto factory = [stub, request, route_to_leader, tracing_enabled, + auto factory = [stub, request, op_context, route_to_leader, tracing_enabled, tracing_options](std::string const& resume_token) mutable { if (!resume_token.empty()) request.set_resume_token(resume_token); auto context = std::make_shared(); auto const& options = internal::CurrentOptions(); internal::ConfigureContext(*context, options); if (route_to_leader) RouteToLeader(*context); - auto stream = stub->ExecuteStreamingSql(context, options, request); + op_context->PreCall(*context); + auto stream = + stub->ExecuteStreamingSql(context, options, request, op_context); std::unique_ptr reader = - std::make_unique(std::move(context), - std::move(stream)); + std::make_unique( + std::move(context), op_context, std::move(stream)); if (tracing_enabled) { reader = std::make_unique(std::move(reader), tracing_options); @@ -911,7 +953,7 @@ ResultType ConnectionImpl::CommonQueryImpl( std::move(factory), Idempotency::kIdempotent, retry_policy_prototype->clone(), backoff_policy_prototype->clone()); - return PartialResultSetSource::Create(std::move(rpc)); + return PartialResultSetSource::Create(std::move(rpc), op_context); }; StatusOr response = @@ -955,27 +997,31 @@ StatusOr ConnectionImpl::CommonDmlImpl( if (!prepare_status.ok()) { return prepare_status; } - // Capture a copy of of these to ensure the `shared_ptr<>` remains valid - // through the lifetime of the lambda. Note that the local variables are a - // reference to avoid increasing refcounts twice, but the capture is by value. - auto stub = GetStubBasedOnSessionMode(*session, ctx); + auto stub_and_channel = GetStubBasedOnSessionMode(*session, ctx); auto current = google::cloud::internal::SaveCurrentOptions(); auto const& retry_policy_prototype = RetryPolicyPrototype(*current); auto const& backoff_policy_prototype = BackoffPolicyPrototype(*current); + auto op_context = + std::make_shared(context_factory_->ExecuteSql()); + op_context->BindChannel(stub_and_channel.channel_id); auto retry_resume_fn = - [function_name, stub, retry_policy_prototype, backoff_policy_prototype, - session, route_to_leader = ctx.route_to_leader, + [function_name, stub = stub_and_channel.stub, op_context, + retry_policy_prototype, backoff_policy_prototype, session, + route_to_leader = ctx.route_to_leader, current](google::spanner::v1::ExecuteSqlRequest& request) mutable -> StatusOr> { StatusOr response = RetryLoop( retry_policy_prototype->clone(), backoff_policy_prototype->clone(), Idempotency::kIdempotent, - [stub, route_to_leader]( + [stub, op_context, route_to_leader]( grpc::ClientContext& context, Options const& options, google::spanner::v1::ExecuteSqlRequest const& request) { if (route_to_leader) RouteToLeader(context); - return stub->ExecuteSql(context, options, request); + op_context->PreCall(context); + auto s = stub->ExecuteSql(context, options, request, *op_context); + op_context->PostCall(context, s.status()); + return s; }, *current, request, function_name); if (!response) { @@ -985,8 +1031,11 @@ StatusOr ConnectionImpl::CommonDmlImpl( } return DmlResultSetSource::Create(std::move(*response)); }; - return ExecuteSqlImpl(session, selector, ctx, std::move(params), - query_mode, std::move(retry_resume_fn)); + auto result = + ExecuteSqlImpl(session, selector, ctx, std::move(params), + query_mode, std::move(retry_resume_fn)); + op_context->OnDone(result.status()); + return result; } StatusOr ConnectionImpl::ExecuteDmlImpl( @@ -1044,22 +1093,30 @@ ConnectionImpl::PartitionQueryImpl( std::move(*sql_statement.mutable_param_types()); *request.mutable_partition_options() = ToProto(params.partition_options); - auto stub = GetStubBasedOnSessionMode(*session, ctx); + auto stub_and_channel = GetStubBasedOnSessionMode(*session, ctx); + auto op_context = context_factory_->PartitionQuery(); + op_context.BindChannel(stub_and_channel.channel_id); + auto const& current = internal::CurrentOptions(); for (;;) { auto response = RetryLoop( RetryPolicyPrototype()->clone(), BackoffPolicyPrototype()->clone(), Idempotency::kIdempotent, - [&stub](grpc::ClientContext& context, Options const& options, - google::spanner::v1::PartitionQueryRequest const& request) { + [&stub = stub_and_channel.stub, &op_context]( + grpc::ClientContext& context, Options const& options, + google::spanner::v1::PartitionQueryRequest const& request) { RouteToLeader(context); // always for PartitionQuery() - return stub->PartitionQuery(context, options, request); + op_context.PreCall(context); + auto s = stub->PartitionQuery(context, options, request, op_context); + op_context.PostCall(context, s.status()); + return s; }, current, request, __func__); if (selector->has_begin()) { if (response.ok()) { if (!response->has_transaction()) { selector = MissingTransactionStatus(__func__); + op_context.OnDone(selector.status()); return selector.status(); } selector->set_id(response->transaction().id()); @@ -1076,6 +1133,7 @@ ConnectionImpl::PartitionQueryImpl( } if (!response.ok()) { auto status = std::move(response).status(); + op_context.OnDone(status); if (IsSessionNotFound(status)) session->set_bad(); return status; } @@ -1087,6 +1145,7 @@ ConnectionImpl::PartitionQueryImpl( session->session_name(), partition.partition_token(), params.partition_options.data_boost, params.statement)); } + op_context.OnDone(Status{}); return query_partitions; } } @@ -1118,16 +1177,23 @@ StatusOr ConnectionImpl::ExecuteBatchDmlImpl( request.mutable_request_options()->set_request_tag(request_tag); request.mutable_request_options()->set_transaction_tag(ctx.tag); - auto stub = GetStubBasedOnSessionMode(*session, ctx); + auto stub_and_channel = GetStubBasedOnSessionMode(*session, ctx); + auto op_context = context_factory_->ExecuteBatchDml(); + op_context.BindChannel(stub_and_channel.channel_id); + auto const& current = internal::CurrentOptions(); for (;;) { auto response = RetryLoop( RetryPolicyPrototype()->clone(), BackoffPolicyPrototype()->clone(), Idempotency::kIdempotent, - [&stub](grpc::ClientContext& context, Options const& options, - google::spanner::v1::ExecuteBatchDmlRequest const& request) { + [&stub = stub_and_channel.stub, &op_context]( + grpc::ClientContext& context, Options const& options, + google::spanner::v1::ExecuteBatchDmlRequest const& request) { RouteToLeader(context); // always for ExecuteBatchDml() - return stub->ExecuteBatchDml(context, options, request); + op_context.PreCall(context); + auto s = stub->ExecuteBatchDml(context, options, request, op_context); + op_context.PostCall(context, s.status()); + return s; }, current, request, __func__); if (response.ok() && response->has_precommit_token()) { @@ -1137,6 +1203,7 @@ StatusOr ConnectionImpl::ExecuteBatchDmlImpl( if (response.ok() && response->result_sets_size() > 0) { if (!response->result_sets(0).metadata().has_transaction()) { selector = MissingTransactionStatus(__func__); + op_context.OnDone(selector.status()); return selector.status(); } selector->set_id( @@ -1154,6 +1221,7 @@ StatusOr ConnectionImpl::ExecuteBatchDmlImpl( } if (!response) { auto status = std::move(response).status(); + op_context.OnDone(status); if (IsSessionNotFound(status)) session->set_bad(); return status; } @@ -1162,6 +1230,7 @@ StatusOr ConnectionImpl::ExecuteBatchDmlImpl( for (auto const& result_set : response->result_sets()) { result.stats.push_back({result_set.stats().row_count_exact()}); } + op_context.OnDone(result.status); return result; } } @@ -1273,7 +1342,9 @@ StatusOr ConnectionImpl::CommitImpl( GCP_ERROR_INFO()); } - auto stub = GetStubBasedOnSessionMode(*session, ctx); + auto stub_and_channel = GetStubBasedOnSessionMode(*session, ctx); + auto op_context = context_factory_->Commit(); + op_context.BindChannel(stub_and_channel.channel_id); auto const& current = internal::CurrentOptions(); char const* calling_func = __func__; @@ -1286,15 +1357,20 @@ StatusOr ConnectionImpl::CommitImpl( *request.mutable_precommit_token() = *token; } - return RetryLoop( + auto sor = RetryLoop( RetryPolicyPrototype(current)->clone(), BackoffPolicyPrototype(current)->clone(), Idempotency::kIdempotent, - [&stub](grpc::ClientContext& context, Options const& options, - google::spanner::v1::CommitRequest const& request) { + [&stub = stub_and_channel.stub, &op_context]( + grpc::ClientContext& context, Options const& options, + google::spanner::v1::CommitRequest const& request) { RouteToLeader(context); // always for Commit() - return stub->Commit(context, options, request); + op_context.PreCall(context); + auto s = stub->Commit(context, options, request, op_context); + op_context.PostCall(context, s.status()); + return s; }, current, request, func); + return sor; }; // If the CommitResponse contains a precommit token, it's a signal from the @@ -1307,6 +1383,7 @@ StatusOr ConnectionImpl::CommitImpl( response = retry_loop_fn(ctx.precommit_token); if (!response) { auto status = std::move(response).status(); + op_context.OnDone(status); if (IsSessionNotFound(status)) session->set_bad(); return status; } @@ -1316,6 +1393,7 @@ StatusOr ConnectionImpl::CommitImpl( } } while (response->has_precommit_token()); + op_context.OnDone(response.status()); spanner::CommitResult r; r.commit_timestamp = MakeTimestamp(response->commit_timestamp()); if (response->has_commit_stats()) { @@ -1353,17 +1431,24 @@ Status ConnectionImpl::RollbackImpl( google::spanner::v1::RollbackRequest request; request.set_session(session->session_name()); request.set_transaction_id(selector->id()); - auto stub = GetStubBasedOnSessionMode(*session, ctx); + auto stub_and_channel = GetStubBasedOnSessionMode(*session, ctx); + auto op_context = context_factory_->Rollback(); + op_context.BindChannel(stub_and_channel.channel_id); auto const& current = internal::CurrentOptions(); auto status = RetryLoop( RetryPolicyPrototype(current)->clone(), BackoffPolicyPrototype(current)->clone(), Idempotency::kIdempotent, - [&stub](grpc::ClientContext& context, Options const& options, - google::spanner::v1::RollbackRequest const& request) { + [&stub = stub_and_channel.stub, &op_context]( + grpc::ClientContext& context, Options const& options, + google::spanner::v1::RollbackRequest const& request) { RouteToLeader(context); // always for Rollback() - return stub->Rollback(context, options, request); + op_context.PreCall(context); + auto s = stub->Rollback(context, options, request, op_context); + op_context.PostCall(context, s); + return s; }, current, request, __func__); + op_context.OnDone(status); if (IsSessionNotFound(status)) session->set_bad(); return status; } @@ -1404,14 +1489,19 @@ spanner::BatchedCommitResultStream ConnectionImpl::BatchWriteImpl( // There's no client-side transaction involved with BatchWrite, so no need // to store the resulting stub in the case of a Multiplexed Session. - auto stub = session_pool_->GetStub(*session); - auto factory = [stub = std::move(stub)]( + auto stub_and_channel = session_pool_->GetStub(*session); + auto op_context = + std::make_shared(context_factory_->BatchWrite()); + op_context->BindChannel(stub_and_channel.channel_id); + + auto factory = [stub = stub_and_channel.stub, op_context]( google::spanner::v1::BatchWriteRequest const& request) { auto context = std::make_shared(); auto const& options = internal::CurrentOptions(); internal::ConfigureContext(*context, options); RouteToLeader(*context); // always for BatchWrite() - return stub->BatchWrite(std::move(context), options, request); + op_context->PreCall(*context); + return stub->BatchWrite(std::move(context), options, request, op_context); }; auto updater = [](google::spanner::v1::BatchWriteResponse const&, google::spanner::v1::BatchWriteRequest&) { @@ -1422,15 +1512,38 @@ spanner::BatchedCommitResultStream ConnectionImpl::BatchWriteImpl( google::spanner::v1::BatchWriteRequest>( RetryPolicyPrototype()->clone(), BackoffPolicyPrototype()->clone(), std::move(factory), std::move(updater), std::move(request)); + struct BatchWriteGuard { + explicit BatchWriteGuard(std::shared_ptr context) + : op_context(std::move(context)), + final_status( + internal::CancelledError("Stream cancelled", GCP_ERROR_INFO())) {} + + ~BatchWriteGuard() { + if (!called && op_context) { + op_context->OnDone(final_status); + } + } + + std::shared_ptr op_context; + Status final_status; + bool called = false; + }; + auto guard = std::make_shared(op_context); // Because there is no enclosing client-side transaction, we move the // session into the stream range so that it is not returned to the pool // until the stream is exhausted. return internal::MakeStreamRange( - [reader = std::move(reader), session = std::move(session)]() - -> std::variant { + [reader = std::move(reader), session = std::move(session), + guard = std::move( + guard)]() -> std::variant { google::spanner::v1::BatchWriteResponse response; auto result = reader->Read(&response); if (result.has_value()) { + guard->final_status = *result; + guard->called = true; + if (guard->op_context) { + guard->op_context->OnDone(*result); + } // "Session not found" can really only happen on the first // response, but, rather than tracking that, we just check // on non-first failures too. diff --git a/google/cloud/spanner/internal/connection_impl.h b/google/cloud/spanner/internal/connection_impl.h index 16b9218f9b26a..254123b89acdf 100644 --- a/google/cloud/spanner/internal/connection_impl.h +++ b/google/cloud/spanner/internal/connection_impl.h @@ -69,11 +69,11 @@ class ConnectionImpl : public spanner::Connection { spanner::BatchedCommitResultStream BatchWrite(BatchWriteParams) override; private: - Status PrepareSession(SessionHolder& session, - Session::Mode mode = Session::Mode::kPooled); + Status PrepareSession(SessionHolder& session); + Status PrepareSession(SessionHolder& session, Session::Mode mode); - std::shared_ptr GetStubBasedOnSessionMode( - Session& session, TransactionContext& ctx); + StubAndChannel GetStubBasedOnSessionMode(Session& session, + TransactionContext& ctx); StatusOr BeginTransaction( SessionHolder& session, google::spanner::v1::TransactionOptions options, @@ -173,6 +173,7 @@ class ConnectionImpl : public spanner::Connection { spanner::Database db_; std::unique_ptr background_threads_; Options opts_; + std::shared_ptr context_factory_; std::shared_ptr session_pool_; }; diff --git a/google/cloud/spanner/internal/connection_impl_test.cc b/google/cloud/spanner/internal/connection_impl_test.cc index c3c372013a05d..81f741b86cf88 100644 --- a/google/cloud/spanner/internal/connection_impl_test.cc +++ b/google/cloud/spanner/internal/connection_impl_test.cc @@ -28,6 +28,7 @@ #include "google/cloud/log.h" #include "google/cloud/testing_util/is_proto_equal.h" #include "google/cloud/testing_util/status_matchers.h" +#include "google/cloud/testing_util/validate_metadata.h" #include "absl/time/time.h" #include #include @@ -397,7 +398,7 @@ TEST(ConnectionImplTest, ReadCreateSessionFailure) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillRepeatedly(Return( Status(StatusCode::kPermissionDenied, "uh-oh in CreateSession"))); @@ -418,7 +419,7 @@ TEST(ConnectionImplTest, ReadStreamingReadFailure) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); auto finish_status = internal::PermissionDeniedError("uh-oh in GrpcReader::Finish"); @@ -443,7 +444,7 @@ TEST(ConnectionImplTest, ReadSuccess) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); auto retry_status = internal::UnavailableError("try-again"); std::vector responses = { @@ -476,17 +477,20 @@ TEST(ConnectionImplTest, ReadSuccess) { EXPECT_CALL( *mock, StreamingRead( - _, _, HasPriority(google::spanner::v1::RequestOptions::PRIORITY_LOW))) - .WillOnce([&retry_status]( - std::shared_ptr const&, Options const&, - google::spanner::v1::ReadRequest const& request) { + _, _, HasPriority(google::spanner::v1::RequestOptions::PRIORITY_LOW), + _)) + .WillOnce([&retry_status](std::shared_ptr const&, + Options const&, + google::spanner::v1::ReadRequest const& request, + auto const&) { // The beginning of the row stream, but immediately fail. EXPECT_THAT(request.resume_token(), IsEmpty()); return MakeReader({}, retry_status); }) .WillOnce([&responses, &retry_status]( std::shared_ptr const&, Options const&, - google::spanner::v1::ReadRequest const& request) { + google::spanner::v1::ReadRequest const& request, + auto const&) { // Restart from the beginning, but return the metadata and first // row before failing again. EXPECT_THAT(request.resume_token(), IsEmpty()); @@ -494,7 +498,8 @@ TEST(ConnectionImplTest, ReadSuccess) { }) .WillOnce([&responses, &retry_status]( std::shared_ptr const&, Options const&, - google::spanner::v1::ReadRequest const& request) { + google::spanner::v1::ReadRequest const& request, + auto const&) { // Restart from the second row, but only return part of it before // failing once more. EXPECT_THAT(request.resume_token(), Eq("restart-row-2")); @@ -502,7 +507,8 @@ TEST(ConnectionImplTest, ReadSuccess) { }) .WillOnce([&responses](std::shared_ptr const&, Options const&, - google::spanner::v1::ReadRequest const& request) { + google::spanner::v1::ReadRequest const& request, + auto const&) { // Restart from the second row, but now deliver it all in two chunks. EXPECT_THAT(request.resume_token(), Eq("restart-row-2")); return MakeReader({responses[1], responses[2]}); @@ -530,12 +536,13 @@ TEST(ConnectionImplTest, ReadDirectedRead) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, BeginTransaction).Times(0); EXPECT_CALL(*mock, StreamingRead) .WillOnce([](std::shared_ptr const&, Options const&, - google::spanner::v1::ReadRequest const& request) { + google::spanner::v1::ReadRequest const& request, + auto const&) { EXPECT_EQ("multiplexed", request.session()); EXPECT_TRUE(request.has_directed_read_options()); auto const& directed_read_options = request.directed_read_options(); @@ -575,7 +582,7 @@ TEST(ConnectionImplTest, ReadPermanentFailure) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, StreamingRead) .WillOnce(Return(ByMove(MakeReader( @@ -598,7 +605,7 @@ TEST(ConnectionImplTest, ReadTooManyTransientFailures) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, StreamingRead) .Times(AtLeast(2)) @@ -626,7 +633,7 @@ TEST(ConnectionImplTest, ReadImplicitBeginTransaction) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, BeginTransaction).Times(0); auto constexpr kText = R"pb( @@ -676,17 +683,19 @@ TEST(ConnectionImplTest, ReadImplicitBeginTransactionOneTransientFailure) { // n.b. these calls are explicitly sequenced because using the scoped // `InSequence` object causes gMock to get confused by the reader calls. Sequence s; - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .InSequence(s) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL( - *mock, StreamingRead( - _, _, AllOf(HasSession("multiplexed"), HasBeginTransaction()))) + *mock, + StreamingRead(_, _, + AllOf(HasSession("multiplexed"), HasBeginTransaction()), _)) .InSequence(s) .WillOnce(Return(ByMove(std::move(failing_reader)))); EXPECT_CALL( - *mock, StreamingRead( - _, _, AllOf(HasSession("multiplexed"), HasBeginTransaction()))) + *mock, + StreamingRead(_, _, + AllOf(HasSession("multiplexed"), HasBeginTransaction()), _)) .InSequence(s) .WillOnce(Return(ByMove(std::move(ok_reader)))); @@ -735,12 +744,13 @@ TEST(ConnectionImplTest, ReadImplicitBeginTransactionOnePermanentFailure) { // n.b. these calls are explicitly sequenced because using the scoped // `InSequence` object causes gMock to get confused by the reader calls. Sequence s; - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .InSequence(s) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL( - *mock, StreamingRead( - _, _, AllOf(HasSession("multiplexed"), HasBeginTransaction()))) + *mock, + StreamingRead(_, _, + AllOf(HasSession("multiplexed"), HasBeginTransaction()), _)) .InSequence(s) .WillOnce(Return(ByMove(std::move(failing_reader)))); EXPECT_CALL(*mock, BeginTransaction) @@ -748,7 +758,8 @@ TEST(ConnectionImplTest, ReadImplicitBeginTransactionOnePermanentFailure) { .WillOnce(Return(MakeTestTransaction("FEDCBA98"))); EXPECT_CALL(*mock, StreamingRead(_, _, AllOf(HasSession("multiplexed"), - HasTransactionId("FEDCBA98")))) + HasTransactionId("FEDCBA98")), + _)) .InSequence(s) .WillOnce(Return(ByMove(std::move(ok_reader)))); @@ -779,12 +790,13 @@ TEST(ConnectionImplTest, ReadImplicitBeginTransactionPermanentFailure) { // n.b. these calls are explicitly sequenced because using the scoped // `InSequence` object causes gMock to get confused by the reader calls. Sequence s; - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .InSequence(s) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL( - *mock, StreamingRead( - _, _, AllOf(HasSession("multiplexed"), HasBeginTransaction()))) + *mock, + StreamingRead(_, _, + AllOf(HasSession("multiplexed"), HasBeginTransaction()), _)) .InSequence(s) .WillOnce(Return(ByMove(std::move(reader1)))); EXPECT_CALL(*mock, BeginTransaction) @@ -792,7 +804,8 @@ TEST(ConnectionImplTest, ReadImplicitBeginTransactionPermanentFailure) { .WillOnce(Return(MakeTestTransaction("FEDCBA98"))); EXPECT_CALL(*mock, StreamingRead(_, _, AllOf(HasSession("multiplexed"), - HasTransactionId("FEDCBA98")))) + HasTransactionId("FEDCBA98")), + _)) .InSequence(s) .WillOnce(Return(ByMove(std::move(reader2)))); @@ -813,7 +826,7 @@ TEST(ConnectionImplTest, ExecuteQueryCreateSessionFailure) { auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillRepeatedly(Return( Status(StatusCode::kPermissionDenied, "uh-oh in CreateSession"))); @@ -832,7 +845,7 @@ TEST(ConnectionImplTest, ExecuteQueryStreamingReadFailure) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, ExecuteStreamingSql) .WillOnce(Return(ByMove(MakeReader( @@ -854,7 +867,7 @@ TEST(ConnectionImplTest, ExecuteQueryReadSuccess) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); auto constexpr kText = R"pb( metadata: { @@ -904,12 +917,13 @@ TEST(ConnectionImplTest, ExecuteQueryDirectedRead) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, BeginTransaction).Times(0); EXPECT_CALL(*mock, ExecuteStreamingSql) .WillOnce([](std::shared_ptr const&, Options const&, - google::spanner::v1::ExecuteSqlRequest const& request) { + google::spanner::v1::ExecuteSqlRequest const& request, + auto const&) { EXPECT_EQ("multiplexed", request.session()); EXPECT_TRUE(request.has_directed_read_options()); auto const& directed_read_options = request.directed_read_options(); @@ -945,7 +959,7 @@ TEST(ConnectionImplTest, ExecuteQueryPgNumericResult) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); auto constexpr kText = R"pb( metadata: { @@ -991,7 +1005,7 @@ TEST(ConnectionImplTest, ExecuteQueryJsonBResult) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); auto constexpr kText = R"pb( metadata: { @@ -1035,7 +1049,7 @@ TEST(ConnectionImplTest, ExecuteQueryNumericParameter) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); auto constexpr kResponseNumeric = R"pb( metadata: { @@ -1061,7 +1075,8 @@ TEST(ConnectionImplTest, ExecuteQueryNumericParameter) { )pb"; EXPECT_CALL(*mock, ExecuteStreamingSql) .WillOnce([&](std::shared_ptr const&, Options const&, - google::spanner::v1::ExecuteSqlRequest const& request) { + google::spanner::v1::ExecuteSqlRequest const& request, + auto const&) { EXPECT_EQ(request.params().fields().at("value").string_value(), "998"); EXPECT_EQ(request.param_types().at("value").code(), google::spanner::v1::TypeCode::NUMERIC); @@ -1071,7 +1086,8 @@ TEST(ConnectionImplTest, ExecuteQueryNumericParameter) { return MakeReader({kResponseNumeric}); }) .WillOnce([&](std::shared_ptr const&, Options const&, - google::spanner::v1::ExecuteSqlRequest const& request) { + google::spanner::v1::ExecuteSqlRequest const& request, + auto const&) { EXPECT_EQ(request.params().fields().at("value").string_value(), "999"); EXPECT_EQ(request.param_types().at("value").code(), google::spanner::v1::TypeCode::NUMERIC); @@ -1080,7 +1096,8 @@ TEST(ConnectionImplTest, ExecuteQueryNumericParameter) { return MakeReader({kResponsePgNumeric}); }) .WillOnce([&](std::shared_ptr const&, Options const&, - google::spanner::v1::ExecuteSqlRequest const& request) { + google::spanner::v1::ExecuteSqlRequest const& request, + auto const&) { EXPECT_EQ(request.params().fields().at("value").string_value(), "NaN"); EXPECT_EQ(request.param_types().at("value").code(), google::spanner::v1::TypeCode::NUMERIC); @@ -1118,7 +1135,7 @@ TEST(ConnectionImplTest, ExecuteQueryPgOidResult) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); auto constexpr kText = R"pb( metadata: { @@ -1162,7 +1179,7 @@ TEST(ConnectionImplTest, ExecuteQueryImplicitBeginTransaction) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, BeginTransaction).Times(0); auto constexpr kText = R"pb( @@ -1381,55 +1398,57 @@ TEST(ConnectionImplTest, QueryOptions) { .Build())); // ExecuteQuery(). - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))) .RetiresOnSaturation(); - EXPECT_CALL(*mock, ExecuteStreamingSql(_, _, execute_sql_request_matcher)) + EXPECT_CALL(*mock, + ExecuteStreamingSql(_, _, execute_sql_request_matcher, _)) .WillOnce(Return(ByMove(std::move(stream)))) .RetiresOnSaturation(); // ProfileQuery(). - EXPECT_CALL(*mock, ExecuteStreamingSql(_, _, execute_sql_request_matcher)) + EXPECT_CALL(*mock, + ExecuteStreamingSql(_, _, execute_sql_request_matcher, _)) .WillOnce(Return( ByMove(std::make_unique< NiceMock>>()))) .RetiresOnSaturation(); // ExecutePartitionedDml(). - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))) .RetiresOnSaturation(); EXPECT_CALL(*mock, - BeginTransaction(_, _, begin_transaction_request_matcher)) + BeginTransaction(_, _, begin_transaction_request_matcher, _)) .WillOnce(Return(MakeTestTransaction("2468ACE"))) .RetiresOnSaturation(); EXPECT_CALL(*mock, ExecuteStreamingSql( - _, _, untagged_execute_sql_request_matcher)) + _, _, untagged_execute_sql_request_matcher, _)) .WillOnce(Return( ByMove(std::make_unique< NiceMock>>()))) .RetiresOnSaturation(); // ExecuteDml(), ProfileDml(), AnalyzeSql(). - EXPECT_CALL(*mock, ExecuteSql(_, _, execute_sql_request_matcher)) + EXPECT_CALL(*mock, ExecuteSql(_, _, execute_sql_request_matcher, _)) .Times(3) .RetiresOnSaturation(); // ExecuteBatchDml(). EXPECT_CALL(*mock, - ExecuteBatchDml(_, _, execute_batch_dml_request_matcher)) + ExecuteBatchDml(_, _, execute_batch_dml_request_matcher, _)) .Times(1) .RetiresOnSaturation(); // Read(). - EXPECT_CALL(*mock, StreamingRead(_, _, read_request_matcher)) + EXPECT_CALL(*mock, StreamingRead(_, _, read_request_matcher, _)) .WillOnce(Return( ByMove(std::make_unique< NiceMock>>()))) .RetiresOnSaturation(); // Commit(). - EXPECT_CALL(*mock, Commit(_, _, commit_request_matcher)) + EXPECT_CALL(*mock, Commit(_, _, commit_request_matcher, _)) .Times(1) .RetiresOnSaturation(); } @@ -1452,7 +1471,7 @@ TEST(ConnectionImplTest, ExecuteDmlCreateSessionFailure) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, HasDatabase(db))) + EXPECT_CALL(*mock, CreateSession(_, _, HasDatabase(db), _)) .WillRepeatedly(Return( Status(StatusCode::kPermissionDenied, "uh-oh in CreateSession"))); @@ -1471,7 +1490,7 @@ TEST(ConnectionImplTest, ExecuteDmlDeleteSuccess) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); auto constexpr kText = R"pb( metadata: { transaction: { id: "1234567890" } } @@ -1500,7 +1519,7 @@ TEST(ConnectionImplTest, ExecuteDmlDeletePermanentFailure) { "placeholder_database_id"); { InSequence seq; - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); Status status(StatusCode::kPermissionDenied, "uh-oh in ExecuteDml"); EXPECT_CALL(*mock, ExecuteSql).WillOnce(Return(status)); @@ -1525,7 +1544,7 @@ TEST(ConnectionImplTest, ExecuteDmlDeleteTooManyTransientFailures) { "placeholder_database_id"); { InSequence seq; - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); Status status(StatusCode::kUnavailable, "try-again in ExecuteDml"); EXPECT_CALL(*mock, ExecuteSql) @@ -1558,7 +1577,7 @@ TEST(ConnectionImplTest, ExecuteDmlTransactionAtomicity) { Status begin_status(StatusCode::kInvalidArgument, "BeginTransaction status"); { InSequence seq; - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); // The first `ExecuteDml` call tries to implicitly begin the transaction @@ -1589,7 +1608,7 @@ TEST(ConnectionImplTest, ExecuteDmlTransactionMissing) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); // Return an otherwise valid response that does not contain a transaction. @@ -1612,7 +1631,7 @@ TEST(ConnectionImplTest, ProfileQuerySuccess) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); auto constexpr kText = R"pb( metadata: { @@ -1676,7 +1695,7 @@ TEST(ConnectionImplTest, ProfileQueryCreateSessionFailure) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, HasDatabase(db))) + EXPECT_CALL(*mock, CreateSession(_, _, HasDatabase(db), _)) .WillRepeatedly(Return( Status(StatusCode::kPermissionDenied, "uh-oh in CreateSession"))); @@ -1695,7 +1714,7 @@ TEST(ConnectionImplTest, ProfileQueryStreamingReadFailure) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); auto finish_status = internal::PermissionDeniedError("uh-oh in GrpcReader::Finish"); @@ -1718,7 +1737,7 @@ TEST(ConnectionImplTest, ProfileDmlCreateSessionFailure) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, HasDatabase(db))) + EXPECT_CALL(*mock, CreateSession(_, _, HasDatabase(db), _)) .WillRepeatedly(Return( Status(StatusCode::kPermissionDenied, "uh-oh in CreateSession"))); EXPECT_CALL(*mock, AsyncDeleteSession).Times(0); @@ -1737,7 +1756,7 @@ TEST(ConnectionImplTest, ProfileDmlDeleteSuccess) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); auto constexpr kText = R"pb( metadata: { transaction: { id: "1234567890" } } @@ -1790,7 +1809,7 @@ TEST(ConnectionImplTest, ProfileDmlDeletePermanentFailure) { "placeholder_database_id"); { InSequence seq; - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); Status status(StatusCode::kPermissionDenied, "uh-oh in ExecuteDml"); EXPECT_CALL(*mock, ExecuteSql).WillOnce(Return(status)); @@ -1815,7 +1834,7 @@ TEST(ConnectionImplTest, ProfileDmlDeleteTooManyTransientFailures) { "placeholder_database_id"); { InSequence seq; - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); Status status(StatusCode::kUnavailable, "try-again in ExecuteDml"); EXPECT_CALL(*mock, ExecuteSql) @@ -1843,7 +1862,7 @@ TEST(ConnectionImplTest, AnalyzeSqlSuccess) { auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); auto constexpr kText = R"pb( metadata: {} @@ -1874,7 +1893,7 @@ TEST(ConnectionImplTest, AnalyzeSqlCreateSessionFailure) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, HasDatabase(db))) + EXPECT_CALL(*mock, CreateSession(_, _, HasDatabase(db), _)) .WillRepeatedly(Return( Status(StatusCode::kPermissionDenied, "uh-oh in CreateSession"))); @@ -1894,7 +1913,7 @@ TEST(ConnectionImplTest, AnalyzeSqlDeletePermanentFailure) { "placeholder_database_id"); { InSequence seq; - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); Status status(StatusCode::kPermissionDenied, "uh-oh in ExecuteDml"); EXPECT_CALL(*mock, ExecuteSql).WillOnce(Return(status)); @@ -1919,7 +1938,7 @@ TEST(ConnectionImplTest, AnalyzeSqlDeleteTooManyTransientFailures) { "placeholder_database_id"); { InSequence seq; - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); Status status(StatusCode::kUnavailable, "try-again in ExecuteDml"); EXPECT_CALL(*mock, ExecuteSql) @@ -1946,7 +1965,7 @@ TEST(ConnectionImplTest, ExecuteBatchDmlSuccess) { auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); auto mock = std::make_shared(); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); auto constexpr kText = R"pb( result_sets: { @@ -1962,7 +1981,7 @@ TEST(ConnectionImplTest, ExecuteBatchDmlSuccess) { *mock, ExecuteBatchDml( _, _, - HasPriority(google::spanner::v1::RequestOptions::PRIORITY_MEDIUM))) + HasPriority(google::spanner::v1::RequestOptions::PRIORITY_MEDIUM), _)) .WillOnce(Return(Status(StatusCode::kUnavailable, "try-again"))) .WillOnce(Return(response)); @@ -1995,7 +2014,7 @@ TEST(ConnectionImplTest, MultiplexedExecuteBatchDmlSuccess) { auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); auto mock = std::make_shared(); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); auto constexpr kText = R"pb( @@ -2013,20 +2032,21 @@ TEST(ConnectionImplTest, MultiplexedExecuteBatchDmlSuccess) { *mock, ExecuteBatchDml( _, _, - HasPriority(google::spanner::v1::RequestOptions::PRIORITY_MEDIUM))) + HasPriority(google::spanner::v1::RequestOptions::PRIORITY_MEDIUM), _)) .WillOnce(Return(Status(StatusCode::kUnavailable, "try-again"))) .WillOnce(Return(response)); EXPECT_CALL(*mock, Commit) - .WillOnce([&](grpc::ClientContext&, Options const&, - google::spanner::v1::CommitRequest const& request) - -> StatusOr { - google::spanner::v1::CommitResponse response; - EXPECT_THAT(request.precommit_token().precommit_token(), - Eq("test-precommit-token-1")); - EXPECT_THAT(request.precommit_token().seq_num(), Eq(1)); - return response; - }); + .WillOnce( + [&](grpc::ClientContext&, Options const&, + google::spanner::v1::CommitRequest const& request, + auto const&) -> StatusOr { + google::spanner::v1::CommitResponse response; + EXPECT_THAT(request.precommit_token().precommit_token(), + Eq("test-precommit-token-1")); + EXPECT_THAT(request.precommit_token().seq_num(), Eq(1)); + return response; + }); auto request = { spanner::SqlStatement("UPDATE ..."), @@ -2059,7 +2079,7 @@ TEST(ConnectionImplTest, ExecuteBatchDmlPartialFailure) { auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); auto mock = std::make_shared(); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); auto constexpr kText = R"pb( result_sets: { @@ -2100,7 +2120,7 @@ TEST(ConnectionImplTest, ExecuteBatchDmlPermanentFailure) { auto mock = std::make_shared(); { InSequence seq; - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); Status status(StatusCode::kPermissionDenied, "uh-oh in ExecuteBatchDml"); @@ -2130,7 +2150,7 @@ TEST(ConnectionImplTest, ExecuteBatchDmlTooManyTransientFailures) { auto mock = std::make_shared(); { InSequence seq; - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); Status status(StatusCode::kUnavailable, "try-again in ExecuteBatchDml"); @@ -2164,7 +2184,7 @@ TEST(ConnectionImplTest, ExecuteBatchDmlNoResultSets) { auto mock = std::make_shared(); { InSequence seq; - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); // The `ExecuteBatchDml` call can succeed, but with no `ResultSet`s and an // error status in the response. @@ -2173,15 +2193,17 @@ TEST(ConnectionImplTest, ExecuteBatchDmlNoResultSets) { )pb"; google::spanner::v1::ExecuteBatchDmlResponse response; ASSERT_TRUE(TextFormat::ParseFromString(kText, &response)); - EXPECT_CALL(*mock, ExecuteBatchDml(_, _, - AllOf(HasSession("multiplexed"), - HasBeginTransaction()))) + EXPECT_CALL( + *mock, + ExecuteBatchDml( + _, _, AllOf(HasSession("multiplexed"), HasBeginTransaction()), _)) .WillOnce(Return(response)); EXPECT_CALL(*mock, BeginTransaction) .WillOnce(Return(MakeTestTransaction("BD000001"))); EXPECT_CALL(*mock, ExecuteBatchDml(_, _, AllOf(HasSession("multiplexed"), - HasTransactionId("BD000001")))) + HasTransactionId("BD000001")), + _)) .WillOnce(Return(response)); } @@ -2200,24 +2222,24 @@ TEST(ConnectionImplTest, ExecutePartitionedDmlDeleteSuccess) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, BeginTransaction) .WillOnce(Return(Status(StatusCode::kUnavailable, "try-again"))) - .WillOnce( - [](grpc::ClientContext&, Options const&, - google::spanner::v1::BeginTransactionRequest const& request) { - EXPECT_TRUE(request.options().has_partitioned_dml()); - EXPECT_FALSE(request.options().exclude_txn_from_change_streams()); - return MakeTestTransaction(); - }); + .WillOnce([](grpc::ClientContext&, Options const&, + google::spanner::v1::BeginTransactionRequest const& request, + auto const&) { + EXPECT_TRUE(request.options().has_partitioned_dml()); + EXPECT_FALSE(request.options().exclude_txn_from_change_streams()); + return MakeTestTransaction(); + }); auto constexpr kTextResponse = R"pb( metadata: {} stats: { row_count_lower_bound: 42 } )pb"; EXPECT_CALL(*mock, ExecuteStreamingSql( - _, _, AllOf(HasRequestTag("tag"), HasTransactionTag("")))) + _, _, AllOf(HasRequestTag("tag"), HasTransactionTag("")), _)) .WillOnce(Return(ByMove(MakeReader( {}, internal::UnavailableError("try-again in ExecutePartitionedDml"))))) @@ -2236,23 +2258,23 @@ TEST(ConnectionImplTest, ExecutePartitionedDmlExcludeFromChangeStreams) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, BeginTransaction) - .WillOnce( - [](grpc::ClientContext&, Options const&, - google::spanner::v1::BeginTransactionRequest const& request) { - EXPECT_TRUE(request.options().has_partitioned_dml()); - EXPECT_TRUE(request.options().exclude_txn_from_change_streams()); - return MakeTestTransaction(); - }); + .WillOnce([](grpc::ClientContext&, Options const&, + google::spanner::v1::BeginTransactionRequest const& request, + auto const&) { + EXPECT_TRUE(request.options().has_partitioned_dml()); + EXPECT_TRUE(request.options().exclude_txn_from_change_streams()); + return MakeTestTransaction(); + }); auto constexpr kTextResponse = R"pb( metadata: {} stats: { row_count_lower_bound: 42 } )pb"; EXPECT_CALL(*mock, ExecuteStreamingSql( - _, _, AllOf(HasRequestTag("tag"), HasTransactionTag("")))) + _, _, AllOf(HasRequestTag("tag"), HasTransactionTag("")), _)) .WillOnce(Return(ByMove(MakeReader({kTextResponse})))); auto conn = MakeConnectionImpl(db, mock); @@ -2270,7 +2292,7 @@ TEST(ConnectionImplTest, ExecutePartitionedDmlCreateSessionFailure) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, HasDatabase(db))) + EXPECT_CALL(*mock, CreateSession(_, _, HasDatabase(db), _)) .WillRepeatedly(Return( Status(StatusCode::kPermissionDenied, "uh-oh in CreateSession"))); @@ -2286,7 +2308,7 @@ TEST(ConnectionImplTest, ExecutePartitionedDmlDeletePermanentFailure) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, BeginTransaction) .WillOnce(Return(Status(StatusCode::kUnavailable, "try-again"))) @@ -2310,7 +2332,7 @@ TEST(ConnectionImplTest, ExecutePartitionedDmlDeleteTooManyTransientFailures) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, BeginTransaction) .WillOnce(Return(Status(StatusCode::kUnavailable, "try-again"))) @@ -2337,7 +2359,7 @@ TEST(ConnectionImplTest, ExecutePartitionedDmlRetryableInternalErrors) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, BeginTransaction) .WillOnce(Return(MakeTestTransaction("2345678901"))); @@ -2371,7 +2393,7 @@ TEST(ConnectionImplTest, auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, BeginTransaction) .WillOnce(Return(Status(StatusCode::kPermissionDenied, @@ -2390,7 +2412,7 @@ TEST(ConnectionImplTest, auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, BeginTransaction) .Times(AtLeast(2)) @@ -2410,7 +2432,7 @@ TEST(ConnectionImplTest, CommitCreateSessionPermanentFailure) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, HasDatabase(db))) + EXPECT_CALL(*mock, CreateSession(_, _, HasDatabase(db), _)) .WillRepeatedly(Return( Status(StatusCode::kPermissionDenied, "uh-oh in CreateSession"))); EXPECT_CALL(*mock, AsyncDeleteSession).Times(0); @@ -2426,7 +2448,7 @@ TEST(ConnectionImplTest, CommitCreateSessionTooManyTransientFailures) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .Times(AtLeast(2)) .WillRepeatedly(Return( Status(StatusCode::kUnavailable, "try-again in CreateSession"))); @@ -2442,7 +2464,7 @@ TEST(ConnectionImplTest, CommitCreateSessionRetry) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return( Status(StatusCode::kUnavailable, "try-again in CreateSession"))) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); @@ -2451,7 +2473,8 @@ TEST(ConnectionImplTest, CommitCreateSessionRetry) { EXPECT_CALL(*mock, BeginTransaction).WillOnce(Return(txn)); EXPECT_CALL(*mock, Commit(_, _, AllOf(HasSession("multiplexed"), - HasNakedTransactionId(txn.id())))) + HasNakedTransactionId(txn.id())), + _)) .WillOnce( Return(Status(StatusCode::kPermissionDenied, "uh-oh in Commit"))); @@ -2466,7 +2489,7 @@ TEST(ConnectionImplTest, CommitBeginTransactionRetry) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); google::spanner::v1::Transaction txn = MakeTestTransaction(); EXPECT_CALL(*mock, BeginTransaction) @@ -2477,7 +2500,8 @@ TEST(ConnectionImplTest, CommitBeginTransactionRetry) { .value(); EXPECT_CALL(*mock, Commit(_, _, AllOf(HasSession("multiplexed"), - HasNakedTransactionId(txn.id())))) + HasNakedTransactionId(txn.id())), + _)) .WillOnce(Return(MakeCommitResponse(commit_timestamp))); auto conn = MakeConnectionImpl(db, mock); @@ -2491,7 +2515,7 @@ TEST(ConnectionImplTest, CommitBeginTransactionSessionNotFound) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, BeginTransaction) .WillOnce(Return(SessionNotFoundError("multiplexed"))); @@ -2510,7 +2534,7 @@ TEST(ConnectionImplTest, CommitBeginTransactionPermanentFailure) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, BeginTransaction) .WillOnce(Return( @@ -2534,13 +2558,14 @@ TEST(ConnectionImplTest, CommitCommitPermanentFailure) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); google::spanner::v1::Transaction txn = MakeTestTransaction(); EXPECT_CALL(*mock, BeginTransaction).WillOnce(Return(txn)); EXPECT_CALL(*mock, Commit(_, _, AllOf(HasSession("multiplexed"), - HasNakedTransactionId(txn.id())))) + HasNakedTransactionId(txn.id())), + _)) .WillOnce( Return(Status(StatusCode::kPermissionDenied, "uh-oh in Commit"))); @@ -2555,13 +2580,14 @@ TEST(ConnectionImplTest, CommitCommitTooManyTransientFailures) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); google::spanner::v1::Transaction txn = MakeTestTransaction(); EXPECT_CALL(*mock, BeginTransaction).WillOnce(Return(txn)); EXPECT_CALL(*mock, Commit(_, _, AllOf(HasSession("multiplexed"), - HasNakedTransactionId(txn.id())))) + HasNakedTransactionId(txn.id())), + _)) .WillOnce( Return(Status(StatusCode::kPermissionDenied, "uh-oh in Commit"))); @@ -2576,7 +2602,7 @@ TEST(ConnectionImplTest, CommitCommitInvalidatedTransaction) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, BeginTransaction).Times(0); EXPECT_CALL(*mock, Commit).Times(0); @@ -2598,14 +2624,15 @@ TEST(ConnectionImplTest, CommitCommitIdempotentTransientSuccess) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); auto const commit_timestamp = spanner::MakeTimestamp(std::chrono::system_clock::from_time_t(123)) .value(); EXPECT_CALL(*mock, Commit(_, _, AllOf(HasSession("multiplexed"), - HasNakedTransactionId("test-txn-id")))) + HasNakedTransactionId("test-txn-id")), + _)) .WillOnce(Return(Status(StatusCode::kUnavailable, "try-again"))) .WillOnce(Return(MakeCommitResponse(commit_timestamp))); @@ -2625,7 +2652,7 @@ TEST(ConnectionImplTest, CommitSuccessWithTransactionId) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL( *mock, @@ -2633,7 +2660,8 @@ TEST(ConnectionImplTest, CommitSuccessWithTransactionId) { _, _, AllOf( HasSession("multiplexed"), HasNakedTransactionId("test-txn-id"), - HasPriority(google::spanner::v1::RequestOptions::PRIORITY_HIGH)))) + HasPriority(google::spanner::v1::RequestOptions::PRIORITY_HIGH)), + _)) .WillOnce(Return(MakeCommitResponse( spanner::MakeTimestamp(std::chrono::system_clock::from_time_t(123)) .value()))); @@ -2656,20 +2684,20 @@ TEST(ConnectionImplTest, CommitSuccessWithStats) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, BeginTransaction) - .WillOnce( - [](grpc::ClientContext&, Options const&, - google::spanner::v1::BeginTransactionRequest const& request) { - EXPECT_TRUE(request.options().has_read_write()); - EXPECT_FALSE(request.options().exclude_txn_from_change_streams()); - EXPECT_FALSE(request.has_mutation_key()); - return MakeTestTransaction(); - }); + .WillOnce([](grpc::ClientContext&, Options const&, + google::spanner::v1::BeginTransactionRequest const& request, + auto const&) { + EXPECT_TRUE(request.options().has_read_write()); + EXPECT_FALSE(request.options().exclude_txn_from_change_streams()); + EXPECT_FALSE(request.has_mutation_key()); + return MakeTestTransaction(); + }); EXPECT_CALL( *mock, - Commit(_, _, AllOf(HasSession("multiplexed"), HasReturnStats(true)))) + Commit(_, _, AllOf(HasSession("multiplexed"), HasReturnStats(true)), _)) .WillOnce(Return(MakeCommitResponse( spanner::MakeTimestamp(std::chrono::system_clock::from_time_t(123)) .value(), @@ -2702,18 +2730,18 @@ TEST(ConnectionImplTest, MutationCommitSuccess) { token.set_precommit_token("test-precommit-token"); token.set_seq_num(1); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, BeginTransaction) - .WillOnce( - [&](grpc::ClientContext&, Options const&, - google::spanner::v1::BeginTransactionRequest const& request) { - EXPECT_TRUE(request.options().has_read_write()); - EXPECT_FALSE(request.options().exclude_txn_from_change_streams()); - EXPECT_TRUE(request.has_mutation_key()); - return MakeTestTransaction(token); - }); + .WillOnce([&](grpc::ClientContext&, Options const&, + google::spanner::v1::BeginTransactionRequest const& request, + auto const&) { + EXPECT_TRUE(request.options().has_read_write()); + EXPECT_FALSE(request.options().exclude_txn_from_change_streams()); + EXPECT_TRUE(request.has_mutation_key()); + return MakeTestTransaction(token); + }); auto const commit_timestamp = spanner::MakeTimestamp(std::chrono::system_clock::from_time_t(123)) @@ -2721,7 +2749,8 @@ TEST(ConnectionImplTest, MutationCommitSuccess) { EXPECT_CALL(*mock, Commit) .WillOnce([commit_timestamp]( grpc::ClientContext&, Options const&, - google::spanner::v1::CommitRequest const& request) { + google::spanner::v1::CommitRequest const& request, + auto const&) { EXPECT_EQ("multiplexed", request.session()); EXPECT_FALSE(request.has_single_use_transaction()); EXPECT_EQ(3, request.mutations_size()); @@ -2757,18 +2786,18 @@ TEST(ConnectionImplTest, MutationCommitRetryOnceSuccess) { token.set_precommit_token("test-precommit-token"); token.set_seq_num(1); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, BeginTransaction) - .WillOnce( - [&](grpc::ClientContext&, Options const&, - google::spanner::v1::BeginTransactionRequest const& request) { - EXPECT_TRUE(request.options().has_read_write()); - EXPECT_FALSE(request.options().exclude_txn_from_change_streams()); - EXPECT_TRUE(request.has_mutation_key()); - return MakeTestTransaction(token); - }); + .WillOnce([&](grpc::ClientContext&, Options const&, + google::spanner::v1::BeginTransactionRequest const& request, + auto const&) { + EXPECT_TRUE(request.options().has_read_write()); + EXPECT_FALSE(request.options().exclude_txn_from_change_streams()); + EXPECT_TRUE(request.has_mutation_key()); + return MakeTestTransaction(token); + }); auto const commit_timestamp = spanner::MakeTimestamp(std::chrono::system_clock::from_time_t(123)) @@ -2780,7 +2809,8 @@ TEST(ConnectionImplTest, MutationCommitRetryOnceSuccess) { std::int64_t original_mutations_size = mutations.size(); EXPECT_CALL(*mock, Commit) .WillOnce([&](grpc::ClientContext&, Options const&, - google::spanner::v1::CommitRequest const& request) { + google::spanner::v1::CommitRequest const& request, + auto const&) { EXPECT_EQ("multiplexed", request.session()); EXPECT_FALSE(request.has_single_use_transaction()); EXPECT_EQ(3, request.mutations_size()); @@ -2791,7 +2821,8 @@ TEST(ConnectionImplTest, MutationCommitRetryOnceSuccess) { retry_token); }) .WillOnce([&](grpc::ClientContext&, Options const&, - google::spanner::v1::CommitRequest const& request) { + google::spanner::v1::CommitRequest const& request, + auto const&) { EXPECT_EQ("multiplexed", request.session()); EXPECT_FALSE(request.has_single_use_transaction()); EXPECT_EQ(0, request.mutations_size()); @@ -2827,18 +2858,18 @@ TEST(ConnectionImplTest, MutationCommitRetryMoreThanOnceSuccess) { token.set_precommit_token("test-precommit-token"); token.set_seq_num(1); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, BeginTransaction) - .WillOnce( - [&](grpc::ClientContext&, Options const&, - google::spanner::v1::BeginTransactionRequest const& request) { - EXPECT_TRUE(request.options().has_read_write()); - EXPECT_FALSE(request.options().exclude_txn_from_change_streams()); - EXPECT_TRUE(request.has_mutation_key()); - return MakeTestTransaction(token); - }); + .WillOnce([&](grpc::ClientContext&, Options const&, + google::spanner::v1::BeginTransactionRequest const& request, + auto const&) { + EXPECT_TRUE(request.options().has_read_write()); + EXPECT_FALSE(request.options().exclude_txn_from_change_streams()); + EXPECT_TRUE(request.has_mutation_key()); + return MakeTestTransaction(token); + }); auto const commit_timestamp = spanner::MakeTimestamp(std::chrono::system_clock::from_time_t(123)) @@ -2853,7 +2884,8 @@ TEST(ConnectionImplTest, MutationCommitRetryMoreThanOnceSuccess) { std::int64_t original_mutations_size = mutations.size(); EXPECT_CALL(*mock, Commit) .WillOnce([&](grpc::ClientContext&, Options const&, - google::spanner::v1::CommitRequest const& request) { + google::spanner::v1::CommitRequest const& request, + auto const&) { EXPECT_EQ("multiplexed", request.session()); EXPECT_FALSE(request.has_single_use_transaction()); EXPECT_EQ(3, request.mutations_size()); @@ -2864,7 +2896,8 @@ TEST(ConnectionImplTest, MutationCommitRetryMoreThanOnceSuccess) { retry_token_1); }) .WillOnce([&](grpc::ClientContext&, Options const&, - google::spanner::v1::CommitRequest const& request) { + google::spanner::v1::CommitRequest const& request, + auto const&) { EXPECT_EQ("multiplexed", request.session()); EXPECT_FALSE(request.has_single_use_transaction()); EXPECT_EQ(0, request.mutations_size()); @@ -2875,7 +2908,8 @@ TEST(ConnectionImplTest, MutationCommitRetryMoreThanOnceSuccess) { retry_token_2); }) .WillOnce([&](grpc::ClientContext&, Options const&, - google::spanner::v1::CommitRequest const& request) { + google::spanner::v1::CommitRequest const& request, + auto const&) { EXPECT_EQ("multiplexed", request.session()); EXPECT_FALSE(request.has_single_use_transaction()); EXPECT_EQ(0, request.mutations_size()); @@ -2901,7 +2935,7 @@ TEST(ConnectionImplTest, MultiplexedPrecommitUpdated) { { InSequence seq; - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); auto constexpr kText = R"pb( @@ -2956,15 +2990,16 @@ TEST(ConnectionImplTest, MultiplexedPrecommitUpdated) { EXPECT_CALL(*mock, ExecuteSql).WillOnce(Return(sql_response2)); EXPECT_CALL(*mock, Commit) - .WillOnce([&](grpc::ClientContext&, Options const&, - google::spanner::v1::CommitRequest const& request) - -> StatusOr { - google::spanner::v1::CommitResponse response; - EXPECT_THAT(request.precommit_token().precommit_token(), - Eq("test-precommit-token-4")); - EXPECT_THAT(request.precommit_token().seq_num(), Eq(4)); - return response; - }); + .WillOnce( + [&](grpc::ClientContext&, Options const&, + google::spanner::v1::CommitRequest const& request, + auto const&) -> StatusOr { + google::spanner::v1::CommitResponse response; + EXPECT_THAT(request.precommit_token().precommit_token(), + Eq("test-precommit-token-4")); + EXPECT_THAT(request.precommit_token().seq_num(), Eq(4)); + return response; + }); } auto conn = MakeConnectionImpl(db, mock, Options{}); @@ -2985,17 +3020,17 @@ TEST(ConnectionImplTest, CommitSuccessExcludeFromChangeStreams) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, BeginTransaction) - .WillOnce( - [](grpc::ClientContext&, Options const&, - google::spanner::v1::BeginTransactionRequest const& request) { - EXPECT_TRUE(request.options().has_read_write()); - EXPECT_TRUE(request.options().exclude_txn_from_change_streams()); - return MakeTestTransaction(); - }); - EXPECT_CALL(*mock, Commit(_, _, HasSession("multiplexed"))) + .WillOnce([](grpc::ClientContext&, Options const&, + google::spanner::v1::BeginTransactionRequest const& request, + auto const&) { + EXPECT_TRUE(request.options().has_read_write()); + EXPECT_TRUE(request.options().exclude_txn_from_change_streams()); + return MakeTestTransaction(); + }); + EXPECT_CALL(*mock, Commit(_, _, HasSession("multiplexed"), _)) .WillOnce(Return(MakeCommitResponse( spanner::MakeTimestamp(std::chrono::system_clock::from_time_t(123)) .value()))); @@ -3016,14 +3051,15 @@ TEST(ConnectionImplTest, CommitSuccessWithMaxCommitDelay) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); google::spanner::v1::Transaction txn = MakeTestTransaction(); EXPECT_CALL(*mock, BeginTransaction).WillOnce(Return(txn)); EXPECT_CALL(*mock, Commit(_, _, AllOf(HasSession("multiplexed"), - HasMaxCommitDelay(std::chrono::milliseconds(100))))) + HasMaxCommitDelay(std::chrono::milliseconds(100))), + _)) .WillOnce(Return(MakeCommitResponse( spanner::MakeTimestamp(std::chrono::system_clock::from_time_t(123)) .value()))); @@ -3041,14 +3077,14 @@ TEST(ConnectionImplTest, CommitSuccessWithCompression) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); google::spanner::v1::Transaction txn = MakeTestTransaction(); EXPECT_CALL(*mock, BeginTransaction).WillOnce(Return(txn)); EXPECT_CALL(*mock, Commit(HasCompressionAlgorithm(GRPC_COMPRESS_GZIP), _, - HasSession("multiplexed"))) + HasSession("multiplexed"), _)) .WillOnce([](grpc::ClientContext&, Options const&, - google::spanner::v1::CommitRequest const&) { + google::spanner::v1::CommitRequest const&, auto const&) { return MakeCommitResponse( spanner::MakeTimestamp(std::chrono::system_clock::from_time_t(123)) .value()); @@ -3070,26 +3106,27 @@ TEST(ConnectionImplTest, CommitAtLeastOnce) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, BeginTransaction).Times(0); // The whole point! auto const commit_timestamp = spanner::MakeTimestamp(std::chrono::system_clock::from_time_t(123)) .value(); EXPECT_CALL(*mock, Commit) - .WillOnce([commit_timestamp]( - grpc::ClientContext&, Options const&, - google::spanner::v1::CommitRequest const& request) { - EXPECT_EQ("multiplexed", request.session()); - EXPECT_TRUE(request.has_single_use_transaction()); - EXPECT_EQ(0, request.mutations_size()); - EXPECT_FALSE(request.return_commit_stats()); - EXPECT_EQ(google::spanner::v1::RequestOptions::PRIORITY_UNSPECIFIED, - request.request_options().priority()); - EXPECT_THAT(request.request_options().request_tag(), IsEmpty()); - EXPECT_THAT(request.request_options().transaction_tag(), IsEmpty()); - return MakeCommitResponse(commit_timestamp); - }); + .WillOnce( + [commit_timestamp](grpc::ClientContext&, Options const&, + google::spanner::v1::CommitRequest const& request, + auto const&) { + EXPECT_EQ("multiplexed", request.session()); + EXPECT_TRUE(request.has_single_use_transaction()); + EXPECT_EQ(0, request.mutations_size()); + EXPECT_FALSE(request.return_commit_stats()); + EXPECT_EQ(google::spanner::v1::RequestOptions::PRIORITY_UNSPECIFIED, + request.request_options().priority()); + EXPECT_THAT(request.request_options().request_tag(), IsEmpty()); + EXPECT_THAT(request.request_options().transaction_tag(), IsEmpty()); + return MakeCommitResponse(commit_timestamp); + }); auto conn = MakeConnectionImpl(db, mock); internal::OptionsSpan span(MakeLimitedTimeOptions()); @@ -3104,12 +3141,12 @@ TEST(ConnectionImplTest, CommitAtLeastOnceBatched) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); using BatchWriteRequest = google::spanner::v1::BatchWriteRequest; EXPECT_CALL(*mock, BatchWrite) .WillOnce([](std::shared_ptr const&, Options const&, - BatchWriteRequest const& request) { + BatchWriteRequest const& request, auto const&) { EXPECT_EQ("multiplexed", request.session()); EXPECT_EQ(google::spanner::v1::RequestOptions::PRIORITY_UNSPECIFIED, request.request_options().priority()); @@ -3124,7 +3161,7 @@ TEST(ConnectionImplTest, CommitAtLeastOnceBatched) { {}, Status(StatusCode::kUnavailable, "try-again in BatchWrite")); }) .WillOnce([&](std::shared_ptr const&, Options const&, - BatchWriteRequest const& request) { + BatchWriteRequest const& request, auto const&) { EXPECT_EQ("multiplexed", request.session()); EXPECT_EQ(google::spanner::v1::RequestOptions::PRIORITY_UNSPECIFIED, request.request_options().priority()); @@ -3163,12 +3200,12 @@ TEST(ConnectionImplTest, CommitAtLeastOnceBatchedExcludeFromChangeStreams) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); using BatchWriteRequest = google::spanner::v1::BatchWriteRequest; EXPECT_CALL(*mock, BatchWrite) .WillOnce([&](std::shared_ptr const&, Options const&, - BatchWriteRequest const& request) { + BatchWriteRequest const& request, auto const&) { EXPECT_EQ("multiplexed", request.session()); EXPECT_EQ(google::spanner::v1::RequestOptions::PRIORITY_UNSPECIFIED, request.request_options().priority()); @@ -3208,7 +3245,7 @@ TEST(ConnectionImplTest, CommitAtLeastOnceBatchedExcludeFromChangeStreams) { TEST(ConnectionImplTest, RollbackCreateSessionFailure) { auto mock = std::make_shared(); auto db = spanner::Database("project", "instance", "database"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillRepeatedly(Return( Status(StatusCode::kPermissionDenied, "uh-oh in CreateSession"))); EXPECT_CALL(*mock, Rollback).Times(0); @@ -3227,14 +3264,15 @@ TEST(ConnectionImplTest, RollbackBeginTransaction) { auto mock = std::make_shared(); auto db = spanner::Database("project", "instance", "database"); std::string const session_name = "multiplexed"; - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); std::string const transaction_id = "RollbackBeginTransaction"; EXPECT_CALL(*mock, BeginTransaction) .WillOnce(Return(MakeTestTransaction(transaction_id))); EXPECT_CALL(*mock, Rollback(_, _, AllOf(HasSession(session_name), - HasNakedTransactionId(transaction_id)))) + HasNakedTransactionId(transaction_id)), + _)) .WillOnce(Return(Status())); auto conn = MakeConnectionImpl(db, mock); @@ -3247,7 +3285,7 @@ TEST(ConnectionImplTest, RollbackBeginTransaction) { TEST(ConnectionImplTest, RollbackSingleUseTransaction) { auto mock = std::make_shared(); auto db = spanner::Database("project", "instance", "database"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, Rollback).Times(0); @@ -3265,12 +3303,13 @@ TEST(ConnectionImplTest, RollbackPermanentFailure) { auto mock = std::make_shared(); auto db = spanner::Database("project", "instance", "database"); std::string const session_name = "multiplexed"; - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); std::string const transaction_id = "test-txn-id"; EXPECT_CALL(*mock, Rollback(_, _, AllOf(HasSession("multiplexed"), - HasNakedTransactionId(transaction_id)))) + HasNakedTransactionId(transaction_id)), + _)) .WillOnce( Return(Status(StatusCode::kPermissionDenied, "uh-oh in Rollback"))); @@ -3287,12 +3326,13 @@ TEST(ConnectionImplTest, RollbackTooManyTransientFailures) { auto mock = std::make_shared(); auto db = spanner::Database("project", "instance", "database"); std::string const session_name = "multiplexed"; - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); std::string const transaction_id = "test-txn-id"; EXPECT_CALL(*mock, Rollback(_, _, AllOf(HasSession(session_name), - HasNakedTransactionId(transaction_id)))) + HasNakedTransactionId(transaction_id)), + _)) .Times(AtLeast(2)) .WillRepeatedly( Return(Status(StatusCode::kUnavailable, "try-again in Rollback"))); @@ -3310,12 +3350,13 @@ TEST(ConnectionImplTest, RollbackSuccess) { auto mock = std::make_shared(); auto db = spanner::Database("project", "instance", "database"); std::string const session_name = "multiplexed"; - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); std::string const transaction_id = "test-txn-id"; EXPECT_CALL(*mock, Rollback(_, _, AllOf(HasSession("multiplexed"), - HasNakedTransactionId(transaction_id)))) + HasNakedTransactionId(transaction_id)), + _)) .WillOnce(Return(Status(StatusCode::kUnavailable, "try-again"))) .WillOnce(Return(Status())); @@ -3331,7 +3372,7 @@ TEST(ConnectionImplTest, RollbackInvalidatedTransaction) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, Rollback).Times(0); @@ -3353,12 +3394,13 @@ TEST(ConnectionImplTest, ReadPartition) { auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, BeginTransaction).Times(0); EXPECT_CALL(*mock, StreamingRead) .WillOnce([](std::shared_ptr const&, Options const&, - google::spanner::v1::ReadRequest const& request) { + google::spanner::v1::ReadRequest const& request, + auto const&) { EXPECT_EQ("multiplexed", request.session()); EXPECT_EQ("Table", request.table()); EXPECT_EQ("DEADBEEF", request.partition_token()); @@ -3383,7 +3425,7 @@ TEST(ConnectionImplTest, PartitionReadSuccess) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); auto constexpr kTextPartitionResponse = R"pb( partitions: { partition_token: "BADDECAF" } @@ -3410,7 +3452,7 @@ TEST(ConnectionImplTest, PartitionReadSuccess) { ASSERT_TRUE( TextFormat::ParseFromString(kTextPartitionRequest, &partition_request)); - EXPECT_CALL(*mock, PartitionRead(_, _, IsProtoEqual(partition_request))) + EXPECT_CALL(*mock, PartitionRead(_, _, IsProtoEqual(partition_request), _)) .WillOnce(Return(Status(StatusCode::kUnavailable, "try-again"))) .WillOnce(Return(partition_response)); @@ -3453,7 +3495,7 @@ TEST(ConnectionImplTest, PartitionReadPermanentFailure) { "placeholder_database_id"); { InSequence seq; - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); Status status(StatusCode::kPermissionDenied, "uh-oh"); EXPECT_CALL(*mock, PartitionRead).WillOnce(Return(status)); @@ -3479,7 +3521,7 @@ TEST(ConnectionImplTest, PartitionReadTooManyTransientFailures) { "placeholder_database_id"); { InSequence seq; - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); Status status(StatusCode::kUnavailable, "try-again"); EXPECT_CALL(*mock, PartitionRead) @@ -3508,12 +3550,13 @@ TEST(ConnectionImplTest, QueryPartition) { auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, BeginTransaction).Times(0); EXPECT_CALL(*mock, ExecuteStreamingSql) .WillOnce([](std::shared_ptr const&, Options const&, - google::spanner::v1::ExecuteSqlRequest const& request) { + google::spanner::v1::ExecuteSqlRequest const& request, + auto const&) { EXPECT_EQ("multiplexed", request.session()); EXPECT_EQ("SELECT * FROM Table", request.sql()); EXPECT_EQ("DEADBEEF", request.partition_token()); @@ -3536,7 +3579,7 @@ TEST(ConnectionImplTest, PartitionQuerySuccess) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); auto constexpr kTextPartitionResponse = R"pb( partitions: { partition_token: "BADDECAF" } @@ -3559,7 +3602,7 @@ TEST(ConnectionImplTest, PartitionQuerySuccess) { google::spanner::v1::PartitionQueryRequest partition_request; ASSERT_TRUE( TextFormat::ParseFromString(kTextPartitionRequest, &partition_request)); - EXPECT_CALL(*mock, PartitionQuery(_, _, IsProtoEqual(partition_request))) + EXPECT_CALL(*mock, PartitionQuery(_, _, IsProtoEqual(partition_request), _)) .WillOnce(Return(Status(StatusCode::kUnavailable, "try-again"))) .WillOnce(Return(partition_response)); @@ -3591,7 +3634,7 @@ TEST(ConnectionImplTest, PartitionQueryPermanentFailure) { Status failed_status = Status(StatusCode::kPermissionDenied, "End of line."); { InSequence seq; - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, PartitionQuery).WillOnce(Return(failed_status)); EXPECT_CALL(*mock, BeginTransaction) @@ -3617,7 +3660,7 @@ TEST(ConnectionImplTest, PartitionQueryTooManyTransientFailures) { Status(StatusCode::kUnavailable, "try-again in PartitionQuery"); { InSequence seq; - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, PartitionQuery) .Times(AtLeast(2)) @@ -3644,15 +3687,16 @@ TEST(ConnectionImplTest, MultipleThreads) { auto db = spanner::Database("project", "instance", "database"); std::string const session_prefix = "test-session-prefix-"; std::string const role = "TestRole"; - EXPECT_CALL(*mock, CreateSession(_, _, AllOf(IsMultiplexed()))) + EXPECT_CALL(*mock, CreateSession(_, _, AllOf(IsMultiplexed()), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, Rollback) - .WillRepeatedly([session_prefix]( - grpc::ClientContext&, Options const&, - google::spanner::v1::RollbackRequest const& request) { - EXPECT_THAT(request.session(), Eq("multiplexed")); - return Status(); - }); + .WillRepeatedly( + [session_prefix](grpc::ClientContext&, Options const&, + google::spanner::v1::RollbackRequest const& request, + auto const&) { + EXPECT_THAT(request.session(), Eq("multiplexed")); + return Status(); + }); int const per_thread_iterations = 1000; auto const thread_count = []() -> unsigned { @@ -3694,7 +3738,7 @@ TEST(ConnectionImplTest, TransactionOutlivesConnection) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, BeginTransaction).Times(0); @@ -3723,7 +3767,7 @@ TEST(ConnectionImplTest, TransactionOutlivesConnection) { TEST(ConnectionImplTest, PartitionReadSessionNotFound) { auto mock = std::make_shared(); auto db = spanner::Database("project", "instance", "database"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, PartitionRead) .WillOnce(Return(SessionNotFoundError("multiplexed"))); @@ -3743,7 +3787,7 @@ TEST(ConnectionImplTest, PartitionReadSessionNotFound) { TEST(ConnectionImplTest, ExecuteQuerySessionNotFound) { auto mock = std::make_shared(); auto db = spanner::Database("project", "instance", "database"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); auto finish_status = SessionNotFoundError("multiplexed"); EXPECT_CALL(*mock, ExecuteStreamingSql) @@ -3764,7 +3808,7 @@ TEST(ConnectionImplTest, ExecuteQuerySessionNotFound) { TEST(ConnectionImplTest, ProfileQuerySessionNotFound) { auto mock = std::make_shared(); auto db = spanner::Database("project", "instance", "database"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); auto finish_status = SessionNotFoundError("multiplexed"); EXPECT_CALL(*mock, ExecuteStreamingSql) @@ -3785,7 +3829,7 @@ TEST(ConnectionImplTest, ProfileQuerySessionNotFound) { TEST(ConnectionImplTest, ExecuteDmlSessionNotFound) { auto mock = std::make_shared(); auto db = spanner::Database("project", "instance", "database"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, ExecuteSql) .WillOnce(Return(SessionNotFoundError("multiplexed"))); @@ -3804,7 +3848,7 @@ TEST(ConnectionImplTest, ExecuteDmlSessionNotFound) { TEST(ConnectionImplTest, ProfileDmlSessionNotFound) { auto mock = std::make_shared(); auto db = spanner::Database("project", "instance", "database"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, ExecuteSql) .WillOnce(Return(SessionNotFoundError("multiplexed"))); @@ -3823,7 +3867,7 @@ TEST(ConnectionImplTest, ProfileDmlSessionNotFound) { TEST(ConnectionImplTest, AnalyzeSqlSessionNotFound) { auto mock = std::make_shared(); auto db = spanner::Database("project", "instance", "database"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, ExecuteSql) .WillOnce(Return(SessionNotFoundError("multiplexed"))); @@ -3842,7 +3886,7 @@ TEST(ConnectionImplTest, AnalyzeSqlSessionNotFound) { TEST(ConnectionImplTest, PartitionQuerySessionNotFound) { auto mock = std::make_shared(); auto db = spanner::Database("project", "instance", "database"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, PartitionQuery) .WillOnce(Return(SessionNotFoundError("multiplexed"))); @@ -3861,7 +3905,7 @@ TEST(ConnectionImplTest, PartitionQuerySessionNotFound) { TEST(ConnectionImplTest, ExecuteBatchDmlSessionNotFound) { auto mock = std::make_shared(); auto db = spanner::Database("project", "instance", "database"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, ExecuteBatchDml) .WillOnce(Return(SessionNotFoundError("multiplexed"))); @@ -3887,7 +3931,7 @@ TEST(ConnectionImplTest, ExecutePartitionedDmlSessionNotFound) { TEST(ConnectionImplTest, CommitSessionNotFound) { auto mock = std::make_shared(); auto db = spanner::Database("project", "instance", "database"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, Commit) .WillOnce(Return(SessionNotFoundError("multiplexed"))); @@ -3906,7 +3950,7 @@ TEST(ConnectionImplTest, CommitSessionNotFound) { TEST(ConnectionImplTest, RollbackSessionNotFound) { auto mock = std::make_shared(); auto db = spanner::Database("project", "instance", "database"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); EXPECT_CALL(*mock, Rollback) .WillOnce(Return(SessionNotFoundError("multiplexed"))); @@ -3923,7 +3967,7 @@ TEST(ConnectionImplTest, RollbackSessionNotFound) { TEST(ConnectionImplTest, ReadRequestOrderByParameterUnspecified) { auto mock = std::make_shared(); auto db = spanner::Database("project", "instance", "database"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); Sequence s; EXPECT_CALL( @@ -3932,7 +3976,8 @@ TEST(ConnectionImplTest, ReadRequestOrderByParameterUnspecified) { _, _, AllOf(HasSession("multiplexed"), HasOrderBy( - google::spanner::v1::ReadRequest::ORDER_BY_UNSPECIFIED)))) + google::spanner::v1::ReadRequest::ORDER_BY_UNSPECIFIED)), + _)) .InSequence(s) .WillOnce(Return(ByMove(MakeReader( {R"pb(metadata: { transaction: { id: "txn1" } })pb"})))); @@ -3955,7 +4000,7 @@ TEST(ConnectionImplTest, ReadRequestOrderByParameterUnspecified) { TEST(ConnectionImplTest, ReadRequestOrderByParameterNoOrder) { auto mock = std::make_shared(); auto db = spanner::Database("project", "instance", "database"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); Sequence s; EXPECT_CALL( @@ -3964,7 +4009,8 @@ TEST(ConnectionImplTest, ReadRequestOrderByParameterNoOrder) { _, _, AllOf( HasSession("multiplexed"), - HasOrderBy(google::spanner::v1::ReadRequest::ORDER_BY_NO_ORDER)))) + HasOrderBy(google::spanner::v1::ReadRequest::ORDER_BY_NO_ORDER)), + _)) .InSequence(s) .WillOnce(Return(ByMove(MakeReader( {R"pb(metadata: { transaction: { id: "txn1" } })pb"})))); @@ -3994,7 +4040,7 @@ TEST(ConnectionImplTest, ReadRequestOrderByParameterNoOrder) { TEST(ConnectionImplTest, ReadRequestLockHintParameterUnspecified) { auto mock = std::make_shared(); auto db = spanner::Database("project", "instance", "database"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); Sequence s; EXPECT_CALL( @@ -4003,7 +4049,8 @@ TEST(ConnectionImplTest, ReadRequestLockHintParameterUnspecified) { _, _, AllOf(HasSession("multiplexed"), HasLockHint( - google::spanner::v1::ReadRequest::LOCK_HINT_UNSPECIFIED)))) + google::spanner::v1::ReadRequest::LOCK_HINT_UNSPECIFIED)), + _)) .InSequence(s) .WillOnce(Return(ByMove(MakeReader( {R"pb(metadata: { transaction: { id: "txn1" } })pb"})))); @@ -4026,16 +4073,16 @@ TEST(ConnectionImplTest, ReadRequestLockHintParameterUnspecified) { TEST(ConnectionImplTest, ReadRequestLockHintShared) { auto mock = std::make_shared(); auto db = spanner::Database("project", "instance", "database"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); Sequence s; EXPECT_CALL( - *mock, - StreamingRead( - _, _, - AllOf( - HasSession("multiplexed"), - HasLockHint(google::spanner::v1::ReadRequest::LOCK_HINT_SHARED)))) + *mock, StreamingRead( + _, _, + AllOf(HasSession("multiplexed"), + HasLockHint( + google::spanner::v1::ReadRequest::LOCK_HINT_SHARED)), + _)) .InSequence(s) .WillOnce(Return(ByMove(MakeReader( {R"pb(metadata: { transaction: { id: "txn1" } })pb"})))); @@ -4067,7 +4114,7 @@ TEST(ConnectionImplTest, OperationsFailOnInvalidatedTransaction) { auto mock = std::make_shared(); auto db = spanner::Database("placeholder_project", "placeholder_instance", "placeholder_database_id"); - EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed())) + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); auto conn = MakeConnectionImpl(db, mock); @@ -4139,6 +4186,63 @@ TEST(ConnectionImplTest, OperationsFailOnInvalidatedTransaction) { HasSubstr("BeginTransaction failed"))); } +TEST(ConnectionImplTest, RequestIdHeaderInjectionAndAttemptIncrement) { + auto mock = std::make_shared(); + auto db = spanner::Database("placeholder_project", "placeholder_instance", + "placeholder_database_id"); + EXPECT_CALL(*mock, CreateSession(_, _, IsMultiplexed(), _)) + .WillRepeatedly(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); + + std::vector request_ids; + std::mutex mu; + testing_util::ValidateMetadataFixture fixture; + + EXPECT_CALL(*mock, ExecuteSql) + .WillOnce([&](grpc::ClientContext& context, Options const&, + google::spanner::v1::ExecuteSqlRequest const&, + auto const&) { + auto metadata = fixture.GetMetadata(context); + auto it = metadata.find("x-goog-spanner-request-id"); + if (it != metadata.end()) { + std::scoped_lock lock(mu); + request_ids.push_back(it->second); + } + return internal::UnavailableError("try-again"); + }) + .WillOnce([&](grpc::ClientContext& context, Options const&, + google::spanner::v1::ExecuteSqlRequest const&, + auto const&) { + auto metadata = fixture.GetMetadata(context); + auto it = metadata.find("x-goog-spanner-request-id"); + if (it != metadata.end()) { + std::scoped_lock lock(mu); + request_ids.push_back(it->second); + } + google::spanner::v1::ResultSet result; + result.mutable_metadata()->mutable_transaction()->set_id("1234567890"); + return result; + }); + + auto conn = MakeConnectionImpl(db, mock); + internal::OptionsSpan span(MakeLimitedRetryOptions()); + auto result = conn->ExecuteDml( + {spanner::MakeReadWriteTransaction(), + spanner::SqlStatement("UPDATE Table SET v = 1 WHERE k = 1")}); + EXPECT_THAT(result, IsOk()); + + ASSERT_THAT(request_ids, testing::SizeIs(2)); + // Request ID format: .... + EXPECT_THAT(request_ids[0], testing::EndsWith(".1")); + EXPECT_THAT(request_ids[1], testing::EndsWith(".2")); + // Both attempts should share the prefix + // ... + std::string const prefix0 = + request_ids[0].substr(0, request_ids[0].rfind('.')); + std::string const prefix1 = + request_ids[1].substr(0, request_ids[1].rfind('.')); + EXPECT_THAT(prefix0, testing::Eq(prefix1)); +} + #if defined(__GNUC__) || defined(__clang__) #pragma GCC diagnostic pop #endif diff --git a/google/cloud/spanner/internal/operation_context.cc b/google/cloud/spanner/internal/operation_context.cc new file mode 100644 index 0000000000000..770945943f4fa --- /dev/null +++ b/google/cloud/spanner/internal/operation_context.cc @@ -0,0 +1,103 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "google/cloud/spanner/internal/operation_context.h" +#include "google/cloud/spanner/internal/spanner_request_id.h" +#include +#include + +namespace google { +namespace cloud { +namespace spanner_internal { +GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN + +OperationContext::OperationContext( + std::shared_ptr static_prefix, + std::uint64_t request_index, std::string_view rpc_name) + : static_prefix_(std::move(static_prefix)), + request_index_(request_index), + rpc_name_(rpc_name) {} + +OperationContext::OperationContext(OperationContext&& other) noexcept { + std::scoped_lock lock(other.mu_); + static_prefix_ = std::move(other.static_prefix_); + request_index_ = other.request_index_; + rpc_name_ = other.rpc_name_; + channel_id_ = other.channel_id_; + attempt_index_ = other.attempt_index_; + current_request_id_ = std::move(other.current_request_id_); +} + +OperationContext& OperationContext::operator=( + OperationContext&& other) noexcept { + if (this != &other) { + std::scoped_lock lock(mu_, other.mu_); + static_prefix_ = std::move(other.static_prefix_); + request_index_ = other.request_index_; + rpc_name_ = other.rpc_name_; + channel_id_ = other.channel_id_; + attempt_index_ = other.attempt_index_; + current_request_id_ = std::move(other.current_request_id_); + } + return *this; +} + +void OperationContext::BindChannel(std::uint32_t channel_id) { + std::scoped_lock lock(mu_); + channel_id_ = channel_id; +} + +void OperationContext::PreCall(grpc::ClientContext& client_context) { + std::scoped_lock lock(mu_); + ++attempt_index_; + if (static_prefix_ != nullptr) { + current_request_id_ = FormatSpannerRequestId( + *static_prefix_, channel_id_, request_index_, attempt_index_); + client_context.AddMetadata("x-goog-spanner-request-id", + current_request_id_); + } +} + +void OperationContext::PostCall(grpc::ClientContext const&, Status const&) { + // Hook for metrics / debugging +} + +void OperationContext::OnDone(Status const&) { + // Hook for metrics / latencies +} + +std::optional OperationContext::RequestId() const { + std::scoped_lock lock(mu_); + if (current_request_id_.empty()) return std::nullopt; + return current_request_id_; +} + +std::uint64_t OperationContext::request_index() const { return request_index_; } + +std::uint32_t OperationContext::attempt_index() const { + std::scoped_lock lock(mu_); + return attempt_index_; +} + +std::uint32_t OperationContext::channel_id() const { + std::scoped_lock lock(mu_); + return channel_id_; +} + +std::string_view OperationContext::rpc_name() const { return rpc_name_; } + +GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END +} // namespace spanner_internal +} // namespace cloud +} // namespace google diff --git a/google/cloud/spanner/internal/operation_context.h b/google/cloud/spanner/internal/operation_context.h new file mode 100644 index 0000000000000..43e307c3d731b --- /dev/null +++ b/google/cloud/spanner/internal/operation_context.h @@ -0,0 +1,82 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_SPANNER_INTERNAL_OPERATION_CONTEXT_H +#define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_SPANNER_INTERNAL_OPERATION_CONTEXT_H + +#include "google/cloud/spanner/version.h" +#include "google/cloud/internal/operation_context.h" +#include +#include +#include +#include +#include +#include + +namespace google { +namespace cloud { +namespace spanner_internal { +GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN + +class OperationContext : public google::cloud::internal::OperationContext { + public: + OperationContext(std::shared_ptr static_prefix, + std::uint64_t request_index, std::string_view rpc_name); + + // Move operations transfer context state across threads. + OperationContext(OperationContext&& other) noexcept; + OperationContext& operator=(OperationContext&& other) noexcept; + + OperationContext(OperationContext const&) = delete; + OperationContext& operator=(OperationContext const&) = delete; + + // Binds the physical gRPC channel index (0..N-1) for subsequent attempts. + void BindChannel(std::uint32_t channel_id); + + // Advances attempt_index_, formats header, and injects + // "x-goog-spanner-request-id" into context. + void PreCall(grpc::ClientContext& client_context) override; + + // Called immediately after an attempt returns (hook for metrics / debugging). + void PostCall(grpc::ClientContext const& client_context, + Status const& status) override; + + // Called when the overall logical operation completes (hook for metrics / + // latencies). + void OnDone(Status const& status) override; + + // Returns the active formatted request ID string for the current attempt. + std::optional RequestId() const; + + std::uint64_t request_index() const; + std::uint32_t attempt_index() const; + std::uint32_t channel_id() const; + std::string_view rpc_name() const; + + private: + std::shared_ptr static_prefix_; + std::uint64_t request_index_; + std::string_view rpc_name_; + mutable std::mutex mu_; + std::uint32_t channel_id_ = 0; + std::uint32_t attempt_index_ = 0; + std::string current_request_id_; +}; + +GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END +} // namespace spanner_internal +} // namespace cloud +} // namespace google + +#endif // GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_SPANNER_INTERNAL_OPERATION_CONTEXT_H diff --git a/google/cloud/spanner/internal/operation_context_test.cc b/google/cloud/spanner/internal/operation_context_test.cc new file mode 100644 index 0000000000000..80cb82a720a8b --- /dev/null +++ b/google/cloud/spanner/internal/operation_context_test.cc @@ -0,0 +1,128 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "google/cloud/spanner/internal/operation_context.h" +#include "google/cloud/testing_util/validate_metadata.h" +#include +#include +#include +#include + +namespace google { +namespace cloud { +namespace spanner_internal { +GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN +namespace { + +using ::testing::Contains; +using ::testing::Eq; +using ::testing::Optional; +using ::testing::Pair; +using ::testing::StrEq; + +class OperationContextTest : public ::testing::Test { + protected: + std::multimap GetMetadata( + grpc::ClientContext& context) { + return validate_metadata_fixture_.GetMetadata(context); + } + + testing_util::ValidateMetadataFixture validate_metadata_fixture_; +}; + +TEST_F(OperationContextTest, PreCallSetsRequestIdHeader) { + auto static_prefix = + std::make_shared("1.0123456789abcdef.1."); + OperationContext op_context(static_prefix, 42, "ExecuteSql"); + + EXPECT_THAT(op_context.request_index(), Eq(42ULL)); + EXPECT_THAT(op_context.attempt_index(), Eq(0U)); + EXPECT_THAT(op_context.channel_id(), Eq(0U)); + EXPECT_THAT(op_context.rpc_name(), StrEq("ExecuteSql")); + EXPECT_THAT(op_context.RequestId(), Eq(std::nullopt)); + + op_context.BindChannel(3); + EXPECT_THAT(op_context.channel_id(), Eq(3U)); + + grpc::ClientContext context1; + op_context.PreCall(context1); + + EXPECT_THAT(op_context.attempt_index(), Eq(1U)); + EXPECT_THAT(op_context.RequestId(), + Optional(Eq("1.0123456789abcdef.1.3.42.1"))); + + auto metadata = GetMetadata(context1); + EXPECT_THAT(metadata, Contains(Pair("x-goog-spanner-request-id", + "1.0123456789abcdef.1.3.42.1"))); + + // Next attempt increments attempt_index to 2 + grpc::ClientContext context2; + op_context.PreCall(context2); + + EXPECT_THAT(op_context.attempt_index(), Eq(2U)); + EXPECT_THAT(op_context.RequestId(), + Optional(Eq("1.0123456789abcdef.1.3.42.2"))); + + metadata = GetMetadata(context2); + EXPECT_THAT(metadata, Contains(Pair("x-goog-spanner-request-id", + "1.0123456789abcdef.1.3.42.2"))); +} + +TEST_F(OperationContextTest, MoveConstructAndAssign) { + auto static_prefix = + std::make_shared("1.0123456789abcdef.1."); + OperationContext op_context(static_prefix, 10, "Commit"); + op_context.BindChannel(2); + + grpc::ClientContext context1; + op_context.PreCall(context1); + EXPECT_THAT(op_context.attempt_index(), Eq(1U)); + + // Move construct + OperationContext moved_context(std::move(op_context)); + EXPECT_THAT(moved_context.request_index(), Eq(10ULL)); + EXPECT_THAT(moved_context.channel_id(), Eq(2U)); + EXPECT_THAT(moved_context.attempt_index(), Eq(1U)); + EXPECT_THAT(moved_context.rpc_name(), StrEq("Commit")); + EXPECT_THAT(moved_context.RequestId(), + Optional(Eq("1.0123456789abcdef.1.2.10.1"))); + + // Move assign + OperationContext target_context(static_prefix, 99, "Rollback"); + target_context = std::move(moved_context); + EXPECT_THAT(target_context.request_index(), Eq(10ULL)); + EXPECT_THAT(target_context.channel_id(), Eq(2U)); + EXPECT_THAT(target_context.attempt_index(), Eq(1U)); + EXPECT_THAT(target_context.rpc_name(), StrEq("Commit")); + EXPECT_THAT(target_context.RequestId(), + Optional(Eq("1.0123456789abcdef.1.2.10.1"))); +} + +TEST_F(OperationContextTest, PostCallAndOnDoneHooks) { + auto static_prefix = + std::make_shared("1.0123456789abcdef.1."); + OperationContext op_context(static_prefix, 1, "ExecuteSql"); + + grpc::ClientContext context; + op_context.PreCall(context); + op_context.PostCall(context, Status{}); + op_context.OnDone(Status{}); + EXPECT_THAT(op_context.attempt_index(), Eq(1U)); +} + +} // namespace +GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END +} // namespace spanner_internal +} // namespace cloud +} // namespace google diff --git a/google/cloud/spanner/internal/partial_result_set_resume_test.cc b/google/cloud/spanner/internal/partial_result_set_resume_test.cc index b1ed22bdd9786..d39c0f8cf7d30 100644 --- a/google/cloud/spanner/internal/partial_result_set_resume_test.cc +++ b/google/cloud/spanner/internal/partial_result_set_resume_test.cc @@ -65,7 +65,7 @@ CreatePartialResultSetSource(std::unique_ptr reader, Options opts = {}) { internal::OptionsSpan span( internal::MergeOptions(std::move(opts), internal::CurrentOptions())); - return PartialResultSetSource::Create(std::move(reader)); + return PartialResultSetSource::Create(std::move(reader), nullptr); } MATCHER_P(IsValidAndEquals, expected, diff --git a/google/cloud/spanner/internal/partial_result_set_source.cc b/google/cloud/spanner/internal/partial_result_set_source.cc index fbd929c41af17..7bcfb2bbc4229 100644 --- a/google/cloud/spanner/internal/partial_result_set_source.cc +++ b/google/cloud/spanner/internal/partial_result_set_source.cc @@ -58,32 +58,41 @@ void ExtractSubrangeAndAppend(Values& src, int start, Values& dst) { } // namespace StatusOr> -PartialResultSetSource::Create(std::unique_ptr reader) { - std::unique_ptr source( - new PartialResultSetSource(std::move(reader))); +PartialResultSetSource::Create( + std::unique_ptr reader, + std::shared_ptr operation_context) { + std::unique_ptr source(new PartialResultSetSource( + std::move(reader), std::move(operation_context))); // Do an initial read from the stream to determine the fate of the factory. auto status = source->ReadFromStream(); // If the initial read finished the stream, and `Finish()` failed, then // creating the `PartialResultSetSource` should fail with the same error. - if (source->state_ == kFinished && !status.ok()) return status; + if (source->state_ == kFinished && !status.ok()) { + source->NotifyOnDone(status); + return status; + } // Otherwise we require that the first response contains the metadata. // Without it, creating the `PartialResultSetSource` should fail. if (!source->metadata_) { - return internal::InternalError( + auto err = internal::InternalError( "PartialResultSetSource response contained no metadata", GCP_ERROR_INFO()); + source->NotifyOnDone(err); + return err; } return {std::move(source)}; } PartialResultSetSource::PartialResultSetSource( - std::unique_ptr reader) + std::unique_ptr reader, + std::shared_ptr operation_context) : options_(internal::CurrentOptions()), reader_(std::move(reader)), + operation_context_(std::move(operation_context)), values_(std::make_optional( google::protobuf::Arena::Create< google::protobuf::RepeatedPtrField>( @@ -94,6 +103,14 @@ PartialResultSetSource::PartialResultSetSource( } } +void PartialResultSetSource::NotifyOnDone(Status const& status) { + if (on_done_called_) return; + on_done_called_ = true; + if (operation_context_) { + operation_context_->OnDone(status); + } +} + PartialResultSetSource::~PartialResultSetSource() { internal::OptionsSpan span(options_); if (state_ == kReading) { @@ -112,7 +129,9 @@ PartialResultSetSource::~PartialResultSetSource() { << status; } state_ = kFinished; + last_status_ = std::move(status); } + NotifyOnDone(last_status_); } StatusOr PartialResultSetSource::NextRow() { @@ -143,10 +162,16 @@ StatusOr PartialResultSetSource::NextRow() { rows_returned_ = 0; } while (usable_rows_ == 0) { - if (state_ == kFinished) return spanner::Row(); + if (state_ == kFinished) { + NotifyOnDone(Status{}); + return spanner::Row(); + } internal::OptionsSpan span(options_); auto status = ReadFromStream(); - if (!status.ok()) return status; + if (!status.ok()) { + NotifyOnDone(status); + return status; + } } auto value_it = (*values_)->begin() + rows_returned_ * columns_->size(); ++rows_returned_; diff --git a/google/cloud/spanner/internal/partial_result_set_source.h b/google/cloud/spanner/internal/partial_result_set_source.h index bd8046b0a67b0..d4100e3b76b83 100644 --- a/google/cloud/spanner/internal/partial_result_set_source.h +++ b/google/cloud/spanner/internal/partial_result_set_source.h @@ -15,6 +15,7 @@ #ifndef GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_SPANNER_INTERNAL_PARTIAL_RESULT_SET_SOURCE_H #define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_SPANNER_INTERNAL_PARTIAL_RESULT_SET_SOURCE_H +#include "google/cloud/spanner/internal/operation_context.h" #include "google/cloud/spanner/internal/partial_result_set_reader.h" #include "google/cloud/spanner/results.h" #include "google/cloud/spanner/value.h" @@ -61,7 +62,8 @@ class PartialResultSetSource : public PartialResultSourceInterface { public: /// Factory method to create a PartialResultSetSource. static StatusOr> Create( - std::unique_ptr reader); + std::unique_ptr reader, + std::shared_ptr operation_context); ~PartialResultSetSource() override; @@ -82,15 +84,21 @@ class PartialResultSetSource : public PartialResultSourceInterface { private: explicit PartialResultSetSource( - std::unique_ptr reader); + std::unique_ptr reader, + std::shared_ptr operation_context); Status ReadFromStream(); + void NotifyOnDone(Status const& status); + // Arena for the values_ field. google::protobuf::Arena arena_; Options options_; std::unique_ptr reader_; + std::shared_ptr operation_context_; + bool on_done_called_ = false; + Status last_status_; // The `PartialResultSet.metadata` we received in the first response, and // the column names it contained (which will be shared between rows). diff --git a/google/cloud/spanner/internal/partial_result_set_source_test.cc b/google/cloud/spanner/internal/partial_result_set_source_test.cc index 7e98f58f33e9e..61c8ed0bb9556 100644 --- a/google/cloud/spanner/internal/partial_result_set_source_test.cc +++ b/google/cloud/spanner/internal/partial_result_set_source_test.cc @@ -59,7 +59,7 @@ CreatePartialResultSetSource(std::unique_ptr reader, internal::OptionsSpan span(internal::MergeOptions( std::move(opts.set(CurrentTestName())), internal::CurrentOptions())); - return PartialResultSetSource::Create(std::move(reader)); + return PartialResultSetSource::Create(std::move(reader), nullptr); } // Returns a functor that expects the current `StringOption` to match the test diff --git a/google/cloud/spanner/internal/session_pool.cc b/google/cloud/spanner/internal/session_pool.cc index ca064dad712c0..21a8ab01541cb 100644 --- a/google/cloud/spanner/internal/session_pool.cc +++ b/google/cloud/spanner/internal/session_pool.cc @@ -41,16 +41,21 @@ using ::google::cloud::Idempotency; std::shared_ptr MakeSessionPool( spanner::Database db, std::vector> stubs, - google::cloud::CompletionQueue cq, Options opts) { - auto pool = std::shared_ptr(new SessionPool( - std::move(db), std::move(stubs), std::move(cq), std::move(opts))); + google::cloud::CompletionQueue cq, + std::shared_ptr context_factory, + Options opts) { + auto pool = std::shared_ptr( + new SessionPool(std::move(db), std::move(stubs), std::move(cq), + std::move(context_factory), std::move(opts))); pool->Initialize(); return pool; } -SessionPool::SessionPool(spanner::Database db, - std::vector> stubs, - google::cloud::CompletionQueue cq, Options opts) +SessionPool::SessionPool( + spanner::Database db, std::vector> stubs, + google::cloud::CompletionQueue cq, + std::shared_ptr context_factory, + Options opts) : db_(std::move(db)), cq_(std::move(cq)), opts_(std::move(opts)), @@ -75,7 +80,8 @@ SessionPool::SessionPool(spanner::Database db, ? opts_.get() : std::chrono::minutes(10)), - channels_(stubs.size()) { + channels_(stubs.size()), + context_factory_(std::move(context_factory)) { if (stubs.empty()) { google::cloud::internal::ThrowInvalidArgument( "SessionPool requires a non-empty set of stubs"); @@ -86,7 +92,8 @@ SessionPool::SessionPool(spanner::Database db, GCP_ERROR_INFO()); for (auto i = 0U; i < stubs.size(); ++i) { - channels_[i] = std::make_shared(std::move(stubs[i])); + channels_[i] = std::make_shared(std::move(stubs[i]), + static_cast(i)); } // `channels_` is never resized after this point. next_dissociated_stub_channel_ = channels_.begin(); @@ -130,7 +137,9 @@ SessionPool::~SessionPool() { // Multiplexed Sessions do not require an explicit Delete call. for (auto const& session : sessions_) { if (session->is_bad()) continue; - AsyncDeleteSession(cq_, GetStub(*session), session->session_name()) + auto stub_and_channel = GetStub(*session); + AsyncDeleteSession(cq_, std::move(stub_and_channel.stub), + stub_and_channel.channel_id, session->session_name()) .then([](auto result) { auto status = result.get(); }); } } @@ -157,9 +166,9 @@ void SessionPool::ReplaceMultiplexedSession() { if (create_calls_in_progress_ == 0 && (*multiplexed_session_)->creation_time() <= refresh_limit) { ++create_calls_in_progress_; - auto stub = GetStub(std::move(lk)); + auto stub_and_channel = GetStub(std::move(lk)); std::weak_ptr pool = shared_from_this(); - CreateMultiplexedSessionAsync(std::move(stub)) + CreateMultiplexedSessionAsync(std::move(stub_and_channel)) .then([pool](future> response) { if (auto shared_pool = pool.lock()) { shared_pool->HandleMultiplexedCreateSessionDone( @@ -207,8 +216,12 @@ void SessionPool::MaintainPoolSize() { // Refresh all sessions whose last-use time is older than the keep-alive // interval. Issues asynchronous RPCs, so this method does not block. void SessionPool::RefreshExpiringSessions() { - std::vector, std::string>> - sessions_to_refresh; + struct RefreshInfo { + std::shared_ptr stub; + std::uint32_t channel_id; + std::string session_name; + }; + std::vector sessions_to_refresh; auto now = clock_->Now(); auto refresh_limit = now - opts_.get(); @@ -219,8 +232,11 @@ void SessionPool::RefreshExpiringSessions() { for (auto const& session : sessions_) { auto last_use_time = session->last_use_time(); if (last_use_time <= refresh_limit) { - sessions_to_refresh.emplace_back(session->channel()->stub, - session->session_name()); + auto const& channel = session->channel(); + if (channel) { + sessions_to_refresh.push_back( + {channel->stub, channel->channel_id, session->session_name()}); + } session->update_last_use_time(); } else if (last_use_time < last_use_time_lower_bound_) { last_use_time_lower_bound_ = last_use_time; @@ -231,7 +247,7 @@ void SessionPool::RefreshExpiringSessions() { std::weak_ptr pool = shared_from_this(); for (auto& refresh : sessions_to_refresh) { auto handler = - [pool, session_name = refresh.second]( + [pool, session_name = refresh.session_name]( future> result) { auto response = result.get(); if (!response && IsSessionNotFound(response.status())) { @@ -245,8 +261,8 @@ void SessionPool::RefreshExpiringSessions() { } } }; - AsyncRefreshSession(cq_, std::move(refresh.first), - std::move(refresh.second)) + AsyncRefreshSession(cq_, std::move(refresh.stub), refresh.channel_id, + std::move(refresh.session_name)) .then(std::move(handler)); } } @@ -281,14 +297,14 @@ Status SessionPool::HandleMultiplexedCreateSessionDone( Status SessionPool::CreateMultiplexedSession(std::unique_lock& lk) { if (create_calls_in_progress_ == 0) { create_calls_in_progress_++; - auto stub = GetStub(std::move(lk)); - return CreateMultiplexedSessionSync(std::move(stub)); + auto stub_and_channel = GetStub(std::move(lk)); + return CreateMultiplexedSessionSync(std::move(stub_and_channel)); } return Status{}; } Status SessionPool::CreateMultiplexedSessionSync( - std::shared_ptr stub) { + StubAndChannel const& stub_and_channel) { google::spanner::v1::CreateSessionRequest request; request.set_database(db_.FullName()); auto* session = request.mutable_session(); @@ -300,20 +316,26 @@ Status SessionPool::CreateMultiplexedSessionSync( if (!role.empty()) session->set_creator_role(role); session->set_multiplexed(true); + auto op_context = context_factory_->BackgroundCreateSession(); + op_context.BindChannel(stub_and_channel.channel_id); + auto const& stub = stub_and_channel.stub; auto response = RetryLoop( retry_policy_prototype_->clone(), backoff_policy_prototype_->clone(), google::cloud::Idempotency::kIdempotent, - [&stub](grpc::ClientContext& context, Options const& options, - google::spanner::v1::CreateSessionRequest const& request) { + [&stub, &op_context]( + grpc::ClientContext& context, Options const& options, + google::spanner::v1::CreateSessionRequest const& request) { RouteToLeader(context); // always for CreateSession() - return stub->CreateSession(context, options, request); + return stub->CreateSession(context, options, request, op_context); }, opts_, request, __func__); + op_context.OnDone(response.status()); return HandleMultiplexedCreateSessionDone(std::move(response)); } future> -SessionPool::CreateMultiplexedSessionAsync(std::shared_ptr stub) { +SessionPool::CreateMultiplexedSessionAsync( + StubAndChannel const& stub_and_channel) { google::spanner::v1::CreateSessionRequest request; request.set_database(db_.FullName()); auto* session = request.mutable_session(); @@ -325,17 +347,31 @@ SessionPool::CreateMultiplexedSessionAsync(std::shared_ptr stub) { if (!role.empty()) session->set_creator_role(role); session->set_multiplexed(true); + auto op_context = std::make_shared( + context_factory_->BackgroundCreateSession()); + op_context->BindChannel(stub_and_channel.channel_id); + auto const& stub = stub_and_channel.stub; + return google::cloud::internal::AsyncRetryLoop( - retry_policy_prototype_->clone(), backoff_policy_prototype_->clone(), - google::cloud::Idempotency::kIdempotent, cq_, - [&stub](CompletionQueue cq, std::shared_ptr context, - internal::ImmutableOptions options, - google::spanner::v1::CreateSessionRequest const& request) { - RouteToLeader(*context); // always for CreateSession() - return stub->AsyncCreateSession(cq, std::move(context), - std::move(options), request); - }, - internal::SaveCurrentOptions(), std::move(request), __func__); + retry_policy_prototype_->clone(), + backoff_policy_prototype_->clone(), + google::cloud::Idempotency::kIdempotent, cq_, + [stub, op_context]( + CompletionQueue cq, + std::shared_ptr context, + internal::ImmutableOptions options, + google::spanner::v1::CreateSessionRequest const& request) { + RouteToLeader(*context); // always for CreateSession() + return stub->AsyncCreateSession(cq, std::move(context), + std::move(options), request, + op_context); + }, + internal::SaveCurrentOptions(), std::move(request), __func__) + .then([op_context](future> f) { + auto res = f.get(); + op_context->OnDone(res.status()); + return res; + }); } /* @@ -457,10 +493,10 @@ StatusOr SessionPool::Multiplexed(Session::Mode mode) { "multiplexed sessions are not enabled", GCP_ERROR_INFO()); } -std::shared_ptr SessionPool::GetStub(Session const& session) { +StubAndChannel SessionPool::GetStub(Session const& session) { if (!session.is_disassociated() && !session.is_multiplexed()) { auto const& channel = session.channel(); - if (channel) return channel->stub; + if (channel) return {channel->stub, channel->channel_id}; } // Multiplexed sessions, or sessions that were created for partitioned @@ -472,15 +508,19 @@ std::shared_ptr SessionPool::GetStub(Session const& session) { // In order to maintain transaction/channel affinity, we leverage the // TransactionContext to make sure we use the same stub for the life of the // transaction. -std::shared_ptr SessionPool::GetStub(Session const& session, - TransactionContext& context) { - if (context.stub.has_value()) return *(context.stub); - context.stub = GetStub(session); - return *context.stub; +StubAndChannel SessionPool::GetStub(Session const& session, + TransactionContext& context) { + if (context.stub.has_value() && context.channel_id.has_value()) { + return {*context.stub, *context.channel_id}; + } + auto stub_and_channel = GetStub(session); + context.stub = stub_and_channel.stub; + context.channel_id = stub_and_channel.channel_id; + return stub_and_channel; } int SessionPool::total_sessions() const { - std::lock_guard lk(mu_); + std::scoped_lock lk(mu_); return total_sessions_; } @@ -548,13 +588,13 @@ StatusOr SessionPool::Allocate(std::unique_lock lk, } } -std::shared_ptr SessionPool::GetStub( - std::unique_lock) { - auto stub = (*next_dissociated_stub_channel_)->stub; +StubAndChannel SessionPool::GetStub(std::unique_lock) { + auto const& channel = *next_dissociated_stub_channel_; + auto stub_and_channel = StubAndChannel{channel->stub, channel->channel_id}; if (++next_dissociated_stub_channel_ == channels_.end()) { next_dissociated_stub_channel_ = channels_.begin(); } - return stub; + return stub_and_channel; } void SessionPool::Release(std::unique_ptr session) { @@ -590,15 +630,19 @@ Status SessionPool::CreateSessionsSync( request.set_session_count(std::int32_t{num_sessions}); auto const& stub = channel->stub; auto const& current = internal::CurrentOptions(); + auto op_context = context_factory_->BackgroundBatchCreateSessions(); + op_context.BindChannel(channel->channel_id); auto response = RetryLoop( retry_policy_prototype_->clone(), backoff_policy_prototype_->clone(), google::cloud::Idempotency::kIdempotent, - [&stub](grpc::ClientContext& context, Options const& options, - google::spanner::v1::BatchCreateSessionsRequest const& request) { + [&stub, &op_context]( + grpc::ClientContext& context, Options const& options, + google::spanner::v1::BatchCreateSessionsRequest const& request) { RouteToLeader(context); // always for BatchCreateSessions() - return stub->BatchCreateSessions(context, options, request); + return stub->BatchCreateSessions(context, options, request, op_context); }, current, request, __func__); + op_context.OnDone(response.status()); return HandleBatchCreateSessionsDone(channel, std::move(response)); } @@ -607,7 +651,8 @@ void SessionPool::CreateSessionsAsync( std::map const& labels, std::string const& role, int num_sessions) { std::weak_ptr pool = shared_from_this(); - AsyncBatchCreateSessions(cq_, channel->stub, labels, role, num_sessions) + AsyncBatchCreateSessions(cq_, channel->stub, channel->channel_id, labels, + role, num_sessions) .then( [pool, channel]( future> @@ -638,8 +683,8 @@ SessionHolder SessionPool::MakeSessionHolder(std::unique_ptr session, future> SessionPool::AsyncBatchCreateSessions( CompletionQueue& cq, std::shared_ptr const& stub, - std::map const& labels, std::string const& role, - int num_sessions) { + std::uint32_t channel_id, std::map const& labels, + std::string const& role, int num_sessions) { google::spanner::v1::BatchCreateSessionsRequest request; request.set_database(db_.FullName()); if (!labels.empty()) { @@ -650,40 +695,67 @@ SessionPool::AsyncBatchCreateSessions( request.mutable_session_template()->set_creator_role(role); } request.set_session_count(std::int32_t{num_sessions}); + auto op_context = std::make_shared( + context_factory_->BackgroundBatchCreateSessions()); + op_context->BindChannel(channel_id); return google::cloud::internal::AsyncRetryLoop( - retry_policy_prototype_->clone(), backoff_policy_prototype_->clone(), - Idempotency::kIdempotent, cq, - [stub](CompletionQueue& cq, std::shared_ptr context, - internal::ImmutableOptions options, - google::spanner::v1::BatchCreateSessionsRequest const& request) { - RouteToLeader(*context); // always for BatchCreateSessions() - return stub->AsyncBatchCreateSessions(cq, std::move(context), - std::move(options), request); - }, - internal::SaveCurrentOptions(), std::move(request), __func__); + retry_policy_prototype_->clone(), + backoff_policy_prototype_->clone(), Idempotency::kIdempotent, cq, + [stub, op_context]( + CompletionQueue& cq, + std::shared_ptr context, + internal::ImmutableOptions options, + google::spanner::v1::BatchCreateSessionsRequest const& + request) { + RouteToLeader(*context); // always for BatchCreateSessions() + return stub->AsyncBatchCreateSessions(cq, std::move(context), + std::move(options), + request, op_context); + }, + internal::SaveCurrentOptions(), std::move(request), __func__) + .then( + [op_context]( + future> + f) { + auto res = f.get(); + op_context->OnDone(res.status()); + return res; + }); } future SessionPool::AsyncDeleteSession( CompletionQueue& cq, std::shared_ptr const& stub, - std::string session_name) { + std::uint32_t channel_id, std::string session_name) { google::spanner::v1::DeleteSessionRequest request; request.set_name(std::move(session_name)); + auto op_context = std::make_shared( + context_factory_->BackgroundDeleteSession()); + op_context->BindChannel(channel_id); return google::cloud::internal::AsyncRetryLoop( - retry_policy_prototype_->clone(), backoff_policy_prototype_->clone(), - Idempotency::kIdempotent, cq, - [stub](CompletionQueue& cq, std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::spanner::v1::DeleteSessionRequest const& request) { - return stub->AsyncDeleteSession(cq, std::move(context), - std::move(options), request); - }, - internal::SaveCurrentOptions(), std::move(request), __func__); + retry_policy_prototype_->clone(), + backoff_policy_prototype_->clone(), Idempotency::kIdempotent, cq, + [stub, op_context]( + CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::spanner::v1::DeleteSessionRequest const& request) { + return stub->AsyncDeleteSession(cq, std::move(context), + std::move(options), request, + op_context); + }, + internal::SaveCurrentOptions(), std::move(request), __func__) + .then([op_context](future f) { + auto res = f.get(); + op_context->OnDone(res); + return res; + }); } /// Refresh the session `session_name` by executing a `SELECT 1` query on it. future> SessionPool::AsyncRefreshSession(CompletionQueue& cq, std::shared_ptr const& stub, + std::uint32_t channel_id, std::string session_name) { google::spanner::v1::ExecuteSqlRequest request; request.set_session(std::move(session_name)); @@ -691,17 +763,28 @@ SessionPool::AsyncRefreshSession(CompletionQueue& cq, request.set_sql("SELECT 1;"); request.mutable_request_options()->set_priority( google::spanner::v1::RequestOptions::PRIORITY_LOW); + auto op_context = std::make_shared( + context_factory_->BackgroundRefreshSession()); + op_context->BindChannel(channel_id); return google::cloud::internal::AsyncRetryLoop( - retry_policy_prototype_->clone(), backoff_policy_prototype_->clone(), - Idempotency::kIdempotent, cq, - [stub](CompletionQueue& cq, std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::spanner::v1::ExecuteSqlRequest const& request) { - // Read-only transaction, so no route-to-leader. - return stub->AsyncExecuteSql(cq, std::move(context), std::move(options), - request); - }, - internal::SaveCurrentOptions(), std::move(request), __func__); + retry_policy_prototype_->clone(), + backoff_policy_prototype_->clone(), Idempotency::kIdempotent, cq, + [stub, op_context]( + CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::spanner::v1::ExecuteSqlRequest const& request) { + // Read-only transaction, so no route-to-leader. + return stub->AsyncExecuteSql(cq, std::move(context), + std::move(options), request, + op_context); + }, + internal::SaveCurrentOptions(), std::move(request), __func__) + .then([op_context](future> f) { + auto res = f.get(); + op_context->OnDone(res.status()); + return res; + }); } Status SessionPool::HandleBatchCreateSessionsDone( diff --git a/google/cloud/spanner/internal/session_pool.h b/google/cloud/spanner/internal/session_pool.h index 82ce975b059f4..d29f8b8bda519 100644 --- a/google/cloud/spanner/internal/session_pool.h +++ b/google/cloud/spanner/internal/session_pool.h @@ -19,6 +19,7 @@ #include "google/cloud/spanner/database.h" #include "google/cloud/spanner/internal/channel.h" #include "google/cloud/spanner/internal/session.h" +#include "google/cloud/spanner/internal/spanner_operation_context_factory.h" #include "google/cloud/spanner/internal/spanner_stub.h" #include "google/cloud/spanner/internal/transaction_impl.h" #include "google/cloud/spanner/retry_policy.h" @@ -32,6 +33,7 @@ #include #include #include +#include #include #include #include @@ -63,6 +65,11 @@ struct MultiplexedSessionBackgroundWorkIntervalOption { using Type = std::chrono::minutes; }; +struct StubAndChannel { + std::shared_ptr stub; + std::uint32_t channel_id; +}; + class SessionPool; /** @@ -74,7 +81,9 @@ class SessionPool; */ std::shared_ptr MakeSessionPool( spanner::Database db, std::vector> stubs, - google::cloud::CompletionQueue cq, Options opts); + google::cloud::CompletionQueue cq, + std::shared_ptr context_factory, + Options opts); /** * Maintains a pool of `Session` objects. @@ -124,11 +133,11 @@ class SessionPool : public std::enable_shared_from_this { Session::Mode mode = Session::Mode::kMultiplexed); /** - * Return a `SpannerStub` to be used when making calls using `session`. + * Return a `SpannerStub` and `channel_id` to be used when making calls using + * `session`. */ - std::shared_ptr GetStub(Session const& session); - std::shared_ptr GetStub(Session const& session, - TransactionContext& context); + StubAndChannel GetStub(Session const& session); + StubAndChannel GetStub(Session const& session, TransactionContext& context); /** * Returns the number of sessions in the session pool plus the number of @@ -142,7 +151,8 @@ class SessionPool : public std::enable_shared_from_this { private: friend std::shared_ptr MakeSessionPool( spanner::Database, std::vector>, - google::cloud::CompletionQueue, Options); + google::cloud::CompletionQueue, + std::shared_ptr, Options); /** * Construct a `SessionPool`. @@ -151,7 +161,9 @@ class SessionPool : public std::enable_shared_from_this { */ SessionPool(spanner::Database db, std::vector> stubs, - google::cloud::CompletionQueue cq, Options opts); + google::cloud::CompletionQueue cq, + std::shared_ptr context_factory, + Options opts); void Initialize(); @@ -167,8 +179,8 @@ class SessionPool : public std::enable_shared_from_this { StatusOr Allocate(std::unique_lock, Session::Mode mode); - // Returns a stub to use by round-robining between the channels. - std::shared_ptr GetStub(std::unique_lock); + // Returns a stub and channel to use by round-robining between the channels. + StubAndChannel GetStub(std::unique_lock); // Release session back to the pool. void Release(std::unique_ptr session); @@ -184,9 +196,9 @@ class SessionPool : public std::enable_shared_from_this { Status CreateMultiplexedSession( std::unique_lock& lk); // EXCLUSIVE_LOCKS_REQUIRED(mu_) - Status CreateMultiplexedSessionSync(std::shared_ptr); + Status CreateMultiplexedSessionSync(StubAndChannel const& stub_and_channel); future> CreateMultiplexedSessionAsync( - std::shared_ptr); + StubAndChannel const& stub_and_channel); Status HandleMultiplexedCreateSessionDone( StatusOr response); bool HasValidMultiplexedSession(std::unique_lock const&) const; @@ -214,14 +226,16 @@ class SessionPool : public std::enable_shared_from_this { future> AsyncBatchCreateSessions(CompletionQueue& cq, std::shared_ptr const& stub, + std::uint32_t channel_id, std::map const& labels, std::string const& role, int num_sessions); future AsyncDeleteSession(CompletionQueue& cq, std::shared_ptr const& stub, + std::uint32_t channel_id, std::string session_name); future> AsyncRefreshSession( CompletionQueue& cq, std::shared_ptr const& stub, - std::string session_name); + std::uint32_t channel_id, std::string session_name); Status HandleBatchCreateSessionsDone( std::shared_ptr const& channel, @@ -277,6 +291,7 @@ class SessionPool : public std::enable_shared_from_this { using ChannelVec = absl::FixedArray>; ChannelVec channels_; // GUARDED_BY(mu_) ChannelVec::iterator next_dissociated_stub_channel_; // GUARDED_BY(mu_) + std::shared_ptr context_factory_; }; GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END diff --git a/google/cloud/spanner/internal/session_pool_test.cc b/google/cloud/spanner/internal/session_pool_test.cc index 9f5cc8e619f6a..07f669c319419 100644 --- a/google/cloud/spanner/internal/session_pool_test.cc +++ b/google/cloud/spanner/internal/session_pool_test.cc @@ -15,6 +15,8 @@ #include "google/cloud/spanner/internal/session_pool.h" #include "google/cloud/spanner/internal/defaults.h" #include "google/cloud/spanner/internal/session.h" +#include "google/cloud/spanner/internal/spanner_operation_context_factory.h" +#include "google/cloud/spanner/internal/spanner_request_id.h" #include "google/cloud/spanner/options.h" #include "google/cloud/spanner/testing/mock_spanner_stub.h" #include "google/cloud/spanner/testing/status_utils.h" @@ -52,6 +54,7 @@ using ::testing::AllOf; using ::testing::AnyOf; using ::testing::ByMove; using ::testing::Contains; +using ::testing::Eq; using ::testing::HasSubstr; using ::testing::Not; using ::testing::Pair; @@ -124,7 +127,13 @@ google::spanner::v1::Session MakeMultiplexedSession(std::string name, std::shared_ptr MakeTestSessionPool( spanner::Database db, std::vector> stubs, - CompletionQueue cq, Options opts = {}) { + CompletionQueue cq, Options opts = {}, + std::shared_ptr context_factory = {}) { + if (!context_factory) { + context_factory = std::make_shared( + /*client_id=*/1, + std::make_shared(ProcessRandomId())); + } opts.set( std::make_shared( std::chrono::minutes(10))); @@ -133,7 +142,7 @@ std::shared_ptr MakeTestSessionPool( std::chrono::milliseconds(100), std::chrono::minutes(1), 2.0)); opts = DefaultOptions(std::move(opts)); return MakeSessionPool(std::move(db), std::move(stubs), std::move(cq), - std::move(opts)); + std::move(context_factory), std::move(opts)); } TEST_F(SessionPoolTest, Multiplexed) { @@ -141,14 +150,14 @@ TEST_F(SessionPoolTest, Multiplexed) { auto db = spanner::Database("project", "instance", "database"); EXPECT_CALL( *mock, - CreateSession(_, _, AllOf(DatabaseIs(db.FullName()), IsMultiplexed()))) + CreateSession(_, _, AllOf(DatabaseIs(db.FullName()), IsMultiplexed()), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); google::cloud::internal::AutomaticallyCreatedBackgroundThreads threads; auto pool = MakeTestSessionPool(db, {mock}, threads.cq(), {}); auto session = pool->Multiplexed(); ASSERT_STATUS_OK(session); - EXPECT_EQ((*session)->session_name(), "multiplexed"); + EXPECT_THAT((*session)->session_name(), Eq("multiplexed")); } TEST_F(SessionPoolTest, MultiplexedAllocateRouteToLeader) { @@ -156,9 +165,10 @@ TEST_F(SessionPoolTest, MultiplexedAllocateRouteToLeader) { auto db = spanner::Database("project", "instance", "database"); EXPECT_CALL( *mock, - CreateSession(_, _, AllOf(DatabaseIs(db.FullName()), IsMultiplexed()))) + CreateSession(_, _, AllOf(DatabaseIs(db.FullName()), IsMultiplexed()), _)) .WillOnce([this](grpc::ClientContext& context, Options const&, - google::spanner::v1::CreateSessionRequest const&) { + google::spanner::v1::CreateSessionRequest const&, + spanner_internal::OperationContext&) { EXPECT_THAT(GetMetadata(context), Contains(Pair(kRouteToLeader, "true"))); return MakeMultiplexedSession("multiplexed"); @@ -171,8 +181,10 @@ TEST_F(SessionPoolTest, MultiplexedAllocateRouteToLeader) { auto session = pool->Multiplexed(); ASSERT_STATUS_OK(session); - EXPECT_EQ((*session)->session_name(), "multiplexed"); - EXPECT_EQ(pool->GetStub(**session), mock); + EXPECT_THAT((*session)->session_name(), Eq("multiplexed")); + auto stub_and_channel = pool->GetStub(**session); + EXPECT_THAT(stub_and_channel.stub, Eq(mock)); + EXPECT_THAT(stub_and_channel.channel_id, Eq(0)); } TEST_F(SessionPoolTest, AllocateRouteToLeader) { @@ -181,9 +193,10 @@ TEST_F(SessionPoolTest, AllocateRouteToLeader) { EXPECT_CALL( *mock, - CreateSession(_, _, AllOf(DatabaseIs(db.FullName()), IsMultiplexed()))) + CreateSession(_, _, AllOf(DatabaseIs(db.FullName()), IsMultiplexed()), _)) .WillOnce([this](grpc::ClientContext& context, Options const&, - google::spanner::v1::CreateSessionRequest const&) { + google::spanner::v1::CreateSessionRequest const&, + spanner_internal::OperationContext&) { EXPECT_THAT(GetMetadata(context), Contains(Pair(kRouteToLeader, "true"))); return MakeMultiplexedSession("multiplexed"); @@ -197,8 +210,10 @@ TEST_F(SessionPoolTest, AllocateRouteToLeader) { .set(42)); auto session = pool->Multiplexed(); ASSERT_STATUS_OK(session); - EXPECT_EQ((*session)->session_name(), "multiplexed"); - EXPECT_EQ(pool->GetStub(**session), mock); + EXPECT_THAT((*session)->session_name(), Eq("multiplexed")); + auto stub_and_channel = pool->GetStub(**session); + EXPECT_THAT(stub_and_channel.stub, Eq(mock)); + EXPECT_THAT(stub_and_channel.channel_id, Eq(0)); } TEST_F(SessionPoolTest, MultiplexedAllocateNoRouteToLeader) { @@ -207,9 +222,10 @@ TEST_F(SessionPoolTest, MultiplexedAllocateNoRouteToLeader) { EXPECT_CALL( *mock, - CreateSession(_, _, AllOf(DatabaseIs(db.FullName()), IsMultiplexed()))) + CreateSession(_, _, AllOf(DatabaseIs(db.FullName()), IsMultiplexed()), _)) .WillOnce([this](grpc::ClientContext& context, Options const&, - google::spanner::v1::CreateSessionRequest const&) { + google::spanner::v1::CreateSessionRequest const&, + spanner_internal::OperationContext&) { EXPECT_THAT(GetMetadata(context), AnyOf(Contains(Pair(kRouteToLeader, "false")), Not(Contains(Pair(kRouteToLeader, _))))); @@ -223,8 +239,10 @@ TEST_F(SessionPoolTest, MultiplexedAllocateNoRouteToLeader) { auto session = pool->Multiplexed(); ASSERT_STATUS_OK(session); - EXPECT_EQ((*session)->session_name(), "multiplexed"); - EXPECT_EQ(pool->GetStub(**session), mock); + EXPECT_THAT((*session)->session_name(), Eq("multiplexed")); + auto stub_and_channel = pool->GetStub(**session); + EXPECT_THAT(stub_and_channel.stub, Eq(mock)); + EXPECT_THAT(stub_and_channel.channel_id, Eq(0)); } TEST_F(SessionPoolTest, AllocateNoRouteToLeader) { @@ -233,9 +251,10 @@ TEST_F(SessionPoolTest, AllocateNoRouteToLeader) { EXPECT_CALL( *mock, - CreateSession(_, _, AllOf(DatabaseIs(db.FullName()), IsMultiplexed()))) + CreateSession(_, _, AllOf(DatabaseIs(db.FullName()), IsMultiplexed()), _)) .WillOnce([this](grpc::ClientContext& context, Options const&, - google::spanner::v1::CreateSessionRequest const&) { + google::spanner::v1::CreateSessionRequest const&, + spanner_internal::OperationContext&) { EXPECT_THAT(GetMetadata(context), AnyOf(Contains(Pair(kRouteToLeader, "false")), Not(Contains(Pair(kRouteToLeader, _))))); @@ -250,14 +269,16 @@ TEST_F(SessionPoolTest, AllocateNoRouteToLeader) { .set(42)); auto session = pool->Multiplexed(); ASSERT_STATUS_OK(session); - EXPECT_EQ((*session)->session_name(), "multiplexed"); - EXPECT_EQ(pool->GetStub(**session), mock); + EXPECT_THAT((*session)->session_name(), Eq("multiplexed")); + auto stub_and_channel = pool->GetStub(**session); + EXPECT_THAT(stub_and_channel.stub, Eq(mock)); + EXPECT_THAT(stub_and_channel.channel_id, Eq(0)); } TEST_F(SessionPoolTest, MultiplexedCreateError) { auto mock = std::make_shared(); auto db = spanner::Database("project", "instance", "database"); - EXPECT_CALL(*mock, CreateSession) + EXPECT_CALL(*mock, CreateSession(_, _, _, _)) .WillRepeatedly( Return(ByMove(Status(StatusCode::kInternal, "init failure")))); @@ -274,9 +295,10 @@ TEST_F(SessionPoolTest, ReuseSession) { auto db = spanner::Database("project", "instance", "database"); EXPECT_CALL( *mock, - CreateSession(_, _, AllOf(DatabaseIs(db.FullName()), IsMultiplexed()))) + CreateSession(_, _, AllOf(DatabaseIs(db.FullName()), IsMultiplexed()), _)) .WillOnce([](grpc::ClientContext&, Options const&, - google::spanner::v1::CreateSessionRequest const&) { + google::spanner::v1::CreateSessionRequest const&, + spanner_internal::OperationContext&) { return MakeMultiplexedSession("multiplexed"); }); @@ -284,12 +306,12 @@ TEST_F(SessionPoolTest, ReuseSession) { auto pool = MakeTestSessionPool(db, {mock}, threads.cq()); auto session = pool->Multiplexed(); ASSERT_STATUS_OK(session); - EXPECT_EQ((*session)->session_name(), "multiplexed"); + EXPECT_THAT((*session)->session_name(), Eq("multiplexed")); session->reset(); auto session2 = pool->Multiplexed(); ASSERT_STATUS_OK(session2); - EXPECT_EQ((*session2)->session_name(), "multiplexed"); + EXPECT_THAT((*session2)->session_name(), Eq("multiplexed")); } TEST_F(SessionPoolTest, MultiplexedLabels) { @@ -299,14 +321,15 @@ TEST_F(SessionPoolTest, MultiplexedLabels) { {"k1", "v1"}, {"k2", "v2"}, {"k3", "v3"}}; EXPECT_CALL( *mock, - CreateSession(_, _, AllOf(DatabaseIs(db.FullName()), IsMultiplexed()))) + CreateSession(_, _, AllOf(DatabaseIs(db.FullName()), IsMultiplexed()), _)) .WillOnce( [labels](grpc::ClientContext&, Options const&, - google::spanner::v1::CreateSessionRequest const& request) { + google::spanner::v1::CreateSessionRequest const& request, + spanner_internal::OperationContext&) { auto const& request_labels = request.session().labels(); - EXPECT_EQ((std::map( - request_labels.begin(), request_labels.end())), - labels); + EXPECT_THAT((std::map( + request_labels.begin(), request_labels.end())), + Eq(labels)); return MakeMultiplexedSession("multiplexed"); }); @@ -316,7 +339,7 @@ TEST_F(SessionPoolTest, MultiplexedLabels) { Options{}.set(std::move(labels))); auto session = pool->Multiplexed(); ASSERT_STATUS_OK(session); - EXPECT_EQ((*session)->session_name(), "multiplexed"); + EXPECT_THAT((*session)->session_name(), Eq("multiplexed")); } TEST_F(SessionPoolTest, MultiplexedCreatorRole) { @@ -325,13 +348,13 @@ TEST_F(SessionPoolTest, MultiplexedCreatorRole) { std::string const role = "public"; EXPECT_CALL( *mock, - CreateSession(_, _, AllOf(DatabaseIs(db.FullName()), IsMultiplexed()))) - .WillOnce( - [role](grpc::ClientContext&, Options const&, - google::spanner::v1::CreateSessionRequest const& request) { - EXPECT_EQ(request.session().creator_role(), role); - return MakeMultiplexedSession("multiplexed"); - }); + CreateSession(_, _, AllOf(DatabaseIs(db.FullName()), IsMultiplexed()), _)) + .WillOnce([role](grpc::ClientContext&, Options const&, + google::spanner::v1::CreateSessionRequest const& request, + spanner_internal::OperationContext&) { + EXPECT_THAT(request.session().creator_role(), Eq(role)); + return MakeMultiplexedSession("multiplexed"); + }); google::cloud::internal::AutomaticallyCreatedBackgroundThreads threads; auto pool = MakeTestSessionPool( @@ -340,13 +363,13 @@ TEST_F(SessionPoolTest, MultiplexedCreatorRole) { auto session = pool->Multiplexed(); ASSERT_STATUS_OK(session); - EXPECT_EQ((*session)->session_name(), "multiplexed"); + EXPECT_THAT((*session)->session_name(), Eq("multiplexed")); } TEST_F(SessionPoolTest, GetStubForStublessSession) { auto mock = std::make_shared(); auto db = spanner::Database("project", "instance", "database"); - EXPECT_CALL(*mock, CreateSession) + EXPECT_CALL(*mock, CreateSession(_, _, _, _)) .WillRepeatedly( Return(ByMove(Status(StatusCode::kInternal, "init failure")))); google::cloud::internal::AutomaticallyCreatedBackgroundThreads threads; @@ -355,7 +378,7 @@ TEST_F(SessionPoolTest, GetStubForStublessSession) { Options{}.set(0)); // ensure we get a stub even if we didn't allocate from the pool. auto session = MakeDissociatedSessionHolder("session_id"); - EXPECT_EQ(pool->GetStub(*session), mock); + EXPECT_THAT(pool->GetStub(*session).stub, Eq(mock)); } TEST_F(SessionPoolTest, MultilpexedSessionReplacementSuccess) { @@ -363,9 +386,9 @@ TEST_F(SessionPoolTest, MultilpexedSessionReplacementSuccess) { auto db = spanner::Database("project", "instance", "database"); EXPECT_CALL( *mock, - CreateSession(_, _, AllOf(DatabaseIs(db.FullName()), IsMultiplexed()))) + CreateSession(_, _, AllOf(DatabaseIs(db.FullName()), IsMultiplexed()), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed1"})))); - EXPECT_CALL(*mock, AsyncCreateSession(_, _, _, _)) + EXPECT_CALL(*mock, AsyncCreateSession(_, _, _, _, _)) .WillOnce(Return(make_ready_future(StatusOr( MakeMultiplexedSession({"multiplexed2"}))))); @@ -386,14 +409,14 @@ TEST_F(SessionPoolTest, MultilpexedSessionReplacementSuccess) { auto s1 = pool->Multiplexed(); ASSERT_STATUS_OK(s1); - EXPECT_EQ((*s1)->session_name(), "multiplexed1"); + EXPECT_THAT((*s1)->session_name(), Eq("multiplexed1")); clock->AdvanceTime(background_interval); impl->SimulateCompletion(true); auto s2 = pool->Multiplexed(); ASSERT_STATUS_OK(s2); - EXPECT_EQ((*s2)->session_name(), "multiplexed2"); + EXPECT_THAT((*s2)->session_name(), Eq("multiplexed2")); // Cancel all pending operations, satisfying any remaining futures. impl->SimulateCompletion(false); @@ -404,9 +427,9 @@ TEST_F(SessionPoolTest, MultilpexedSessionReplacementRpcPermanentFailure) { auto db = spanner::Database("project", "instance", "database"); EXPECT_CALL( *mock, - CreateSession(_, _, AllOf(DatabaseIs(db.FullName()), IsMultiplexed()))) + CreateSession(_, _, AllOf(DatabaseIs(db.FullName()), IsMultiplexed()), _)) .WillOnce(Return(ByMove(MakeMultiplexedSession({"multiplexed"})))); - EXPECT_CALL(*mock, AsyncCreateSession(_, _, _, _)) + EXPECT_CALL(*mock, AsyncCreateSession(_, _, _, _, _)) .WillOnce(Return(make_ready_future(StatusOr( Status(StatusCode::kResourceExhausted, "retry policy exhausted"))))); @@ -427,19 +450,75 @@ TEST_F(SessionPoolTest, MultilpexedSessionReplacementRpcPermanentFailure) { auto s1 = pool->Multiplexed(); ASSERT_STATUS_OK(s1); - EXPECT_EQ((*s1)->session_name(), "multiplexed"); + EXPECT_THAT((*s1)->session_name(), Eq("multiplexed")); clock->AdvanceTime(background_interval); impl->SimulateCompletion(true); auto s2 = pool->Multiplexed(); ASSERT_STATUS_OK(s2); - EXPECT_EQ((*s2)->session_name(), "multiplexed"); + EXPECT_THAT((*s2)->session_name(), Eq("multiplexed")); // Cancel all pending operations, satisfying any remaining futures. impl->SimulateCompletion(false); } +TEST_F(SessionPoolTest, ChannelIdRoundRobinAndAffinity) { + auto mock1 = std::make_shared(); + auto mock2 = std::make_shared(); + auto mock3 = std::make_shared(); + auto db = spanner::Database("project", "instance", "database"); + + google::cloud::internal::AutomaticallyCreatedBackgroundThreads threads; + auto opts = + DefaultOptions(Options{}.set(0)); + opts.unset(); + auto context_factory = + std::make_shared( + /*client_id=*/1, + std::make_shared(ProcessRandomId())); + auto pool = MakeSessionPool(db, {mock1, mock2, mock3}, threads.cq(), + context_factory, std::move(opts)); + + auto session = MakeDissociatedSessionHolder("session_id"); + + // Round-robin for dissociated sessions + auto stub_and_channel1 = pool->GetStub(*session); + EXPECT_THAT(stub_and_channel1.stub, Eq(mock1)); + EXPECT_THAT(stub_and_channel1.channel_id, Eq(0)); + + auto stub_and_channel2 = pool->GetStub(*session); + EXPECT_THAT(stub_and_channel2.stub, Eq(mock2)); + EXPECT_THAT(stub_and_channel2.channel_id, Eq(1)); + + auto stub_and_channel3 = pool->GetStub(*session); + EXPECT_THAT(stub_and_channel3.stub, Eq(mock3)); + EXPECT_THAT(stub_and_channel3.channel_id, Eq(2)); + + auto stub_and_channel4 = pool->GetStub(*session); + EXPECT_THAT(stub_and_channel4.stub, Eq(mock1)); + EXPECT_THAT(stub_and_channel4.channel_id, Eq(0)); + + // TransactionContext pins stub and channel_id + std::string tag = "test_tag"; + TransactionContext ctx{/*route_to_leader=*/false, + tag, + /*seqno=*/1, + /*stub=*/std::nullopt, + /*channel_id=*/std::nullopt, + /*precommit_token=*/std::nullopt}; + auto pinned1 = pool->GetStub(*session, ctx); + EXPECT_THAT(pinned1.stub, Eq(mock2)); + EXPECT_THAT(pinned1.channel_id, Eq(1)); + EXPECT_THAT(ctx.channel_id, Eq(1)); + EXPECT_THAT(*ctx.stub, Eq(mock2)); + + // Second call with the same ctx returns the cached channel + auto pinned2 = pool->GetStub(*session, ctx); + EXPECT_THAT(pinned2.stub, Eq(mock2)); + EXPECT_THAT(pinned2.channel_id, Eq(1)); +} + } // namespace GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END } // namespace spanner_internal diff --git a/google/cloud/spanner/internal/spanner_auth_decorator.cc b/google/cloud/spanner/internal/spanner_auth_decorator.cc index 4c7a4b9605a06..2782603feeec3 100644 --- a/google/cloud/spanner/internal/spanner_auth_decorator.cc +++ b/google/cloud/spanner/internal/spanner_auth_decorator.cc @@ -36,150 +36,172 @@ SpannerAuth::SpannerAuth( StatusOr SpannerAuth::CreateSession( grpc::ClientContext& context, Options const& options, - google::spanner::v1::CreateSessionRequest const& request) { + google::spanner::v1::CreateSessionRequest const& request, + spanner_internal::OperationContext& operation_context) { auto status = auth_->ConfigureContext(context); if (!status.ok()) return status; - return child_->CreateSession(context, options, request); + return child_->CreateSession(context, options, request, operation_context); } StatusOr SpannerAuth::BatchCreateSessions( grpc::ClientContext& context, Options const& options, - google::spanner::v1::BatchCreateSessionsRequest const& request) { + google::spanner::v1::BatchCreateSessionsRequest const& request, + spanner_internal::OperationContext& operation_context) { auto status = auth_->ConfigureContext(context); if (!status.ok()) return status; - return child_->BatchCreateSessions(context, options, request); + return child_->BatchCreateSessions(context, options, request, + operation_context); } Status SpannerAuth::DeleteSession( grpc::ClientContext& context, Options const& options, - google::spanner::v1::DeleteSessionRequest const& request) { + google::spanner::v1::DeleteSessionRequest const& request, + spanner_internal::OperationContext& operation_context) { auto status = auth_->ConfigureContext(context); if (!status.ok()) return status; - return child_->DeleteSession(context, options, request); + return child_->DeleteSession(context, options, request, operation_context); } StatusOr SpannerAuth::ExecuteSql( grpc::ClientContext& context, Options const& options, - google::spanner::v1::ExecuteSqlRequest const& request) { + google::spanner::v1::ExecuteSqlRequest const& request, + spanner_internal::OperationContext& operation_context) { auto status = auth_->ConfigureContext(context); if (!status.ok()) return status; - return child_->ExecuteSql(context, options, request); + return child_->ExecuteSql(context, options, request, operation_context); } std::unique_ptr> SpannerAuth::ExecuteStreamingSql( std::shared_ptr context, Options const& options, - google::spanner::v1::ExecuteSqlRequest const& request) { + google::spanner::v1::ExecuteSqlRequest const& request, + std::shared_ptr operation_context) { using ErrorStream = ::google::cloud::internal::StreamingReadRpcError< google::spanner::v1::PartialResultSet>; auto status = auth_->ConfigureContext(*context); if (!status.ok()) return std::make_unique(std::move(status)); - return child_->ExecuteStreamingSql(std::move(context), options, request); + return child_->ExecuteStreamingSql(std::move(context), options, request, + std::move(operation_context)); } StatusOr SpannerAuth::ExecuteBatchDml( grpc::ClientContext& context, Options const& options, - google::spanner::v1::ExecuteBatchDmlRequest const& request) { + google::spanner::v1::ExecuteBatchDmlRequest const& request, + spanner_internal::OperationContext& operation_context) { auto status = auth_->ConfigureContext(context); if (!status.ok()) return status; - return child_->ExecuteBatchDml(context, options, request); + return child_->ExecuteBatchDml(context, options, request, operation_context); } std::unique_ptr> -SpannerAuth::StreamingRead(std::shared_ptr context, - Options const& options, - google::spanner::v1::ReadRequest const& request) { +SpannerAuth::StreamingRead( + std::shared_ptr context, Options const& options, + google::spanner::v1::ReadRequest const& request, + std::shared_ptr operation_context) { using ErrorStream = ::google::cloud::internal::StreamingReadRpcError< google::spanner::v1::PartialResultSet>; auto status = auth_->ConfigureContext(*context); if (!status.ok()) return std::make_unique(std::move(status)); - return child_->StreamingRead(std::move(context), options, request); + return child_->StreamingRead(std::move(context), options, request, + std::move(operation_context)); } StatusOr SpannerAuth::BeginTransaction( grpc::ClientContext& context, Options const& options, - google::spanner::v1::BeginTransactionRequest const& request) { + google::spanner::v1::BeginTransactionRequest const& request, + spanner_internal::OperationContext& operation_context) { auto status = auth_->ConfigureContext(context); if (!status.ok()) return status; - return child_->BeginTransaction(context, options, request); + return child_->BeginTransaction(context, options, request, operation_context); } StatusOr SpannerAuth::Commit( grpc::ClientContext& context, Options const& options, - google::spanner::v1::CommitRequest const& request) { + google::spanner::v1::CommitRequest const& request, + spanner_internal::OperationContext& operation_context) { auto status = auth_->ConfigureContext(context); if (!status.ok()) return status; - return child_->Commit(context, options, request); + return child_->Commit(context, options, request, operation_context); } Status SpannerAuth::Rollback( grpc::ClientContext& context, Options const& options, - google::spanner::v1::RollbackRequest const& request) { + google::spanner::v1::RollbackRequest const& request, + spanner_internal::OperationContext& operation_context) { auto status = auth_->ConfigureContext(context); if (!status.ok()) return status; - return child_->Rollback(context, options, request); + return child_->Rollback(context, options, request, operation_context); } StatusOr SpannerAuth::PartitionQuery( grpc::ClientContext& context, Options const& options, - google::spanner::v1::PartitionQueryRequest const& request) { + google::spanner::v1::PartitionQueryRequest const& request, + spanner_internal::OperationContext& operation_context) { auto status = auth_->ConfigureContext(context); if (!status.ok()) return status; - return child_->PartitionQuery(context, options, request); + return child_->PartitionQuery(context, options, request, operation_context); } StatusOr SpannerAuth::PartitionRead( grpc::ClientContext& context, Options const& options, - google::spanner::v1::PartitionReadRequest const& request) { + google::spanner::v1::PartitionReadRequest const& request, + spanner_internal::OperationContext& operation_context) { auto status = auth_->ConfigureContext(context); if (!status.ok()) return status; - return child_->PartitionRead(context, options, request); + return child_->PartitionRead(context, options, request, operation_context); } std::unique_ptr> -SpannerAuth::BatchWrite(std::shared_ptr context, - Options const& options, - google::spanner::v1::BatchWriteRequest const& request) { +SpannerAuth::BatchWrite( + std::shared_ptr context, Options const& options, + google::spanner::v1::BatchWriteRequest const& request, + std::shared_ptr operation_context) { using ErrorStream = ::google::cloud::internal::StreamingReadRpcError< google::spanner::v1::BatchWriteResponse>; auto status = auth_->ConfigureContext(*context); if (!status.ok()) return std::make_unique(std::move(status)); - return child_->BatchWrite(std::move(context), options, request); + return child_->BatchWrite(std::move(context), options, request, + std::move(operation_context)); } std::unique_ptr< google::cloud::internal::StreamingReadRpc> SpannerAuth::FetchCacheUpdate( std::shared_ptr context, Options const& options, - google::spanner::v1::FetchCacheUpdateRequest const& request) { + google::spanner::v1::FetchCacheUpdateRequest const& request, + std::shared_ptr operation_context) { using ErrorStream = ::google::cloud::internal::StreamingReadRpcError< google::spanner::v1::CacheUpdate>; auto status = auth_->ConfigureContext(*context); if (!status.ok()) return std::make_unique(std::move(status)); - return child_->FetchCacheUpdate(std::move(context), options, request); + return child_->FetchCacheUpdate(std::move(context), options, request, + std::move(operation_context)); } future> SpannerAuth::AsyncCreateSession( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::CreateSessionRequest const& request) { + google::spanner::v1::CreateSessionRequest const& request, + std::shared_ptr operation_context) { return auth_->AsyncConfigureContext(std::move(context)) - .then([cq, child = child_, options = std::move(options), - request](future>> - f) mutable { + .then([cq, child = child_, options = std::move(options), request, + operation_context = std::move(operation_context)]( + future>> + f) mutable { auto context = f.get(); if (!context) { return make_ready_future(StatusOr( std::move(context).status())); } return child->AsyncCreateSession(cq, *std::move(context), - std::move(options), request); + std::move(options), request, + std::move(operation_context)); }); } @@ -188,11 +210,13 @@ SpannerAuth::AsyncBatchCreateSessions( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::BatchCreateSessionsRequest const& request) { + google::spanner::v1::BatchCreateSessionsRequest const& request, + std::shared_ptr operation_context) { return auth_->AsyncConfigureContext(std::move(context)) - .then([cq, child = child_, options = std::move(options), - request](future>> - f) mutable { + .then([cq, child = child_, options = std::move(options), request, + operation_context = std::move(operation_context)]( + future>> + f) mutable { auto context = f.get(); if (!context) { return make_ready_future( @@ -200,7 +224,8 @@ SpannerAuth::AsyncBatchCreateSessions( std::move(context).status())); } return child->AsyncBatchCreateSessions(cq, *std::move(context), - std::move(options), request); + std::move(options), request, + std::move(operation_context)); }); } @@ -208,15 +233,18 @@ future SpannerAuth::AsyncDeleteSession( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::DeleteSessionRequest const& request) { + google::spanner::v1::DeleteSessionRequest const& request, + std::shared_ptr operation_context) { return auth_->AsyncConfigureContext(std::move(context)) - .then([cq, child = child_, options = std::move(options), - request](future>> - f) mutable { + .then([cq, child = child_, options = std::move(options), request, + operation_context = std::move(operation_context)]( + future>> + f) mutable { auto context = f.get(); if (!context) return make_ready_future(std::move(context).status()); return child->AsyncDeleteSession(cq, *std::move(context), - std::move(options), request); + std::move(options), request, + std::move(operation_context)); }); } @@ -224,18 +252,21 @@ future> SpannerAuth::AsyncExecuteSql( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::ExecuteSqlRequest const& request) { + google::spanner::v1::ExecuteSqlRequest const& request, + std::shared_ptr operation_context) { return auth_->AsyncConfigureContext(std::move(context)) - .then([cq, child = child_, options = std::move(options), - request](future>> - f) mutable { + .then([cq, child = child_, options = std::move(options), request, + operation_context = std::move(operation_context)]( + future>> + f) mutable { auto context = f.get(); if (!context) { return make_ready_future(StatusOr( std::move(context).status())); } return child->AsyncExecuteSql(cq, *std::move(context), - std::move(options), request); + std::move(options), request, + std::move(operation_context)); }); } diff --git a/google/cloud/spanner/internal/spanner_auth_decorator.h b/google/cloud/spanner/internal/spanner_auth_decorator.h index 2647b4c7afc6e..3beae6332d790 100644 --- a/google/cloud/spanner/internal/spanner_auth_decorator.h +++ b/google/cloud/spanner/internal/spanner_auth_decorator.h @@ -43,92 +43,119 @@ class SpannerAuth : public SpannerStub { StatusOr CreateSession( grpc::ClientContext& context, Options const& options, - google::spanner::v1::CreateSessionRequest const& request) override; + google::spanner::v1::CreateSessionRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr BatchCreateSessions( grpc::ClientContext& context, Options const& options, - google::spanner::v1::BatchCreateSessionsRequest const& request) override; + google::spanner::v1::BatchCreateSessionsRequest const& request, + spanner_internal::OperationContext& operation_context) override; Status DeleteSession( grpc::ClientContext& context, Options const& options, - google::spanner::v1::DeleteSessionRequest const& request) override; + google::spanner::v1::DeleteSessionRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr ExecuteSql( grpc::ClientContext& context, Options const& options, - google::spanner::v1::ExecuteSqlRequest const& request) override; + google::spanner::v1::ExecuteSqlRequest const& request, + spanner_internal::OperationContext& operation_context) override; std::unique_ptr> - ExecuteStreamingSql( - std::shared_ptr context, Options const& options, - google::spanner::v1::ExecuteSqlRequest const& request) override; + ExecuteStreamingSql(std::shared_ptr context, + Options const& options, + google::spanner::v1::ExecuteSqlRequest const& request, + std::shared_ptr + operation_context) override; StatusOr ExecuteBatchDml( grpc::ClientContext& context, Options const& options, - google::spanner::v1::ExecuteBatchDmlRequest const& request) override; + google::spanner::v1::ExecuteBatchDmlRequest const& request, + spanner_internal::OperationContext& operation_context) override; std::unique_ptr> StreamingRead(std::shared_ptr context, Options const& options, - google::spanner::v1::ReadRequest const& request) override; + google::spanner::v1::ReadRequest const& request, + std::shared_ptr + operation_context) override; StatusOr BeginTransaction( grpc::ClientContext& context, Options const& options, - google::spanner::v1::BeginTransactionRequest const& request) override; + google::spanner::v1::BeginTransactionRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr Commit( grpc::ClientContext& context, Options const& options, - google::spanner::v1::CommitRequest const& request) override; + google::spanner::v1::CommitRequest const& request, + spanner_internal::OperationContext& operation_context) override; - Status Rollback(grpc::ClientContext& context, Options const& options, - google::spanner::v1::RollbackRequest const& request) override; + Status Rollback( + grpc::ClientContext& context, Options const& options, + google::spanner::v1::RollbackRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr PartitionQuery( grpc::ClientContext& context, Options const& options, - google::spanner::v1::PartitionQueryRequest const& request) override; + google::spanner::v1::PartitionQueryRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr PartitionRead( grpc::ClientContext& context, Options const& options, - google::spanner::v1::PartitionReadRequest const& request) override; + google::spanner::v1::PartitionReadRequest const& request, + spanner_internal::OperationContext& operation_context) override; std::unique_ptr> BatchWrite(std::shared_ptr context, Options const& options, - google::spanner::v1::BatchWriteRequest const& request) override; + google::spanner::v1::BatchWriteRequest const& request, + std::shared_ptr + operation_context) override; std::unique_ptr> - FetchCacheUpdate( - std::shared_ptr context, Options const& options, - google::spanner::v1::FetchCacheUpdateRequest const& request) override; + FetchCacheUpdate(std::shared_ptr context, + Options const& options, + google::spanner::v1::FetchCacheUpdateRequest const& request, + std::shared_ptr + operation_context) override; future> AsyncCreateSession( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::CreateSessionRequest const& request) override; + google::spanner::v1::CreateSessionRequest const& request, + std::shared_ptr operation_context) + override; future> AsyncBatchCreateSessions( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::BatchCreateSessionsRequest const& request) override; + google::spanner::v1::BatchCreateSessionsRequest const& request, + std::shared_ptr operation_context) + override; future AsyncDeleteSession( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::DeleteSessionRequest const& request) override; + google::spanner::v1::DeleteSessionRequest const& request, + std::shared_ptr operation_context) + override; future> AsyncExecuteSql( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::ExecuteSqlRequest const& request) override; + google::spanner::v1::ExecuteSqlRequest const& request, + std::shared_ptr operation_context) + override; private: std::shared_ptr auth_; diff --git a/google/cloud/spanner/internal/spanner_logging_decorator.cc b/google/cloud/spanner/internal/spanner_logging_decorator.cc index e690344aacf43..545da5a927449 100644 --- a/google/cloud/spanner/internal/spanner_logging_decorator.cc +++ b/google/cloud/spanner/internal/spanner_logging_decorator.cc @@ -43,11 +43,14 @@ SpannerLogging::SpannerLogging(std::shared_ptr child, StatusOr SpannerLogging::CreateSession( grpc::ClientContext& context, Options const& options, - google::spanner::v1::CreateSessionRequest const& request) { + google::spanner::v1::CreateSessionRequest const& request, + spanner_internal::OperationContext& operation_context) { return google::cloud::internal::LogWrapper( - [this](grpc::ClientContext& context, Options const& options, - google::spanner::v1::CreateSessionRequest const& request) { - return child_->CreateSession(context, options, request); + [this, &operation_context]( + grpc::ClientContext& context, Options const& options, + google::spanner::v1::CreateSessionRequest const& request) { + return child_->CreateSession(context, options, request, + operation_context); }, context, options, request, __func__, tracing_options_); } @@ -55,33 +58,41 @@ StatusOr SpannerLogging::CreateSession( StatusOr SpannerLogging::BatchCreateSessions( grpc::ClientContext& context, Options const& options, - google::spanner::v1::BatchCreateSessionsRequest const& request) { + google::spanner::v1::BatchCreateSessionsRequest const& request, + spanner_internal::OperationContext& operation_context) { return google::cloud::internal::LogWrapper( - [this](grpc::ClientContext& context, Options const& options, - google::spanner::v1::BatchCreateSessionsRequest const& request) { - return child_->BatchCreateSessions(context, options, request); + [this, &operation_context]( + grpc::ClientContext& context, Options const& options, + google::spanner::v1::BatchCreateSessionsRequest const& request) { + return child_->BatchCreateSessions(context, options, request, + operation_context); }, context, options, request, __func__, tracing_options_); } Status SpannerLogging::DeleteSession( grpc::ClientContext& context, Options const& options, - google::spanner::v1::DeleteSessionRequest const& request) { + google::spanner::v1::DeleteSessionRequest const& request, + spanner_internal::OperationContext& operation_context) { return google::cloud::internal::LogWrapper( - [this](grpc::ClientContext& context, Options const& options, - google::spanner::v1::DeleteSessionRequest const& request) { - return child_->DeleteSession(context, options, request); + [this, &operation_context]( + grpc::ClientContext& context, Options const& options, + google::spanner::v1::DeleteSessionRequest const& request) { + return child_->DeleteSession(context, options, request, + operation_context); }, context, options, request, __func__, tracing_options_); } StatusOr SpannerLogging::ExecuteSql( grpc::ClientContext& context, Options const& options, - google::spanner::v1::ExecuteSqlRequest const& request) { + google::spanner::v1::ExecuteSqlRequest const& request, + spanner_internal::OperationContext& operation_context) { return google::cloud::internal::LogWrapper( - [this](grpc::ClientContext& context, Options const& options, - google::spanner::v1::ExecuteSqlRequest const& request) { - return child_->ExecuteSql(context, options, request); + [this, &operation_context]( + grpc::ClientContext& context, Options const& options, + google::spanner::v1::ExecuteSqlRequest const& request) { + return child_->ExecuteSql(context, options, request, operation_context); }, context, options, request, __func__, tracing_options_); } @@ -90,15 +101,16 @@ std::unique_ptr> SpannerLogging::ExecuteStreamingSql( std::shared_ptr context, Options const& options, - google::spanner::v1::ExecuteSqlRequest const& request) { + google::spanner::v1::ExecuteSqlRequest const& request, + std::shared_ptr operation_context) { return google::cloud::internal::LogWrapper( - [this](std::shared_ptr context, - Options const& options, - google::spanner::v1::ExecuteSqlRequest const& request) + [this, operation_context = std::move(operation_context)]( + std::shared_ptr context, Options const& options, + google::spanner::v1::ExecuteSqlRequest const& request) -> std::unique_ptr> { - auto stream = - child_->ExecuteStreamingSql(std::move(context), options, request); + auto stream = child_->ExecuteStreamingSql( + std::move(context), options, request, std::move(operation_context)); if (stream_logging_) { stream = std::make_unique SpannerLogging::ExecuteBatchDml( grpc::ClientContext& context, Options const& options, - google::spanner::v1::ExecuteBatchDmlRequest const& request) { + google::spanner::v1::ExecuteBatchDmlRequest const& request, + spanner_internal::OperationContext& operation_context) { return google::cloud::internal::LogWrapper( - [this](grpc::ClientContext& context, Options const& options, - google::spanner::v1::ExecuteBatchDmlRequest const& request) { - return child_->ExecuteBatchDml(context, options, request); + [this, &operation_context]( + grpc::ClientContext& context, Options const& options, + google::spanner::v1::ExecuteBatchDmlRequest const& request) { + return child_->ExecuteBatchDml(context, options, request, + operation_context); }, context, options, request, __func__, tracing_options_); } std::unique_ptr> -SpannerLogging::StreamingRead(std::shared_ptr context, - Options const& options, - google::spanner::v1::ReadRequest const& request) { +SpannerLogging::StreamingRead( + std::shared_ptr context, Options const& options, + google::spanner::v1::ReadRequest const& request, + std::shared_ptr operation_context) { return google::cloud::internal::LogWrapper( - [this](std::shared_ptr context, - Options const& options, - google::spanner::v1::ReadRequest const& request) + [this, operation_context = std::move(operation_context)]( + std::shared_ptr context, Options const& options, + google::spanner::v1::ReadRequest const& request) -> std::unique_ptr> { - auto stream = - child_->StreamingRead(std::move(context), options, request); + auto stream = child_->StreamingRead( + std::move(context), options, request, std::move(operation_context)); if (stream_logging_) { stream = std::make_unique context, StatusOr SpannerLogging::BeginTransaction( grpc::ClientContext& context, Options const& options, - google::spanner::v1::BeginTransactionRequest const& request) { + google::spanner::v1::BeginTransactionRequest const& request, + spanner_internal::OperationContext& operation_context) { return google::cloud::internal::LogWrapper( - [this](grpc::ClientContext& context, Options const& options, - google::spanner::v1::BeginTransactionRequest const& request) { - return child_->BeginTransaction(context, options, request); + [this, &operation_context]( + grpc::ClientContext& context, Options const& options, + google::spanner::v1::BeginTransactionRequest const& request) { + return child_->BeginTransaction(context, options, request, + operation_context); }, context, options, request, __func__, tracing_options_); } StatusOr SpannerLogging::Commit( grpc::ClientContext& context, Options const& options, - google::spanner::v1::CommitRequest const& request) { + google::spanner::v1::CommitRequest const& request, + spanner_internal::OperationContext& operation_context) { return google::cloud::internal::LogWrapper( - [this](grpc::ClientContext& context, Options const& options, - google::spanner::v1::CommitRequest const& request) { - return child_->Commit(context, options, request); + [this, &operation_context]( + grpc::ClientContext& context, Options const& options, + google::spanner::v1::CommitRequest const& request) { + return child_->Commit(context, options, request, operation_context); }, context, options, request, __func__, tracing_options_); } Status SpannerLogging::Rollback( grpc::ClientContext& context, Options const& options, - google::spanner::v1::RollbackRequest const& request) { + google::spanner::v1::RollbackRequest const& request, + spanner_internal::OperationContext& operation_context) { return google::cloud::internal::LogWrapper( - [this](grpc::ClientContext& context, Options const& options, - google::spanner::v1::RollbackRequest const& request) { - return child_->Rollback(context, options, request); + [this, &operation_context]( + grpc::ClientContext& context, Options const& options, + google::spanner::v1::RollbackRequest const& request) { + return child_->Rollback(context, options, request, operation_context); }, context, options, request, __func__, tracing_options_); } StatusOr SpannerLogging::PartitionQuery( grpc::ClientContext& context, Options const& options, - google::spanner::v1::PartitionQueryRequest const& request) { + google::spanner::v1::PartitionQueryRequest const& request, + spanner_internal::OperationContext& operation_context) { return google::cloud::internal::LogWrapper( - [this](grpc::ClientContext& context, Options const& options, - google::spanner::v1::PartitionQueryRequest const& request) { - return child_->PartitionQuery(context, options, request); + [this, &operation_context]( + grpc::ClientContext& context, Options const& options, + google::spanner::v1::PartitionQueryRequest const& request) { + return child_->PartitionQuery(context, options, request, + operation_context); }, context, options, request, __func__, tracing_options_); } StatusOr SpannerLogging::PartitionRead( grpc::ClientContext& context, Options const& options, - google::spanner::v1::PartitionReadRequest const& request) { + google::spanner::v1::PartitionReadRequest const& request, + spanner_internal::OperationContext& operation_context) { return google::cloud::internal::LogWrapper( - [this](grpc::ClientContext& context, Options const& options, - google::spanner::v1::PartitionReadRequest const& request) { - return child_->PartitionRead(context, options, request); + [this, &operation_context]( + grpc::ClientContext& context, Options const& options, + google::spanner::v1::PartitionReadRequest const& request) { + return child_->PartitionRead(context, options, request, + operation_context); }, context, options, request, __func__, tracing_options_); } @@ -207,14 +236,16 @@ std::unique_ptr> SpannerLogging::BatchWrite( std::shared_ptr context, Options const& options, - google::spanner::v1::BatchWriteRequest const& request) { + google::spanner::v1::BatchWriteRequest const& request, + std::shared_ptr operation_context) { return google::cloud::internal::LogWrapper( - [this](std::shared_ptr context, - Options const& options, - google::spanner::v1::BatchWriteRequest const& request) + [this, operation_context = std::move(operation_context)]( + std::shared_ptr context, Options const& options, + google::spanner::v1::BatchWriteRequest const& request) -> std::unique_ptr> { - auto stream = child_->BatchWrite(std::move(context), options, request); + auto stream = child_->BatchWrite(std::move(context), options, request, + std::move(operation_context)); if (stream_logging_) { stream = std::make_unique> SpannerLogging::FetchCacheUpdate( std::shared_ptr context, Options const& options, - google::spanner::v1::FetchCacheUpdateRequest const& request) { + google::spanner::v1::FetchCacheUpdateRequest const& request, + std::shared_ptr operation_context) { return google::cloud::internal::LogWrapper( - [this](std::shared_ptr context, - Options const& options, - google::spanner::v1::FetchCacheUpdateRequest const& request) + [this, operation_context = std::move(operation_context)]( + std::shared_ptr context, Options const& options, + google::spanner::v1::FetchCacheUpdateRequest const& request) -> std::unique_ptr> { - auto stream = - child_->FetchCacheUpdate(std::move(context), options, request); + auto stream = child_->FetchCacheUpdate( + std::move(context), options, request, std::move(operation_context)); if (stream_logging_) { stream = std::make_unique context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::CreateSessionRequest const& request) { + google::spanner::v1::CreateSessionRequest const& request, + std::shared_ptr operation_context) { return google::cloud::internal::LogWrapper( - [this](google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::spanner::v1::CreateSessionRequest const& request) { + [this, operation_context = std::move(operation_context)]( + google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::spanner::v1::CreateSessionRequest const& request) { return child_->AsyncCreateSession(cq, std::move(context), - std::move(options), request); + std::move(options), request, + std::move(operation_context)); }, cq, std::move(context), std::move(options), request, __func__, tracing_options_); @@ -275,14 +310,17 @@ SpannerLogging::AsyncBatchCreateSessions( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::BatchCreateSessionsRequest const& request) { + google::spanner::v1::BatchCreateSessionsRequest const& request, + std::shared_ptr operation_context) { return google::cloud::internal::LogWrapper( - [this](google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::spanner::v1::BatchCreateSessionsRequest const& request) { + [this, operation_context = std::move(operation_context)]( + google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::spanner::v1::BatchCreateSessionsRequest const& request) { return child_->AsyncBatchCreateSessions(cq, std::move(context), - std::move(options), request); + std::move(options), request, + std::move(operation_context)); }, cq, std::move(context), std::move(options), request, __func__, tracing_options_); @@ -292,14 +330,17 @@ future SpannerLogging::AsyncDeleteSession( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::DeleteSessionRequest const& request) { + google::spanner::v1::DeleteSessionRequest const& request, + std::shared_ptr operation_context) { return google::cloud::internal::LogWrapper( - [this](google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::spanner::v1::DeleteSessionRequest const& request) { + [this, operation_context = std::move(operation_context)]( + google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::spanner::v1::DeleteSessionRequest const& request) { return child_->AsyncDeleteSession(cq, std::move(context), - std::move(options), request); + std::move(options), request, + std::move(operation_context)); }, cq, std::move(context), std::move(options), request, __func__, tracing_options_); @@ -310,14 +351,17 @@ SpannerLogging::AsyncExecuteSql( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::ExecuteSqlRequest const& request) { + google::spanner::v1::ExecuteSqlRequest const& request, + std::shared_ptr operation_context) { return google::cloud::internal::LogWrapper( - [this](google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::spanner::v1::ExecuteSqlRequest const& request) { + [this, operation_context = std::move(operation_context)]( + google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::spanner::v1::ExecuteSqlRequest const& request) { return child_->AsyncExecuteSql(cq, std::move(context), - std::move(options), request); + std::move(options), request, + std::move(operation_context)); }, cq, std::move(context), std::move(options), request, __func__, tracing_options_); diff --git a/google/cloud/spanner/internal/spanner_logging_decorator.h b/google/cloud/spanner/internal/spanner_logging_decorator.h index da409b8f523a7..5d21880e3024e 100644 --- a/google/cloud/spanner/internal/spanner_logging_decorator.h +++ b/google/cloud/spanner/internal/spanner_logging_decorator.h @@ -43,92 +43,119 @@ class SpannerLogging : public SpannerStub { StatusOr CreateSession( grpc::ClientContext& context, Options const& options, - google::spanner::v1::CreateSessionRequest const& request) override; + google::spanner::v1::CreateSessionRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr BatchCreateSessions( grpc::ClientContext& context, Options const& options, - google::spanner::v1::BatchCreateSessionsRequest const& request) override; + google::spanner::v1::BatchCreateSessionsRequest const& request, + spanner_internal::OperationContext& operation_context) override; Status DeleteSession( grpc::ClientContext& context, Options const& options, - google::spanner::v1::DeleteSessionRequest const& request) override; + google::spanner::v1::DeleteSessionRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr ExecuteSql( grpc::ClientContext& context, Options const& options, - google::spanner::v1::ExecuteSqlRequest const& request) override; + google::spanner::v1::ExecuteSqlRequest const& request, + spanner_internal::OperationContext& operation_context) override; std::unique_ptr> - ExecuteStreamingSql( - std::shared_ptr context, Options const& options, - google::spanner::v1::ExecuteSqlRequest const& request) override; + ExecuteStreamingSql(std::shared_ptr context, + Options const& options, + google::spanner::v1::ExecuteSqlRequest const& request, + std::shared_ptr + operation_context) override; StatusOr ExecuteBatchDml( grpc::ClientContext& context, Options const& options, - google::spanner::v1::ExecuteBatchDmlRequest const& request) override; + google::spanner::v1::ExecuteBatchDmlRequest const& request, + spanner_internal::OperationContext& operation_context) override; std::unique_ptr> StreamingRead(std::shared_ptr context, Options const& options, - google::spanner::v1::ReadRequest const& request) override; + google::spanner::v1::ReadRequest const& request, + std::shared_ptr + operation_context) override; StatusOr BeginTransaction( grpc::ClientContext& context, Options const& options, - google::spanner::v1::BeginTransactionRequest const& request) override; + google::spanner::v1::BeginTransactionRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr Commit( grpc::ClientContext& context, Options const& options, - google::spanner::v1::CommitRequest const& request) override; + google::spanner::v1::CommitRequest const& request, + spanner_internal::OperationContext& operation_context) override; - Status Rollback(grpc::ClientContext& context, Options const& options, - google::spanner::v1::RollbackRequest const& request) override; + Status Rollback( + grpc::ClientContext& context, Options const& options, + google::spanner::v1::RollbackRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr PartitionQuery( grpc::ClientContext& context, Options const& options, - google::spanner::v1::PartitionQueryRequest const& request) override; + google::spanner::v1::PartitionQueryRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr PartitionRead( grpc::ClientContext& context, Options const& options, - google::spanner::v1::PartitionReadRequest const& request) override; + google::spanner::v1::PartitionReadRequest const& request, + spanner_internal::OperationContext& operation_context) override; std::unique_ptr> BatchWrite(std::shared_ptr context, Options const& options, - google::spanner::v1::BatchWriteRequest const& request) override; + google::spanner::v1::BatchWriteRequest const& request, + std::shared_ptr + operation_context) override; std::unique_ptr> - FetchCacheUpdate( - std::shared_ptr context, Options const& options, - google::spanner::v1::FetchCacheUpdateRequest const& request) override; + FetchCacheUpdate(std::shared_ptr context, + Options const& options, + google::spanner::v1::FetchCacheUpdateRequest const& request, + std::shared_ptr + operation_context) override; future> AsyncCreateSession( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::CreateSessionRequest const& request) override; + google::spanner::v1::CreateSessionRequest const& request, + std::shared_ptr operation_context) + override; future> AsyncBatchCreateSessions( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::BatchCreateSessionsRequest const& request) override; + google::spanner::v1::BatchCreateSessionsRequest const& request, + std::shared_ptr operation_context) + override; future AsyncDeleteSession( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::DeleteSessionRequest const& request) override; + google::spanner::v1::DeleteSessionRequest const& request, + std::shared_ptr operation_context) + override; future> AsyncExecuteSql( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::ExecuteSqlRequest const& request) override; + google::spanner::v1::ExecuteSqlRequest const& request, + std::shared_ptr operation_context) + override; private: std::shared_ptr child_; diff --git a/google/cloud/spanner/internal/spanner_metadata_decorator.cc b/google/cloud/spanner/internal/spanner_metadata_decorator.cc index 25235d811a581..92d2dcc6e9cba 100644 --- a/google/cloud/spanner/internal/spanner_metadata_decorator.cc +++ b/google/cloud/spanner/internal/spanner_metadata_decorator.cc @@ -50,128 +50,147 @@ SpannerMetadata::SpannerMetadata( StatusOr SpannerMetadata::CreateSession( grpc::ClientContext& context, Options const& options, - google::spanner::v1::CreateSessionRequest const& request) { + google::spanner::v1::CreateSessionRequest const& request, + spanner_internal::OperationContext& operation_context) { SetMetadata( context, options, absl::StrCat("database=", internal::UrlEncode(request.database()))); - return child_->CreateSession(context, options, request); + return child_->CreateSession(context, options, request, operation_context); } StatusOr SpannerMetadata::BatchCreateSessions( grpc::ClientContext& context, Options const& options, - google::spanner::v1::BatchCreateSessionsRequest const& request) { + google::spanner::v1::BatchCreateSessionsRequest const& request, + spanner_internal::OperationContext& operation_context) { SetMetadata( context, options, absl::StrCat("database=", internal::UrlEncode(request.database()))); - return child_->BatchCreateSessions(context, options, request); + return child_->BatchCreateSessions(context, options, request, + operation_context); } Status SpannerMetadata::DeleteSession( grpc::ClientContext& context, Options const& options, - google::spanner::v1::DeleteSessionRequest const& request) { + google::spanner::v1::DeleteSessionRequest const& request, + spanner_internal::OperationContext& operation_context) { SetMetadata(context, options, absl::StrCat("name=", internal::UrlEncode(request.name()))); - return child_->DeleteSession(context, options, request); + return child_->DeleteSession(context, options, request, operation_context); } StatusOr SpannerMetadata::ExecuteSql( grpc::ClientContext& context, Options const& options, - google::spanner::v1::ExecuteSqlRequest const& request) { + google::spanner::v1::ExecuteSqlRequest const& request, + spanner_internal::OperationContext& operation_context) { SetMetadata(context, options, absl::StrCat("session=", internal::UrlEncode(request.session()))); - return child_->ExecuteSql(context, options, request); + return child_->ExecuteSql(context, options, request, operation_context); } std::unique_ptr> SpannerMetadata::ExecuteStreamingSql( std::shared_ptr context, Options const& options, - google::spanner::v1::ExecuteSqlRequest const& request) { + google::spanner::v1::ExecuteSqlRequest const& request, + std::shared_ptr operation_context) { SetMetadata(*context, options, absl::StrCat("session=", internal::UrlEncode(request.session()))); - return child_->ExecuteStreamingSql(std::move(context), options, request); + return child_->ExecuteStreamingSql(std::move(context), options, request, + std::move(operation_context)); } StatusOr SpannerMetadata::ExecuteBatchDml( grpc::ClientContext& context, Options const& options, - google::spanner::v1::ExecuteBatchDmlRequest const& request) { + google::spanner::v1::ExecuteBatchDmlRequest const& request, + spanner_internal::OperationContext& operation_context) { SetMetadata(context, options, absl::StrCat("session=", internal::UrlEncode(request.session()))); - return child_->ExecuteBatchDml(context, options, request); + return child_->ExecuteBatchDml(context, options, request, operation_context); } std::unique_ptr> SpannerMetadata::StreamingRead( std::shared_ptr context, Options const& options, - google::spanner::v1::ReadRequest const& request) { + google::spanner::v1::ReadRequest const& request, + std::shared_ptr operation_context) { SetMetadata(*context, options, absl::StrCat("session=", internal::UrlEncode(request.session()))); - return child_->StreamingRead(std::move(context), options, request); + return child_->StreamingRead(std::move(context), options, request, + std::move(operation_context)); } StatusOr SpannerMetadata::BeginTransaction( grpc::ClientContext& context, Options const& options, - google::spanner::v1::BeginTransactionRequest const& request) { + google::spanner::v1::BeginTransactionRequest const& request, + spanner_internal::OperationContext& operation_context) { SetMetadata(context, options, absl::StrCat("session=", internal::UrlEncode(request.session()))); - return child_->BeginTransaction(context, options, request); + return child_->BeginTransaction(context, options, request, operation_context); } StatusOr SpannerMetadata::Commit( grpc::ClientContext& context, Options const& options, - google::spanner::v1::CommitRequest const& request) { + google::spanner::v1::CommitRequest const& request, + spanner_internal::OperationContext& operation_context) { SetMetadata(context, options, absl::StrCat("session=", internal::UrlEncode(request.session()))); - return child_->Commit(context, options, request); + return child_->Commit(context, options, request, operation_context); } Status SpannerMetadata::Rollback( grpc::ClientContext& context, Options const& options, - google::spanner::v1::RollbackRequest const& request) { + google::spanner::v1::RollbackRequest const& request, + spanner_internal::OperationContext& operation_context) { SetMetadata(context, options, absl::StrCat("session=", internal::UrlEncode(request.session()))); - return child_->Rollback(context, options, request); + return child_->Rollback(context, options, request, operation_context); } StatusOr SpannerMetadata::PartitionQuery( grpc::ClientContext& context, Options const& options, - google::spanner::v1::PartitionQueryRequest const& request) { + google::spanner::v1::PartitionQueryRequest const& request, + spanner_internal::OperationContext& operation_context) { SetMetadata(context, options, absl::StrCat("session=", internal::UrlEncode(request.session()))); - return child_->PartitionQuery(context, options, request); + return child_->PartitionQuery(context, options, request, operation_context); } StatusOr SpannerMetadata::PartitionRead( grpc::ClientContext& context, Options const& options, - google::spanner::v1::PartitionReadRequest const& request) { + google::spanner::v1::PartitionReadRequest const& request, + spanner_internal::OperationContext& operation_context) { SetMetadata(context, options, absl::StrCat("session=", internal::UrlEncode(request.session()))); - return child_->PartitionRead(context, options, request); + return child_->PartitionRead(context, options, request, operation_context); } std::unique_ptr> SpannerMetadata::BatchWrite( std::shared_ptr context, Options const& options, - google::spanner::v1::BatchWriteRequest const& request) { + google::spanner::v1::BatchWriteRequest const& request, + std::shared_ptr operation_context) { SetMetadata(*context, options, absl::StrCat("session=", internal::UrlEncode(request.session()))); - return child_->BatchWrite(std::move(context), options, request); + return child_->BatchWrite(std::move(context), options, request, + std::move(operation_context)); } std::unique_ptr< google::cloud::internal::StreamingReadRpc> SpannerMetadata::FetchCacheUpdate( std::shared_ptr context, Options const& options, - google::spanner::v1::FetchCacheUpdateRequest const& request) { + google::spanner::v1::FetchCacheUpdateRequest const& request, + std::shared_ptr operation_context) { SetMetadata( *context, options, absl::StrCat("database=", internal::UrlEncode(request.database()))); - return child_->FetchCacheUpdate(std::move(context), options, request); + return child_->FetchCacheUpdate(std::move(context), options, request, + std::move(operation_context)); } future> @@ -179,12 +198,13 @@ SpannerMetadata::AsyncCreateSession( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::CreateSessionRequest const& request) { + google::spanner::v1::CreateSessionRequest const& request, + std::shared_ptr operation_context) { SetMetadata( *context, *options, absl::StrCat("database=", internal::UrlEncode(request.database()))); return child_->AsyncCreateSession(cq, std::move(context), std::move(options), - request); + request, std::move(operation_context)); } future> @@ -192,23 +212,26 @@ SpannerMetadata::AsyncBatchCreateSessions( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::BatchCreateSessionsRequest const& request) { + google::spanner::v1::BatchCreateSessionsRequest const& request, + std::shared_ptr operation_context) { SetMetadata( *context, *options, absl::StrCat("database=", internal::UrlEncode(request.database()))); return child_->AsyncBatchCreateSessions(cq, std::move(context), - std::move(options), request); + std::move(options), request, + std::move(operation_context)); } future SpannerMetadata::AsyncDeleteSession( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::DeleteSessionRequest const& request) { + google::spanner::v1::DeleteSessionRequest const& request, + std::shared_ptr operation_context) { SetMetadata(*context, *options, absl::StrCat("name=", internal::UrlEncode(request.name()))); return child_->AsyncDeleteSession(cq, std::move(context), std::move(options), - request); + request, std::move(operation_context)); } future> @@ -216,11 +239,12 @@ SpannerMetadata::AsyncExecuteSql( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::ExecuteSqlRequest const& request) { + google::spanner::v1::ExecuteSqlRequest const& request, + std::shared_ptr operation_context) { SetMetadata(*context, *options, absl::StrCat("session=", internal::UrlEncode(request.session()))); return child_->AsyncExecuteSql(cq, std::move(context), std::move(options), - request); + request, std::move(operation_context)); } void SpannerMetadata::SetMetadata(grpc::ClientContext& context, diff --git a/google/cloud/spanner/internal/spanner_metadata_decorator.h b/google/cloud/spanner/internal/spanner_metadata_decorator.h index 33769ae8bcb5b..ad560488330f2 100644 --- a/google/cloud/spanner/internal/spanner_metadata_decorator.h +++ b/google/cloud/spanner/internal/spanner_metadata_decorator.h @@ -43,92 +43,119 @@ class SpannerMetadata : public SpannerStub { StatusOr CreateSession( grpc::ClientContext& context, Options const& options, - google::spanner::v1::CreateSessionRequest const& request) override; + google::spanner::v1::CreateSessionRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr BatchCreateSessions( grpc::ClientContext& context, Options const& options, - google::spanner::v1::BatchCreateSessionsRequest const& request) override; + google::spanner::v1::BatchCreateSessionsRequest const& request, + spanner_internal::OperationContext& operation_context) override; Status DeleteSession( grpc::ClientContext& context, Options const& options, - google::spanner::v1::DeleteSessionRequest const& request) override; + google::spanner::v1::DeleteSessionRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr ExecuteSql( grpc::ClientContext& context, Options const& options, - google::spanner::v1::ExecuteSqlRequest const& request) override; + google::spanner::v1::ExecuteSqlRequest const& request, + spanner_internal::OperationContext& operation_context) override; std::unique_ptr> - ExecuteStreamingSql( - std::shared_ptr context, Options const& options, - google::spanner::v1::ExecuteSqlRequest const& request) override; + ExecuteStreamingSql(std::shared_ptr context, + Options const& options, + google::spanner::v1::ExecuteSqlRequest const& request, + std::shared_ptr + operation_context) override; StatusOr ExecuteBatchDml( grpc::ClientContext& context, Options const& options, - google::spanner::v1::ExecuteBatchDmlRequest const& request) override; + google::spanner::v1::ExecuteBatchDmlRequest const& request, + spanner_internal::OperationContext& operation_context) override; std::unique_ptr> StreamingRead(std::shared_ptr context, Options const& options, - google::spanner::v1::ReadRequest const& request) override; + google::spanner::v1::ReadRequest const& request, + std::shared_ptr + operation_context) override; StatusOr BeginTransaction( grpc::ClientContext& context, Options const& options, - google::spanner::v1::BeginTransactionRequest const& request) override; + google::spanner::v1::BeginTransactionRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr Commit( grpc::ClientContext& context, Options const& options, - google::spanner::v1::CommitRequest const& request) override; + google::spanner::v1::CommitRequest const& request, + spanner_internal::OperationContext& operation_context) override; - Status Rollback(grpc::ClientContext& context, Options const& options, - google::spanner::v1::RollbackRequest const& request) override; + Status Rollback( + grpc::ClientContext& context, Options const& options, + google::spanner::v1::RollbackRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr PartitionQuery( grpc::ClientContext& context, Options const& options, - google::spanner::v1::PartitionQueryRequest const& request) override; + google::spanner::v1::PartitionQueryRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr PartitionRead( grpc::ClientContext& context, Options const& options, - google::spanner::v1::PartitionReadRequest const& request) override; + google::spanner::v1::PartitionReadRequest const& request, + spanner_internal::OperationContext& operation_context) override; std::unique_ptr> BatchWrite(std::shared_ptr context, Options const& options, - google::spanner::v1::BatchWriteRequest const& request) override; + google::spanner::v1::BatchWriteRequest const& request, + std::shared_ptr + operation_context) override; std::unique_ptr> - FetchCacheUpdate( - std::shared_ptr context, Options const& options, - google::spanner::v1::FetchCacheUpdateRequest const& request) override; + FetchCacheUpdate(std::shared_ptr context, + Options const& options, + google::spanner::v1::FetchCacheUpdateRequest const& request, + std::shared_ptr + operation_context) override; future> AsyncCreateSession( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::CreateSessionRequest const& request) override; + google::spanner::v1::CreateSessionRequest const& request, + std::shared_ptr operation_context) + override; future> AsyncBatchCreateSessions( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::BatchCreateSessionsRequest const& request) override; + google::spanner::v1::BatchCreateSessionsRequest const& request, + std::shared_ptr operation_context) + override; future AsyncDeleteSession( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::DeleteSessionRequest const& request) override; + google::spanner::v1::DeleteSessionRequest const& request, + std::shared_ptr operation_context) + override; future> AsyncExecuteSql( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::ExecuteSqlRequest const& request) override; + google::spanner::v1::ExecuteSqlRequest const& request, + std::shared_ptr operation_context) + override; private: void SetMetadata(grpc::ClientContext& context, Options const& options, diff --git a/google/cloud/spanner/internal/spanner_operation_context_factory.cc b/google/cloud/spanner/internal/spanner_operation_context_factory.cc new file mode 100644 index 0000000000000..a1f11b7bb4512 --- /dev/null +++ b/google/cloud/spanner/internal/spanner_operation_context_factory.cc @@ -0,0 +1,160 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "google/cloud/spanner/internal/spanner_operation_context_factory.h" +#include "google/cloud/spanner/internal/spanner_request_id.h" + +namespace google { +namespace cloud { +namespace spanner_internal { +GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN + +DefaultSpannerOperationContextFactory::DefaultSpannerOperationContextFactory( + std::uint64_t client_id, + std::shared_ptr process_random_id) + : client_id_(client_id), + process_random_id_(std::move(process_random_id)), + static_prefix_(std::make_shared( + FormatSpannerRequestStaticPrefix(1, *process_random_id_, client_id_))) +#ifndef _WIN32 + , + cached_pid_(getpid()) +#endif +{ +} + +std::shared_ptr +DefaultSpannerOperationContextFactory::StaticPrefix() { +#ifndef _WIN32 + pid_t const current_pid = getpid(); + std::scoped_lock lock(mu_); + if (current_pid != cached_pid_) { + process_random_id_ = std::make_shared(ProcessRandomId()); + static_prefix_ = std::make_shared( + FormatSpannerRequestStaticPrefix(1, *process_random_id_, client_id_)); + cached_pid_ = current_pid; + } + return static_prefix_; +#else + return static_prefix_; +#endif +} + +std::uint64_t DefaultSpannerOperationContextFactory::NextUserRequestIndex() { + return next_user_request_index_.fetch_add(1, std::memory_order_relaxed); +} + +std::uint64_t +DefaultSpannerOperationContextFactory::NextBackgroundRequestIndex() { + return next_background_request_index_.fetch_add(1, std::memory_order_relaxed); +} + +OperationContext DefaultSpannerOperationContextFactory::CreateSession() { + return OperationContext(StaticPrefix(), NextUserRequestIndex(), + "CreateSession"); +} + +OperationContext DefaultSpannerOperationContextFactory::BatchCreateSessions() { + return OperationContext(StaticPrefix(), NextUserRequestIndex(), + "BatchCreateSessions"); +} + +OperationContext DefaultSpannerOperationContextFactory::GetSession() { + return OperationContext(StaticPrefix(), NextUserRequestIndex(), "GetSession"); +} + +OperationContext DefaultSpannerOperationContextFactory::ListSessions() { + return OperationContext(StaticPrefix(), NextUserRequestIndex(), + "ListSessions"); +} + +OperationContext DefaultSpannerOperationContextFactory::DeleteSession() { + return OperationContext(StaticPrefix(), NextUserRequestIndex(), + "DeleteSession"); +} + +OperationContext DefaultSpannerOperationContextFactory::ExecuteSql() { + return OperationContext(StaticPrefix(), NextUserRequestIndex(), "ExecuteSql"); +} + +OperationContext DefaultSpannerOperationContextFactory::ExecuteStreamingSql() { + return OperationContext(StaticPrefix(), NextUserRequestIndex(), + "ExecuteStreamingSql"); +} + +OperationContext DefaultSpannerOperationContextFactory::ExecuteBatchDml() { + return OperationContext(StaticPrefix(), NextUserRequestIndex(), + "ExecuteBatchDml"); +} + +OperationContext DefaultSpannerOperationContextFactory::StreamingRead() { + return OperationContext(StaticPrefix(), NextUserRequestIndex(), + "StreamingRead"); +} + +OperationContext DefaultSpannerOperationContextFactory::BeginTransaction() { + return OperationContext(StaticPrefix(), NextUserRequestIndex(), + "BeginTransaction"); +} + +OperationContext DefaultSpannerOperationContextFactory::Commit() { + return OperationContext(StaticPrefix(), NextUserRequestIndex(), "Commit"); +} + +OperationContext DefaultSpannerOperationContextFactory::Rollback() { + return OperationContext(StaticPrefix(), NextUserRequestIndex(), "Rollback"); +} + +OperationContext DefaultSpannerOperationContextFactory::PartitionQuery() { + return OperationContext(StaticPrefix(), NextUserRequestIndex(), + "PartitionQuery"); +} + +OperationContext DefaultSpannerOperationContextFactory::PartitionRead() { + return OperationContext(StaticPrefix(), NextUserRequestIndex(), + "PartitionRead"); +} + +OperationContext DefaultSpannerOperationContextFactory::BatchWrite() { + return OperationContext(StaticPrefix(), NextUserRequestIndex(), "BatchWrite"); +} + +OperationContext +DefaultSpannerOperationContextFactory::BackgroundCreateSession() { + return OperationContext(StaticPrefix(), NextBackgroundRequestIndex(), + "BackgroundCreateSession"); +} + +OperationContext +DefaultSpannerOperationContextFactory::BackgroundBatchCreateSessions() { + return OperationContext(StaticPrefix(), NextBackgroundRequestIndex(), + "BackgroundBatchCreateSessions"); +} + +OperationContext +DefaultSpannerOperationContextFactory::BackgroundDeleteSession() { + return OperationContext(StaticPrefix(), NextBackgroundRequestIndex(), + "BackgroundDeleteSession"); +} + +OperationContext +DefaultSpannerOperationContextFactory::BackgroundRefreshSession() { + return OperationContext(StaticPrefix(), NextBackgroundRequestIndex(), + "BackgroundRefreshSession"); +} + +GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END +} // namespace spanner_internal +} // namespace cloud +} // namespace google diff --git a/google/cloud/spanner/internal/spanner_operation_context_factory.h b/google/cloud/spanner/internal/spanner_operation_context_factory.h new file mode 100644 index 0000000000000..253c7b9c4a54b --- /dev/null +++ b/google/cloud/spanner/internal/spanner_operation_context_factory.h @@ -0,0 +1,113 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_SPANNER_INTERNAL_SPANNER_OPERATION_CONTEXT_FACTORY_H +#define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_SPANNER_INTERNAL_SPANNER_OPERATION_CONTEXT_FACTORY_H + +#include "google/cloud/spanner/internal/operation_context.h" +#include "google/cloud/spanner/version.h" +#include +#include +#include +#include +#include +#ifndef _WIN32 +#include +#endif + +namespace google { +namespace cloud { +namespace spanner_internal { +GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN + +class SpannerOperationContextFactory { + public: + virtual ~SpannerOperationContextFactory() = default; + + // User-facing RPC context factories (use next_user_request_index_) + virtual OperationContext CreateSession() = 0; + virtual OperationContext BatchCreateSessions() = 0; + virtual OperationContext GetSession() = 0; + virtual OperationContext ListSessions() = 0; + virtual OperationContext DeleteSession() = 0; + virtual OperationContext ExecuteSql() = 0; + virtual OperationContext ExecuteStreamingSql() = 0; + virtual OperationContext ExecuteBatchDml() = 0; + virtual OperationContext StreamingRead() = 0; + virtual OperationContext BeginTransaction() = 0; + virtual OperationContext Commit() = 0; + virtual OperationContext Rollback() = 0; + virtual OperationContext PartitionQuery() = 0; + virtual OperationContext PartitionRead() = 0; + virtual OperationContext BatchWrite() = 0; + + // Background maintenance context factories (use + // next_background_request_index_) + virtual OperationContext BackgroundCreateSession() = 0; + virtual OperationContext BackgroundBatchCreateSessions() = 0; + virtual OperationContext BackgroundDeleteSession() = 0; + virtual OperationContext BackgroundRefreshSession() = 0; +}; + +class DefaultSpannerOperationContextFactory + : public SpannerOperationContextFactory { + public: + // Explicit constructor without default arguments per internal guidelines. + DefaultSpannerOperationContextFactory( + std::uint64_t client_id, + std::shared_ptr process_random_id); + + OperationContext CreateSession() override; + OperationContext BatchCreateSessions() override; + OperationContext GetSession() override; + OperationContext ListSessions() override; + OperationContext DeleteSession() override; + OperationContext ExecuteSql() override; + OperationContext ExecuteStreamingSql() override; + OperationContext ExecuteBatchDml() override; + OperationContext StreamingRead() override; + OperationContext BeginTransaction() override; + OperationContext Commit() override; + OperationContext Rollback() override; + OperationContext PartitionQuery() override; + OperationContext PartitionRead() override; + OperationContext BatchWrite() override; + + OperationContext BackgroundCreateSession() override; + OperationContext BackgroundBatchCreateSessions() override; + OperationContext BackgroundDeleteSession() override; + OperationContext BackgroundRefreshSession() override; + + private: + std::shared_ptr StaticPrefix(); + std::uint64_t NextUserRequestIndex(); + std::uint64_t NextBackgroundRequestIndex(); + + std::uint64_t client_id_; + std::shared_ptr process_random_id_; + std::shared_ptr static_prefix_; +#ifndef _WIN32 + pid_t cached_pid_; +#endif + std::mutex mu_; + std::atomic next_user_request_index_{1}; + std::atomic next_background_request_index_{1}; +}; + +GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END +} // namespace spanner_internal +} // namespace cloud +} // namespace google + +#endif // GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_SPANNER_INTERNAL_SPANNER_OPERATION_CONTEXT_FACTORY_H diff --git a/google/cloud/spanner/internal/spanner_operation_context_factory_test.cc b/google/cloud/spanner/internal/spanner_operation_context_factory_test.cc new file mode 100644 index 0000000000000..43679ee40b247 --- /dev/null +++ b/google/cloud/spanner/internal/spanner_operation_context_factory_test.cc @@ -0,0 +1,156 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "google/cloud/spanner/internal/spanner_operation_context_factory.h" +#include +#include +#include +#include +#include +#ifndef _WIN32 +#include +#include +#endif + +namespace google { +namespace cloud { +namespace spanner_internal { +GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN +namespace { + +using ::testing::Eq; +using ::testing::Ge; +using ::testing::Ne; +using ::testing::Optional; +using ::testing::StrEq; + +MATCHER_P(MatchesStdRegex, pattern, "") { + if (std::regex_match(arg, std::regex(pattern))) { + return true; + } + *result_listener << "which does not match regex \"" << pattern << "\""; + return false; +} + +TEST(SpannerOperationContextFactoryTest, CounterIsolationAndRpcNames) { + auto process_random_id = + std::make_shared("0123456789abcdef"); + DefaultSpannerOperationContextFactory factory(100, process_random_id); + + // User RPCs + auto ctx_create_session = factory.CreateSession(); + EXPECT_THAT(ctx_create_session.request_index(), Eq(1ULL)); + EXPECT_THAT(ctx_create_session.rpc_name(), StrEq("CreateSession")); + + auto ctx_exec_sql = factory.ExecuteSql(); + EXPECT_THAT(ctx_exec_sql.request_index(), Eq(2ULL)); + EXPECT_THAT(ctx_exec_sql.rpc_name(), StrEq("ExecuteSql")); + + auto ctx_commit = factory.Commit(); + EXPECT_THAT(ctx_commit.request_index(), Eq(3ULL)); + EXPECT_THAT(ctx_commit.rpc_name(), StrEq("Commit")); + + // Background RPCs + auto bg_create_session = factory.BackgroundCreateSession(); + EXPECT_THAT(bg_create_session.request_index(), Eq(1ULL)); + EXPECT_THAT(bg_create_session.rpc_name(), StrEq("BackgroundCreateSession")); + + auto bg_batch_create = factory.BackgroundBatchCreateSessions(); + EXPECT_THAT(bg_batch_create.request_index(), Eq(2ULL)); + EXPECT_THAT(bg_batch_create.rpc_name(), + StrEq("BackgroundBatchCreateSessions")); + + auto bg_delete = factory.BackgroundDeleteSession(); + EXPECT_THAT(bg_delete.request_index(), Eq(3ULL)); + EXPECT_THAT(bg_delete.rpc_name(), StrEq("BackgroundDeleteSession")); + + auto bg_refresh = factory.BackgroundRefreshSession(); + EXPECT_THAT(bg_refresh.request_index(), Eq(4ULL)); + EXPECT_THAT(bg_refresh.rpc_name(), StrEq("BackgroundRefreshSession")); + + // Next user RPC resumes user counter without interference from background + // counter + auto ctx_rollback = factory.Rollback(); + EXPECT_THAT(ctx_rollback.request_index(), Eq(4ULL)); + EXPECT_THAT(ctx_rollback.rpc_name(), StrEq("Rollback")); +} + +TEST(SpannerOperationContextFactoryTest, HeaderGenerationFromFactory) { + auto process_random_id = + std::make_shared("0123456789abcdef"); + DefaultSpannerOperationContextFactory factory(5, process_random_id); + + auto op_context = factory.ExecuteSql(); + op_context.BindChannel(2); + + grpc::ClientContext client_context; + op_context.PreCall(client_context); + + EXPECT_THAT(op_context.RequestId(), + Optional(Eq("1.0123456789abcdef.5.2.1.1"))); +} + +#ifndef _WIN32 +TEST(SpannerOperationContextFactoryTest, + ForkDetectsAndRegeneratesStaticPrefix) { + auto process_random_id = + std::make_shared("0123456789abcdef"); + DefaultSpannerOperationContextFactory factory(7, process_random_id); + + auto parent_ctx = factory.ExecuteSql(); + grpc::ClientContext parent_client_context; + parent_ctx.PreCall(parent_client_context); + EXPECT_THAT(parent_ctx.RequestId(), + Optional(Eq("1.0123456789abcdef.7.0.1.1"))); + + int pipe_fds[2]; + ASSERT_THAT(pipe(pipe_fds), Eq(0)); + + pid_t const pid = fork(); + ASSERT_THAT(pid, Ge(0)); + + if (pid == 0) { + close(pipe_fds[0]); + auto child_ctx = factory.ExecuteSql(); + grpc::ClientContext child_client_context; + child_ctx.PreCall(child_client_context); + auto const child_request_id = child_ctx.RequestId().value_or(""); + write(pipe_fds[1], child_request_id.data(), child_request_id.size()); + close(pipe_fds[1]); + _exit(0); + } + + close(pipe_fds[1]); + char buffer[128] = {0}; + ssize_t const bytes_read = read(pipe_fds[0], buffer, sizeof(buffer) - 1); + close(pipe_fds[0]); + + int status = 0; + waitpid(pid, &status, 0); + ASSERT_THAT(WIFEXITED(status), Eq(true)); + ASSERT_THAT(WEXITSTATUS(status), Eq(0)); + + std::string const child_request_id(buffer, + static_cast(bytes_read)); + EXPECT_THAT(child_request_id, + MatchesStdRegex(R"(^1\.[0-9a-f]{16}\.7\.0\.2\.1$)")); + EXPECT_THAT(child_request_id, Ne("1.0123456789abcdef.7.0.2.1")); +} +#endif + +} // namespace +GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END +} // namespace spanner_internal +} // namespace cloud +} // namespace google diff --git a/google/cloud/spanner/internal/spanner_request_id.cc b/google/cloud/spanner/internal/spanner_request_id.cc new file mode 100644 index 0000000000000..7a2f845e0f7a5 --- /dev/null +++ b/google/cloud/spanner/internal/spanner_request_id.cc @@ -0,0 +1,87 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "google/cloud/spanner/internal/spanner_request_id.h" +#include "google/cloud/internal/random.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/str_format.h" +#include +#include +#include +#ifndef _WIN32 +#include +#endif + +namespace google { +namespace cloud { +namespace spanner_internal { +GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN + +std::string ProcessRandomId() { +#ifndef _WIN32 + static std::mutex mu; + static pid_t pid = 0; + static std::string random_id; + pid_t const current_pid = getpid(); + std::scoped_lock lock(mu); + if (pid != current_pid) { + auto generator = google::cloud::internal::MakeDefaultPRNG(); + std::uniform_int_distribution dist; + random_id = absl::StrFormat("%016x", dist(generator)); + pid = current_pid; + } + return random_id; +#else + static std::string const random_id = [] { + auto generator = google::cloud::internal::MakeDefaultPRNG(); + std::uniform_int_distribution dist; + return absl::StrFormat("%016x", dist(generator)); + }(); + return random_id; +#endif +} + +std::uint64_t NextClientId() { + static std::atomic counter{0}; + return ++counter; +} + +std::string FormatSpannerRequestStaticPrefix(std::uint32_t version, + std::string_view process_random_id, + std::uint64_t client_id) { + return absl::StrCat(version, ".", process_random_id, ".", client_id, "."); +} + +std::string FormatSpannerRequestId(std::string_view static_prefix, + std::uint32_t channel_id, + std::uint64_t request_index, + std::uint32_t attempt_index) { + return absl::StrCat(static_prefix, channel_id, ".", request_index, ".", + attempt_index); +} + +std::string FormatSpannerRequestId(std::uint32_t version, + std::string_view process_random_id, + std::uint64_t client_id, + std::uint32_t channel_id, + std::uint64_t request_index, + std::uint32_t attempt_index) { + return absl::StrCat(version, ".", process_random_id, ".", client_id, ".", + channel_id, ".", request_index, ".", attempt_index); +} + +GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END +} // namespace spanner_internal +} // namespace cloud +} // namespace google diff --git a/google/cloud/spanner/internal/spanner_request_id.h b/google/cloud/spanner/internal/spanner_request_id.h new file mode 100644 index 0000000000000..e8fa29a16b3c1 --- /dev/null +++ b/google/cloud/spanner/internal/spanner_request_id.h @@ -0,0 +1,63 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_SPANNER_INTERNAL_SPANNER_REQUEST_ID_H +#define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_SPANNER_INTERNAL_SPANNER_REQUEST_ID_H + +#include "google/cloud/spanner/version.h" +#include +#include +#include + +namespace google { +namespace cloud { +namespace spanner_internal { +GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN + +/// Generates a 16-character zero-padded lowercase hex random ID from a 64-bit +/// random integer. On POSIX systems, re-generates the ID if a process fork is +/// detected. On non-POSIX systems (e.g. Windows), generates once per process +/// lifecycle. +std::string ProcessRandomId(); + +/// Returns the next sequential process-wide client ID (thread-safe). +std::uint64_t NextClientId(); + +/// Formats the static 3-field prefix: "..." +std::string FormatSpannerRequestStaticPrefix(std::uint32_t version, + std::string_view process_random_id, + std::uint64_t client_id); + +/// Formats full request ID using the cached static prefix: +/// ".." +std::string FormatSpannerRequestId(std::string_view static_prefix, + std::uint32_t channel_id, + std::uint64_t request_index, + std::uint32_t attempt_index); + +/// Convenience overload formatting all 6 fields: +/// "....." +std::string FormatSpannerRequestId(std::uint32_t version, + std::string_view process_random_id, + std::uint64_t client_id, + std::uint32_t channel_id, + std::uint64_t request_index, + std::uint32_t attempt_index); + +GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END +} // namespace spanner_internal +} // namespace cloud +} // namespace google + +#endif // GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_SPANNER_INTERNAL_SPANNER_REQUEST_ID_H diff --git a/google/cloud/spanner/internal/spanner_request_id_test.cc b/google/cloud/spanner/internal/spanner_request_id_test.cc new file mode 100644 index 0000000000000..52725a56c42c9 --- /dev/null +++ b/google/cloud/spanner/internal/spanner_request_id_test.cc @@ -0,0 +1,120 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "google/cloud/spanner/internal/spanner_request_id.h" +#include +#include +#include +#ifndef _WIN32 +#include +#include +#endif + +namespace google { +namespace cloud { +namespace spanner_internal { +GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN +namespace { + +using ::testing::Eq; +using ::testing::Ge; +using ::testing::Gt; +using ::testing::Ne; + +MATCHER_P(MatchesStdRegex, pattern, "") { + if (std::regex_match(arg, std::regex(pattern))) { + return true; + } + *result_listener << "which does not match regex \"" << pattern << "\""; + return false; +} + +TEST(SpannerRequestIdTest, ProcessRandomIdFormat) { + std::string const id1 = ProcessRandomId(); + EXPECT_THAT(id1, MatchesStdRegex("^[0-9a-f]{16}$")); + std::string const id2 = ProcessRandomId(); + EXPECT_THAT(id2, Eq(id1)); +} + +TEST(SpannerRequestIdTest, NextClientIdMonotonic) { + std::uint64_t const c1 = NextClientId(); + std::uint64_t const c2 = NextClientId(); + std::uint64_t const c3 = NextClientId(); + EXPECT_THAT(c1, Gt(0ULL)); + EXPECT_THAT(c2, Eq(c1 + 1)); + EXPECT_THAT(c3, Eq(c2 + 1)); +} + +TEST(SpannerRequestIdTest, FormatSpannerRequestStaticPrefix) { + std::string const prefix = + FormatSpannerRequestStaticPrefix(1, "0123456789abcdef", 42); + EXPECT_THAT(prefix, Eq("1.0123456789abcdef.42.")); +} + +TEST(SpannerRequestIdTest, FormatSpannerRequestIdWithPrefix) { + std::string const prefix = "1.0123456789abcdef.42."; + std::string const request_id = FormatSpannerRequestId(prefix, 1, 100, 2); + EXPECT_THAT(request_id, Eq("1.0123456789abcdef.42.1.100.2")); +} + +TEST(SpannerRequestIdTest, FormatSpannerRequestIdDirect) { + std::string const request_id = + FormatSpannerRequestId(1, "0123456789abcdef", 42, 3, 200, 1); + EXPECT_THAT(request_id, Eq("1.0123456789abcdef.42.3.200.1")); +} + +#ifndef _WIN32 +TEST(SpannerRequestIdTest, ProcessRandomIdForkRegeneration) { + // Ensure the parent has already initialized ProcessRandomId + std::string const parent_id = ProcessRandomId(); + ASSERT_THAT(parent_id, MatchesStdRegex("^[0-9a-f]{16}$")); + + int pipe_fds[2]; + ASSERT_THAT(pipe(pipe_fds), Eq(0)); + + pid_t const pid = fork(); + ASSERT_THAT(pid, Ge(0)); + + if (pid == 0) { + // Child process: read ProcessRandomId and write to pipe + close(pipe_fds[0]); + std::string const child_id = ProcessRandomId(); + write(pipe_fds[1], child_id.data(), child_id.size()); + close(pipe_fds[1]); + _exit(0); + } + + // Parent process: read child's ID from pipe + close(pipe_fds[1]); + char buffer[32] = {0}; + ssize_t const bytes_read = read(pipe_fds[0], buffer, sizeof(buffer) - 1); + close(pipe_fds[0]); + + int status = 0; + waitpid(pid, &status, 0); + ASSERT_THAT(WIFEXITED(status), Eq(true)); + ASSERT_THAT(WEXITSTATUS(status), Eq(0)); + + ASSERT_THAT(bytes_read, Eq(16)); + std::string const child_id(buffer, static_cast(bytes_read)); + EXPECT_THAT(child_id, MatchesStdRegex("^[0-9a-f]{16}$")); + EXPECT_THAT(child_id, Ne(parent_id)); +} +#endif + +} // namespace +GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END +} // namespace spanner_internal +} // namespace cloud +} // namespace google diff --git a/google/cloud/spanner/internal/spanner_stub.cc b/google/cloud/spanner/internal/spanner_stub.cc index a7b47a87b4bd5..37ec0b4177af8 100644 --- a/google/cloud/spanner/internal/spanner_stub.cc +++ b/google/cloud/spanner/internal/spanner_stub.cc @@ -35,7 +35,8 @@ SpannerStub::~SpannerStub() = default; StatusOr DefaultSpannerStub::CreateSession( grpc::ClientContext& context, Options const&, - google::spanner::v1::CreateSessionRequest const& request) { + google::spanner::v1::CreateSessionRequest const& request, + spanner_internal::OperationContext&) { google::spanner::v1::Session response; auto status = grpc_stub_->CreateSession(&context, request, &response); if (!status.ok()) { @@ -47,7 +48,8 @@ StatusOr DefaultSpannerStub::CreateSession( StatusOr DefaultSpannerStub::BatchCreateSessions( grpc::ClientContext& context, Options const&, - google::spanner::v1::BatchCreateSessionsRequest const& request) { + google::spanner::v1::BatchCreateSessionsRequest const& request, + spanner_internal::OperationContext&) { google::spanner::v1::BatchCreateSessionsResponse response; auto status = grpc_stub_->BatchCreateSessions(&context, request, &response); if (!status.ok()) { @@ -58,7 +60,8 @@ DefaultSpannerStub::BatchCreateSessions( Status DefaultSpannerStub::DeleteSession( grpc::ClientContext& context, Options const&, - google::spanner::v1::DeleteSessionRequest const& request) { + google::spanner::v1::DeleteSessionRequest const& request, + spanner_internal::OperationContext&) { google::protobuf::Empty response; auto status = grpc_stub_->DeleteSession(&context, request, &response); if (!status.ok()) { @@ -69,7 +72,8 @@ Status DefaultSpannerStub::DeleteSession( StatusOr DefaultSpannerStub::ExecuteSql( grpc::ClientContext& context, Options const&, - google::spanner::v1::ExecuteSqlRequest const& request) { + google::spanner::v1::ExecuteSqlRequest const& request, + spanner_internal::OperationContext&) { google::spanner::v1::ResultSet response; auto status = grpc_stub_->ExecuteSql(&context, request, &response); if (!status.ok()) { @@ -82,7 +86,8 @@ std::unique_ptr> DefaultSpannerStub::ExecuteStreamingSql( std::shared_ptr context, Options const&, - google::spanner::v1::ExecuteSqlRequest const& request) { + google::spanner::v1::ExecuteSqlRequest const& request, + std::shared_ptr) { auto stream = grpc_stub_->ExecuteStreamingSql(context.get(), request); return std::make_unique>(std::move(context), @@ -92,7 +97,8 @@ DefaultSpannerStub::ExecuteStreamingSql( StatusOr DefaultSpannerStub::ExecuteBatchDml( grpc::ClientContext& context, Options const&, - google::spanner::v1::ExecuteBatchDmlRequest const& request) { + google::spanner::v1::ExecuteBatchDmlRequest const& request, + spanner_internal::OperationContext&) { google::spanner::v1::ExecuteBatchDmlResponse response; auto status = grpc_stub_->ExecuteBatchDml(&context, request, &response); if (!status.ok()) { @@ -105,7 +111,8 @@ std::unique_ptr> DefaultSpannerStub::StreamingRead( std::shared_ptr context, Options const&, - google::spanner::v1::ReadRequest const& request) { + google::spanner::v1::ReadRequest const& request, + std::shared_ptr) { auto stream = grpc_stub_->StreamingRead(context.get(), request); return std::make_unique>(std::move(context), @@ -114,7 +121,8 @@ DefaultSpannerStub::StreamingRead( StatusOr DefaultSpannerStub::BeginTransaction( grpc::ClientContext& context, Options const&, - google::spanner::v1::BeginTransactionRequest const& request) { + google::spanner::v1::BeginTransactionRequest const& request, + spanner_internal::OperationContext&) { google::spanner::v1::Transaction response; auto status = grpc_stub_->BeginTransaction(&context, request, &response); if (!status.ok()) { @@ -125,7 +133,8 @@ StatusOr DefaultSpannerStub::BeginTransaction( StatusOr DefaultSpannerStub::Commit( grpc::ClientContext& context, Options const&, - google::spanner::v1::CommitRequest const& request) { + google::spanner::v1::CommitRequest const& request, + spanner_internal::OperationContext&) { google::spanner::v1::CommitResponse response; auto status = grpc_stub_->Commit(&context, request, &response); if (!status.ok()) { @@ -136,7 +145,8 @@ StatusOr DefaultSpannerStub::Commit( Status DefaultSpannerStub::Rollback( grpc::ClientContext& context, Options const&, - google::spanner::v1::RollbackRequest const& request) { + google::spanner::v1::RollbackRequest const& request, + spanner_internal::OperationContext&) { google::protobuf::Empty response; auto status = grpc_stub_->Rollback(&context, request, &response); if (!status.ok()) { @@ -148,7 +158,8 @@ Status DefaultSpannerStub::Rollback( StatusOr DefaultSpannerStub::PartitionQuery( grpc::ClientContext& context, Options const&, - google::spanner::v1::PartitionQueryRequest const& request) { + google::spanner::v1::PartitionQueryRequest const& request, + spanner_internal::OperationContext&) { google::spanner::v1::PartitionResponse response; auto status = grpc_stub_->PartitionQuery(&context, request, &response); if (!status.ok()) { @@ -160,7 +171,8 @@ DefaultSpannerStub::PartitionQuery( StatusOr DefaultSpannerStub::PartitionRead( grpc::ClientContext& context, Options const&, - google::spanner::v1::PartitionReadRequest const& request) { + google::spanner::v1::PartitionReadRequest const& request, + spanner_internal::OperationContext&) { google::spanner::v1::PartitionResponse response; auto status = grpc_stub_->PartitionRead(&context, request, &response); if (!status.ok()) { @@ -173,7 +185,8 @@ std::unique_ptr> DefaultSpannerStub::BatchWrite( std::shared_ptr context, Options const&, - google::spanner::v1::BatchWriteRequest const& request) { + google::spanner::v1::BatchWriteRequest const& request, + std::shared_ptr) { auto stream = grpc_stub_->BatchWrite(context.get(), request); return std::make_unique>(std::move(context), @@ -184,7 +197,8 @@ std::unique_ptr< google::cloud::internal::StreamingReadRpc> DefaultSpannerStub::FetchCacheUpdate( std::shared_ptr context, Options const&, - google::spanner::v1::FetchCacheUpdateRequest const& request) { + google::spanner::v1::FetchCacheUpdateRequest const& request, + std::shared_ptr) { auto stream = grpc_stub_->FetchCacheUpdate(context.get(), request); return std::make_unique>(std::move(context), std::move(stream)); @@ -196,7 +210,8 @@ DefaultSpannerStub::AsyncCreateSession( std::shared_ptr context, // NOLINTNEXTLINE(performance-unnecessary-value-param) google::cloud::internal::ImmutableOptions, - google::spanner::v1::CreateSessionRequest const& request) { + google::spanner::v1::CreateSessionRequest const& request, + std::shared_ptr) { return internal::MakeUnaryRpcImpl( cq, @@ -214,7 +229,8 @@ DefaultSpannerStub::AsyncBatchCreateSessions( std::shared_ptr context, // NOLINTNEXTLINE(performance-unnecessary-value-param) google::cloud::internal::ImmutableOptions, - google::spanner::v1::BatchCreateSessionsRequest const& request) { + google::spanner::v1::BatchCreateSessionsRequest const& request, + std::shared_ptr) { return internal::MakeUnaryRpcImpl< google::spanner::v1::BatchCreateSessionsRequest, google::spanner::v1::BatchCreateSessionsResponse>( @@ -232,7 +248,8 @@ future DefaultSpannerStub::AsyncDeleteSession( std::shared_ptr context, // NOLINTNEXTLINE(performance-unnecessary-value-param) google::cloud::internal::ImmutableOptions, - google::spanner::v1::DeleteSessionRequest const& request) { + google::spanner::v1::DeleteSessionRequest const& request, + std::shared_ptr) { return internal::MakeUnaryRpcImpl( cq, @@ -253,7 +270,8 @@ DefaultSpannerStub::AsyncExecuteSql( std::shared_ptr context, // NOLINTNEXTLINE(performance-unnecessary-value-param) google::cloud::internal::ImmutableOptions, - google::spanner::v1::ExecuteSqlRequest const& request) { + google::spanner::v1::ExecuteSqlRequest const& request, + std::shared_ptr) { return internal::MakeUnaryRpcImpl( cq, diff --git a/google/cloud/spanner/internal/spanner_stub.h b/google/cloud/spanner/internal/spanner_stub.h index 7908a18a352c6..c63dad2482e47 100644 --- a/google/cloud/spanner/internal/spanner_stub.h +++ b/google/cloud/spanner/internal/spanner_stub.h @@ -19,6 +19,7 @@ #ifndef GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_SPANNER_INTERNAL_SPANNER_STUB_H #define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_SPANNER_INTERNAL_SPANNER_STUB_H +#include "google/cloud/spanner/internal/operation_context.h" #include "google/cloud/completion_queue.h" #include "google/cloud/future.h" #include "google/cloud/internal/streaming_read_rpc.h" @@ -43,94 +44,119 @@ class SpannerStub { virtual StatusOr CreateSession( grpc::ClientContext& context, Options const& options, - google::spanner::v1::CreateSessionRequest const& request) = 0; + google::spanner::v1::CreateSessionRequest const& request, + spanner_internal::OperationContext& operation_context) = 0; virtual StatusOr BatchCreateSessions( grpc::ClientContext& context, Options const& options, - google::spanner::v1::BatchCreateSessionsRequest const& request) = 0; + google::spanner::v1::BatchCreateSessionsRequest const& request, + spanner_internal::OperationContext& operation_context) = 0; virtual Status DeleteSession( grpc::ClientContext& context, Options const& options, - google::spanner::v1::DeleteSessionRequest const& request) = 0; + google::spanner::v1::DeleteSessionRequest const& request, + spanner_internal::OperationContext& operation_context) = 0; virtual StatusOr ExecuteSql( grpc::ClientContext& context, Options const& options, - google::spanner::v1::ExecuteSqlRequest const& request) = 0; + google::spanner::v1::ExecuteSqlRequest const& request, + spanner_internal::OperationContext& operation_context) = 0; virtual std::unique_ptr> - ExecuteStreamingSql( - std::shared_ptr context, Options const& options, - google::spanner::v1::ExecuteSqlRequest const& request) = 0; + ExecuteStreamingSql(std::shared_ptr context, + Options const& options, + google::spanner::v1::ExecuteSqlRequest const& request, + std::shared_ptr + operation_context) = 0; virtual StatusOr - ExecuteBatchDml( - grpc::ClientContext& context, Options const& options, - google::spanner::v1::ExecuteBatchDmlRequest const& request) = 0; + ExecuteBatchDml(grpc::ClientContext& context, Options const& options, + google::spanner::v1::ExecuteBatchDmlRequest const& request, + spanner_internal::OperationContext& operation_context) = 0; virtual std::unique_ptr> StreamingRead(std::shared_ptr context, Options const& options, - google::spanner::v1::ReadRequest const& request) = 0; + google::spanner::v1::ReadRequest const& request, + std::shared_ptr + operation_context) = 0; virtual StatusOr BeginTransaction( grpc::ClientContext& context, Options const& options, - google::spanner::v1::BeginTransactionRequest const& request) = 0; + google::spanner::v1::BeginTransactionRequest const& request, + spanner_internal::OperationContext& operation_context) = 0; virtual StatusOr Commit( grpc::ClientContext& context, Options const& options, - google::spanner::v1::CommitRequest const& request) = 0; + google::spanner::v1::CommitRequest const& request, + spanner_internal::OperationContext& operation_context) = 0; virtual Status Rollback( grpc::ClientContext& context, Options const& options, - google::spanner::v1::RollbackRequest const& request) = 0; + google::spanner::v1::RollbackRequest const& request, + spanner_internal::OperationContext& operation_context) = 0; virtual StatusOr PartitionQuery( grpc::ClientContext& context, Options const& options, - google::spanner::v1::PartitionQueryRequest const& request) = 0; + google::spanner::v1::PartitionQueryRequest const& request, + spanner_internal::OperationContext& operation_context) = 0; virtual StatusOr PartitionRead( grpc::ClientContext& context, Options const& options, - google::spanner::v1::PartitionReadRequest const& request) = 0; + google::spanner::v1::PartitionReadRequest const& request, + spanner_internal::OperationContext& operation_context) = 0; virtual std::unique_ptr> BatchWrite(std::shared_ptr context, Options const& options, - google::spanner::v1::BatchWriteRequest const& request) = 0; + google::spanner::v1::BatchWriteRequest const& request, + std::shared_ptr + operation_context) = 0; virtual std::unique_ptr> - FetchCacheUpdate( - std::shared_ptr context, Options const& options, - google::spanner::v1::FetchCacheUpdateRequest const& request) = 0; + FetchCacheUpdate(std::shared_ptr context, + Options const& options, + google::spanner::v1::FetchCacheUpdateRequest const& request, + std::shared_ptr + operation_context) = 0; virtual future> AsyncCreateSession( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::CreateSessionRequest const& request) = 0; + google::spanner::v1::CreateSessionRequest const& request, + std::shared_ptr + operation_context) = 0; virtual future> AsyncBatchCreateSessions( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::BatchCreateSessionsRequest const& request) = 0; + google::spanner::v1::BatchCreateSessionsRequest const& request, + std::shared_ptr + operation_context) = 0; virtual future AsyncDeleteSession( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::DeleteSessionRequest const& request) = 0; + google::spanner::v1::DeleteSessionRequest const& request, + std::shared_ptr + operation_context) = 0; virtual future> AsyncExecuteSql( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::ExecuteSqlRequest const& request) = 0; + google::spanner::v1::ExecuteSqlRequest const& request, + std::shared_ptr + operation_context) = 0; }; class DefaultSpannerStub : public SpannerStub { @@ -141,92 +167,119 @@ class DefaultSpannerStub : public SpannerStub { StatusOr CreateSession( grpc::ClientContext& context, Options const& options, - google::spanner::v1::CreateSessionRequest const& request) override; + google::spanner::v1::CreateSessionRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr BatchCreateSessions( grpc::ClientContext& context, Options const& options, - google::spanner::v1::BatchCreateSessionsRequest const& request) override; + google::spanner::v1::BatchCreateSessionsRequest const& request, + spanner_internal::OperationContext& operation_context) override; Status DeleteSession( grpc::ClientContext& context, Options const& options, - google::spanner::v1::DeleteSessionRequest const& request) override; + google::spanner::v1::DeleteSessionRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr ExecuteSql( grpc::ClientContext& context, Options const& options, - google::spanner::v1::ExecuteSqlRequest const& request) override; + google::spanner::v1::ExecuteSqlRequest const& request, + spanner_internal::OperationContext& operation_context) override; std::unique_ptr> - ExecuteStreamingSql( - std::shared_ptr context, Options const& options, - google::spanner::v1::ExecuteSqlRequest const& request) override; + ExecuteStreamingSql(std::shared_ptr context, + Options const& options, + google::spanner::v1::ExecuteSqlRequest const& request, + std::shared_ptr + operation_context) override; StatusOr ExecuteBatchDml( grpc::ClientContext& context, Options const& options, - google::spanner::v1::ExecuteBatchDmlRequest const& request) override; + google::spanner::v1::ExecuteBatchDmlRequest const& request, + spanner_internal::OperationContext& operation_context) override; std::unique_ptr> StreamingRead(std::shared_ptr context, Options const& options, - google::spanner::v1::ReadRequest const& request) override; + google::spanner::v1::ReadRequest const& request, + std::shared_ptr + operation_context) override; StatusOr BeginTransaction( grpc::ClientContext& context, Options const& options, - google::spanner::v1::BeginTransactionRequest const& request) override; + google::spanner::v1::BeginTransactionRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr Commit( grpc::ClientContext& context, Options const& options, - google::spanner::v1::CommitRequest const& request) override; + google::spanner::v1::CommitRequest const& request, + spanner_internal::OperationContext& operation_context) override; - Status Rollback(grpc::ClientContext& context, Options const& options, - google::spanner::v1::RollbackRequest const& request) override; + Status Rollback( + grpc::ClientContext& context, Options const& options, + google::spanner::v1::RollbackRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr PartitionQuery( grpc::ClientContext& context, Options const& options, - google::spanner::v1::PartitionQueryRequest const& request) override; + google::spanner::v1::PartitionQueryRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr PartitionRead( grpc::ClientContext& context, Options const& options, - google::spanner::v1::PartitionReadRequest const& request) override; + google::spanner::v1::PartitionReadRequest const& request, + spanner_internal::OperationContext& operation_context) override; std::unique_ptr> BatchWrite(std::shared_ptr context, Options const& options, - google::spanner::v1::BatchWriteRequest const& request) override; + google::spanner::v1::BatchWriteRequest const& request, + std::shared_ptr + operation_context) override; std::unique_ptr> - FetchCacheUpdate( - std::shared_ptr context, Options const& options, - google::spanner::v1::FetchCacheUpdateRequest const& request) override; + FetchCacheUpdate(std::shared_ptr context, + Options const& options, + google::spanner::v1::FetchCacheUpdateRequest const& request, + std::shared_ptr + operation_context) override; future> AsyncCreateSession( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::CreateSessionRequest const& request) override; + google::spanner::v1::CreateSessionRequest const& request, + std::shared_ptr operation_context) + override; future> AsyncBatchCreateSessions( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::BatchCreateSessionsRequest const& request) override; + google::spanner::v1::BatchCreateSessionsRequest const& request, + std::shared_ptr operation_context) + override; future AsyncDeleteSession( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::DeleteSessionRequest const& request) override; + google::spanner::v1::DeleteSessionRequest const& request, + std::shared_ptr operation_context) + override; future> AsyncExecuteSql( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::ExecuteSqlRequest const& request) override; + google::spanner::v1::ExecuteSqlRequest const& request, + std::shared_ptr operation_context) + override; private: std::unique_ptr grpc_stub_; diff --git a/google/cloud/spanner/internal/spanner_stub_factory_test.cc b/google/cloud/spanner/internal/spanner_stub_factory_test.cc index b6315d1de8c26..334b0fa094b93 100644 --- a/google/cloud/spanner/internal/spanner_stub_factory_test.cc +++ b/google/cloud/spanner/internal/spanner_stub_factory_test.cc @@ -13,6 +13,7 @@ // limitations under the License. #include "google/cloud/spanner/internal/spanner_stub_factory.h" +#include "google/cloud/spanner/internal/operation_context.h" #include "google/cloud/spanner/testing/mock_spanner_stub.h" #include "google/cloud/common_options.h" #include "google/cloud/grpc_options.h" @@ -44,7 +45,8 @@ TEST(DecorateSpannerStub, Auth) { auto mock = std::make_shared(); EXPECT_CALL(*mock, CreateSession) .WillOnce([](grpc::ClientContext& context, Options const&, - google::spanner::v1::CreateSessionRequest const&) { + google::spanner::v1::CreateSessionRequest const&, + internal::OperationContext&) { EXPECT_THAT(context.credentials(), NotNull()); return internal::AbortedError("fail"); }); @@ -61,7 +63,8 @@ TEST(DecorateSpannerStub, Auth) { ASSERT_NE(stub, nullptr); grpc::ClientContext context; - auto session = stub->CreateSession(context, Options{}, {}); + OperationContext op_context(nullptr, 1, "CreateSession"); + auto session = stub->CreateSession(context, Options{}, {}, op_context); EXPECT_THAT(session, StatusIs(StatusCode::kAborted)); } @@ -70,7 +73,8 @@ TEST(DecorateSpannerStub, Metadata) { auto const db = spanner::Database("foo", "bar", "baz"); EXPECT_CALL(*mock, CreateSession) .WillOnce([&db](grpc::ClientContext& context, Options const&, - google::spanner::v1::CreateSessionRequest const&) { + google::spanner::v1::CreateSessionRequest const&, + OperationContext&) { testing_util::ValidateMetadataFixture fixture; auto metadata = fixture.GetMetadata(context); EXPECT_THAT(metadata, Contains(Pair("google-cloud-resource-prefix", @@ -85,7 +89,8 @@ TEST(DecorateSpannerStub, Metadata) { ASSERT_NE(stub, nullptr); grpc::ClientContext context; - auto session = stub->CreateSession(context, Options{}, {}); + OperationContext op_context(nullptr, 1, "CreateSession"); + auto session = stub->CreateSession(context, Options{}, {}, op_context); EXPECT_THAT(session, StatusIs(StatusCode::kAborted)); } @@ -105,7 +110,8 @@ TEST(DecorateSpannerStub, Logging) { ASSERT_NE(stub, nullptr); grpc::ClientContext context; - auto session = stub->CreateSession(context, Options{}, {}); + OperationContext op_context(nullptr, 1, "CreateSession"); + auto session = stub->CreateSession(context, Options{}, {}, op_context); EXPECT_THAT(session, StatusIs(StatusCode::kAborted)); EXPECT_THAT(log.ExtractLines(), @@ -123,7 +129,7 @@ TEST(DecorateSpannerStub, TracingEnabled) { auto mock = std::make_shared(); EXPECT_CALL(*mock, CreateSession) - .WillOnce([](auto& context, auto const&, auto const&) { + .WillOnce([](auto& context, auto const&, auto const&, auto&) { testing_util::ValidatePropagator(context); return internal::AbortedError("fail"); }); @@ -136,7 +142,8 @@ TEST(DecorateSpannerStub, TracingEnabled) { ASSERT_NE(stub, nullptr); grpc::ClientContext context; - auto session = stub->CreateSession(context, Options{}, {}); + OperationContext op_context(nullptr, 1, "CreateSession"); + auto session = stub->CreateSession(context, Options{}, {}, op_context); EXPECT_THAT(session, StatusIs(StatusCode::kAborted)); EXPECT_THAT( @@ -149,7 +156,7 @@ TEST(DecorateSpannerStub, TracingDisabled) { auto mock = std::make_shared(); EXPECT_CALL(*mock, CreateSession) - .WillOnce([](auto& context, auto const&, auto const&) { + .WillOnce([](auto& context, auto const&, auto const&, auto&) { testing_util::ValidateNoPropagator(context); return internal::AbortedError("fail"); }); @@ -162,7 +169,8 @@ TEST(DecorateSpannerStub, TracingDisabled) { EXPECT_NE(stub, nullptr); grpc::ClientContext context; - auto session = stub->CreateSession(context, Options{}, {}); + OperationContext op_context(nullptr, 1, "CreateSession"); + auto session = stub->CreateSession(context, Options{}, {}, op_context); EXPECT_THAT(session, StatusIs(StatusCode::kAborted)); EXPECT_THAT(span_catcher->GetSpans(), IsEmpty()); diff --git a/google/cloud/spanner/internal/spanner_tracing_stub.cc b/google/cloud/spanner/internal/spanner_tracing_stub.cc index 3b895f7d3e7d9..19c35308e5dc6 100644 --- a/google/cloud/spanner/internal/spanner_tracing_stub.cc +++ b/google/cloud/spanner/internal/spanner_tracing_stub.cc @@ -35,58 +35,68 @@ SpannerTracingStub::SpannerTracingStub(std::shared_ptr child) StatusOr SpannerTracingStub::CreateSession( grpc::ClientContext& context, Options const& options, - google::spanner::v1::CreateSessionRequest const& request) { + google::spanner::v1::CreateSessionRequest const& request, + spanner_internal::OperationContext& operation_context) { auto span = internal::MakeSpanGrpc("google.spanner.v1.Spanner", "CreateSession"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(context, *propagator_); - return internal::EndSpan(context, *span, - child_->CreateSession(context, options, request)); + return internal::EndSpan( + context, *span, + child_->CreateSession(context, options, request, operation_context)); } StatusOr SpannerTracingStub::BatchCreateSessions( grpc::ClientContext& context, Options const& options, - google::spanner::v1::BatchCreateSessionsRequest const& request) { + google::spanner::v1::BatchCreateSessionsRequest const& request, + spanner_internal::OperationContext& operation_context) { auto span = internal::MakeSpanGrpc("google.spanner.v1.Spanner", "BatchCreateSessions"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(context, *propagator_); - return internal::EndSpan( - context, *span, child_->BatchCreateSessions(context, options, request)); + return internal::EndSpan(context, *span, + child_->BatchCreateSessions( + context, options, request, operation_context)); } Status SpannerTracingStub::DeleteSession( grpc::ClientContext& context, Options const& options, - google::spanner::v1::DeleteSessionRequest const& request) { + google::spanner::v1::DeleteSessionRequest const& request, + spanner_internal::OperationContext& operation_context) { auto span = internal::MakeSpanGrpc("google.spanner.v1.Spanner", "DeleteSession"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(context, *propagator_); - return internal::EndSpan(context, *span, - child_->DeleteSession(context, options, request)); + return internal::EndSpan( + context, *span, + child_->DeleteSession(context, options, request, operation_context)); } StatusOr SpannerTracingStub::ExecuteSql( grpc::ClientContext& context, Options const& options, - google::spanner::v1::ExecuteSqlRequest const& request) { + google::spanner::v1::ExecuteSqlRequest const& request, + spanner_internal::OperationContext& operation_context) { auto span = internal::MakeSpanGrpc("google.spanner.v1.Spanner", "ExecuteSql"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(context, *propagator_); - return internal::EndSpan(context, *span, - child_->ExecuteSql(context, options, request)); + return internal::EndSpan( + context, *span, + child_->ExecuteSql(context, options, request, operation_context)); } std::unique_ptr> SpannerTracingStub::ExecuteStreamingSql( std::shared_ptr context, Options const& options, - google::spanner::v1::ExecuteSqlRequest const& request) { + google::spanner::v1::ExecuteSqlRequest const& request, + std::shared_ptr operation_context) { auto span = internal::MakeSpanGrpc("google.spanner.v1.Spanner", "ExecuteStreamingSql"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(*context, *propagator_); - auto stream = child_->ExecuteStreamingSql(context, options, request); + auto stream = child_->ExecuteStreamingSql(context, options, request, + std::move(operation_context)); return std::make_unique< internal::StreamingReadRpcTracing>( std::move(context), std::move(stream), std::move(span)); @@ -95,25 +105,29 @@ SpannerTracingStub::ExecuteStreamingSql( StatusOr SpannerTracingStub::ExecuteBatchDml( grpc::ClientContext& context, Options const& options, - google::spanner::v1::ExecuteBatchDmlRequest const& request) { + google::spanner::v1::ExecuteBatchDmlRequest const& request, + spanner_internal::OperationContext& operation_context) { auto span = internal::MakeSpanGrpc("google.spanner.v1.Spanner", "ExecuteBatchDml"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(context, *propagator_); - return internal::EndSpan(context, *span, - child_->ExecuteBatchDml(context, options, request)); + return internal::EndSpan( + context, *span, + child_->ExecuteBatchDml(context, options, request, operation_context)); } std::unique_ptr> SpannerTracingStub::StreamingRead( std::shared_ptr context, Options const& options, - google::spanner::v1::ReadRequest const& request) { + google::spanner::v1::ReadRequest const& request, + std::shared_ptr operation_context) { auto span = internal::MakeSpanGrpc("google.spanner.v1.Spanner", "StreamingRead"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(*context, *propagator_); - auto stream = child_->StreamingRead(context, options, request); + auto stream = child_->StreamingRead(context, options, request, + std::move(operation_context)); return std::make_unique< internal::StreamingReadRpcTracing>( std::move(context), std::move(stream), std::move(span)); @@ -121,68 +135,80 @@ SpannerTracingStub::StreamingRead( StatusOr SpannerTracingStub::BeginTransaction( grpc::ClientContext& context, Options const& options, - google::spanner::v1::BeginTransactionRequest const& request) { + google::spanner::v1::BeginTransactionRequest const& request, + spanner_internal::OperationContext& operation_context) { auto span = internal::MakeSpanGrpc("google.spanner.v1.Spanner", "BeginTransaction"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(context, *propagator_); - return internal::EndSpan(context, *span, - child_->BeginTransaction(context, options, request)); + return internal::EndSpan( + context, *span, + child_->BeginTransaction(context, options, request, operation_context)); } StatusOr SpannerTracingStub::Commit( grpc::ClientContext& context, Options const& options, - google::spanner::v1::CommitRequest const& request) { + google::spanner::v1::CommitRequest const& request, + spanner_internal::OperationContext& operation_context) { auto span = internal::MakeSpanGrpc("google.spanner.v1.Spanner", "Commit"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(context, *propagator_); - return internal::EndSpan(context, *span, - child_->Commit(context, options, request)); + return internal::EndSpan( + context, *span, + child_->Commit(context, options, request, operation_context)); } Status SpannerTracingStub::Rollback( grpc::ClientContext& context, Options const& options, - google::spanner::v1::RollbackRequest const& request) { + google::spanner::v1::RollbackRequest const& request, + spanner_internal::OperationContext& operation_context) { auto span = internal::MakeSpanGrpc("google.spanner.v1.Spanner", "Rollback"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(context, *propagator_); - return internal::EndSpan(context, *span, - child_->Rollback(context, options, request)); + return internal::EndSpan( + context, *span, + child_->Rollback(context, options, request, operation_context)); } StatusOr SpannerTracingStub::PartitionQuery( grpc::ClientContext& context, Options const& options, - google::spanner::v1::PartitionQueryRequest const& request) { + google::spanner::v1::PartitionQueryRequest const& request, + spanner_internal::OperationContext& operation_context) { auto span = internal::MakeSpanGrpc("google.spanner.v1.Spanner", "PartitionQuery"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(context, *propagator_); - return internal::EndSpan(context, *span, - child_->PartitionQuery(context, options, request)); + return internal::EndSpan( + context, *span, + child_->PartitionQuery(context, options, request, operation_context)); } StatusOr SpannerTracingStub::PartitionRead( grpc::ClientContext& context, Options const& options, - google::spanner::v1::PartitionReadRequest const& request) { + google::spanner::v1::PartitionReadRequest const& request, + spanner_internal::OperationContext& operation_context) { auto span = internal::MakeSpanGrpc("google.spanner.v1.Spanner", "PartitionRead"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(context, *propagator_); - return internal::EndSpan(context, *span, - child_->PartitionRead(context, options, request)); + return internal::EndSpan( + context, *span, + child_->PartitionRead(context, options, request, operation_context)); } std::unique_ptr> SpannerTracingStub::BatchWrite( std::shared_ptr context, Options const& options, - google::spanner::v1::BatchWriteRequest const& request) { + google::spanner::v1::BatchWriteRequest const& request, + std::shared_ptr operation_context) { auto span = internal::MakeSpanGrpc("google.spanner.v1.Spanner", "BatchWrite"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(*context, *propagator_); - auto stream = child_->BatchWrite(context, options, request); + auto stream = child_->BatchWrite(context, options, request, + std::move(operation_context)); return std::make_unique>( std::move(context), std::move(stream), std::move(span)); @@ -192,12 +218,14 @@ std::unique_ptr< google::cloud::internal::StreamingReadRpc> SpannerTracingStub::FetchCacheUpdate( std::shared_ptr context, Options const& options, - google::spanner::v1::FetchCacheUpdateRequest const& request) { + google::spanner::v1::FetchCacheUpdateRequest const& request, + std::shared_ptr operation_context) { auto span = internal::MakeSpanGrpc("google.spanner.v1.Spanner", "FetchCacheUpdate"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(*context, *propagator_); - auto stream = child_->FetchCacheUpdate(context, options, request); + auto stream = child_->FetchCacheUpdate(context, options, request, + std::move(operation_context)); return std::make_unique< internal::StreamingReadRpcTracing>( std::move(context), std::move(stream), std::move(span)); @@ -208,12 +236,14 @@ SpannerTracingStub::AsyncCreateSession( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::CreateSessionRequest const& request) { + google::spanner::v1::CreateSessionRequest const& request, + std::shared_ptr operation_context) { auto span = internal::MakeSpanGrpc("google.spanner.v1.Spanner", "CreateSession"); internal::OTelScope scope(span); internal::InjectTraceContext(*context, *propagator_); - auto f = child_->AsyncCreateSession(cq, context, std::move(options), request); + auto f = child_->AsyncCreateSession(cq, context, std::move(options), request, + std::move(operation_context)); return internal::EndSpan(std::move(context), std::move(span), std::move(f)); } @@ -222,13 +252,14 @@ SpannerTracingStub::AsyncBatchCreateSessions( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::BatchCreateSessionsRequest const& request) { + google::spanner::v1::BatchCreateSessionsRequest const& request, + std::shared_ptr operation_context) { auto span = internal::MakeSpanGrpc("google.spanner.v1.Spanner", "BatchCreateSessions"); internal::OTelScope scope(span); internal::InjectTraceContext(*context, *propagator_); - auto f = child_->AsyncBatchCreateSessions(cq, context, std::move(options), - request); + auto f = child_->AsyncBatchCreateSessions( + cq, context, std::move(options), request, std::move(operation_context)); return internal::EndSpan(std::move(context), std::move(span), std::move(f)); } @@ -236,12 +267,14 @@ future SpannerTracingStub::AsyncDeleteSession( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::DeleteSessionRequest const& request) { + google::spanner::v1::DeleteSessionRequest const& request, + std::shared_ptr operation_context) { auto span = internal::MakeSpanGrpc("google.spanner.v1.Spanner", "DeleteSession"); internal::OTelScope scope(span); internal::InjectTraceContext(*context, *propagator_); - auto f = child_->AsyncDeleteSession(cq, context, std::move(options), request); + auto f = child_->AsyncDeleteSession(cq, context, std::move(options), request, + std::move(operation_context)); return internal::EndSpan(std::move(context), std::move(span), std::move(f)); } @@ -250,11 +283,13 @@ SpannerTracingStub::AsyncExecuteSql( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::ExecuteSqlRequest const& request) { + google::spanner::v1::ExecuteSqlRequest const& request, + std::shared_ptr operation_context) { auto span = internal::MakeSpanGrpc("google.spanner.v1.Spanner", "ExecuteSql"); internal::OTelScope scope(span); internal::InjectTraceContext(*context, *propagator_); - auto f = child_->AsyncExecuteSql(cq, context, std::move(options), request); + auto f = child_->AsyncExecuteSql(cq, context, std::move(options), request, + std::move(operation_context)); return internal::EndSpan(std::move(context), std::move(span), std::move(f)); } diff --git a/google/cloud/spanner/internal/spanner_tracing_stub.h b/google/cloud/spanner/internal/spanner_tracing_stub.h index ca8678470c3d5..29c7302655951 100644 --- a/google/cloud/spanner/internal/spanner_tracing_stub.h +++ b/google/cloud/spanner/internal/spanner_tracing_stub.h @@ -41,92 +41,119 @@ class SpannerTracingStub : public SpannerStub { StatusOr CreateSession( grpc::ClientContext& context, Options const& options, - google::spanner::v1::CreateSessionRequest const& request) override; + google::spanner::v1::CreateSessionRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr BatchCreateSessions( grpc::ClientContext& context, Options const& options, - google::spanner::v1::BatchCreateSessionsRequest const& request) override; + google::spanner::v1::BatchCreateSessionsRequest const& request, + spanner_internal::OperationContext& operation_context) override; Status DeleteSession( grpc::ClientContext& context, Options const& options, - google::spanner::v1::DeleteSessionRequest const& request) override; + google::spanner::v1::DeleteSessionRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr ExecuteSql( grpc::ClientContext& context, Options const& options, - google::spanner::v1::ExecuteSqlRequest const& request) override; + google::spanner::v1::ExecuteSqlRequest const& request, + spanner_internal::OperationContext& operation_context) override; std::unique_ptr> - ExecuteStreamingSql( - std::shared_ptr context, Options const& options, - google::spanner::v1::ExecuteSqlRequest const& request) override; + ExecuteStreamingSql(std::shared_ptr context, + Options const& options, + google::spanner::v1::ExecuteSqlRequest const& request, + std::shared_ptr + operation_context) override; StatusOr ExecuteBatchDml( grpc::ClientContext& context, Options const& options, - google::spanner::v1::ExecuteBatchDmlRequest const& request) override; + google::spanner::v1::ExecuteBatchDmlRequest const& request, + spanner_internal::OperationContext& operation_context) override; std::unique_ptr> StreamingRead(std::shared_ptr context, Options const& options, - google::spanner::v1::ReadRequest const& request) override; + google::spanner::v1::ReadRequest const& request, + std::shared_ptr + operation_context) override; StatusOr BeginTransaction( grpc::ClientContext& context, Options const& options, - google::spanner::v1::BeginTransactionRequest const& request) override; + google::spanner::v1::BeginTransactionRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr Commit( grpc::ClientContext& context, Options const& options, - google::spanner::v1::CommitRequest const& request) override; + google::spanner::v1::CommitRequest const& request, + spanner_internal::OperationContext& operation_context) override; - Status Rollback(grpc::ClientContext& context, Options const& options, - google::spanner::v1::RollbackRequest const& request) override; + Status Rollback( + grpc::ClientContext& context, Options const& options, + google::spanner::v1::RollbackRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr PartitionQuery( grpc::ClientContext& context, Options const& options, - google::spanner::v1::PartitionQueryRequest const& request) override; + google::spanner::v1::PartitionQueryRequest const& request, + spanner_internal::OperationContext& operation_context) override; StatusOr PartitionRead( grpc::ClientContext& context, Options const& options, - google::spanner::v1::PartitionReadRequest const& request) override; + google::spanner::v1::PartitionReadRequest const& request, + spanner_internal::OperationContext& operation_context) override; std::unique_ptr> BatchWrite(std::shared_ptr context, Options const& options, - google::spanner::v1::BatchWriteRequest const& request) override; + google::spanner::v1::BatchWriteRequest const& request, + std::shared_ptr + operation_context) override; std::unique_ptr> - FetchCacheUpdate( - std::shared_ptr context, Options const& options, - google::spanner::v1::FetchCacheUpdateRequest const& request) override; + FetchCacheUpdate(std::shared_ptr context, + Options const& options, + google::spanner::v1::FetchCacheUpdateRequest const& request, + std::shared_ptr + operation_context) override; future> AsyncCreateSession( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::CreateSessionRequest const& request) override; + google::spanner::v1::CreateSessionRequest const& request, + std::shared_ptr operation_context) + override; future> AsyncBatchCreateSessions( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::BatchCreateSessionsRequest const& request) override; + google::spanner::v1::BatchCreateSessionsRequest const& request, + std::shared_ptr operation_context) + override; future AsyncDeleteSession( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::DeleteSessionRequest const& request) override; + google::spanner::v1::DeleteSessionRequest const& request, + std::shared_ptr operation_context) + override; future> AsyncExecuteSql( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::spanner::v1::ExecuteSqlRequest const& request) override; + google::spanner::v1::ExecuteSqlRequest const& request, + std::shared_ptr operation_context) + override; private: std::shared_ptr child_; diff --git a/google/cloud/spanner/internal/transaction_impl.h b/google/cloud/spanner/internal/transaction_impl.h index 48623c336b74b..f8df8d89324c4 100644 --- a/google/cloud/spanner/internal/transaction_impl.h +++ b/google/cloud/spanner/internal/transaction_impl.h @@ -37,6 +37,7 @@ struct TransactionContext { std::string const& tag; std::int64_t seqno; std::optional> stub; + std::optional channel_id; std::optional precommit_token; }; @@ -91,13 +92,14 @@ class TransactionImpl { StatusOr&, TransactionContext&>::value, "TransactionImpl::Visit() functor has incompatible type."); - TransactionContext ctx{route_to_leader_, tag_, 0, std::nullopt, - std::nullopt}; + TransactionContext ctx{route_to_leader_, tag_, 0, + std::nullopt, std::nullopt, std::nullopt}; { std::unique_lock lock(mu_); ctx.seqno = ++seqno_; // what about overflow? cond_.wait(lock, [this] { return state_ != State::kPending; }); ctx.stub = stub_; + ctx.channel_id = channel_id_; ctx.precommit_token = precommit_token_; if (state_ == State::kDone) { lock.unlock(); @@ -117,6 +119,7 @@ class TransactionImpl { { std::unique_lock lock(mu_); stub_ = ctx.stub; + channel_id_ = ctx.channel_id; UpdatePrecommitToken(lock, ctx.precommit_token); state_ = selector_ && selector_->has_begin() ? State::kBegin : State::kDone; @@ -131,7 +134,7 @@ class TransactionImpl { #if GOOGLE_CLOUD_CPP_HAVE_EXCEPTIONS } catch (...) { { - std::lock_guard lock(mu_); + std::scoped_lock lock(mu_); state_ = State::kBegin; } cond_.notify_one(); @@ -161,6 +164,7 @@ class TransactionImpl { std::string tag_; std::int64_t seqno_; std::optional> stub_ = std::nullopt; + std::optional channel_id_ = std::nullopt; std::optional precommit_token_ = std::nullopt; }; diff --git a/google/cloud/spanner/spanner_client_unit_tests.bzl b/google/cloud/spanner/spanner_client_unit_tests.bzl index 6a501b33fdeb2..3c246a806ed5a 100644 --- a/google/cloud/spanner/spanner_client_unit_tests.bzl +++ b/google/cloud/spanner/spanner_client_unit_tests.bzl @@ -33,10 +33,13 @@ spanner_client_unit_tests = [ "internal/instance_admin_metadata_test.cc", "internal/logging_result_set_reader_test.cc", "internal/merge_chunk_test.cc", + "internal/operation_context_test.cc", "internal/partial_result_set_resume_test.cc", "internal/partial_result_set_source_test.cc", "internal/route_to_leader_test.cc", "internal/session_pool_test.cc", + "internal/spanner_operation_context_factory_test.cc", + "internal/spanner_request_id_test.cc", "internal/spanner_stub_factory_test.cc", "internal/status_utils_test.cc", "internal/transaction_impl_test.cc", diff --git a/google/cloud/spanner/testing/mock_spanner_stub.h b/google/cloud/spanner/testing/mock_spanner_stub.h index b2d8385010780..5c1bc8dca0f21 100644 --- a/google/cloud/spanner/testing/mock_spanner_stub.h +++ b/google/cloud/spanner/testing/mock_spanner_stub.h @@ -15,6 +15,7 @@ #ifndef GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_SPANNER_TESTING_MOCK_SPANNER_STUB_H #define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_SPANNER_TESTING_MOCK_SPANNER_STUB_H +#include "google/cloud/spanner/internal/operation_context.h" #include "google/cloud/spanner/internal/spanner_stub.h" #include "google/cloud/spanner/version.h" #include @@ -29,21 +30,24 @@ class MockSpannerStub : public google::cloud::spanner_internal::SpannerStub { public: MOCK_METHOD(StatusOr, CreateSession, (grpc::ClientContext&, Options const&, - google::spanner::v1::CreateSessionRequest const&), + google::spanner::v1::CreateSessionRequest const&, + spanner_internal::OperationContext&), (override)); - MOCK_METHOD(future>, - AsyncCreateSession, - (google::cloud::CompletionQueue & cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::spanner::v1::CreateSessionRequest const& request), - (override)); + MOCK_METHOD( + future>, AsyncCreateSession, + (google::cloud::CompletionQueue & cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::spanner::v1::CreateSessionRequest const& request, + std::shared_ptr operation_context), + (override)); MOCK_METHOD(StatusOr, BatchCreateSessions, (grpc::ClientContext&, Options const&, - google::spanner::v1::BatchCreateSessionsRequest const&), + google::spanner::v1::BatchCreateSessionsRequest const&, + spanner_internal::OperationContext&), (override)); MOCK_METHOD( @@ -51,29 +55,34 @@ class MockSpannerStub : public google::cloud::spanner_internal::SpannerStub { AsyncBatchCreateSessions, (CompletionQueue&, std::shared_ptr, google::cloud::internal::ImmutableOptions, - google::spanner::v1::BatchCreateSessionsRequest const&), + google::spanner::v1::BatchCreateSessionsRequest const&, + std::shared_ptr), (override)); MOCK_METHOD(Status, DeleteSession, (grpc::ClientContext&, Options const&, - google::spanner::v1::DeleteSessionRequest const&), + google::spanner::v1::DeleteSessionRequest const&, + spanner_internal::OperationContext&), (override)); MOCK_METHOD(future, AsyncDeleteSession, (CompletionQueue&, std::shared_ptr, google::cloud::internal::ImmutableOptions, - google::spanner::v1::DeleteSessionRequest const&), + google::spanner::v1::DeleteSessionRequest const&, + std::shared_ptr), (override)); MOCK_METHOD(StatusOr, ExecuteSql, (grpc::ClientContext&, Options const&, - google::spanner::v1::ExecuteSqlRequest const&), + google::spanner::v1::ExecuteSqlRequest const&, + spanner_internal::OperationContext&), (override)); MOCK_METHOD(future>, AsyncExecuteSql, (CompletionQueue&, std::shared_ptr, google::cloud::internal::ImmutableOptions, - google::spanner::v1::ExecuteSqlRequest const&), + google::spanner::v1::ExecuteSqlRequest const&, + std::shared_ptr), (override)); MOCK_METHOD( @@ -81,13 +90,15 @@ class MockSpannerStub : public google::cloud::spanner_internal::SpannerStub { internal::StreamingReadRpc>, ExecuteStreamingSql, (std::shared_ptr, Options const&, - google::spanner::v1::ExecuteSqlRequest const&), + google::spanner::v1::ExecuteSqlRequest const&, + std::shared_ptr), (override)); MOCK_METHOD(StatusOr, ExecuteBatchDml, (grpc::ClientContext&, Options const&, - google::spanner::v1::ExecuteBatchDmlRequest const&), + google::spanner::v1::ExecuteBatchDmlRequest const&, + spanner_internal::OperationContext&), (override)); MOCK_METHOD( @@ -95,32 +106,38 @@ class MockSpannerStub : public google::cloud::spanner_internal::SpannerStub { internal::StreamingReadRpc>, StreamingRead, (std::shared_ptr, Options const&, - google::spanner::v1::ReadRequest const&), + google::spanner::v1::ReadRequest const&, + std::shared_ptr), (override)); MOCK_METHOD(StatusOr, BeginTransaction, (grpc::ClientContext&, Options const&, - google::spanner::v1::BeginTransactionRequest const&), + google::spanner::v1::BeginTransactionRequest const&, + spanner_internal::OperationContext&), (override)); MOCK_METHOD(StatusOr, Commit, (grpc::ClientContext&, Options const&, - google::spanner::v1::CommitRequest const&), + google::spanner::v1::CommitRequest const&, + spanner_internal::OperationContext&), (override)); MOCK_METHOD(Status, Rollback, (grpc::ClientContext&, Options const&, - google::spanner::v1::RollbackRequest const&), + google::spanner::v1::RollbackRequest const&, + spanner_internal::OperationContext&), (override)); MOCK_METHOD(StatusOr, PartitionQuery, (grpc::ClientContext&, Options const&, - google::spanner::v1::PartitionQueryRequest const&), + google::spanner::v1::PartitionQueryRequest const&, + spanner_internal::OperationContext&), (override)); MOCK_METHOD(StatusOr, PartitionRead, (grpc::ClientContext&, Options const&, - google::spanner::v1::PartitionReadRequest const&), + google::spanner::v1::PartitionReadRequest const&, + spanner_internal::OperationContext&), (override)); MOCK_METHOD( @@ -128,14 +145,16 @@ class MockSpannerStub : public google::cloud::spanner_internal::SpannerStub { internal::StreamingReadRpc>, BatchWrite, (std::shared_ptr, Options const&, - google::spanner::v1::BatchWriteRequest const&), + google::spanner::v1::BatchWriteRequest const&, + std::shared_ptr), (override)); MOCK_METHOD(std::unique_ptr>, FetchCacheUpdate, (std::shared_ptr, Options const&, - google::spanner::v1::FetchCacheUpdateRequest const&), + google::spanner::v1::FetchCacheUpdateRequest const&, + std::shared_ptr), (override)); };