diff --git a/google/cloud/bigtable/internal/async_bulk_apply.cc b/google/cloud/bigtable/internal/async_bulk_apply.cc index 1039aa862b30b..88dd567e7c568 100644 --- a/google/cloud/bigtable/internal/async_bulk_apply.cc +++ b/google/cloud/bigtable/internal/async_bulk_apply.cc @@ -83,7 +83,7 @@ void AsyncBulkApplier::MakeRequest() { auto self = this->shared_from_this(); PerformAsyncStreamingRead( stub_->AsyncMutateRows(cq_, client_context_, options_, - state_.BeforeStart()), + state_.BeforeStart(), *operation_context_), [self](google::bigtable::v2::MutateRowsResponse r) { self->OnRead(std::move(r)); return make_ready_future(self->keep_reading_.load()); diff --git a/google/cloud/bigtable/internal/async_bulk_apply_test.cc b/google/cloud/bigtable/internal/async_bulk_apply_test.cc index afddfdc964e86..78ed55b8864f1 100644 --- a/google/cloud/bigtable/internal/async_bulk_apply_test.cc +++ b/google/cloud/bigtable/internal/async_bulk_apply_test.cc @@ -163,7 +163,8 @@ class MockMetric : public Metric { ElementDeliveryParams const&), (override)); MOCK_METHOD(std::unique_ptr, clone, - (ResourceLabels resource_labels, DataLabels data_labels), + (TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels), (const, override)); }; @@ -172,7 +173,8 @@ class CloningMetric : public Metric { public: explicit CloningMetric(std::unique_ptr metric) : metric_(std::move(metric)) {} - std::unique_ptr clone(ResourceLabels, DataLabels) const override { + std::unique_ptr clone(TableResourceLabels const&, + TableDataLabels const&) const override { return std::move(metric_); } @@ -207,7 +209,8 @@ TEST_F(AsyncBulkApplyTest, Success) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncMutateRows) .WillOnce([this](CompletionQueue const&, auto client_context, auto, - v2::MutateRowsRequest const& request) { + v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -281,7 +284,8 @@ TEST_F(AsyncBulkApplyTest, PartialStreamIsRetried) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncMutateRows) .WillOnce([this](CompletionQueue const&, auto context, auto, - v2::MutateRowsRequest const& request) { + v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -306,7 +310,8 @@ TEST_F(AsyncBulkApplyTest, PartialStreamIsRetried) { return stream; }) .WillOnce([this](CompletionQueue const&, auto client_context, auto, - v2::MutateRowsRequest const& request) { + v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -383,7 +388,8 @@ TEST_F(AsyncBulkApplyTest, IdempotentMutationPolicy) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncMutateRows) .WillOnce([this](CompletionQueue const&, auto client_context, auto, - v2::MutateRowsRequest const& request) { + v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -408,7 +414,8 @@ TEST_F(AsyncBulkApplyTest, IdempotentMutationPolicy) { return stream; }) .WillOnce([this](CompletionQueue const&, auto client_context, auto, - v2::MutateRowsRequest const& request) { + v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -464,7 +471,8 @@ TEST_F(AsyncBulkApplyTest, TooManyStreamFailures) { EXPECT_CALL(*mock, AsyncMutateRows) .Times(kNumRetries + 1) .WillRepeatedly([this](CompletionQueue const&, auto context, auto, - v2::MutateRowsRequest const& request) { + v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -530,7 +538,8 @@ TEST_F(AsyncBulkApplyTest, RetryInfoHeeded) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncMutateRows) .WillOnce([this](CompletionQueue const&, auto context, auto, - v2::MutateRowsRequest const&) { + v2::MutateRowsRequest const&, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); auto stream = std::make_unique(); EXPECT_CALL(*stream, Start).WillOnce([] { @@ -544,7 +553,8 @@ TEST_F(AsyncBulkApplyTest, RetryInfoHeeded) { return stream; }) .WillOnce([this](CompletionQueue const&, auto context, auto, - v2::MutateRowsRequest const&) { + v2::MutateRowsRequest const&, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); auto stream = std::make_unique(); EXPECT_CALL(*stream, Start).WillOnce([] { @@ -603,7 +613,8 @@ TEST_F(AsyncBulkApplyTest, RetryInfoIgnored) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncMutateRows) .WillOnce([this](CompletionQueue const&, auto context, auto, - v2::MutateRowsRequest const&) { + v2::MutateRowsRequest const&, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); auto stream = std::make_unique(); EXPECT_CALL(*stream, Start).WillOnce([] { @@ -640,7 +651,8 @@ TEST_F(AsyncBulkApplyTest, TimerError) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncMutateRows) .WillOnce([this](CompletionQueue const&, auto context, auto, - v2::MutateRowsRequest const& request) { + v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -705,7 +717,8 @@ TEST_F(AsyncBulkApplyTest, CancelAfterSuccess) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncMutateRows) .WillOnce([&p, this](CompletionQueue const&, auto client_context, auto, - v2::MutateRowsRequest const& request) { + v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -781,7 +794,8 @@ TEST_F(AsyncBulkApplyTest, CancelMidStream) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncMutateRows) .WillOnce([&p, this](CompletionQueue const&, auto client_context, auto, - v2::MutateRowsRequest const& request) { + v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -849,7 +863,8 @@ TEST_F(AsyncBulkApplyTest, CurrentOptionsContinuedOnRetries) { EXPECT_CALL(*mock, AsyncMutateRows) .Times(2) .WillRepeatedly([this](CompletionQueue const&, auto context, auto, - v2::MutateRowsRequest const&) { + v2::MutateRowsRequest const&, + bigtable_internal::OperationContext&) { EXPECT_EQ(5, internal::CurrentOptions().get()); metadata_fixture_.SetServerMetadata(*context, {}); auto stream = std::make_unique(); @@ -902,7 +917,8 @@ TEST_F(AsyncBulkApplyTest, RetriesOkStreamWithFailedMutations) { EXPECT_CALL(*mock, AsyncMutateRows) .Times(kNumRetries + 1) .WillRepeatedly([this](CompletionQueue const&, auto context, auto, - v2::MutateRowsRequest const& request) { + v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -974,7 +990,8 @@ TEST_F(AsyncBulkApplyTest, Throttling) { }); EXPECT_CALL(*mock, AsyncMutateRows) .WillOnce([&limiter, this](CompletionQueue const&, auto client_context, - auto, v2::MutateRowsRequest const&) { + auto, v2::MutateRowsRequest const&, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); ::testing::InSequence seq2; auto stream = std::make_unique(); @@ -1011,7 +1028,8 @@ TEST_F(AsyncBulkApplyTest, ThrottlingBeforeEachRetry) { EXPECT_CALL(*mock, AsyncMutateRows) .Times(kNumRetries + 1) .WillRepeatedly([this](CompletionQueue const&, auto context, auto, - v2::MutateRowsRequest const&) { + v2::MutateRowsRequest const&, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); auto stream = std::make_unique(); EXPECT_CALL(*stream, Start).WillOnce([] { @@ -1065,7 +1083,8 @@ TEST_F(AsyncBulkApplyTest, BigtableCookie) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncMutateRows) .WillOnce([this](CompletionQueue const&, auto context, auto, - v2::MutateRowsRequest const&) { + v2::MutateRowsRequest const&, + bigtable_internal::OperationContext&) { // Return a bigtable cookie in the first request. metadata_fixture_.SetServerMetadata( *context, {{}, {{"x-goog-cbt-cookie-routing", "routing"}}}); @@ -1079,7 +1098,8 @@ TEST_F(AsyncBulkApplyTest, BigtableCookie) { return stream; }) .WillOnce([this](CompletionQueue const&, auto context, auto, - v2::MutateRowsRequest const&) { + v2::MutateRowsRequest const&, + bigtable_internal::OperationContext&) { // Verify that the next request includes the bigtable cookie from above. auto headers = metadata_fixture_.GetMetadata(*context); EXPECT_THAT(headers, @@ -1128,7 +1148,7 @@ TEST_F(AsyncBulkApplyTest, TracedBackoff) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncMutateRows) .Times(kNumRetries + 1) - .WillRepeatedly([this](auto&, auto context, auto, auto const&) { + .WillRepeatedly([this](auto&, auto context, auto, auto const&, auto&) { metadata_fixture_.SetServerMetadata(*context, {}); return std::make_unique(TransientError()); }); @@ -1160,11 +1180,12 @@ TEST_F(AsyncBulkApplyTest, CallSpanActiveThroughout) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncMutateRows) .Times(kNumRetries + 1) - .WillRepeatedly([this, span](auto&, auto context, auto, auto const&) { - metadata_fixture_.SetServerMetadata(*context, {}); - EXPECT_THAT(span, IsActive()); - return std::make_unique(TransientError()); - }); + .WillRepeatedly( + [this, span](auto&, auto context, auto, auto const&, auto&) { + metadata_fixture_.SetServerMetadata(*context, {}); + EXPECT_THAT(span, IsActive()); + return std::make_unique(TransientError()); + }); internal::AutomaticallyCreatedBackgroundThreads background; auto retry = DataLimitedErrorCountRetryPolicy(kNumRetries).clone(); diff --git a/google/cloud/bigtable/internal/async_row_reader.cc b/google/cloud/bigtable/internal/async_row_reader.cc index 23a5bfdf5f8ee..3c7b0973c70cd 100644 --- a/google/cloud/bigtable/internal/async_row_reader.cc +++ b/google/cloud/bigtable/internal/async_row_reader.cc @@ -51,7 +51,8 @@ void AsyncRowReader::MakeRequest() { auto self = this->shared_from_this(); PerformAsyncStreamingRead( - stub_->AsyncReadRows(cq_, client_context_, options_, request), + stub_->AsyncReadRows(cq_, client_context_, options_, request, + *operation_context_), [self](v2::ReadRowsResponse r) { return self->OnDataReceived(std::move(r)); }, diff --git a/google/cloud/bigtable/internal/async_row_reader_test.cc b/google/cloud/bigtable/internal/async_row_reader_test.cc index 4bf4654dd9e20..92e0d29274760 100644 --- a/google/cloud/bigtable/internal/async_row_reader_test.cc +++ b/google/cloud/bigtable/internal/async_row_reader_test.cc @@ -13,6 +13,7 @@ // limitations under the License. #include "google/cloud/bigtable/internal/async_row_reader.h" +#include "google/cloud/bigtable/internal/operation_context.h" #include "google/cloud/bigtable/row_reader.h" #include "google/cloud/bigtable/testing/mock_bigtable_stub.h" #include "google/cloud/internal/async_streaming_read_rpc_impl.h" @@ -115,7 +116,8 @@ class MockMetric : public Metric { ElementDeliveryParams const&), (override)); MOCK_METHOD(std::unique_ptr, clone, - (ResourceLabels resource_labels, DataLabels data_labels), + (TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels), (const, override)); }; @@ -124,7 +126,8 @@ class CloningMetric : public Metric { public: explicit CloningMetric(std::unique_ptr metric) : metric_(std::move(metric)) {} - std::unique_ptr clone(ResourceLabels, DataLabels) const override { + std::unique_ptr clone(TableResourceLabels const&, + TableDataLabels const&) const override { return std::move(metric_); } @@ -160,7 +163,8 @@ TEST_F(AsyncRowReaderTest, Success) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) .WillOnce([this](Unused, auto client_context, Unused, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -259,7 +263,8 @@ TEST_F(AsyncRowReaderTest, SuccessDelayedFuture) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) .WillOnce([this](Unused, auto client_context, Unused, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -352,7 +357,8 @@ TEST_F(AsyncRowReaderTest, ResponseInMultipleChunks) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) .WillOnce([this](Unused, auto client_context, Unused, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -425,7 +431,8 @@ TEST_F(AsyncRowReaderTest, ParserEofFailsOnUnfinishedRow) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) .WillOnce([this](Unused, auto client_context, Unused, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -497,7 +504,8 @@ TEST_F(AsyncRowReaderTest, ParserEofDoesntFailOnUnfinishedRowIfRowLimit) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) .WillOnce([this](Unused, auto client_context, Unused, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -575,7 +583,8 @@ TEST_F(AsyncRowReaderTest, PermanentFailure) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) .WillOnce([this](Unused, auto client_context, Unused, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -648,7 +657,8 @@ TEST_F(AsyncRowReaderTest, RetryPolicyExhausted) { EXPECT_CALL(*mock, AsyncReadRows) .Times(kNumRetries + 1) .WillRepeatedly([this](Unused, auto context, Unused, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -698,21 +708,22 @@ TEST_F(AsyncRowReaderTest, RetryInfoHeeded) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) - .WillOnce( - [this](Unused, auto context, Unused, v2::ReadRowsRequest const&) { - metadata_fixture_.SetServerMetadata(*context, {}); - auto stream = std::make_unique(); - EXPECT_CALL(*stream, Start) - .WillOnce(Return(ByMove(make_ready_future(false)))); - EXPECT_CALL(*stream, Finish).WillOnce([] { - auto status = internal::PermissionDeniedError("try again"); - internal::SetRetryInfo(status, internal::RetryInfo{ms(0)}); - return make_ready_future(status); - }); - return stream; - }) + .WillOnce([this](Unused, auto context, Unused, v2::ReadRowsRequest const&, + bigtable_internal::OperationContext&) { + metadata_fixture_.SetServerMetadata(*context, {}); + auto stream = std::make_unique(); + EXPECT_CALL(*stream, Start) + .WillOnce(Return(ByMove(make_ready_future(false)))); + EXPECT_CALL(*stream, Finish).WillOnce([] { + auto status = internal::PermissionDeniedError("try again"); + internal::SetRetryInfo(status, internal::RetryInfo{ms(0)}); + return make_ready_future(status); + }); + return stream; + }) .WillOnce([this](Unused, auto client_context, Unused, - v2::ReadRowsRequest const&) { + v2::ReadRowsRequest const&, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); auto stream = std::make_unique(); EXPECT_CALL(*stream, Start) @@ -749,19 +760,19 @@ TEST_F(AsyncRowReaderTest, RetryInfoIgnored) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) - .WillOnce( - [this](Unused, auto context, Unused, v2::ReadRowsRequest const&) { - metadata_fixture_.SetServerMetadata(*context, {}); - auto stream = std::make_unique(); - EXPECT_CALL(*stream, Start) - .WillOnce(Return(ByMove(make_ready_future(false)))); - EXPECT_CALL(*stream, Finish).WillOnce([] { - auto status = internal::PermissionDeniedError("try again"); - internal::SetRetryInfo(status, internal::RetryInfo{ms(0)}); - return make_ready_future(status); - }); - return stream; - }); + .WillOnce([this](Unused, auto context, Unused, v2::ReadRowsRequest const&, + bigtable_internal::OperationContext&) { + metadata_fixture_.SetServerMetadata(*context, {}); + auto stream = std::make_unique(); + EXPECT_CALL(*stream, Start) + .WillOnce(Return(ByMove(make_ready_future(false)))); + EXPECT_CALL(*stream, Finish).WillOnce([] { + auto status = internal::PermissionDeniedError("try again"); + internal::SetRetryInfo(status, internal::RetryInfo{ms(0)}); + return make_ready_future(status); + }); + return stream; + }); MockFunction(bigtable::Row const&)> on_row; EXPECT_CALL(on_row, Call).Times(0); @@ -813,7 +824,8 @@ TEST_F(AsyncRowReaderTest, RetrySkipsReadRows) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) .WillOnce([this](Unused, auto context, Unused, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -834,7 +846,8 @@ TEST_F(AsyncRowReaderTest, RetrySkipsReadRows) { return stream; }) .WillOnce([this](Unused, auto client_context, Unused, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -895,7 +908,8 @@ TEST_F(AsyncRowReaderTest, NoRetryIfRowSetIsEmpty) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) .WillOnce([this](Unused, auto client_context, Unused, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -957,7 +971,8 @@ TEST_F(AsyncRowReaderTest, LastScannedRowKeyIsRespected) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) .WillOnce([this](Unused, auto context, Unused, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -984,7 +999,8 @@ TEST_F(AsyncRowReaderTest, LastScannedRowKeyIsRespected) { return stream; }) .WillOnce([this](Unused, auto client_context, Unused, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -1045,7 +1061,8 @@ TEST_F(AsyncRowReaderTest, ParserFailsOnOutOfOrderRowKeys) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) .WillOnce([this](Unused, auto client_context, Unused, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -1110,7 +1127,8 @@ TEST_P(AsyncRowReaderExceptionTest, CancelMidStream) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) .WillOnce([this](Unused, auto client_context, Unused, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -1216,7 +1234,8 @@ TEST_F(AsyncRowReaderTest, CancelAfterStreamFinish) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) .WillOnce([this](Unused, auto client_context, Unused, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -1290,7 +1309,8 @@ TEST_F(AsyncRowReaderTest, DeepStack) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) .WillOnce([this](Unused, auto client_context, Unused, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -1358,7 +1378,8 @@ TEST_F(AsyncRowReaderTest, TimerErrorEndsLoop) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) .WillOnce([this](Unused, auto context, Unused, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -1422,7 +1443,8 @@ TEST_F(AsyncRowReaderTest, CurrentOptionsContinuedOnRetries) { EXPECT_CALL(*mock, AsyncReadRows) .Times(2) .WillRepeatedly([this](CompletionQueue const&, auto context, auto, - v2::ReadRowsRequest const&) { + v2::ReadRowsRequest const&, + bigtable_internal::OperationContext&) { EXPECT_EQ(5, internal::CurrentOptions().get()); metadata_fixture_.SetServerMetadata(*context, {}); auto stream = std::make_unique(); @@ -1474,7 +1496,8 @@ TEST_F(AsyncRowReaderTest, ReverseScanSuccess) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) .WillOnce([this](Unused, auto client_context, Unused, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_TRUE(request.reversed()); auto stream = std::make_unique(); @@ -1533,7 +1556,8 @@ TEST_F(AsyncRowReaderTest, ReverseScanFailsOnIncreasingRowKeyOrder) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) .WillOnce([this](Unused, auto client_context, Unused, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_TRUE(request.reversed()); auto stream = std::make_unique(); @@ -1593,7 +1617,8 @@ TEST_F(AsyncRowReaderTest, ReverseScanResumption) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) .WillOnce([this](Unused, auto context, Unused, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_TRUE(request.reversed()); // The initial row set contains three rows: "r1", "r2", and "r3". @@ -1619,7 +1644,8 @@ TEST_F(AsyncRowReaderTest, ReverseScanResumption) { return stream; }) .WillOnce([this](Unused, auto client_context, Unused, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -1682,36 +1708,36 @@ TEST_F(AsyncRowReaderTest, BigtableCookie) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) - .WillOnce( - [this](Unused, auto context, Unused, v2::ReadRowsRequest const&) { - // Return a bigtable cookie in the first request. - metadata_fixture_.SetServerMetadata( - *context, {{}, {{"x-goog-cbt-cookie-routing", "routing"}}}); - auto stream = std::make_unique(); - EXPECT_CALL(*stream, Start).WillOnce([] { - return make_ready_future(false); - }); - EXPECT_CALL(*stream, Finish).WillOnce([] { - return make_ready_future(TransientError()); - }); - return stream; - }) - .WillOnce( - [this](Unused, auto context, Unused, v2::ReadRowsRequest const&) { - // Verify that the next request includes the bigtable cookie from - // above. - auto headers = metadata_fixture_.GetMetadata(*context); - EXPECT_THAT(headers, - Contains(Pair("x-goog-cbt-cookie-routing", "routing"))); - auto stream = std::make_unique(); - EXPECT_CALL(*stream, Start).WillOnce([] { - return make_ready_future(false); - }); - EXPECT_CALL(*stream, Finish).WillOnce([] { - return make_ready_future(internal::PermissionDeniedError("fail")); - }); - return stream; - }); + .WillOnce([this](Unused, auto context, Unused, v2::ReadRowsRequest const&, + bigtable_internal::OperationContext&) { + // Return a bigtable cookie in the first request. + metadata_fixture_.SetServerMetadata( + *context, {{}, {{"x-goog-cbt-cookie-routing", "routing"}}}); + auto stream = std::make_unique(); + EXPECT_CALL(*stream, Start).WillOnce([] { + return make_ready_future(false); + }); + EXPECT_CALL(*stream, Finish).WillOnce([] { + return make_ready_future(TransientError()); + }); + return stream; + }) + .WillOnce([this](Unused, auto context, Unused, v2::ReadRowsRequest const&, + bigtable_internal::OperationContext&) { + // Verify that the next request includes the bigtable cookie from + // above. + auto headers = metadata_fixture_.GetMetadata(*context); + EXPECT_THAT(headers, + Contains(Pair("x-goog-cbt-cookie-routing", "routing"))); + auto stream = std::make_unique(); + EXPECT_CALL(*stream, Start).WillOnce([] { + return make_ready_future(false); + }); + EXPECT_CALL(*stream, Finish).WillOnce([] { + return make_ready_future(internal::PermissionDeniedError("fail")); + }); + return stream; + }); MockFunction(bigtable::Row const&)> on_row; EXPECT_CALL(on_row, Call).Times(0); @@ -1747,7 +1773,7 @@ TEST_F(AsyncRowReaderTest, TracedBackoff) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) .Times(kNumRetries + 1) - .WillRepeatedly([this](auto&, auto context, auto, auto const&) { + .WillRepeatedly([this](auto&, auto context, auto, auto const&, auto&) { metadata_fixture_.SetServerMetadata(*context); return std::make_unique(TransientError()); }); @@ -1784,11 +1810,12 @@ TEST_F(AsyncRowReaderTest, CallSpanActiveThroughout) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) .Times(kNumRetries + 1) - .WillRepeatedly([this, span](auto&, auto context, auto, auto const&) { - metadata_fixture_.SetServerMetadata(*context); - EXPECT_THAT(span, IsActive()); - return std::make_unique(TransientError()); - }); + .WillRepeatedly( + [this, span](auto&, auto context, auto, auto const&, auto&) { + metadata_fixture_.SetServerMetadata(*context); + EXPECT_THAT(span, IsActive()); + return std::make_unique(TransientError()); + }); promise p; internal::AutomaticallyCreatedBackgroundThreads background; diff --git a/google/cloud/bigtable/internal/async_row_sampler.cc b/google/cloud/bigtable/internal/async_row_sampler.cc index 66593522f0cc6..a40966d7bb433 100644 --- a/google/cloud/bigtable/internal/async_row_sampler.cc +++ b/google/cloud/bigtable/internal/async_row_sampler.cc @@ -71,7 +71,8 @@ void AsyncRowSampler::StartIteration() { auto self = this->shared_from_this(); PerformAsyncStreamingRead( - stub_->AsyncSampleRowKeys(cq_, client_context_, options_, request), + stub_->AsyncSampleRowKeys(cq_, client_context_, options_, request, + *operation_context_), [self](v2::SampleRowKeysResponse response) { return self->OnRead(std::move(response)); }, diff --git a/google/cloud/bigtable/internal/async_row_sampler_test.cc b/google/cloud/bigtable/internal/async_row_sampler_test.cc index 8b89b356b63c4..2b585dd2aebfd 100644 --- a/google/cloud/bigtable/internal/async_row_sampler_test.cc +++ b/google/cloud/bigtable/internal/async_row_sampler_test.cc @@ -13,6 +13,7 @@ // limitations under the License. #include "google/cloud/bigtable/internal/async_row_sampler.h" +#include "google/cloud/bigtable/internal/operation_context.h" #include "google/cloud/bigtable/testing/mock_bigtable_stub.h" #include "google/cloud/grpc_options.h" #include "google/cloud/internal/async_streaming_read_rpc_impl.h" @@ -103,7 +104,8 @@ class MockMetric : public Metric { ElementDeliveryParams const&), (override)); MOCK_METHOD(std::unique_ptr, clone, - (ResourceLabels resource_labels, DataLabels data_labels), + (TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels), (const, override)); }; @@ -112,7 +114,8 @@ class CloningMetric : public Metric { public: explicit CloningMetric(std::unique_ptr metric) : metric_(std::move(metric)) {} - std::unique_ptr clone(ResourceLabels, DataLabels) const override { + std::unique_ptr clone(TableResourceLabels const&, + TableDataLabels const&) const override { return std::move(metric_); } @@ -148,7 +151,8 @@ TEST_F(AsyncSampleRowKeysTest, Simple) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncSampleRowKeys) .WillOnce([this](CompletionQueue const&, auto client_context, auto, - v2::SampleRowKeysRequest const& request) { + v2::SampleRowKeysRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -216,7 +220,8 @@ TEST_F(AsyncSampleRowKeysTest, RetryResetsSamples) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncSampleRowKeys) .WillOnce([this](CompletionQueue const&, auto context, auto, - v2::SampleRowKeysRequest const& request) { + v2::SampleRowKeysRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -238,7 +243,8 @@ TEST_F(AsyncSampleRowKeysTest, RetryResetsSamples) { return stream; }) .WillOnce([this](CompletionQueue const&, auto client_context, auto, - v2::SampleRowKeysRequest const& request) { + v2::SampleRowKeysRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -308,7 +314,8 @@ TEST_F(AsyncSampleRowKeysTest, TooManyFailures) { EXPECT_CALL(*mock, AsyncSampleRowKeys) .Times(kNumRetries + 1) .WillRepeatedly([this](CompletionQueue const&, auto context, auto, - v2::SampleRowKeysRequest const& request) { + v2::SampleRowKeysRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -371,7 +378,8 @@ TEST_F(AsyncSampleRowKeysTest, RetryInfoHeeded) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncSampleRowKeys) .WillOnce([this](CompletionQueue const&, auto context, auto, - v2::SampleRowKeysRequest const&) { + v2::SampleRowKeysRequest const&, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); auto stream = std::make_unique(); EXPECT_CALL(*stream, Start).WillOnce([] { @@ -385,7 +393,8 @@ TEST_F(AsyncSampleRowKeysTest, RetryInfoHeeded) { return stream; }) .WillOnce([this](CompletionQueue const&, auto context, auto, - v2::SampleRowKeysRequest const&) { + v2::SampleRowKeysRequest const&, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); auto stream = std::make_unique(); EXPECT_CALL(*stream, Start).WillOnce([] { @@ -443,7 +452,8 @@ TEST_F(AsyncSampleRowKeysTest, RetryInfoIgnored) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncSampleRowKeys) .WillOnce([this](CompletionQueue const&, auto context, auto, - v2::SampleRowKeysRequest const&) { + v2::SampleRowKeysRequest const&, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); auto stream = std::make_unique(); EXPECT_CALL(*stream, Start).WillOnce([] { @@ -476,7 +486,8 @@ TEST_F(AsyncSampleRowKeysTest, TimerError) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncSampleRowKeys) .WillOnce([this](CompletionQueue const&, auto context, auto, - v2::SampleRowKeysRequest const& request) { + v2::SampleRowKeysRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -543,7 +554,8 @@ TEST_F(AsyncSampleRowKeysTest, CancelAfterSuccess) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncSampleRowKeys) .WillOnce([&p, this](CompletionQueue const&, auto client_context, auto, - v2::SampleRowKeysRequest const& request) { + v2::SampleRowKeysRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -613,7 +625,8 @@ TEST_F(AsyncSampleRowKeysTest, CancelMidStream) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncSampleRowKeys) .WillOnce([&p, this](CompletionQueue const&, auto client_context, auto, - v2::SampleRowKeysRequest const& request) { + v2::SampleRowKeysRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -679,7 +692,8 @@ TEST_F(AsyncSampleRowKeysTest, CurrentOptionsContinuedOnRetries) { EXPECT_CALL(*mock, AsyncSampleRowKeys) .Times(2) .WillRepeatedly([this](CompletionQueue const&, auto context, auto, - v2::SampleRowKeysRequest const&) { + v2::SampleRowKeysRequest const&, + bigtable_internal::OperationContext&) { EXPECT_EQ(5, internal::CurrentOptions().get()); metadata_fixture_.SetServerMetadata(*context); auto stream = std::make_unique(); @@ -725,7 +739,8 @@ TEST_F(AsyncSampleRowKeysTest, BigtableCookie) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncSampleRowKeys) .WillOnce([this](CompletionQueue const&, auto context, auto, - v2::SampleRowKeysRequest const&) { + v2::SampleRowKeysRequest const&, + bigtable_internal::OperationContext&) { // Return a bigtable cookie in the first request. metadata_fixture_.SetServerMetadata( *context, {{}, {{"x-goog-cbt-cookie-routing", "routing"}}}); @@ -739,7 +754,8 @@ TEST_F(AsyncSampleRowKeysTest, BigtableCookie) { return stream; }) .WillOnce([this](CompletionQueue const&, auto context, auto, - v2::SampleRowKeysRequest const&) { + v2::SampleRowKeysRequest const&, + bigtable_internal::OperationContext&) { // Verify that the next request includes the bigtable cookie from above. auto headers = metadata_fixture_.GetMetadata(*context); EXPECT_THAT(headers, @@ -787,7 +803,7 @@ TEST_F(AsyncSampleRowKeysTest, TracedBackoff) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncSampleRowKeys) .Times(kNumRetries + 1) - .WillRepeatedly([this](auto&, auto context, auto, auto const&) { + .WillRepeatedly([this](auto&, auto context, auto, auto const&, auto&) { metadata_fixture_.SetServerMetadata(*context, {}); return std::make_unique( internal::UnavailableError("try again")); @@ -816,12 +832,13 @@ TEST_F(AsyncSampleRowKeysTest, CallSpanActiveThroughout) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncSampleRowKeys) .Times(kNumRetries + 1) - .WillRepeatedly([this, span](auto&, auto context, auto, auto const&) { - metadata_fixture_.SetServerMetadata(*context, {}); - EXPECT_THAT(span, IsActive()); - return std::make_unique( - internal::UnavailableError("try again")); - }); + .WillRepeatedly( + [this, span](auto&, auto context, auto, auto const&, auto&) { + metadata_fixture_.SetServerMetadata(*context, {}); + EXPECT_THAT(span, IsActive()); + return std::make_unique( + internal::UnavailableError("try again")); + }); internal::AutomaticallyCreatedBackgroundThreads background; auto retry = DataLimitedErrorCountRetryPolicy(kNumRetries).clone(); diff --git a/google/cloud/bigtable/internal/bigtable_auth_decorator.cc b/google/cloud/bigtable/internal/bigtable_auth_decorator.cc index e67d547efbfa3..ac3b72446979b 100644 --- a/google/cloud/bigtable/internal/bigtable_auth_decorator.cc +++ b/google/cloud/bigtable/internal/bigtable_auth_decorator.cc @@ -40,99 +40,109 @@ std::unique_ptr> BigtableAuth::ReadRows(std::shared_ptr context, Options const& options, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) { using ErrorStream = ::google::cloud::internal::StreamingReadRpcError< google::bigtable::v2::ReadRowsResponse>; auto status = auth_->ConfigureContext(*context); if (!status.ok()) return std::make_unique(std::move(status)); - return child_->ReadRows(std::move(context), options, request); + return child_->ReadRows(std::move(context), options, request, oc); } std::unique_ptr> BigtableAuth::SampleRowKeys( std::shared_ptr context, Options const& options, - google::bigtable::v2::SampleRowKeysRequest const& request) { + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) { using ErrorStream = ::google::cloud::internal::StreamingReadRpcError< google::bigtable::v2::SampleRowKeysResponse>; auto status = auth_->ConfigureContext(*context); if (!status.ok()) return std::make_unique(std::move(status)); - return child_->SampleRowKeys(std::move(context), options, request); + return child_->SampleRowKeys(std::move(context), options, request, oc); } StatusOr BigtableAuth::MutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::MutateRowRequest const& request) { + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) { auto status = auth_->ConfigureContext(context); if (!status.ok()) return status; - return child_->MutateRow(context, options, request); + return child_->MutateRow(context, options, request, oc); } std::unique_ptr> -BigtableAuth::MutateRows( - std::shared_ptr context, Options const& options, - google::bigtable::v2::MutateRowsRequest const& request) { +BigtableAuth::MutateRows(std::shared_ptr context, + Options const& options, + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) { using ErrorStream = ::google::cloud::internal::StreamingReadRpcError< google::bigtable::v2::MutateRowsResponse>; auto status = auth_->ConfigureContext(*context); if (!status.ok()) return std::make_unique(std::move(status)); - return child_->MutateRows(std::move(context), options, request); + return child_->MutateRows(std::move(context), options, request, oc); } StatusOr BigtableAuth::CheckAndMutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) { + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) { auto status = auth_->ConfigureContext(context); if (!status.ok()) return status; - return child_->CheckAndMutateRow(context, options, request); + return child_->CheckAndMutateRow(context, options, request, oc); } StatusOr BigtableAuth::PingAndWarm( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PingAndWarmRequest const& request) { + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) { auto status = auth_->ConfigureContext(context); if (!status.ok()) return status; - return child_->PingAndWarm(context, options, request); + return child_->PingAndWarm(context, options, request, oc); } StatusOr BigtableAuth::ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) { + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) { auto status = auth_->ConfigureContext(context); if (!status.ok()) return status; - return child_->ReadModifyWriteRow(context, options, request); + return child_->ReadModifyWriteRow(context, options, request, oc); } StatusOr BigtableAuth::PrepareQuery( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PrepareQueryRequest const& request) { + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) { auto status = auth_->ConfigureContext(context); if (!status.ok()) return status; - return child_->PrepareQuery(context, options, request); + return child_->PrepareQuery(context, options, request, oc); } std::unique_ptr> BigtableAuth::ExecuteQuery( std::shared_ptr context, Options const& options, - google::bigtable::v2::ExecuteQueryRequest const& request) { + google::bigtable::v2::ExecuteQueryRequest const& request, + OperationContext& oc) { using ErrorStream = ::google::cloud::internal::StreamingReadRpcError< google::bigtable::v2::ExecuteQueryResponse>; auto status = auth_->ConfigureContext(*context); if (!status.ok()) return std::make_unique(std::move(status)); - return child_->ExecuteQuery(std::move(context), options, request); + return child_->ExecuteQuery(std::move(context), options, request, oc); } StatusOr BigtableAuth::GetClientConfiguration( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::GetClientConfigurationRequest const& request) { + google::bigtable::v2::GetClientConfigurationRequest const& request, + OperationContext& oc) { auto status = auth_->ConfigureContext(context); if (!status.ok()) return status; - return child_->GetClientConfiguration(context, options, request); + return child_->GetClientConfiguration(context, options, request, oc); } std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< @@ -198,14 +208,15 @@ BigtableAuth::AsyncReadRows( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) { using StreamAuth = google::cloud::internal::AsyncStreamingReadRpcAuth< google::bigtable::v2::ReadRowsResponse>; auto& child = child_; - auto call = [child, cq, opts = std::move(options), - request](std::shared_ptr ctx) { - return child->AsyncReadRows(cq, std::move(ctx), opts, request); + auto call = [child, cq, opts = std::move(options), request, + &oc](std::shared_ptr ctx) { + return child->AsyncReadRows(cq, std::move(ctx), opts, request, oc); }; return std::make_unique( std::move(context), auth_, StreamAuth::StreamFactory(std::move(call))); @@ -217,14 +228,15 @@ BigtableAuth::AsyncSampleRowKeys( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::SampleRowKeysRequest const& request) { + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) { using StreamAuth = google::cloud::internal::AsyncStreamingReadRpcAuth< google::bigtable::v2::SampleRowKeysResponse>; auto& child = child_; - auto call = [child, cq, opts = std::move(options), - request](std::shared_ptr ctx) { - return child->AsyncSampleRowKeys(cq, std::move(ctx), opts, request); + auto call = [child, cq, opts = std::move(options), request, + &oc](std::shared_ptr ctx) { + return child->AsyncSampleRowKeys(cq, std::move(ctx), opts, request, oc); }; return std::make_unique( std::move(context), auth_, StreamAuth::StreamFactory(std::move(call))); @@ -235,11 +247,12 @@ BigtableAuth::AsyncMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowRequest const& request) { + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) { 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, + &oc](future>> + f) mutable { auto context = f.get(); if (!context) { return make_ready_future( @@ -247,7 +260,7 @@ BigtableAuth::AsyncMutateRow( std::move(context).status())); } return child->AsyncMutateRow(cq, *std::move(context), - std::move(options), request); + std::move(options), request, oc); }); } @@ -257,14 +270,15 @@ BigtableAuth::AsyncMutateRows( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) { using StreamAuth = google::cloud::internal::AsyncStreamingReadRpcAuth< google::bigtable::v2::MutateRowsResponse>; auto& child = child_; - auto call = [child, cq, opts = std::move(options), - request](std::shared_ptr ctx) { - return child->AsyncMutateRows(cq, std::move(ctx), opts, request); + auto call = [child, cq, opts = std::move(options), request, + &oc](std::shared_ptr ctx) { + return child->AsyncMutateRows(cq, std::move(ctx), opts, request, oc); }; return std::make_unique( std::move(context), auth_, StreamAuth::StreamFactory(std::move(call))); @@ -275,11 +289,12 @@ BigtableAuth::AsyncCheckAndMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) { + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) { 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, + &oc](future>> + f) mutable { auto context = f.get(); if (!context) { return make_ready_future( @@ -287,7 +302,7 @@ BigtableAuth::AsyncCheckAndMutateRow( std::move(context).status())); } return child->AsyncCheckAndMutateRow(cq, *std::move(context), - std::move(options), request); + std::move(options), request, oc); }); } @@ -296,11 +311,12 @@ BigtableAuth::AsyncPingAndWarm( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PingAndWarmRequest const& request) { + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) { 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, + &oc](future>> + f) mutable { auto context = f.get(); if (!context) { return make_ready_future( @@ -308,7 +324,7 @@ BigtableAuth::AsyncPingAndWarm( std::move(context).status())); } return child->AsyncPingAndWarm(cq, *std::move(context), - std::move(options), request); + std::move(options), request, oc); }); } @@ -317,11 +333,12 @@ BigtableAuth::AsyncReadModifyWriteRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) { + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) { 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, + &oc](future>> + f) mutable { auto context = f.get(); if (!context) { return make_ready_future( @@ -329,7 +346,7 @@ BigtableAuth::AsyncReadModifyWriteRow( std::move(context).status())); } return child->AsyncReadModifyWriteRow(cq, *std::move(context), - std::move(options), request); + std::move(options), request, oc); }); } @@ -338,11 +355,12 @@ BigtableAuth::AsyncPrepareQuery( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request) { + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) { 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, + &oc](future>> + f) mutable { auto context = f.get(); if (!context) { return make_ready_future( @@ -350,7 +368,7 @@ BigtableAuth::AsyncPrepareQuery( std::move(context).status())); } return child->AsyncPrepareQuery(cq, *std::move(context), - std::move(options), request); + std::move(options), request, oc); }); } diff --git a/google/cloud/bigtable/internal/bigtable_auth_decorator.h b/google/cloud/bigtable/internal/bigtable_auth_decorator.h index 6c954417e6a70..a8bb8fc012eef 100644 --- a/google/cloud/bigtable/internal/bigtable_auth_decorator.h +++ b/google/cloud/bigtable/internal/bigtable_auth_decorator.h @@ -44,50 +44,59 @@ class BigtableAuth : public BigtableStub { std::unique_ptr> ReadRows(std::shared_ptr context, Options const& options, - google::bigtable::v2::ReadRowsRequest const& request) override; + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) override; std::unique_ptr> - SampleRowKeys( - std::shared_ptr context, Options const& options, - google::bigtable::v2::SampleRowKeysRequest const& request) override; + SampleRowKeys(std::shared_ptr context, + Options const& options, + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) override; StatusOr MutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::MutateRowRequest const& request) override; + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) override; std::unique_ptr> MutateRows(std::shared_ptr context, Options const& options, - google::bigtable::v2::MutateRowsRequest const& request) override; + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) override; StatusOr CheckAndMutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) override; + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) override; StatusOr PingAndWarm( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PingAndWarmRequest const& request) override; + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) override; StatusOr ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) override; + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) override; StatusOr PrepareQuery( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PrepareQueryRequest const& request) override; + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) override; std::unique_ptr> - ExecuteQuery( - std::shared_ptr context, Options const& options, - google::bigtable::v2::ExecuteQueryRequest const& request) override; + ExecuteQuery(std::shared_ptr context, + Options const& options, + google::bigtable::v2::ExecuteQueryRequest const& request, + OperationContext& oc) override; StatusOr GetClientConfiguration( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::GetClientConfigurationRequest const& request) - override; + google::bigtable::v2::GetClientConfigurationRequest const& request, + OperationContext& oc) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, @@ -117,56 +126,61 @@ class BigtableAuth : public BigtableStub { AsyncReadRows(google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadRowsRequest const& request) override; + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) 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) override; + AsyncSampleRowKeys(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) override; future> AsyncMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowRequest const& request) override; + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) 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) override; + AsyncMutateRows(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) override; future> AsyncCheckAndMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) override; + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) override; future> AsyncPingAndWarm( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PingAndWarmRequest const& request) override; + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) override; future> AsyncReadModifyWriteRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) override; + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) override; future> - AsyncPrepareQuery( - google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request) override; + AsyncPrepareQuery(google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) 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 830402857fa8d..06ad8826f1711 100644 --- a/google/cloud/bigtable/internal/bigtable_channel_refresh.cc +++ b/google/cloud/bigtable/internal/bigtable_channel_refresh.cc @@ -23,67 +23,76 @@ std::unique_ptr> BigtableChannelRefresh::ReadRows( std::shared_ptr client_context, Options const& options, - google::bigtable::v2::ReadRowsRequest const& request) { - return child_->ReadRows(std::move(client_context), options, request); + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) { + return child_->ReadRows(std::move(client_context), options, request, oc); } std::unique_ptr> BigtableChannelRefresh::SampleRowKeys( std::shared_ptr client_context, Options const& options, - google::bigtable::v2::SampleRowKeysRequest const& request) { - return child_->SampleRowKeys(std::move(client_context), options, request); + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) { + return child_->SampleRowKeys(std::move(client_context), options, request, oc); } StatusOr BigtableChannelRefresh::MutateRow( grpc::ClientContext& client_context, Options const& options, - google::bigtable::v2::MutateRowRequest const& request) { - return child_->MutateRow(client_context, options, request); + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) { + return child_->MutateRow(client_context, options, request, oc); } std::unique_ptr> BigtableChannelRefresh::MutateRows( std::shared_ptr client_context, Options const& options, - google::bigtable::v2::MutateRowsRequest const& request) { - return child_->MutateRows(std::move(client_context), options, request); + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) { + return child_->MutateRows(std::move(client_context), options, request, oc); } StatusOr BigtableChannelRefresh::CheckAndMutateRow( grpc::ClientContext& client_context, Options const& options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) { - return child_->CheckAndMutateRow(client_context, options, request); + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) { + return child_->CheckAndMutateRow(client_context, options, request, oc); } StatusOr BigtableChannelRefresh::PingAndWarm( grpc::ClientContext& client_context, Options const& options, - google::bigtable::v2::PingAndWarmRequest const& request) { - return child_->PingAndWarm(client_context, options, request); + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) { + return child_->PingAndWarm(client_context, options, request, oc); } StatusOr BigtableChannelRefresh::ReadModifyWriteRow( grpc::ClientContext& client_context, Options const& options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) { - return child_->ReadModifyWriteRow(client_context, options, request); + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) { + return child_->ReadModifyWriteRow(client_context, options, request, oc); } StatusOr BigtableChannelRefresh::PrepareQuery( grpc::ClientContext& client_context, Options const& options, - google::bigtable::v2::PrepareQueryRequest const& request) { - return child_->PrepareQuery(client_context, options, request); + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) { + return child_->PrepareQuery(client_context, options, request, oc); } std::unique_ptr> BigtableChannelRefresh::ExecuteQuery( std::shared_ptr client_context, Options const& options, - google::bigtable::v2::ExecuteQueryRequest const& request) { - return child_->ExecuteQuery(client_context, options, request); + google::bigtable::v2::ExecuteQueryRequest const& request, + OperationContext& oc) { + return child_->ExecuteQuery(client_context, options, request, oc); } std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< @@ -92,9 +101,10 @@ BigtableChannelRefresh::AsyncReadRows( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) { return child_->AsyncReadRows(cq, std::move(context), std::move(options), - request); + request, oc); } std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< @@ -103,9 +113,10 @@ BigtableChannelRefresh::AsyncSampleRowKeys( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::SampleRowKeysRequest const& request) { + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) { return child_->AsyncSampleRowKeys(cq, std::move(context), std::move(options), - request); + request, oc); } future> @@ -113,9 +124,10 @@ BigtableChannelRefresh::AsyncMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowRequest const& request) { + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) { return child_->AsyncMutateRow(cq, std::move(context), std::move(options), - request); + request, oc); } std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< @@ -124,9 +136,10 @@ BigtableChannelRefresh::AsyncMutateRows( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) { return child_->AsyncMutateRows(cq, std::move(context), std::move(options), - request); + request, oc); } future> @@ -134,9 +147,10 @@ BigtableChannelRefresh::AsyncCheckAndMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) { + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) { return child_->AsyncCheckAndMutateRow(cq, std::move(context), - std::move(options), request); + std::move(options), request, oc); } future> @@ -144,9 +158,10 @@ BigtableChannelRefresh::AsyncPingAndWarm( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PingAndWarmRequest const& request) { + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) { return child_->AsyncPingAndWarm(cq, std::move(context), std::move(options), - request); + request, oc); } future> @@ -154,9 +169,10 @@ BigtableChannelRefresh::AsyncReadModifyWriteRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) { + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) { return child_->AsyncReadModifyWriteRow(cq, std::move(context), - std::move(options), request); + std::move(options), request, oc); } future> @@ -164,16 +180,18 @@ BigtableChannelRefresh::AsyncPrepareQuery( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request) { + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) { return child_->AsyncPrepareQuery(cq, std::move(context), std::move(options), - request); + request, oc); } StatusOr BigtableChannelRefresh::GetClientConfiguration( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::GetClientConfigurationRequest const& request) { - return child_->GetClientConfiguration(context, options, request); + google::bigtable::v2::GetClientConfigurationRequest const& request, + OperationContext& oc) { + return child_->GetClientConfiguration(context, options, request, oc); } std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< diff --git a/google/cloud/bigtable/internal/bigtable_channel_refresh.h b/google/cloud/bigtable/internal/bigtable_channel_refresh.h index b80ca5d0a6562..c68ebb7828617 100644 --- a/google/cloud/bigtable/internal/bigtable_channel_refresh.h +++ b/google/cloud/bigtable/internal/bigtable_channel_refresh.h @@ -46,108 +46,120 @@ class BigtableChannelRefresh : public BigtableStub { google::bigtable::v2::ReadRowsResponse>> ReadRows(std::shared_ptr client_context, Options const& options, - google::bigtable::v2::ReadRowsRequest const& request) override; + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) override; std::unique_ptr> - SampleRowKeys( - std::shared_ptr client_context, - Options const& options, - google::bigtable::v2::SampleRowKeysRequest const& request) override; + SampleRowKeys(std::shared_ptr client_context, + Options const& options, + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) override; StatusOr MutateRow( grpc::ClientContext& client_context, Options const& options, - google::bigtable::v2::MutateRowRequest const& request) override; + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) override; std::unique_ptr> MutateRows(std::shared_ptr client_context, Options const& options, - google::bigtable::v2::MutateRowsRequest const& request) override; + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) override; StatusOr CheckAndMutateRow( grpc::ClientContext& client_context, Options const& options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) override; + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) override; StatusOr PingAndWarm( grpc::ClientContext& client_context, Options const& options, - google::bigtable::v2::PingAndWarmRequest const& request) override; + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) override; StatusOr ReadModifyWriteRow( grpc::ClientContext& client_context, Options const& options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) override; + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) override; StatusOr PrepareQuery( grpc::ClientContext& client, Options const& options, - google::bigtable::v2::PrepareQueryRequest const& request) override; + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) override; std::unique_ptr> - ExecuteQuery( - std::shared_ptr client_context, - Options const& options, - google::bigtable::v2::ExecuteQueryRequest const& request) override; + ExecuteQuery(std::shared_ptr client_context, + Options const& options, + google::bigtable::v2::ExecuteQueryRequest const& request, + OperationContext& oc) 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) override; + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) 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) override; + AsyncSampleRowKeys(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) override; future> AsyncMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions, - google::bigtable::v2::MutateRowRequest const& request) override; + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) 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) override; + AsyncMutateRows(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) override; future> AsyncCheckAndMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions, - google::bigtable::v2::CheckAndMutateRowRequest const& request) override; + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) override; future> AsyncPingAndWarm( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PingAndWarmRequest const& request) override; + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) override; future> AsyncReadModifyWriteRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) override; + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) override; future> - AsyncPrepareQuery( - google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request) override; + AsyncPrepareQuery(google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) override; StatusOr GetClientConfiguration( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::GetClientConfigurationRequest const& request) - override; + google::bigtable::v2::GetClientConfigurationRequest const& request, + OperationContext& oc) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, diff --git a/google/cloud/bigtable/internal/bigtable_logging_decorator.cc b/google/cloud/bigtable/internal/bigtable_logging_decorator.cc index 13b9051d7a425..566cfe3667129 100644 --- a/google/cloud/bigtable/internal/bigtable_logging_decorator.cc +++ b/google/cloud/bigtable/internal/bigtable_logging_decorator.cc @@ -45,16 +45,18 @@ BigtableLogging::BigtableLogging(std::shared_ptr child, std::unique_ptr> -BigtableLogging::ReadRows( - std::shared_ptr context, Options const& options, - google::bigtable::v2::ReadRowsRequest const& request) { +BigtableLogging::ReadRows(std::shared_ptr context, + Options const& options, + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) { return google::cloud::internal::LogWrapper( - [this](std::shared_ptr context, - Options const& options, - google::bigtable::v2::ReadRowsRequest const& request) + [this, &oc](std::shared_ptr context, + Options const& options, + google::bigtable::v2::ReadRowsRequest const& request) -> std::unique_ptr> { - auto stream = child_->ReadRows(std::move(context), options, request); + auto stream = + child_->ReadRows(std::move(context), options, request, oc); if (stream_logging_) { stream = std::make_unique> BigtableLogging::SampleRowKeys( std::shared_ptr context, Options const& options, - google::bigtable::v2::SampleRowKeysRequest const& request) { + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) { return google::cloud::internal::LogWrapper( - [this](std::shared_ptr context, - Options const& options, - google::bigtable::v2::SampleRowKeysRequest const& request) + [this, &oc](std::shared_ptr context, + Options const& options, + google::bigtable::v2::SampleRowKeysRequest const& request) -> std::unique_ptr> { auto stream = - child_->SampleRowKeys(std::move(context), options, request); + child_->SampleRowKeys(std::move(context), options, request, oc); if (stream_logging_) { stream = std::make_unique BigtableLogging::MutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::MutateRowRequest const& request) { + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) { return google::cloud::internal::LogWrapper( - [this](grpc::ClientContext& context, Options const& options, - google::bigtable::v2::MutateRowRequest const& request) { - return child_->MutateRow(context, options, request); + [this, &oc](grpc::ClientContext& context, Options const& options, + google::bigtable::v2::MutateRowRequest const& request) { + return child_->MutateRow(context, options, request, oc); }, context, options, request, __func__, tracing_options_); } @@ -107,14 +111,16 @@ std::unique_ptr> BigtableLogging::MutateRows( std::shared_ptr context, Options const& options, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) { return google::cloud::internal::LogWrapper( - [this](std::shared_ptr context, - Options const& options, - google::bigtable::v2::MutateRowsRequest const& request) + [this, &oc](std::shared_ptr context, + Options const& options, + google::bigtable::v2::MutateRowsRequest const& request) -> std::unique_ptr> { - auto stream = child_->MutateRows(std::move(context), options, request); + auto stream = + child_->MutateRows(std::move(context), options, request, oc); if (stream_logging_) { stream = std::make_unique BigtableLogging::CheckAndMutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) { + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) { return google::cloud::internal::LogWrapper( - [this](grpc::ClientContext& context, Options const& options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) { - return child_->CheckAndMutateRow(context, options, request); + [this, &oc]( + grpc::ClientContext& context, Options const& options, + google::bigtable::v2::CheckAndMutateRowRequest const& request) { + return child_->CheckAndMutateRow(context, options, request, oc); }, context, options, request, __func__, tracing_options_); } @@ -142,11 +150,12 @@ BigtableLogging::CheckAndMutateRow( StatusOr BigtableLogging::PingAndWarm( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PingAndWarmRequest const& request) { + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) { return google::cloud::internal::LogWrapper( - [this](grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PingAndWarmRequest const& request) { - return child_->PingAndWarm(context, options, request); + [this, &oc](grpc::ClientContext& context, Options const& options, + google::bigtable::v2::PingAndWarmRequest const& request) { + return child_->PingAndWarm(context, options, request, oc); }, context, options, request, __func__, tracing_options_); } @@ -154,11 +163,13 @@ BigtableLogging::PingAndWarm( StatusOr BigtableLogging::ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) { + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) { return google::cloud::internal::LogWrapper( - [this](grpc::ClientContext& context, Options const& options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) { - return child_->ReadModifyWriteRow(context, options, request); + [this, &oc]( + grpc::ClientContext& context, Options const& options, + google::bigtable::v2::ReadModifyWriteRowRequest const& request) { + return child_->ReadModifyWriteRow(context, options, request, oc); }, context, options, request, __func__, tracing_options_); } @@ -166,11 +177,12 @@ BigtableLogging::ReadModifyWriteRow( StatusOr BigtableLogging::PrepareQuery( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PrepareQueryRequest const& request) { + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) { return google::cloud::internal::LogWrapper( - [this](grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PrepareQueryRequest const& request) { - return child_->PrepareQuery(context, options, request); + [this, &oc](grpc::ClientContext& context, Options const& options, + google::bigtable::v2::PrepareQueryRequest const& request) { + return child_->PrepareQuery(context, options, request, oc); }, context, options, request, __func__, tracing_options_); } @@ -179,15 +191,16 @@ std::unique_ptr> BigtableLogging::ExecuteQuery( std::shared_ptr context, Options const& options, - google::bigtable::v2::ExecuteQueryRequest const& request) { + google::bigtable::v2::ExecuteQueryRequest const& request, + OperationContext& oc) { return google::cloud::internal::LogWrapper( - [this](std::shared_ptr context, - Options const& options, - google::bigtable::v2::ExecuteQueryRequest const& request) + [this, &oc](std::shared_ptr context, + Options const& options, + google::bigtable::v2::ExecuteQueryRequest const& request) -> std::unique_ptr> { auto stream = - child_->ExecuteQuery(std::move(context), options, request); + child_->ExecuteQuery(std::move(context), options, request, oc); if (stream_logging_) { stream = std::make_unique BigtableLogging::GetClientConfiguration( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::GetClientConfigurationRequest const& request) { + google::bigtable::v2::GetClientConfigurationRequest const& request, + OperationContext& oc) { return google::cloud::internal::LogWrapper( - [this]( + [this, &oc]( grpc::ClientContext& context, Options const& options, google::bigtable::v2::GetClientConfigurationRequest const& request) { - return child_->GetClientConfiguration(context, options, request); + return child_->GetClientConfiguration(context, options, request, oc); }, context, options, request, __func__, tracing_options_); } @@ -288,7 +302,8 @@ BigtableLogging::AsyncReadRows( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) { using LoggingStream = ::google::cloud::internal::AsyncStreamingReadRpcLogging< google::bigtable::v2::ReadRowsResponse>; @@ -297,7 +312,7 @@ BigtableLogging::AsyncReadRows( __func__, request_id, google::cloud::internal::DebugString(request, tracing_options_)); auto stream = child_->AsyncReadRows(cq, std::move(context), - std::move(options), request); + std::move(options), request, oc); if (stream_logging_) { stream = std::make_unique( std::move(stream), tracing_options_, std::move(request_id)); @@ -311,7 +326,8 @@ BigtableLogging::AsyncSampleRowKeys( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::SampleRowKeysRequest const& request) { + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) { using LoggingStream = ::google::cloud::internal::AsyncStreamingReadRpcLogging< google::bigtable::v2::SampleRowKeysResponse>; @@ -320,7 +336,7 @@ BigtableLogging::AsyncSampleRowKeys( __func__, request_id, google::cloud::internal::DebugString(request, tracing_options_)); auto stream = child_->AsyncSampleRowKeys(cq, std::move(context), - std::move(options), request); + std::move(options), request, oc); if (stream_logging_) { stream = std::make_unique( std::move(stream), tracing_options_, std::move(request_id)); @@ -333,14 +349,15 @@ BigtableLogging::AsyncMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowRequest const& request) { + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) { return google::cloud::internal::LogWrapper( - [this](google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowRequest const& request) { + [this, &oc](google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::MutateRowRequest const& request) { return child_->AsyncMutateRow(cq, std::move(context), - std::move(options), request); + std::move(options), request, oc); }, cq, std::move(context), std::move(options), request, __func__, tracing_options_); @@ -352,7 +369,8 @@ BigtableLogging::AsyncMutateRows( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) { using LoggingStream = ::google::cloud::internal::AsyncStreamingReadRpcLogging< google::bigtable::v2::MutateRowsResponse>; @@ -361,7 +379,7 @@ BigtableLogging::AsyncMutateRows( __func__, request_id, google::cloud::internal::DebugString(request, tracing_options_)); auto stream = child_->AsyncMutateRows(cq, std::move(context), - std::move(options), request); + std::move(options), request, oc); if (stream_logging_) { stream = std::make_unique( std::move(stream), tracing_options_, std::move(request_id)); @@ -374,14 +392,16 @@ BigtableLogging::AsyncCheckAndMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) { + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) { return google::cloud::internal::LogWrapper( - [this](google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) { + [this, &oc]( + google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::CheckAndMutateRowRequest const& request) { return child_->AsyncCheckAndMutateRow(cq, std::move(context), - std::move(options), request); + std::move(options), request, oc); }, cq, std::move(context), std::move(options), request, __func__, tracing_options_); @@ -392,14 +412,15 @@ BigtableLogging::AsyncPingAndWarm( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PingAndWarmRequest const& request) { + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) { return google::cloud::internal::LogWrapper( - [this](google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PingAndWarmRequest const& request) { + [this, &oc](google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::PingAndWarmRequest const& request) { return child_->AsyncPingAndWarm(cq, std::move(context), - std::move(options), request); + std::move(options), request, oc); }, cq, std::move(context), std::move(options), request, __func__, tracing_options_); @@ -410,14 +431,16 @@ BigtableLogging::AsyncReadModifyWriteRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) { + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) { return google::cloud::internal::LogWrapper( - [this](google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) { + [this, &oc]( + google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::ReadModifyWriteRowRequest const& request) { return child_->AsyncReadModifyWriteRow(cq, std::move(context), - std::move(options), request); + std::move(options), request, oc); }, cq, std::move(context), std::move(options), request, __func__, tracing_options_); @@ -428,14 +451,15 @@ BigtableLogging::AsyncPrepareQuery( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request) { + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) { return google::cloud::internal::LogWrapper( - [this](google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request) { + [this, &oc](google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::PrepareQueryRequest const& request) { return child_->AsyncPrepareQuery(cq, std::move(context), - std::move(options), request); + std::move(options), request, oc); }, cq, std::move(context), std::move(options), request, __func__, tracing_options_); diff --git a/google/cloud/bigtable/internal/bigtable_logging_decorator.h b/google/cloud/bigtable/internal/bigtable_logging_decorator.h index 29f7b8132ddf8..81ef558111bac 100644 --- a/google/cloud/bigtable/internal/bigtable_logging_decorator.h +++ b/google/cloud/bigtable/internal/bigtable_logging_decorator.h @@ -44,50 +44,59 @@ class BigtableLogging : public BigtableStub { std::unique_ptr> ReadRows(std::shared_ptr context, Options const& options, - google::bigtable::v2::ReadRowsRequest const& request) override; + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) override; std::unique_ptr> - SampleRowKeys( - std::shared_ptr context, Options const& options, - google::bigtable::v2::SampleRowKeysRequest const& request) override; + SampleRowKeys(std::shared_ptr context, + Options const& options, + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) override; StatusOr MutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::MutateRowRequest const& request) override; + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) override; std::unique_ptr> MutateRows(std::shared_ptr context, Options const& options, - google::bigtable::v2::MutateRowsRequest const& request) override; + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) override; StatusOr CheckAndMutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) override; + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) override; StatusOr PingAndWarm( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PingAndWarmRequest const& request) override; + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) override; StatusOr ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) override; + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) override; StatusOr PrepareQuery( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PrepareQueryRequest const& request) override; + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) override; std::unique_ptr> - ExecuteQuery( - std::shared_ptr context, Options const& options, - google::bigtable::v2::ExecuteQueryRequest const& request) override; + ExecuteQuery(std::shared_ptr context, + Options const& options, + google::bigtable::v2::ExecuteQueryRequest const& request, + OperationContext& oc) override; StatusOr GetClientConfiguration( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::GetClientConfigurationRequest const& request) - override; + google::bigtable::v2::GetClientConfigurationRequest const& request, + OperationContext& oc) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, @@ -117,56 +126,61 @@ class BigtableLogging : public BigtableStub { AsyncReadRows(google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadRowsRequest const& request) override; + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) 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) override; + AsyncSampleRowKeys(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) override; future> AsyncMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowRequest const& request) override; + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) 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) override; + AsyncMutateRows(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) override; future> AsyncCheckAndMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) override; + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) override; future> AsyncPingAndWarm( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PingAndWarmRequest const& request) override; + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) override; future> AsyncReadModifyWriteRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) override; + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) override; future> - AsyncPrepareQuery( - google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request) override; + AsyncPrepareQuery(google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) 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 1ffca73ff9c70..cf2cecc3f2203 100644 --- a/google/cloud/bigtable/internal/bigtable_metadata_decorator.cc +++ b/google/cloud/bigtable/internal/bigtable_metadata_decorator.cc @@ -51,9 +51,10 @@ BigtableMetadata::BigtableMetadata( std::unique_ptr> -BigtableMetadata::ReadRows( - std::shared_ptr context, Options const& options, - google::bigtable::v2::ReadRowsRequest const& request) { +BigtableMetadata::ReadRows(std::shared_ptr context, + Options const& options, + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) { std::vector params; params.reserve(3); @@ -101,14 +102,15 @@ BigtableMetadata::ReadRows( } else { SetMetadata(*context, options, absl::StrJoin(params, "&")); } - return child_->ReadRows(std::move(context), options, request); + return child_->ReadRows(std::move(context), options, request, oc); } std::unique_ptr> BigtableMetadata::SampleRowKeys( std::shared_ptr context, Options const& options, - google::bigtable::v2::SampleRowKeysRequest const& request) { + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) { std::vector params; params.reserve(3); @@ -156,12 +158,13 @@ BigtableMetadata::SampleRowKeys( } else { SetMetadata(*context, options, absl::StrJoin(params, "&")); } - return child_->SampleRowKeys(std::move(context), options, request); + return child_->SampleRowKeys(std::move(context), options, request, oc); } StatusOr BigtableMetadata::MutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::MutateRowRequest const& request) { + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) { std::vector params; params.reserve(2); @@ -194,14 +197,15 @@ StatusOr BigtableMetadata::MutateRow( } else { SetMetadata(context, options, absl::StrJoin(params, "&")); } - return child_->MutateRow(context, options, request); + return child_->MutateRow(context, options, request, oc); } std::unique_ptr> BigtableMetadata::MutateRows( std::shared_ptr context, Options const& options, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) { std::vector params; params.reserve(2); @@ -234,13 +238,14 @@ BigtableMetadata::MutateRows( } else { SetMetadata(*context, options, absl::StrJoin(params, "&")); } - return child_->MutateRows(std::move(context), options, request); + return child_->MutateRows(std::move(context), options, request, oc); } StatusOr BigtableMetadata::CheckAndMutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) { + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) { std::vector params; params.reserve(2); @@ -273,13 +278,14 @@ BigtableMetadata::CheckAndMutateRow( } else { SetMetadata(context, options, absl::StrJoin(params, "&")); } - return child_->CheckAndMutateRow(context, options, request); + return child_->CheckAndMutateRow(context, options, request, oc); } StatusOr BigtableMetadata::PingAndWarm( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PingAndWarmRequest const& request) { + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) { std::vector params; params.reserve(2); @@ -306,13 +312,14 @@ BigtableMetadata::PingAndWarm( } else { SetMetadata(context, options, absl::StrJoin(params, "&")); } - return child_->PingAndWarm(context, options, request); + return child_->PingAndWarm(context, options, request, oc); } StatusOr BigtableMetadata::ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) { + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) { std::vector params; params.reserve(2); @@ -345,13 +352,14 @@ BigtableMetadata::ReadModifyWriteRow( } else { SetMetadata(context, options, absl::StrJoin(params, "&")); } - return child_->ReadModifyWriteRow(context, options, request); + return child_->ReadModifyWriteRow(context, options, request, oc); } StatusOr BigtableMetadata::PrepareQuery( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PrepareQueryRequest const& request) { + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) { std::vector params; params.reserve(2); @@ -378,14 +386,15 @@ BigtableMetadata::PrepareQuery( } else { SetMetadata(context, options, absl::StrJoin(params, "&")); } - return child_->PrepareQuery(context, options, request); + return child_->PrepareQuery(context, options, request, oc); } std::unique_ptr> BigtableMetadata::ExecuteQuery( std::shared_ptr context, Options const& options, - google::bigtable::v2::ExecuteQueryRequest const& request) { + google::bigtable::v2::ExecuteQueryRequest const& request, + OperationContext& oc) { std::vector params; params.reserve(2); @@ -412,15 +421,16 @@ BigtableMetadata::ExecuteQuery( } else { SetMetadata(*context, options, absl::StrJoin(params, "&")); } - return child_->ExecuteQuery(std::move(context), options, request); + return child_->ExecuteQuery(std::move(context), options, request, oc); } StatusOr BigtableMetadata::GetClientConfiguration( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::GetClientConfigurationRequest const& request) { + google::bigtable::v2::GetClientConfigurationRequest const& request, + OperationContext& oc) { SetMetadata(context, options); - return child_->GetClientConfiguration(context, options, request); + return child_->GetClientConfiguration(context, options, request, oc); } std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< @@ -464,7 +474,8 @@ BigtableMetadata::AsyncReadRows( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) { std::vector params; params.reserve(3); @@ -513,7 +524,7 @@ BigtableMetadata::AsyncReadRows( SetMetadata(*context, *options, absl::StrJoin(params, "&")); } return child_->AsyncReadRows(cq, std::move(context), std::move(options), - request); + request, oc); } std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< @@ -522,7 +533,8 @@ BigtableMetadata::AsyncSampleRowKeys( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::SampleRowKeysRequest const& request) { + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) { std::vector params; params.reserve(3); @@ -571,7 +583,7 @@ BigtableMetadata::AsyncSampleRowKeys( SetMetadata(*context, *options, absl::StrJoin(params, "&")); } return child_->AsyncSampleRowKeys(cq, std::move(context), std::move(options), - request); + request, oc); } future> @@ -579,7 +591,8 @@ BigtableMetadata::AsyncMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowRequest const& request) { + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) { std::vector params; params.reserve(2); @@ -613,7 +626,7 @@ BigtableMetadata::AsyncMutateRow( SetMetadata(*context, *options, absl::StrJoin(params, "&")); } return child_->AsyncMutateRow(cq, std::move(context), std::move(options), - request); + request, oc); } std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< @@ -622,7 +635,8 @@ BigtableMetadata::AsyncMutateRows( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) { std::vector params; params.reserve(2); @@ -656,7 +670,7 @@ BigtableMetadata::AsyncMutateRows( SetMetadata(*context, *options, absl::StrJoin(params, "&")); } return child_->AsyncMutateRows(cq, std::move(context), std::move(options), - request); + request, oc); } future> @@ -664,7 +678,8 @@ BigtableMetadata::AsyncCheckAndMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) { + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) { std::vector params; params.reserve(2); @@ -698,7 +713,7 @@ BigtableMetadata::AsyncCheckAndMutateRow( SetMetadata(*context, *options, absl::StrJoin(params, "&")); } return child_->AsyncCheckAndMutateRow(cq, std::move(context), - std::move(options), request); + std::move(options), request, oc); } future> @@ -706,7 +721,8 @@ BigtableMetadata::AsyncPingAndWarm( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PingAndWarmRequest const& request) { + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) { std::vector params; params.reserve(2); @@ -734,7 +750,7 @@ BigtableMetadata::AsyncPingAndWarm( SetMetadata(*context, *options, absl::StrJoin(params, "&")); } return child_->AsyncPingAndWarm(cq, std::move(context), std::move(options), - request); + request, oc); } future> @@ -742,7 +758,8 @@ BigtableMetadata::AsyncReadModifyWriteRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) { + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) { std::vector params; params.reserve(2); @@ -776,7 +793,7 @@ BigtableMetadata::AsyncReadModifyWriteRow( SetMetadata(*context, *options, absl::StrJoin(params, "&")); } return child_->AsyncReadModifyWriteRow(cq, std::move(context), - std::move(options), request); + std::move(options), request, oc); } future> @@ -784,7 +801,8 @@ BigtableMetadata::AsyncPrepareQuery( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request) { + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) { std::vector params; params.reserve(2); @@ -812,7 +830,7 @@ BigtableMetadata::AsyncPrepareQuery( SetMetadata(*context, *options, absl::StrJoin(params, "&")); } return child_->AsyncPrepareQuery(cq, std::move(context), std::move(options), - request); + request, oc); } void BigtableMetadata::SetMetadata(grpc::ClientContext& context, diff --git a/google/cloud/bigtable/internal/bigtable_metadata_decorator.h b/google/cloud/bigtable/internal/bigtable_metadata_decorator.h index 25d1d946bdd6c..d1d14dae2a96d 100644 --- a/google/cloud/bigtable/internal/bigtable_metadata_decorator.h +++ b/google/cloud/bigtable/internal/bigtable_metadata_decorator.h @@ -44,50 +44,59 @@ class BigtableMetadata : public BigtableStub { std::unique_ptr> ReadRows(std::shared_ptr context, Options const& options, - google::bigtable::v2::ReadRowsRequest const& request) override; + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) override; std::unique_ptr> - SampleRowKeys( - std::shared_ptr context, Options const& options, - google::bigtable::v2::SampleRowKeysRequest const& request) override; + SampleRowKeys(std::shared_ptr context, + Options const& options, + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) override; StatusOr MutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::MutateRowRequest const& request) override; + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) override; std::unique_ptr> MutateRows(std::shared_ptr context, Options const& options, - google::bigtable::v2::MutateRowsRequest const& request) override; + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) override; StatusOr CheckAndMutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) override; + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) override; StatusOr PingAndWarm( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PingAndWarmRequest const& request) override; + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) override; StatusOr ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) override; + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) override; StatusOr PrepareQuery( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PrepareQueryRequest const& request) override; + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) override; std::unique_ptr> - ExecuteQuery( - std::shared_ptr context, Options const& options, - google::bigtable::v2::ExecuteQueryRequest const& request) override; + ExecuteQuery(std::shared_ptr context, + Options const& options, + google::bigtable::v2::ExecuteQueryRequest const& request, + OperationContext& oc) override; StatusOr GetClientConfiguration( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::GetClientConfigurationRequest const& request) - override; + google::bigtable::v2::GetClientConfigurationRequest const& request, + OperationContext& oc) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, @@ -117,56 +126,61 @@ class BigtableMetadata : public BigtableStub { AsyncReadRows(google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadRowsRequest const& request) override; + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) 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) override; + AsyncSampleRowKeys(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) override; future> AsyncMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowRequest const& request) override; + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) 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) override; + AsyncMutateRows(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) override; future> AsyncCheckAndMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) override; + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) override; future> AsyncPingAndWarm( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PingAndWarmRequest const& request) override; + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) override; future> AsyncReadModifyWriteRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) override; + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) override; future> - AsyncPrepareQuery( - google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request) override; + AsyncPrepareQuery(google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) 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 1a98b11c3524f..8d099f6c97662 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,9 +13,12 @@ // limitations under the License. #include "google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.h" +#include "google/cloud/bigtable/internal/metrics.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" +#include "google/cloud/status_or.h" #include #include #include @@ -108,11 +111,15 @@ class AsyncStreamingReadWriteRpcTracking template Response UnaryHelper(std::shared_ptr>& pool, + OperationContext& oc, std::function fn) { - auto child = pool->GetChannelRandomTwoLeastUsed(); - auto stub = child->AcquireStub(); + auto selection = pool->GetChannelRandomTwoLeastUsed(); + oc.StubSelection(StubSelectionParams{ + selection.outstanding_rpcs, ChannelPoolLbPolicy::kRandomTwoLeastUsed, + pool->transport_type(), RpcType::kUnary}); + auto stub = selection.channel->AcquireStub(); auto result = fn(*stub); - child->ReleaseStub(); + selection.channel->ReleaseStub(); return result; } @@ -120,13 +127,17 @@ template std::unique_ptr> StreamingHelper( std::shared_ptr>& pool, + OperationContext& oc, std::function>(BigtableStub&)> fn) { - auto child = pool->GetChannelRandomTwoLeastUsed(); - auto stub = child->AcquireStub(); + auto selection = pool->GetChannelRandomTwoLeastUsed(); + oc.StubSelection(StubSelectionParams{ + selection.outstanding_rpcs, ChannelPoolLbPolicy::kRandomTwoLeastUsed, + pool->transport_type(), RpcType::kStreaming}); + auto stub = selection.channel->AcquireStub(); auto result = fn(*stub); - auto release_fn = [weak = child->MakeWeak()] { + auto release_fn = [weak = selection.channel->MakeWeak()] { auto child = weak.lock(); if (child) child->ReleaseStub(); }; @@ -138,14 +149,18 @@ template std::unique_ptr> AsyncStreamingHelper( std::shared_ptr>& pool, + OperationContext& oc, std::function>( BigtableStub&)> fn) { - auto child = pool->GetChannelRandomTwoLeastUsed(); - auto stub = child->AcquireStub(); + auto selection = pool->GetChannelRandomTwoLeastUsed(); + oc.StubSelection(StubSelectionParams{ + selection.outstanding_rpcs, ChannelPoolLbPolicy::kRandomTwoLeastUsed, + pool->transport_type(), RpcType::kStreaming}); + auto stub = selection.channel->AcquireStub(); auto result = fn(*stub); - auto release_fn = [weak = child->MakeWeak()] { + auto release_fn = [weak = selection.channel->MakeWeak()] { auto child = weak.lock(); if (child) child->ReleaseStub(); }; @@ -160,10 +175,10 @@ AsyncStreamingHelper( std::function>(BigtableStub&)> fn) { - auto child = pool->GetChannelRandomTwoLeastUsed(); - auto stub = child->AcquireStub(); + auto selection = pool->GetChannelRandomTwoLeastUsed(); + auto stub = selection.channel->AcquireStub(); auto result = fn(*stub); - auto release_fn = [weak = child->MakeWeak()] { + auto release_fn = [weak = selection.channel->MakeWeak()] { auto child = weak.lock(); if (child) child->ReleaseStub(); }; @@ -178,10 +193,11 @@ std::unique_ptr> BigtableRandomTwoLeastUsed::ReadRows( std::shared_ptr context, Options const& options, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) { return StreamingHelper( - pool_, [&, context = std::move(context)](BigtableStub& stub) mutable { - return stub.ReadRows(std::move(context), options, request); + pool_, oc, [&, context = std::move(context)](BigtableStub& stub) mutable { + return stub.ReadRows(std::move(context), options, request, oc); }); } @@ -189,20 +205,22 @@ std::unique_ptr> BigtableRandomTwoLeastUsed::SampleRowKeys( std::shared_ptr context, Options const& options, - google::bigtable::v2::SampleRowKeysRequest const& request) { + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) { return StreamingHelper( - pool_, [&, context = std::move(context)](BigtableStub& stub) mutable { - return stub.SampleRowKeys(std::move(context), options, request); + pool_, oc, [&, context = std::move(context)](BigtableStub& stub) mutable { + return stub.SampleRowKeys(std::move(context), options, request, oc); }); } StatusOr BigtableRandomTwoLeastUsed::MutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::MutateRowRequest const& request) { + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) { return UnaryHelper>( - pool_, [&](BigtableStub& stub) { - return stub.MutateRow(context, options, request); + pool_, oc, [&](BigtableStub& stub) { + return stub.MutateRow(context, options, request, oc); }); } @@ -210,51 +228,56 @@ std::unique_ptr> BigtableRandomTwoLeastUsed::MutateRows( std::shared_ptr context, Options const& options, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) { return StreamingHelper( - pool_, [&, context = std::move(context)](BigtableStub& stub) mutable { - return stub.MutateRows(std::move(context), options, request); + pool_, oc, [&, context = std::move(context)](BigtableStub& stub) mutable { + return stub.MutateRows(std::move(context), options, request, oc); }); } StatusOr BigtableRandomTwoLeastUsed::CheckAndMutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) { + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) { return UnaryHelper>( - pool_, [&](BigtableStub& stub) { - return stub.CheckAndMutateRow(context, options, request); + pool_, oc, [&](BigtableStub& stub) { + return stub.CheckAndMutateRow(context, options, request, oc); }); } StatusOr BigtableRandomTwoLeastUsed::PingAndWarm( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PingAndWarmRequest const& request) { + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) { return UnaryHelper>( - pool_, [&](BigtableStub& stub) { - return stub.PingAndWarm(context, options, request); + pool_, oc, [&](BigtableStub& stub) { + return stub.PingAndWarm(context, options, request, oc); }); } StatusOr BigtableRandomTwoLeastUsed::ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) { + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) { return UnaryHelper< StatusOr>( - pool_, [&](BigtableStub& stub) { - return stub.ReadModifyWriteRow(context, options, request); + pool_, oc, [&](BigtableStub& stub) { + return stub.ReadModifyWriteRow(context, options, request, oc); }); } StatusOr BigtableRandomTwoLeastUsed::PrepareQuery( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PrepareQueryRequest const& request) { + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) { return UnaryHelper>( - pool_, [&](BigtableStub& stub) { - return stub.PrepareQuery(context, options, request); + pool_, oc, [&](BigtableStub& stub) { + return stub.PrepareQuery(context, options, request, oc); }); } @@ -262,10 +285,11 @@ std::unique_ptr> BigtableRandomTwoLeastUsed::ExecuteQuery( std::shared_ptr context, Options const& options, - google::bigtable::v2::ExecuteQueryRequest const& request) { + google::bigtable::v2::ExecuteQueryRequest const& request, + OperationContext& oc) { return StreamingHelper( - pool_, [&, context = std::move(context)](BigtableStub& stub) mutable { - return stub.ExecuteQuery(std::move(context), options, request); + pool_, oc, [&, context = std::move(context)](BigtableStub& stub) mutable { + return stub.ExecuteQuery(std::move(context), options, request, oc); }); } @@ -275,12 +299,14 @@ BigtableRandomTwoLeastUsed::AsyncReadRows( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) { return AsyncStreamingHelper( - pool_, [&, context = std::move(context), - options = std::move(options)](BigtableStub& stub) mutable { + pool_, oc, + [&, context = std::move(context), + options = std::move(options)](BigtableStub& stub) mutable { return stub.AsyncReadRows(cq, std::move(context), std::move(options), - request); + request, oc); }); } @@ -290,12 +316,14 @@ BigtableRandomTwoLeastUsed::AsyncSampleRowKeys( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::SampleRowKeysRequest const& request) { + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) { return AsyncStreamingHelper( - pool_, [&, context = std::move(context), - options = std::move(options)](BigtableStub& stub) mutable { + pool_, oc, + [&, context = std::move(context), + options = std::move(options)](BigtableStub& stub) mutable { return stub.AsyncSampleRowKeys(cq, std::move(context), - std::move(options), request); + std::move(options), request, oc); }); } @@ -304,12 +332,14 @@ BigtableRandomTwoLeastUsed::AsyncMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowRequest const& request) { + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) { return UnaryHelper>>( - pool_, [&, context = std::move(context), - options = std::move(options)](BigtableStub& stub) mutable { + pool_, oc, + [&, context = std::move(context), + options = std::move(options)](BigtableStub& stub) mutable { return stub.AsyncMutateRow(cq, std::move(context), std::move(options), - request); + request, oc); }); } @@ -319,12 +349,14 @@ BigtableRandomTwoLeastUsed::AsyncMutateRows( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) { return AsyncStreamingHelper( - pool_, [&, context = std::move(context), - options = std::move(options)](BigtableStub& stub) mutable { + pool_, oc, + [&, context = std::move(context), + options = std::move(options)](BigtableStub& stub) mutable { return stub.AsyncMutateRows(cq, std::move(context), std::move(options), - request); + request, oc); }); } @@ -333,13 +365,15 @@ BigtableRandomTwoLeastUsed::AsyncCheckAndMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) { + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) { return UnaryHelper< future>>( - pool_, [&, context = std::move(context), - options = std::move(options)](BigtableStub& stub) mutable { + pool_, oc, + [&, context = std::move(context), + options = std::move(options)](BigtableStub& stub) mutable { return stub.AsyncCheckAndMutateRow(cq, std::move(context), - std::move(options), request); + std::move(options), request, oc); }); } @@ -348,13 +382,15 @@ BigtableRandomTwoLeastUsed::AsyncPingAndWarm( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PingAndWarmRequest const& request) { + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) { return UnaryHelper< future>>( - pool_, [&, context = std::move(context), - options = std::move(options)](BigtableStub& stub) mutable { + pool_, oc, + [&, context = std::move(context), + options = std::move(options)](BigtableStub& stub) mutable { return stub.AsyncPingAndWarm(cq, std::move(context), std::move(options), - request); + request, oc); }); } @@ -363,13 +399,15 @@ BigtableRandomTwoLeastUsed::AsyncReadModifyWriteRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) { + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) { return UnaryHelper< future>>( - pool_, [&, context = std::move(context), - options = std::move(options)](BigtableStub& stub) mutable { + pool_, oc, + [&, context = std::move(context), + options = std::move(options)](BigtableStub& stub) mutable { return stub.AsyncReadModifyWriteRow(cq, std::move(context), - std::move(options), request); + std::move(options), request, oc); }); } @@ -378,23 +416,26 @@ BigtableRandomTwoLeastUsed::AsyncPrepareQuery( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request) { + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) { return UnaryHelper< future>>( - pool_, [&, context = std::move(context), - options = std::move(options)](BigtableStub& stub) mutable { + pool_, oc, + [&, context = std::move(context), + options = std::move(options)](BigtableStub& stub) mutable { return stub.AsyncPrepareQuery(cq, std::move(context), - std::move(options), request); + std::move(options), request, oc); }); } StatusOr BigtableRandomTwoLeastUsed::GetClientConfiguration( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::GetClientConfigurationRequest const& request) { + google::bigtable::v2::GetClientConfigurationRequest const& request, + OperationContext& oc) { return UnaryHelper>( - pool_, [&](BigtableStub& stub) { - return stub.GetClientConfiguration(context, options, request); + pool_, oc, [&](BigtableStub& stub) { + return stub.GetClientConfiguration(context, options, request, oc); }); } 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 7327896a192f1..d70db9ef496b2 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 @@ -42,106 +42,120 @@ class BigtableRandomTwoLeastUsed : public BigtableStub { std::unique_ptr> ReadRows(std::shared_ptr context, Options const& options, - google::bigtable::v2::ReadRowsRequest const& request) override; + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) override; std::unique_ptr> - SampleRowKeys( - std::shared_ptr context, Options const& options, - google::bigtable::v2::SampleRowKeysRequest const& request) override; + SampleRowKeys(std::shared_ptr context, + Options const& options, + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) override; StatusOr MutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::MutateRowRequest const& request) override; + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) override; std::unique_ptr> MutateRows(std::shared_ptr context, Options const& options, - google::bigtable::v2::MutateRowsRequest const& request) override; + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) override; StatusOr CheckAndMutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) override; + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) override; StatusOr PingAndWarm( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PingAndWarmRequest const& request) override; + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) override; StatusOr ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) override; + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) override; StatusOr PrepareQuery( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PrepareQueryRequest const& request) override; + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) override; std::unique_ptr> - ExecuteQuery( - std::shared_ptr context, Options const& options, - google::bigtable::v2::ExecuteQueryRequest const& request) override; + ExecuteQuery(std::shared_ptr context, + Options const& options, + google::bigtable::v2::ExecuteQueryRequest const& request, + OperationContext& oc) 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) override; + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) 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) override; + AsyncSampleRowKeys(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) override; future> AsyncMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowRequest const& request) override; + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) 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) override; + AsyncMutateRows(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) override; future> AsyncCheckAndMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) override; + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) override; future> AsyncPingAndWarm( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PingAndWarmRequest const& request) override; + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) override; future> AsyncReadModifyWriteRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) override; + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) override; future> - AsyncPrepareQuery( - google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request) override; + AsyncPrepareQuery(google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) override; StatusOr GetClientConfiguration( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::GetClientConfigurationRequest const& request) - override; + google::bigtable::v2::GetClientConfigurationRequest const& request, + OperationContext& oc) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, diff --git a/google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator_test.cc b/google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator_test.cc index a4ae1ba42a6aa..cf69402d6fa46 100644 --- a/google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator_test.cc +++ b/google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator_test.cc @@ -14,6 +14,7 @@ #include "google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.h" #include "google/cloud/bigtable/internal/dynamic_channel_pool.h" +#include "google/cloud/bigtable/internal/operation_context.h" #include "google/cloud/bigtable/testing/mock_bigtable_stub.h" #include "google/cloud/testing_util/fake_completion_queue_impl.h" #include @@ -63,6 +64,7 @@ class BigtableRandomTwoLeastUsedTest : public ::testing::Test { stub_factory_fn_; CompletionQueue cq_; std::shared_ptr fake_cq_impl_; + OperationContext default_oc_; }; TEST_F(BigtableRandomTwoLeastUsedTest, ReadRows) { @@ -70,7 +72,7 @@ TEST_F(BigtableRandomTwoLeastUsedTest, ReadRows) { auto decorator = std::make_shared(pool_); auto client_context = std::make_shared(); google::bigtable::v2::ReadRowsRequest request; - auto response = decorator->ReadRows(client_context, {}, request); + auto response = decorator->ReadRows(client_context, {}, request, default_oc_); } TEST_F(BigtableRandomTwoLeastUsedTest, SampleRowKeys) { @@ -78,7 +80,8 @@ TEST_F(BigtableRandomTwoLeastUsedTest, SampleRowKeys) { auto decorator = std::make_shared(pool_); auto client_context = std::make_shared(); google::bigtable::v2::SampleRowKeysRequest request; - auto response = decorator->SampleRowKeys(client_context, {}, request); + auto response = + decorator->SampleRowKeys(client_context, {}, request, default_oc_); } TEST_F(BigtableRandomTwoLeastUsedTest, MutateRow) { @@ -86,7 +89,8 @@ TEST_F(BigtableRandomTwoLeastUsedTest, MutateRow) { auto decorator = std::make_shared(pool_); grpc::ClientContext client_context; google::bigtable::v2::MutateRowRequest request; - auto response = decorator->MutateRow(client_context, {}, request); + auto response = + decorator->MutateRow(client_context, {}, request, default_oc_); } TEST_F(BigtableRandomTwoLeastUsedTest, MutateRows) { @@ -94,7 +98,8 @@ TEST_F(BigtableRandomTwoLeastUsedTest, MutateRows) { auto decorator = std::make_shared(pool_); auto client_context = std::make_shared(); google::bigtable::v2::MutateRowsRequest request; - auto response = decorator->MutateRows(client_context, {}, request); + auto response = + decorator->MutateRows(client_context, {}, request, default_oc_); } TEST_F(BigtableRandomTwoLeastUsedTest, CheckAndMutateRow) { @@ -102,7 +107,8 @@ TEST_F(BigtableRandomTwoLeastUsedTest, CheckAndMutateRow) { auto decorator = std::make_shared(pool_); grpc::ClientContext client_context; google::bigtable::v2::CheckAndMutateRowRequest request; - auto response = decorator->CheckAndMutateRow(client_context, {}, request); + auto response = + decorator->CheckAndMutateRow(client_context, {}, request, default_oc_); } TEST_F(BigtableRandomTwoLeastUsedTest, PingAndWarm) { @@ -110,7 +116,8 @@ TEST_F(BigtableRandomTwoLeastUsedTest, PingAndWarm) { auto decorator = std::make_shared(pool_); grpc::ClientContext client_context; google::bigtable::v2::PingAndWarmRequest request; - auto response = decorator->PingAndWarm(client_context, {}, request); + auto response = + decorator->PingAndWarm(client_context, {}, request, default_oc_); } TEST_F(BigtableRandomTwoLeastUsedTest, ReadModifyWriteRow) { @@ -118,7 +125,8 @@ TEST_F(BigtableRandomTwoLeastUsedTest, ReadModifyWriteRow) { auto decorator = std::make_shared(pool_); grpc::ClientContext client_context; google::bigtable::v2::ReadModifyWriteRowRequest request; - auto response = decorator->ReadModifyWriteRow(client_context, {}, request); + auto response = + decorator->ReadModifyWriteRow(client_context, {}, request, default_oc_); } TEST_F(BigtableRandomTwoLeastUsedTest, PrepareQuery) { @@ -126,7 +134,8 @@ TEST_F(BigtableRandomTwoLeastUsedTest, PrepareQuery) { auto decorator = std::make_shared(pool_); grpc::ClientContext client_context; google::bigtable::v2::PrepareQueryRequest request; - auto response = decorator->PrepareQuery(client_context, {}, request); + auto response = + decorator->PrepareQuery(client_context, {}, request, default_oc_); } TEST_F(BigtableRandomTwoLeastUsedTest, ExecuteQuery) { @@ -134,7 +143,8 @@ TEST_F(BigtableRandomTwoLeastUsedTest, ExecuteQuery) { auto decorator = std::make_shared(pool_); auto client_context = std::make_shared(); google::bigtable::v2::ExecuteQueryRequest request; - auto response = decorator->ExecuteQuery(client_context, {}, request); + auto response = + decorator->ExecuteQuery(client_context, {}, request, default_oc_); } TEST_F(BigtableRandomTwoLeastUsedTest, AsyncReadRows) { @@ -142,7 +152,8 @@ TEST_F(BigtableRandomTwoLeastUsedTest, AsyncReadRows) { auto decorator = std::make_shared(pool_); auto client_context = std::make_shared(); google::bigtable::v2::ReadRowsRequest request; - auto response = decorator->AsyncReadRows(cq_, client_context, {}, request); + auto response = + decorator->AsyncReadRows(cq_, client_context, {}, request, default_oc_); } TEST_F(BigtableRandomTwoLeastUsedTest, AsyncSampleRowKeys) { @@ -150,8 +161,8 @@ TEST_F(BigtableRandomTwoLeastUsedTest, AsyncSampleRowKeys) { auto decorator = std::make_shared(pool_); auto client_context = std::make_shared(); google::bigtable::v2::SampleRowKeysRequest request; - auto response = - decorator->AsyncSampleRowKeys(cq_, client_context, {}, request); + auto response = decorator->AsyncSampleRowKeys(cq_, client_context, {}, + request, default_oc_); } TEST_F(BigtableRandomTwoLeastUsedTest, AsyncMutateRow) { @@ -159,7 +170,8 @@ TEST_F(BigtableRandomTwoLeastUsedTest, AsyncMutateRow) { auto decorator = std::make_shared(pool_); auto client_context = std::make_shared(); google::bigtable::v2::MutateRowRequest request; - auto response = decorator->AsyncMutateRow(cq_, client_context, {}, request); + auto response = + decorator->AsyncMutateRow(cq_, client_context, {}, request, default_oc_); } TEST_F(BigtableRandomTwoLeastUsedTest, AsyncMutateRows) { @@ -167,7 +179,8 @@ TEST_F(BigtableRandomTwoLeastUsedTest, AsyncMutateRows) { auto decorator = std::make_shared(pool_); auto client_context = std::make_shared(); google::bigtable::v2::MutateRowsRequest request; - auto response = decorator->AsyncMutateRows(cq_, client_context, {}, request); + auto response = + decorator->AsyncMutateRows(cq_, client_context, {}, request, default_oc_); } TEST_F(BigtableRandomTwoLeastUsedTest, AsyncCheckAndMutateRow) { @@ -175,8 +188,8 @@ TEST_F(BigtableRandomTwoLeastUsedTest, AsyncCheckAndMutateRow) { auto decorator = std::make_shared(pool_); auto client_context = std::make_shared(); google::bigtable::v2::CheckAndMutateRowRequest request; - auto response = - decorator->AsyncCheckAndMutateRow(cq_, client_context, {}, request); + auto response = decorator->AsyncCheckAndMutateRow(cq_, client_context, {}, + request, default_oc_); } TEST_F(BigtableRandomTwoLeastUsedTest, AsyncPingAndWarm) { @@ -184,7 +197,8 @@ TEST_F(BigtableRandomTwoLeastUsedTest, AsyncPingAndWarm) { auto decorator = std::make_shared(pool_); auto client_context = std::make_shared(); google::bigtable::v2::PingAndWarmRequest request; - auto response = decorator->AsyncPingAndWarm(cq_, client_context, {}, request); + auto response = decorator->AsyncPingAndWarm(cq_, client_context, {}, request, + default_oc_); } TEST_F(BigtableRandomTwoLeastUsedTest, AsyncReadModifyWriteRow) { @@ -192,8 +206,8 @@ TEST_F(BigtableRandomTwoLeastUsedTest, AsyncReadModifyWriteRow) { auto decorator = std::make_shared(pool_); auto client_context = std::make_shared(); google::bigtable::v2::ReadModifyWriteRowRequest request; - auto response = - decorator->AsyncReadModifyWriteRow(cq_, client_context, {}, request); + auto response = decorator->AsyncReadModifyWriteRow(cq_, client_context, {}, + request, default_oc_); } TEST_F(BigtableRandomTwoLeastUsedTest, AsyncPrepareQuery) { @@ -201,8 +215,8 @@ TEST_F(BigtableRandomTwoLeastUsedTest, AsyncPrepareQuery) { auto decorator = std::make_shared(pool_); auto client_context = std::make_shared(); google::bigtable::v2::PrepareQueryRequest request; - auto response = - decorator->AsyncPrepareQuery(cq_, client_context, {}, request); + auto response = decorator->AsyncPrepareQuery(cq_, client_context, {}, request, + default_oc_); } } // namespace diff --git a/google/cloud/bigtable/internal/bigtable_round_robin_decorator.cc b/google/cloud/bigtable/internal/bigtable_round_robin_decorator.cc index 34a28f6a5f998..e25f35ac0dde6 100644 --- a/google/cloud/bigtable/internal/bigtable_round_robin_decorator.cc +++ b/google/cloud/bigtable/internal/bigtable_round_robin_decorator.cc @@ -17,6 +17,9 @@ // source: google/bigtable/v2/bigtable.proto #include "google/cloud/bigtable/internal/bigtable_round_robin_decorator.h" +#include "google/cloud/bigtable/internal/defaults.h" +#include "google/cloud/bigtable/internal/metrics.h" +#include "google/cloud/bigtable/internal/operation_context.h" #include #include #include @@ -37,73 +40,91 @@ std::unique_ptr> BigtableRoundRobin::ReadRows( std::shared_ptr context, Options const& options, - google::bigtable::v2::ReadRowsRequest const& request) { - return Child()->ReadRows(std::move(context), options, request); + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) { + return Child(oc, RpcType::kStreaming) + ->ReadRows(std::move(context), options, request, oc); } std::unique_ptr> BigtableRoundRobin::SampleRowKeys( std::shared_ptr context, Options const& options, - google::bigtable::v2::SampleRowKeysRequest const& request) { - return Child()->SampleRowKeys(std::move(context), options, request); + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) { + return Child(oc, RpcType::kStreaming) + ->SampleRowKeys(std::move(context), options, request, oc); } StatusOr BigtableRoundRobin::MutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::MutateRowRequest const& request) { - return Child()->MutateRow(context, options, request); + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) { + return Child(oc, RpcType::kUnary)->MutateRow(context, options, request, oc); } std::unique_ptr> BigtableRoundRobin::MutateRows( std::shared_ptr context, Options const& options, - google::bigtable::v2::MutateRowsRequest const& request) { - return Child()->MutateRows(std::move(context), options, request); + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) { + return Child(oc, RpcType::kStreaming) + ->MutateRows(std::move(context), options, request, oc); } StatusOr BigtableRoundRobin::CheckAndMutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) { - return Child()->CheckAndMutateRow(context, options, request); + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) { + return Child(oc, RpcType::kUnary) + ->CheckAndMutateRow(context, options, request, oc); } StatusOr BigtableRoundRobin::PingAndWarm( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PingAndWarmRequest const& request) { - return Child()->PingAndWarm(context, options, request); + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) { + return Child(oc, RpcType::kUnary)->PingAndWarm(context, options, request, oc); } StatusOr BigtableRoundRobin::ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) { - return Child()->ReadModifyWriteRow(context, options, request); + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) { + return Child(oc, RpcType::kUnary) + ->ReadModifyWriteRow(context, options, request, oc); } StatusOr BigtableRoundRobin::PrepareQuery( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PrepareQueryRequest const& request) { - return Child()->PrepareQuery(context, options, request); + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) { + return Child(oc, RpcType::kUnary) + ->PrepareQuery(context, options, request, oc); } std::unique_ptr> BigtableRoundRobin::ExecuteQuery( std::shared_ptr context, Options const& options, - google::bigtable::v2::ExecuteQueryRequest const& request) { - return Child()->ExecuteQuery(std::move(context), options, request); + google::bigtable::v2::ExecuteQueryRequest const& request, + OperationContext& oc) { + return Child(oc, RpcType::kStreaming) + ->ExecuteQuery(std::move(context), options, request, oc); } StatusOr BigtableRoundRobin::GetClientConfiguration( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::GetClientConfigurationRequest const& request) { - return Child()->GetClientConfiguration(context, options, request); + google::bigtable::v2::GetClientConfigurationRequest const& request, + OperationContext& oc) { + return Child(oc, RpcType::kUnary) + ->GetClientConfiguration(context, options, request, oc); } std::unique_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadRowsRequest const& request) { - return Child()->AsyncReadRows(cq, std::move(context), std::move(options), - request); + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) { + return Child(oc, RpcType::kStreaming) + ->AsyncReadRows(cq, std::move(context), std::move(options), request, oc); } std::unique_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::SampleRowKeysRequest const& request) { - return Child()->AsyncSampleRowKeys(cq, std::move(context), std::move(options), - request); + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) { + return Child(oc, RpcType::kStreaming) + ->AsyncSampleRowKeys(cq, std::move(context), std::move(options), request, + oc); } future> @@ -165,9 +189,10 @@ BigtableRoundRobin::AsyncMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowRequest const& request) { - return Child()->AsyncMutateRow(cq, std::move(context), std::move(options), - request); + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) { + return Child(oc, RpcType::kUnary) + ->AsyncMutateRow(cq, std::move(context), std::move(options), request, oc); } std::unique_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowsRequest const& request) { - return Child()->AsyncMutateRows(cq, std::move(context), std::move(options), - request); + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) { + return Child(oc, RpcType::kStreaming) + ->AsyncMutateRows(cq, std::move(context), std::move(options), request, + oc); } future> @@ -186,9 +213,11 @@ BigtableRoundRobin::AsyncCheckAndMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) { - return Child()->AsyncCheckAndMutateRow(cq, std::move(context), - std::move(options), request); + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) { + return Child(oc, RpcType::kUnary) + ->AsyncCheckAndMutateRow(cq, std::move(context), std::move(options), + request, oc); } future> @@ -196,9 +225,11 @@ BigtableRoundRobin::AsyncPingAndWarm( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PingAndWarmRequest const& request) { - return Child()->AsyncPingAndWarm(cq, std::move(context), std::move(options), - request); + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) { + return Child(oc, RpcType::kUnary) + ->AsyncPingAndWarm(cq, std::move(context), std::move(options), request, + oc); } future> @@ -206,9 +237,11 @@ BigtableRoundRobin::AsyncReadModifyWriteRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) { - return Child()->AsyncReadModifyWriteRow(cq, std::move(context), - std::move(options), request); + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) { + return Child(oc, RpcType::kUnary) + ->AsyncReadModifyWriteRow(cq, std::move(context), std::move(options), + request, oc); } future> @@ -216,9 +249,11 @@ BigtableRoundRobin::AsyncPrepareQuery( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request) { - return Child()->AsyncPrepareQuery(cq, std::move(context), std::move(options), - request); + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) { + return Child(oc, RpcType::kUnary) + ->AsyncPrepareQuery(cq, std::move(context), std::move(options), request, + oc); } std::shared_ptr BigtableRoundRobin::Child() { @@ -229,6 +264,16 @@ std::shared_ptr BigtableRoundRobin::Child() { return children_[current]; } +std::shared_ptr BigtableRoundRobin::Child(OperationContext& oc, + RpcType rpc_type) { + oc.StubSelection(StubSelectionParams{0, ChannelPoolLbPolicy::kRoundRobin, + bigtable::internal::IsDirectPath() + ? TransportType::kDirectPath + : TransportType::kCloudPath, + rpc_type}); + return Child(); +} + GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END } // namespace bigtable_internal } // namespace cloud diff --git a/google/cloud/bigtable/internal/bigtable_round_robin_decorator.h b/google/cloud/bigtable/internal/bigtable_round_robin_decorator.h index 58e72d59c1131..8c4387ef792d5 100644 --- a/google/cloud/bigtable/internal/bigtable_round_robin_decorator.h +++ b/google/cloud/bigtable/internal/bigtable_round_robin_decorator.h @@ -20,6 +20,7 @@ #define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_INTERNAL_BIGTABLE_ROUND_ROBIN_DECORATOR_H #include "google/cloud/bigtable/internal/bigtable_stub.h" +#include "google/cloud/bigtable/internal/metrics.h" #include "google/cloud/version.h" #include #include @@ -42,50 +43,59 @@ class BigtableRoundRobin : public BigtableStub { std::unique_ptr> ReadRows(std::shared_ptr context, Options const& options, - google::bigtable::v2::ReadRowsRequest const& request) override; + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) override; std::unique_ptr> - SampleRowKeys( - std::shared_ptr context, Options const& options, - google::bigtable::v2::SampleRowKeysRequest const& request) override; + SampleRowKeys(std::shared_ptr context, + Options const& options, + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) override; StatusOr MutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::MutateRowRequest const& request) override; + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) override; std::unique_ptr> MutateRows(std::shared_ptr context, Options const& options, - google::bigtable::v2::MutateRowsRequest const& request) override; + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) override; StatusOr CheckAndMutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) override; + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) override; StatusOr PingAndWarm( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PingAndWarmRequest const& request) override; + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) override; StatusOr ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) override; + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) override; StatusOr PrepareQuery( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PrepareQueryRequest const& request) override; + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) override; std::unique_ptr> - ExecuteQuery( - std::shared_ptr context, Options const& options, - google::bigtable::v2::ExecuteQueryRequest const& request) override; + ExecuteQuery(std::shared_ptr context, + Options const& options, + google::bigtable::v2::ExecuteQueryRequest const& request, + OperationContext& oc) override; StatusOr GetClientConfiguration( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::GetClientConfigurationRequest const& request) - override; + google::bigtable::v2::GetClientConfigurationRequest const& request, + OperationContext& oc) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, @@ -115,59 +125,65 @@ class BigtableRoundRobin : public BigtableStub { AsyncReadRows(google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadRowsRequest const& request) override; + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) 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) override; + AsyncSampleRowKeys(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) override; future> AsyncMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowRequest const& request) override; + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) 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) override; + AsyncMutateRows(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) override; future> AsyncCheckAndMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) override; + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) override; future> AsyncPingAndWarm( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PingAndWarmRequest const& request) override; + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) override; future> AsyncReadModifyWriteRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) override; + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) override; future> - AsyncPrepareQuery( - google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request) override; + AsyncPrepareQuery(google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) override; private: std::shared_ptr Child(); + std::shared_ptr Child(OperationContext& oc, RpcType rpc_type); std::vector> const children_; std::mutex mu_; diff --git a/google/cloud/bigtable/internal/bigtable_stub.cc b/google/cloud/bigtable/internal/bigtable_stub.cc index 9c8b78fb03a72..4a6eb1f23a25e 100644 --- a/google/cloud/bigtable/internal/bigtable_stub.cc +++ b/google/cloud/bigtable/internal/bigtable_stub.cc @@ -17,6 +17,7 @@ // source: google/bigtable/v2/bigtable.proto #include "google/cloud/bigtable/internal/bigtable_stub.h" +#include "google/cloud/bigtable/internal/operation_context.h" #include "google/cloud/grpc_error_delegate.h" #include "google/cloud/internal/async_read_write_stream_impl.h" #include "google/cloud/internal/async_streaming_read_rpc_impl.h" @@ -39,7 +40,7 @@ std::unique_ptr> DefaultBigtableStub::ReadRows( std::shared_ptr context, Options const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, OperationContext&) { auto stream = grpc_stub_->ReadRows(context.get(), request); return std::make_unique>(std::move(context), @@ -50,7 +51,8 @@ std::unique_ptr> DefaultBigtableStub::SampleRowKeys( std::shared_ptr context, Options const&, - google::bigtable::v2::SampleRowKeysRequest const& request) { + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext&) { auto stream = grpc_stub_->SampleRowKeys(context.get(), request); return std::make_unique>(std::move(context), @@ -60,7 +62,7 @@ DefaultBigtableStub::SampleRowKeys( StatusOr DefaultBigtableStub::MutateRow( grpc::ClientContext& context, Options const&, - google::bigtable::v2::MutateRowRequest const& request) { + google::bigtable::v2::MutateRowRequest const& request, OperationContext&) { google::bigtable::v2::MutateRowResponse response; auto status = grpc_stub_->MutateRow(&context, request, &response); if (!status.ok()) { @@ -73,7 +75,7 @@ std::unique_ptr> DefaultBigtableStub::MutateRows( std::shared_ptr context, Options const&, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, OperationContext&) { auto stream = grpc_stub_->MutateRows(context.get(), request); return std::make_unique>(std::move(context), @@ -83,7 +85,8 @@ DefaultBigtableStub::MutateRows( StatusOr DefaultBigtableStub::CheckAndMutateRow( grpc::ClientContext& context, Options const&, - google::bigtable::v2::CheckAndMutateRowRequest const& request) { + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext&) { google::bigtable::v2::CheckAndMutateRowResponse response; auto status = grpc_stub_->CheckAndMutateRow(&context, request, &response); if (!status.ok()) { @@ -95,7 +98,8 @@ DefaultBigtableStub::CheckAndMutateRow( StatusOr DefaultBigtableStub::PingAndWarm( grpc::ClientContext& context, Options const&, - google::bigtable::v2::PingAndWarmRequest const& request) { + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext&) { google::bigtable::v2::PingAndWarmResponse response; auto status = grpc_stub_->PingAndWarm(&context, request, &response); if (!status.ok()) { @@ -107,7 +111,8 @@ DefaultBigtableStub::PingAndWarm( StatusOr DefaultBigtableStub::ReadModifyWriteRow( grpc::ClientContext& context, Options const&, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) { + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext&) { google::bigtable::v2::ReadModifyWriteRowResponse response; auto status = grpc_stub_->ReadModifyWriteRow(&context, request, &response); if (!status.ok()) { @@ -119,7 +124,8 @@ DefaultBigtableStub::ReadModifyWriteRow( StatusOr DefaultBigtableStub::PrepareQuery( grpc::ClientContext& context, Options const&, - google::bigtable::v2::PrepareQueryRequest const& request) { + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext&) { google::bigtable::v2::PrepareQueryResponse response; auto status = grpc_stub_->PrepareQuery(&context, request, &response); if (!status.ok()) { @@ -132,7 +138,8 @@ std::unique_ptr> DefaultBigtableStub::ExecuteQuery( std::shared_ptr context, Options const&, - google::bigtable::v2::ExecuteQueryRequest const& request) { + google::bigtable::v2::ExecuteQueryRequest const& request, + OperationContext&) { auto stream = grpc_stub_->ExecuteQuery(context.get(), request); return std::make_unique>(std::move(context), @@ -142,7 +149,8 @@ DefaultBigtableStub::ExecuteQuery( StatusOr DefaultBigtableStub::GetClientConfiguration( grpc::ClientContext& context, Options const&, - google::bigtable::v2::GetClientConfigurationRequest const& request) { + google::bigtable::v2::GetClientConfigurationRequest const& request, + OperationContext&) { google::bigtable::v2::ClientConfiguration response; auto status = grpc_stub_->GetClientConfiguration(&context, request, &response); @@ -206,7 +214,7 @@ DefaultBigtableStub::AsyncReadRows( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, OperationContext&) { return google::cloud::internal::MakeStreamingReadRpc< google::bigtable::v2::ReadRowsRequest, google::bigtable::v2::ReadRowsResponse>( @@ -224,7 +232,8 @@ DefaultBigtableStub::AsyncSampleRowKeys( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::SampleRowKeysRequest const& request) { + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext&) { return google::cloud::internal::MakeStreamingReadRpc< google::bigtable::v2::SampleRowKeysRequest, google::bigtable::v2::SampleRowKeysResponse>( @@ -242,7 +251,7 @@ DefaultBigtableStub::AsyncMutateRow( std::shared_ptr context, // NOLINTNEXTLINE(performance-unnecessary-value-param) google::cloud::internal::ImmutableOptions, - google::bigtable::v2::MutateRowRequest const& request) { + google::bigtable::v2::MutateRowRequest const& request, OperationContext&) { return internal::MakeUnaryRpcImpl( cq, @@ -260,7 +269,7 @@ DefaultBigtableStub::AsyncMutateRows( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, OperationContext&) { return google::cloud::internal::MakeStreamingReadRpc< google::bigtable::v2::MutateRowsRequest, google::bigtable::v2::MutateRowsResponse>( @@ -278,7 +287,8 @@ DefaultBigtableStub::AsyncCheckAndMutateRow( std::shared_ptr context, // NOLINTNEXTLINE(performance-unnecessary-value-param) google::cloud::internal::ImmutableOptions, - google::bigtable::v2::CheckAndMutateRowRequest const& request) { + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext&) { return internal::MakeUnaryRpcImpl< google::bigtable::v2::CheckAndMutateRowRequest, google::bigtable::v2::CheckAndMutateRowResponse>( @@ -297,7 +307,8 @@ DefaultBigtableStub::AsyncPingAndWarm( std::shared_ptr context, // NOLINTNEXTLINE(performance-unnecessary-value-param) google::cloud::internal::ImmutableOptions, - google::bigtable::v2::PingAndWarmRequest const& request) { + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext&) { return internal::MakeUnaryRpcImpl( cq, @@ -315,7 +326,8 @@ DefaultBigtableStub::AsyncReadModifyWriteRow( std::shared_ptr context, // NOLINTNEXTLINE(performance-unnecessary-value-param) google::cloud::internal::ImmutableOptions, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) { + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext&) { return internal::MakeUnaryRpcImpl< google::bigtable::v2::ReadModifyWriteRowRequest, google::bigtable::v2::ReadModifyWriteRowResponse>( @@ -334,7 +346,8 @@ DefaultBigtableStub::AsyncPrepareQuery( std::shared_ptr context, // NOLINTNEXTLINE(performance-unnecessary-value-param) google::cloud::internal::ImmutableOptions, - google::bigtable::v2::PrepareQueryRequest const& request) { + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext&) { return internal::MakeUnaryRpcImpl( cq, diff --git a/google/cloud/bigtable/internal/bigtable_stub.h b/google/cloud/bigtable/internal/bigtable_stub.h index dcda73ce2791c..86c620753daa4 100644 --- a/google/cloud/bigtable/internal/bigtable_stub.h +++ b/google/cloud/bigtable/internal/bigtable_stub.h @@ -39,6 +39,8 @@ namespace cloud { namespace bigtable_internal { GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN +class OperationContext; + class BigtableStub { public: virtual ~BigtableStub() = 0; @@ -46,52 +48,62 @@ class BigtableStub { virtual std::unique_ptr> ReadRows(std::shared_ptr context, Options const& options, - google::bigtable::v2::ReadRowsRequest const& request) = 0; + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) = 0; virtual std::unique_ptr> SampleRowKeys(std::shared_ptr context, Options const& options, - google::bigtable::v2::SampleRowKeysRequest const& request) = 0; + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) = 0; virtual StatusOr MutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::MutateRowRequest const& request) = 0; + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) = 0; virtual std::unique_ptr> MutateRows(std::shared_ptr context, Options const& options, - google::bigtable::v2::MutateRowsRequest const& request) = 0; + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) = 0; virtual StatusOr CheckAndMutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) = 0; + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) = 0; virtual StatusOr PingAndWarm( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PingAndWarmRequest const& request) = 0; + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) = 0; virtual StatusOr ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) = 0; + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) = 0; virtual StatusOr PrepareQuery( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PrepareQueryRequest const& request) = 0; + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) = 0; virtual std::unique_ptr> ExecuteQuery(std::shared_ptr context, Options const& options, - google::bigtable::v2::ExecuteQueryRequest const& request) = 0; + google::bigtable::v2::ExecuteQueryRequest const& request, + OperationContext& oc) = 0; virtual StatusOr GetClientConfiguration( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::GetClientConfigurationRequest const& request) = 0; + google::bigtable::v2::GetClientConfigurationRequest const& request, + OperationContext& oc) = 0; virtual std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, @@ -121,55 +133,61 @@ class BigtableStub { AsyncReadRows(google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadRowsRequest const& request) = 0; + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) = 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) = 0; + AsyncSampleRowKeys(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) = 0; virtual future> AsyncMutateRow(google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowRequest const& request) = 0; + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) = 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) = 0; + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) = 0; virtual future> AsyncCheckAndMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) = 0; + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) = 0; virtual future> AsyncPingAndWarm(google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PingAndWarmRequest const& request) = 0; + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) = 0; virtual future> AsyncReadModifyWriteRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) = 0; + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) = 0; virtual future> - AsyncPrepareQuery( - google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request) = 0; + AsyncPrepareQuery(google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) = 0; }; class DefaultBigtableStub : public BigtableStub { @@ -181,50 +199,59 @@ class DefaultBigtableStub : public BigtableStub { std::unique_ptr> ReadRows(std::shared_ptr context, Options const& options, - google::bigtable::v2::ReadRowsRequest const& request) override; + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext& oc) override; std::unique_ptr> - SampleRowKeys( - std::shared_ptr context, Options const& options, - google::bigtable::v2::SampleRowKeysRequest const& request) override; + SampleRowKeys(std::shared_ptr context, + Options const& options, + google::bigtable::v2::SampleRowKeysRequest const& request, + bigtable_internal::OperationContext& oc) override; StatusOr MutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::MutateRowRequest const& request) override; + google::bigtable::v2::MutateRowRequest const& request, + bigtable_internal::OperationContext& oc) override; std::unique_ptr> MutateRows(std::shared_ptr context, Options const& options, - google::bigtable::v2::MutateRowsRequest const& request) override; + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext& oc) override; StatusOr CheckAndMutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) override; + google::bigtable::v2::CheckAndMutateRowRequest const& request, + bigtable_internal::OperationContext& oc) override; StatusOr PingAndWarm( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PingAndWarmRequest const& request) override; + google::bigtable::v2::PingAndWarmRequest const& request, + bigtable_internal::OperationContext& oc) override; StatusOr ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) override; + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + bigtable_internal::OperationContext& oc) override; StatusOr PrepareQuery( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PrepareQueryRequest const& request) override; + google::bigtable::v2::PrepareQueryRequest const& request, + bigtable_internal::OperationContext& oc) override; std::unique_ptr> - ExecuteQuery( - std::shared_ptr context, Options const& options, - google::bigtable::v2::ExecuteQueryRequest const& request) override; + ExecuteQuery(std::shared_ptr context, + Options const& options, + google::bigtable::v2::ExecuteQueryRequest const& request, + bigtable_internal::OperationContext& oc) override; StatusOr GetClientConfiguration( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::GetClientConfigurationRequest const& request) - override; + google::bigtable::v2::GetClientConfigurationRequest const& request, + bigtable_internal::OperationContext& oc) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, @@ -254,56 +281,61 @@ class DefaultBigtableStub : public BigtableStub { AsyncReadRows(google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadRowsRequest const& request) override; + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext& oc) 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) override; + AsyncSampleRowKeys(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::SampleRowKeysRequest const& request, + bigtable_internal::OperationContext& oc) override; future> AsyncMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowRequest const& request) override; + google::bigtable::v2::MutateRowRequest const& request, + bigtable_internal::OperationContext& oc) 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) override; + AsyncMutateRows(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext& oc) override; future> AsyncCheckAndMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) override; + google::bigtable::v2::CheckAndMutateRowRequest const& request, + bigtable_internal::OperationContext& oc) override; future> AsyncPingAndWarm( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PingAndWarmRequest const& request) override; + google::bigtable::v2::PingAndWarmRequest const& request, + bigtable_internal::OperationContext& oc) override; future> AsyncReadModifyWriteRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) override; + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + bigtable_internal::OperationContext& oc) override; future> - AsyncPrepareQuery( - google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request) override; + AsyncPrepareQuery(google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::PrepareQueryRequest const& request, + bigtable_internal::OperationContext& oc) override; private: std::unique_ptr grpc_stub_; diff --git a/google/cloud/bigtable/internal/bigtable_stub_factory.cc b/google/cloud/bigtable/internal/bigtable_stub_factory.cc index b9b2830616371..3dc2397ab821e 100644 --- a/google/cloud/bigtable/internal/bigtable_stub_factory.cc +++ b/google/cloud/bigtable/internal/bigtable_stub_factory.cc @@ -22,6 +22,7 @@ #include "google/cloud/bigtable/internal/bigtable_tracing_stub.h" #include "google/cloud/bigtable/internal/connection_refresh_state.h" #include "google/cloud/bigtable/internal/defaults.h" +#include "google/cloud/bigtable/internal/operation_context.h" #include "google/cloud/bigtable/options.h" #include "google/cloud/common_options.h" #include "google/cloud/grpc_options.h" @@ -147,8 +148,9 @@ std::shared_ptr CreateBigtableStubRandomTwoLeastUsed( grpc::ClientContext client_context; google::bigtable::v2::PingAndWarmRequest request; request.set_name(std::string{instance_name}); - auto response = - stub->PingAndWarm(client_context, options, std::move(request)); + OperationContext default_oc; + auto response = stub->PingAndWarm(client_context, options, + std::move(request), default_oc); if (!response.ok()) return response.status(); } diff --git a/google/cloud/bigtable/internal/bigtable_stub_factory_test.cc b/google/cloud/bigtable/internal/bigtable_stub_factory_test.cc index 439405920b422..aadcf5a54f5c5 100644 --- a/google/cloud/bigtable/internal/bigtable_stub_factory_test.cc +++ b/google/cloud/bigtable/internal/bigtable_stub_factory_test.cc @@ -15,6 +15,7 @@ #include "google/cloud/bigtable/internal/bigtable_stub_factory.h" #include "google/cloud/bigtable/internal/bigtable_random_two_least_used_decorator.h" #include "google/cloud/bigtable/internal/dynamic_channel_pool.h" +#include "google/cloud/bigtable/internal/operation_context.h" #include "google/cloud/bigtable/options.h" #include "google/cloud/bigtable/testing/mock_bigtable_stub.h" #include "google/cloud/common_options.h" @@ -109,7 +110,8 @@ TEST(BigtableStubFactory, RoundRobin) { grpc::ClientContext context; for (int i = 0; i != kTestChannels; ++i) { - auto response = stub->MutateRow(context, Options{}, {}); + OperationContext default_oc; + auto response = stub->MutateRow(context, Options{}, {}, default_oc); EXPECT_THAT(response, StatusIs(StatusCode::kAborted, "fail")); } } @@ -157,7 +159,8 @@ TEST(BigtableStubFactory, RandomTwoLeastUsed) { grpc::ClientContext context; for (int i = 0; i != kTestChannels; ++i) { - auto response = stub->MutateRow(context, Options{}, {}); + OperationContext default_oc; + auto response = stub->MutateRow(context, Options{}, {}, default_oc); EXPECT_THAT(response, StatusIs(StatusCode::kAborted, "fail")); } @@ -247,7 +250,8 @@ TEST(BigtableStubFactory, Auth) { factory.AsStdFunction()); grpc::ClientContext context; - auto response = stub->MutateRow(context, Options{}, {}); + OperationContext default_oc; + auto response = stub->MutateRow(context, Options{}, {}, default_oc); EXPECT_THAT(response, StatusIs(StatusCode::kAborted, "fail")); } @@ -257,15 +261,15 @@ TEST(BigtableStubFactory, Metadata) { .WillOnce([](std::shared_ptr const&) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRow) - .WillOnce( - [](grpc::ClientContext& context, Options const&, - google::bigtable::v2::MutateRowRequest const& request) { - ValidateMetadataFixture fixture; - fixture.IsContextMDValid( - context, "google.bigtable.v2.Bigtable.MutateRow", request, - internal::HandCraftedLibClientHeader()); - return internal::AbortedError("fail"); - }); + .WillOnce([](grpc::ClientContext& context, Options const&, + google::bigtable::v2::MutateRowRequest const& request, + OperationContext&) { + ValidateMetadataFixture fixture; + fixture.IsContextMDValid( + context, "google.bigtable.v2.Bigtable.MutateRow", request, + internal::HandCraftedLibClientHeader()); + return internal::AbortedError("fail"); + }); return mock; }); @@ -280,7 +284,8 @@ TEST(BigtableStubFactory, Metadata) { factory.AsStdFunction()); grpc::ClientContext context; - auto response = stub->MutateRow(context, Options{}, {}); + OperationContext default_oc; + auto response = stub->MutateRow(context, Options{}, {}, default_oc); EXPECT_THAT(response, StatusIs(StatusCode::kAborted, "fail")); } @@ -308,7 +313,8 @@ TEST(BigtableStubFactory, LoggingEnabled) { factory.AsStdFunction()); grpc::ClientContext context; - auto response = stub->MutateRow(context, Options{}, {}); + OperationContext default_oc; + auto response = stub->MutateRow(context, Options{}, {}, default_oc); EXPECT_THAT(response, StatusIs(StatusCode::kAborted, "fail")); EXPECT_THAT(log.ExtractLines(), Contains(HasSubstr("MutateRow"))); @@ -338,7 +344,8 @@ TEST(BigtableStubFactory, LoggingDisabled) { factory.AsStdFunction()); grpc::ClientContext context; - auto response = stub->MutateRow(context, Options{}, {}); + OperationContext default_oc; + auto response = stub->MutateRow(context, Options{}, {}, default_oc); EXPECT_THAT(response, StatusIs(StatusCode::kAborted, "fail")); EXPECT_THAT(log.ExtractLines(), Not(Contains(HasSubstr("MutateRow")))); @@ -351,7 +358,8 @@ TEST(BigtableStubFactory, FeaturesFlagsCloudDirectPath) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRow) .WillOnce([](grpc::ClientContext& context, Options const&, - google::bigtable::v2::MutateRowRequest const&) { + google::bigtable::v2::MutateRowRequest const&, + OperationContext&) { ValidateMetadataFixture fixture; auto headers = fixture.GetMetadata(context); auto it = headers.find("bigtable-features"); @@ -389,7 +397,8 @@ TEST(BigtableStubFactory, FeaturesFlagsCloudDirectPath) { .set(MakeInsecureCredentials()), factory.AsStdFunction()); grpc::ClientContext context; - (void)stub->MutateRow(context, Options{}, {}); + OperationContext default_oc; + (void)stub->MutateRow(context, Options{}, {}, default_oc); testing_util::UnsetEnv("GOOGLE_CLOUD_ENABLE_DIRECT_PATH"); } @@ -400,7 +409,8 @@ TEST(BigtableStubFactory, FeaturesFlagsBigtableDirectPath) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRow) .WillOnce([](grpc::ClientContext& context, Options const&, - google::bigtable::v2::MutateRowRequest const&) { + google::bigtable::v2::MutateRowRequest const&, + OperationContext&) { ValidateMetadataFixture fixture; auto headers = fixture.GetMetadata(context); auto it = headers.find("bigtable-features"); @@ -438,7 +448,8 @@ TEST(BigtableStubFactory, FeaturesFlagsBigtableDirectPath) { .set(MakeInsecureCredentials()), factory.AsStdFunction()); grpc::ClientContext context; - (void)stub->MutateRow(context, Options{}, {}); + OperationContext default_oc; + (void)stub->MutateRow(context, Options{}, {}, default_oc); testing_util::UnsetEnv("CBT_ENABLE_DIRECTPATH"); } @@ -449,7 +460,8 @@ TEST(BigtableStubFactory, FeaturesFlags) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRow) .WillOnce([](grpc::ClientContext& context, Options const&, - google::bigtable::v2::MutateRowRequest const&) { + google::bigtable::v2::MutateRowRequest const&, + OperationContext&) { ValidateMetadataFixture fixture; auto headers = fixture.GetMetadata(context); auto it = headers.find("bigtable-features"); @@ -486,7 +498,8 @@ TEST(BigtableStubFactory, FeaturesFlags) { .set(MakeInsecureCredentials()), factory.AsStdFunction()); grpc::ClientContext context; - (void)stub->MutateRow(context, Options{}, {}); + OperationContext default_oc; + (void)stub->MutateRow(context, Options{}, {}, default_oc); } using ::google::cloud::testing_util::DisableTracing; @@ -505,7 +518,7 @@ TEST(BigtableStubFactory, TracingEnabled) { .WillOnce([](std::shared_ptr const&) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRow) - .WillOnce([](auto& context, auto const&, auto const&) { + .WillOnce([](auto& context, auto const&, auto const&, auto&) { ValidatePropagator(context); return internal::AbortedError("fail"); }); @@ -523,7 +536,8 @@ TEST(BigtableStubFactory, TracingEnabled) { .set(MakeInsecureCredentials())), factory.AsStdFunction()); grpc::ClientContext context; - (void)stub->MutateRow(context, Options{}, {}); + OperationContext default_oc; + (void)stub->MutateRow(context, Options{}, {}, default_oc); EXPECT_THAT(span_catcher->GetSpans(), ElementsAre(SpanNamed("google.bigtable.v2.Bigtable/MutateRow"))); @@ -537,7 +551,7 @@ TEST(BigtableStubFactory, TracingDisabled) { .WillOnce([](std::shared_ptr const&) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRow) - .WillOnce([](auto& context, auto const&, auto const&) { + .WillOnce([](auto& context, auto const&, auto const&, auto&) { ValidateNoPropagator(context); return internal::AbortedError("fail"); }); @@ -555,7 +569,8 @@ TEST(BigtableStubFactory, TracingDisabled) { .set(MakeInsecureCredentials())), factory.AsStdFunction()); grpc::ClientContext context; - (void)stub->MutateRow(context, Options{}, {}); + OperationContext default_oc; + (void)stub->MutateRow(context, Options{}, {}, default_oc); EXPECT_THAT(span_catcher->GetSpans(), IsEmpty()); } diff --git a/google/cloud/bigtable/internal/bigtable_tracing_stub.cc b/google/cloud/bigtable/internal/bigtable_tracing_stub.cc index 7e8bf74be9dbd..a91fe6053f13a 100644 --- a/google/cloud/bigtable/internal/bigtable_tracing_stub.cc +++ b/google/cloud/bigtable/internal/bigtable_tracing_stub.cc @@ -39,11 +39,12 @@ std::unique_ptr> BigtableTracingStub::ReadRows( std::shared_ptr context, Options const& options, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "ReadRows"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(*context, *propagator_); - auto stream = child_->ReadRows(context, options, request); + auto stream = child_->ReadRows(context, options, request, oc); return std::make_unique>( std::move(context), std::move(stream), std::move(span)); @@ -53,12 +54,13 @@ std::unique_ptr> BigtableTracingStub::SampleRowKeys( std::shared_ptr context, Options const& options, - google::bigtable::v2::SampleRowKeysRequest const& request) { + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "SampleRowKeys"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(*context, *propagator_); - auto stream = child_->SampleRowKeys(context, options, request); + auto stream = child_->SampleRowKeys(context, options, request, oc); return std::make_unique>( std::move(context), std::move(stream), std::move(span)); @@ -67,25 +69,27 @@ BigtableTracingStub::SampleRowKeys( StatusOr BigtableTracingStub::MutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::MutateRowRequest const& request) { + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "MutateRow"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(context, *propagator_); return internal::EndSpan(context, *span, - child_->MutateRow(context, options, request)); + child_->MutateRow(context, options, request, oc)); } std::unique_ptr> BigtableTracingStub::MutateRows( std::shared_ptr context, Options const& options, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "MutateRows"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(*context, *propagator_); - auto stream = child_->MutateRows(context, options, request); + auto stream = child_->MutateRows(context, options, request, oc); return std::make_unique>( std::move(context), std::move(stream), std::move(span)); @@ -94,61 +98,67 @@ BigtableTracingStub::MutateRows( StatusOr BigtableTracingStub::CheckAndMutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) { + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "CheckAndMutateRow"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(context, *propagator_); return internal::EndSpan( - context, *span, child_->CheckAndMutateRow(context, options, request)); + context, *span, child_->CheckAndMutateRow(context, options, request, oc)); } StatusOr BigtableTracingStub::PingAndWarm( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PingAndWarmRequest const& request) { + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "PingAndWarm"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(context, *propagator_); return internal::EndSpan(context, *span, - child_->PingAndWarm(context, options, request)); + child_->PingAndWarm(context, options, request, oc)); } StatusOr BigtableTracingStub::ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) { + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "ReadModifyWriteRow"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(context, *propagator_); return internal::EndSpan( - context, *span, child_->ReadModifyWriteRow(context, options, request)); + context, *span, + child_->ReadModifyWriteRow(context, options, request, oc)); } StatusOr BigtableTracingStub::PrepareQuery( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PrepareQueryRequest const& request) { + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "PrepareQuery"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(context, *propagator_); return internal::EndSpan(context, *span, - child_->PrepareQuery(context, options, request)); + child_->PrepareQuery(context, options, request, oc)); } std::unique_ptr> BigtableTracingStub::ExecuteQuery( std::shared_ptr context, Options const& options, - google::bigtable::v2::ExecuteQueryRequest const& request) { + google::bigtable::v2::ExecuteQueryRequest const& request, + OperationContext& oc) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "ExecuteQuery"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(*context, *propagator_); - auto stream = child_->ExecuteQuery(context, options, request); + auto stream = child_->ExecuteQuery(context, options, request, oc); return std::make_unique>( std::move(context), std::move(stream), std::move(span)); @@ -157,14 +167,15 @@ BigtableTracingStub::ExecuteQuery( StatusOr BigtableTracingStub::GetClientConfiguration( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::GetClientConfigurationRequest const& request) { + google::bigtable::v2::GetClientConfigurationRequest const& request, + OperationContext& oc) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "GetClientConfiguration"); auto scope = opentelemetry::trace::Scope(span); internal::InjectTraceContext(context, *propagator_); return internal::EndSpan( context, *span, - child_->GetClientConfiguration(context, options, request)); + child_->GetClientConfiguration(context, options, request, oc)); } std::unique_ptr< @@ -226,11 +237,13 @@ BigtableTracingStub::AsyncReadRows( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "ReadRows"); internal::OTelScope scope(span); internal::InjectTraceContext(*context, *propagator_); - auto stream = child_->AsyncReadRows(cq, context, std::move(options), request); + auto stream = + child_->AsyncReadRows(cq, context, std::move(options), request, oc); return std::make_unique>( std::move(context), std::move(stream), std::move(span)); @@ -242,13 +255,14 @@ BigtableTracingStub::AsyncSampleRowKeys( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::SampleRowKeysRequest const& request) { + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "SampleRowKeys"); internal::OTelScope scope(span); internal::InjectTraceContext(*context, *propagator_); auto stream = - child_->AsyncSampleRowKeys(cq, context, std::move(options), request); + child_->AsyncSampleRowKeys(cq, context, std::move(options), request, oc); return std::make_unique>( std::move(context), std::move(stream), std::move(span)); @@ -259,12 +273,13 @@ BigtableTracingStub::AsyncMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowRequest const& request) { + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "MutateRow"); internal::OTelScope scope(span); internal::InjectTraceContext(*context, *propagator_); - auto f = child_->AsyncMutateRow(cq, context, std::move(options), request); + auto f = child_->AsyncMutateRow(cq, context, std::move(options), request, oc); return internal::EndSpan(std::move(context), std::move(span), std::move(f)); } @@ -274,13 +289,14 @@ BigtableTracingStub::AsyncMutateRows( google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "MutateRows"); internal::OTelScope scope(span); internal::InjectTraceContext(*context, *propagator_); auto stream = - child_->AsyncMutateRows(cq, context, std::move(options), request); + child_->AsyncMutateRows(cq, context, std::move(options), request, oc); return std::make_unique>( std::move(context), std::move(stream), std::move(span)); @@ -291,13 +307,14 @@ BigtableTracingStub::AsyncCheckAndMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) { + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "CheckAndMutateRow"); internal::OTelScope scope(span); internal::InjectTraceContext(*context, *propagator_); - auto f = - child_->AsyncCheckAndMutateRow(cq, context, std::move(options), request); + auto f = child_->AsyncCheckAndMutateRow(cq, context, std::move(options), + request, oc); return internal::EndSpan(std::move(context), std::move(span), std::move(f)); } @@ -306,12 +323,14 @@ BigtableTracingStub::AsyncPingAndWarm( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PingAndWarmRequest const& request) { + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "PingAndWarm"); internal::OTelScope scope(span); internal::InjectTraceContext(*context, *propagator_); - auto f = child_->AsyncPingAndWarm(cq, context, std::move(options), request); + auto f = + child_->AsyncPingAndWarm(cq, context, std::move(options), request, oc); return internal::EndSpan(std::move(context), std::move(span), std::move(f)); } @@ -320,13 +339,14 @@ BigtableTracingStub::AsyncReadModifyWriteRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) { + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "ReadModifyWriteRow"); internal::OTelScope scope(span); internal::InjectTraceContext(*context, *propagator_); - auto f = - child_->AsyncReadModifyWriteRow(cq, context, std::move(options), request); + auto f = child_->AsyncReadModifyWriteRow(cq, context, std::move(options), + request, oc); return internal::EndSpan(std::move(context), std::move(span), std::move(f)); } @@ -335,12 +355,14 @@ BigtableTracingStub::AsyncPrepareQuery( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request) { + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) { auto span = internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "PrepareQuery"); internal::OTelScope scope(span); internal::InjectTraceContext(*context, *propagator_); - auto f = child_->AsyncPrepareQuery(cq, context, std::move(options), request); + auto f = + child_->AsyncPrepareQuery(cq, context, std::move(options), request, oc); return internal::EndSpan(std::move(context), std::move(span), std::move(f)); } diff --git a/google/cloud/bigtable/internal/bigtable_tracing_stub.h b/google/cloud/bigtable/internal/bigtable_tracing_stub.h index 9ddb99b028c3f..5c61ed7d7ce61 100644 --- a/google/cloud/bigtable/internal/bigtable_tracing_stub.h +++ b/google/cloud/bigtable/internal/bigtable_tracing_stub.h @@ -42,50 +42,59 @@ class BigtableTracingStub : public BigtableStub { std::unique_ptr> ReadRows(std::shared_ptr context, Options const& options, - google::bigtable::v2::ReadRowsRequest const& request) override; + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) override; std::unique_ptr> - SampleRowKeys( - std::shared_ptr context, Options const& options, - google::bigtable::v2::SampleRowKeysRequest const& request) override; + SampleRowKeys(std::shared_ptr context, + Options const& options, + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) override; StatusOr MutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::MutateRowRequest const& request) override; + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) override; std::unique_ptr> MutateRows(std::shared_ptr context, Options const& options, - google::bigtable::v2::MutateRowsRequest const& request) override; + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) override; StatusOr CheckAndMutateRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) override; + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) override; StatusOr PingAndWarm( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PingAndWarmRequest const& request) override; + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) override; StatusOr ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) override; + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) override; StatusOr PrepareQuery( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::PrepareQueryRequest const& request) override; + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) override; std::unique_ptr> - ExecuteQuery( - std::shared_ptr context, Options const& options, - google::bigtable::v2::ExecuteQueryRequest const& request) override; + ExecuteQuery(std::shared_ptr context, + Options const& options, + google::bigtable::v2::ExecuteQueryRequest const& request, + OperationContext& oc) override; StatusOr GetClientConfiguration( grpc::ClientContext& context, Options const& options, - google::bigtable::v2::GetClientConfigurationRequest const& request) - override; + google::bigtable::v2::GetClientConfigurationRequest const& request, + OperationContext& oc) override; std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc< google::bigtable::v2::SessionRequest, @@ -115,56 +124,61 @@ class BigtableTracingStub : public BigtableStub { AsyncReadRows(google::cloud::CompletionQueue const& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadRowsRequest const& request) override; + google::bigtable::v2::ReadRowsRequest const& request, + OperationContext& oc) 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) override; + AsyncSampleRowKeys(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::SampleRowKeysRequest const& request, + OperationContext& oc) override; future> AsyncMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::MutateRowRequest const& request) override; + google::bigtable::v2::MutateRowRequest const& request, + OperationContext& oc) 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) override; + AsyncMutateRows(google::cloud::CompletionQueue const& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::MutateRowsRequest const& request, + OperationContext& oc) override; future> AsyncCheckAndMutateRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::CheckAndMutateRowRequest const& request) override; + google::bigtable::v2::CheckAndMutateRowRequest const& request, + OperationContext& oc) override; future> AsyncPingAndWarm( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PingAndWarmRequest const& request) override; + google::bigtable::v2::PingAndWarmRequest const& request, + OperationContext& oc) override; future> AsyncReadModifyWriteRow( google::cloud::CompletionQueue& cq, std::shared_ptr context, google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::ReadModifyWriteRowRequest const& request) override; + google::bigtable::v2::ReadModifyWriteRowRequest const& request, + OperationContext& oc) override; future> - AsyncPrepareQuery( - google::cloud::CompletionQueue& cq, - std::shared_ptr context, - google::cloud::internal::ImmutableOptions options, - google::bigtable::v2::PrepareQueryRequest const& request) override; + AsyncPrepareQuery(google::cloud::CompletionQueue& cq, + std::shared_ptr context, + google::cloud::internal::ImmutableOptions options, + google::bigtable::v2::PrepareQueryRequest const& request, + OperationContext& oc) override; private: std::shared_ptr child_; diff --git a/google/cloud/bigtable/internal/bulk_mutator.cc b/google/cloud/bigtable/internal/bulk_mutator.cc index d2a53a321ed58..413e9b8fe3048 100644 --- a/google/cloud/bigtable/internal/bulk_mutator.cc +++ b/google/cloud/bigtable/internal/bulk_mutator.cc @@ -212,7 +212,8 @@ Status BulkMutator::MakeOneRequest(BigtableStub& stub, limiter.Acquire(); // Read the stream of responses. - auto stream = stub.MutateRows(client_context, options, mutations); + auto stream = + stub.MutateRows(client_context, options, mutations, *operation_context_); std::optional status; while (true) { btproto::MutateRowsResponse response; diff --git a/google/cloud/bigtable/internal/bulk_mutator_test.cc b/google/cloud/bigtable/internal/bulk_mutator_test.cc index adbb69025a245..f034719a2d37f 100644 --- a/google/cloud/bigtable/internal/bulk_mutator_test.cc +++ b/google/cloud/bigtable/internal/bulk_mutator_test.cc @@ -107,8 +107,8 @@ class MockMetric : public bigtable_internal::Metric { bigtable_internal::ElementDeliveryParams const&), (override)); MOCK_METHOD(std::unique_ptr, clone, - (bigtable_internal::ResourceLabels resource_labels, - bigtable_internal::DataLabels data_labels), + (bigtable_internal::TableResourceLabels const& resource_labels, + bigtable_internal::TableDataLabels const& data_labels), (const, override)); }; @@ -118,8 +118,8 @@ class CloningMetric : public bigtable_internal::Metric { explicit CloningMetric(std::unique_ptr metric) : metric_(std::move(metric)) {} std::unique_ptr clone( - bigtable_internal::ResourceLabels, - bigtable_internal::DataLabels) const override { + bigtable_internal::TableResourceLabels const&, + bigtable_internal::TableDataLabels const&) const override { return std::move(metric_); } @@ -156,7 +156,8 @@ TEST_F(BulkMutatorTest, Simple) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRows) .WillOnce([this](auto context, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); @@ -210,7 +211,8 @@ TEST_F(BulkMutatorTest, RetryPartialFailure) { // Prepare the mocks for the request. First create a stream response // which indicates a partial failure. .WillOnce([this](auto context, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); @@ -226,7 +228,8 @@ TEST_F(BulkMutatorTest, RetryPartialFailure) { // Prepare a second stream response, because the client should retry after // the partial failure. .WillOnce([this](auto context, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); @@ -282,7 +285,8 @@ TEST_F(BulkMutatorTest, PermanentFailure) { EXPECT_CALL(*mock, MutateRows) // The first RPC return one recoverable and one unrecoverable failure. .WillOnce([this](auto context, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); @@ -298,7 +302,8 @@ TEST_F(BulkMutatorTest, PermanentFailure) { // The BulkMutator should issue a second request, which will return // success for the remaining mutation. .WillOnce([this](auto context, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); @@ -357,7 +362,8 @@ TEST_F(BulkMutatorTest, PartialStream) { // This will be the stream returned by the first request. It is missing // information about one of the mutations. .WillOnce([this](auto context, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); @@ -373,7 +379,8 @@ TEST_F(BulkMutatorTest, PartialStream) { // returned by the second request, which indicates success for the missed // mutation on r1. .WillOnce([this](auto context, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); @@ -422,7 +429,8 @@ TEST_F(BulkMutatorTest, RetryOnlyIdempotent) { // We will setup the mock to return recoverable transient errors for all // mutations. .WillOnce([this](auto context, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_THAT(request, HasCorrectResourceNames()); EXPECT_EQ(2, request.entries_size()); @@ -440,7 +448,8 @@ TEST_F(BulkMutatorTest, RetryOnlyIdempotent) { // idempotent mutation. Make the mock return success for it. .WillOnce([this, expect_r2]( auto context, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_THAT(request, HasCorrectResourceNames()); expect_r2(request); @@ -480,7 +489,8 @@ TEST_F(BulkMutatorTest, RetryInfoHeeded) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRows) .WillOnce([this](auto context, auto const&, - google::bigtable::v2::MutateRowsRequest const&) { + google::bigtable::v2::MutateRowsRequest const&, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); auto status = google::cloud::internal::ResourceExhaustedError("try again"); @@ -494,7 +504,8 @@ TEST_F(BulkMutatorTest, RetryInfoHeeded) { // us that it is safe to retry the mutation, even though it is not // idempotent. .WillOnce([this](auto context, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); std::vector row_keys; for (auto const& entry : request.entries()) { @@ -533,7 +544,8 @@ TEST_F(BulkMutatorTest, RetryInfoIgnored) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRows) .WillOnce([this](auto context, auto const&, - google::bigtable::v2::MutateRowsRequest const&) { + google::bigtable::v2::MutateRowsRequest const&, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); auto status = google::cloud::internal::ResourceExhaustedError("try again"); @@ -571,7 +583,8 @@ TEST_F(BulkMutatorTest, UnconfirmedAreFailed) { // We will setup the mock to return recoverable failures for idempotent // mutations. .WillOnce([this](auto context, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_THAT(request, HasCorrectResourceNames()); EXPECT_EQ(3, request.entries_size()); @@ -610,7 +623,8 @@ TEST_F(BulkMutatorTest, ConfiguresContext) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRows) .WillOnce([this](auto context, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); @@ -639,7 +653,8 @@ TEST_F(BulkMutatorTest, MutationStatusReportedOnOkStream) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRows) .WillOnce([this](auto context, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); @@ -679,7 +694,8 @@ TEST_F(BulkMutatorTest, ReportEitherRetryableMutationFailOrStreamFail) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRows) .WillOnce([this](auto context, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); @@ -721,7 +737,8 @@ TEST_F(BulkMutatorTest, ReportOnlyLatestMutationStatus) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRows) .WillOnce([this](auto context, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); @@ -734,7 +751,8 @@ TEST_F(BulkMutatorTest, ReportOnlyLatestMutationStatus) { return stream; }) .WillOnce([this](auto context, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); @@ -771,24 +789,24 @@ TEST_F(BulkMutatorTest, Throttling) { ::testing::InSequence seq; EXPECT_CALL(*mock_limiter, Acquire); EXPECT_CALL(*mock_stub, MutateRows) - .WillOnce( - [this](auto context, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { - metadata_fixture_.SetServerMetadata(*context, {}); - EXPECT_THAT(request, HasCorrectResourceNames()); - auto stream = std::make_unique(); - EXPECT_CALL(*stream, Read) - .WillOnce([](google::bigtable::v2::MutateRowsResponse* r) { - *r = MakeResponse({{0, grpc::StatusCode::OK}}); - return std::nullopt; - }) - .WillOnce([](google::bigtable::v2::MutateRowsResponse* r) { - *r = MakeResponse({{1, grpc::StatusCode::OK}}); - return std::nullopt; - }) - .WillOnce(Return(Status())); - return stream; - }); + .WillOnce([this](auto context, auto const&, + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { + metadata_fixture_.SetServerMetadata(*context, {}); + EXPECT_THAT(request, HasCorrectResourceNames()); + auto stream = std::make_unique(); + EXPECT_CALL(*stream, Read) + .WillOnce([](google::bigtable::v2::MutateRowsResponse* r) { + *r = MakeResponse({{0, grpc::StatusCode::OK}}); + return std::nullopt; + }) + .WillOnce([](google::bigtable::v2::MutateRowsResponse* r) { + *r = MakeResponse({{1, grpc::StatusCode::OK}}); + return std::nullopt; + }) + .WillOnce(Return(Status())); + return stream; + }); EXPECT_CALL(*mock_limiter, Update).Times(2); } @@ -811,7 +829,8 @@ TEST_F(BulkMutatorTest, BigtableCookies) { EXPECT_CALL(*mock, MutateRows) .WillOnce([this](auto context, auto const&, - google::bigtable::v2::MutateRowsRequest const&) { + google::bigtable::v2::MutateRowsRequest const&, + bigtable_internal::OperationContext&) { // Return a bigtable cookie in the first request. metadata_fixture_.SetServerMetadata( *context, {{}, {{"x-goog-cbt-cookie-routing", "routing"}}}); @@ -822,7 +841,8 @@ TEST_F(BulkMutatorTest, BigtableCookies) { return stream; }) .WillOnce([this](auto context, auto const&, - google::bigtable::v2::MutateRowsRequest const&) { + google::bigtable::v2::MutateRowsRequest const&, + bigtable_internal::OperationContext&) { // Verify that the next request includes the bigtable cookie from above. auto headers = metadata_fixture_.GetMetadata(*context); EXPECT_THAT(headers, diff --git a/google/cloud/bigtable/internal/connection_refresh_state.cc b/google/cloud/bigtable/internal/connection_refresh_state.cc index 89d0efb6038e0..7793c8dc39652 100644 --- a/google/cloud/bigtable/internal/connection_refresh_state.cc +++ b/google/cloud/bigtable/internal/connection_refresh_state.cc @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and #include "google/cloud/bigtable/internal/connection_refresh_state.h" +#include "google/cloud/bigtable/internal/operation_context.h" #include "google/cloud/log.h" #include @@ -136,8 +137,9 @@ void ScheduleStubRefresh( // AsyncWaitConnectionReady. client_context->set_deadline(std::chrono::system_clock::now() + kConnectionReadyTimeout); + OperationContext default_oc; stub->AsyncPingAndWarm(cq, client_context, std::move(options), - request) + request, default_oc) .then( [weak_stub, weak_cq_impl, state, instance_name, connection_status_fn = std::move(connection_status_fn)]( diff --git a/google/cloud/bigtable/internal/connection_refresh_state_test.cc b/google/cloud/bigtable/internal/connection_refresh_state_test.cc index 8d41d08b28c5c..7f9b22548b19a 100644 --- a/google/cloud/bigtable/internal/connection_refresh_state_test.cc +++ b/google/cloud/bigtable/internal/connection_refresh_state_test.cc @@ -155,7 +155,8 @@ TEST_F(ScheduleStubRefreshTest, RefreshedUsingAsyncPingAndWarm) { .WillRepeatedly( [&](CompletionQueue&, std::shared_ptr const&, internal::ImmutableOptions const&, - google::bigtable::v2::PingAndWarmRequest const& request) + google::bigtable::v2::PingAndWarmRequest const& request, + bigtable_internal::OperationContext&) -> future> { EXPECT_THAT(request.name(), Eq(instance_name)); return p2.get_future(); diff --git a/google/cloud/bigtable/internal/data_connection_impl.cc b/google/cloud/bigtable/internal/data_connection_impl.cc index 5a09bf55acaf1..7366d034088fe 100644 --- a/google/cloud/bigtable/internal/data_connection_impl.cc +++ b/google/cloud/bigtable/internal/data_connection_impl.cc @@ -326,7 +326,7 @@ Status DataConnectionImpl::Apply(std::string const& table_name, grpc::ClientContext& context, Options const& options, google::bigtable::v2::MutateRowRequest const& request) { operation_context->PreCall(context); - auto s = stub->MutateRow(context, options, request); + auto s = stub->MutateRow(context, options, request, *operation_context); operation_context->PostCall(context, s.status()); return s; }, @@ -369,7 +369,7 @@ future DataConnectionImpl::AsyncApply(std::string const& table_name, google::bigtable::v2::MutateRowRequest const& request) { operation_context->PreCall(*context); auto f = stub->AsyncMutateRow(cq, context, std::move(options), - request); + request, *operation_context); return f.then( [operation_context, context = std::move(context)](auto f) { auto s = f.get(); @@ -504,7 +504,8 @@ StatusOr DataConnectionImpl::CheckAndMutateRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::CheckAndMutateRowRequest const& request) { operation_context->PreCall(context); - auto s = stub->CheckAndMutateRow(context, options, request); + auto s = stub->CheckAndMutateRow(context, options, request, + *operation_context); operation_context->PostCall(context, s.status()); return s; }, @@ -553,8 +554,9 @@ DataConnectionImpl::AsyncCheckAndMutateRow( google::bigtable::v2::CheckAndMutateRowRequest const& request) { operation_context->PreCall(*context); - auto f = stub->AsyncCheckAndMutateRow( - cq, context, std::move(options), request); + auto f = + stub->AsyncCheckAndMutateRow(cq, context, std::move(options), + request, *operation_context); return f.then( [operation_context, context = std::move(context)](auto f) { auto s = f.get(); @@ -594,7 +596,8 @@ StatusOr> DataConnectionImpl::SampleRows( auto context = std::make_shared(); internal::ConfigureContext(*context, internal::CurrentOptions()); operation_context->PreCall(*context); - auto stream = stub->SampleRowKeys(context, Options{}, request); + auto stream = + stub->SampleRowKeys(context, Options{}, request, *operation_context); std::optional status; while (true) { @@ -656,7 +659,8 @@ StatusOr DataConnectionImpl::ReadModifyWriteRow( grpc::ClientContext& context, Options const& options, google::bigtable::v2::ReadModifyWriteRowRequest const& request) { operation_context->PreCall(context); - auto result = stub->ReadModifyWriteRow(context, options, request); + auto result = stub->ReadModifyWriteRow(context, options, request, + *operation_context); operation_context->PostCall(context, result.status()); return result; }, @@ -686,7 +690,8 @@ future> DataConnectionImpl::AsyncReadModifyWriteRow( request) { operation_context->PreCall(*context); auto f = stub->AsyncReadModifyWriteRow( - cq, context, std::move(options), request); + cq, context, std::move(options), request, + *operation_context); return f.then( [operation_context, context = std::move(context)](auto f) { auto s = f.get(); @@ -819,7 +824,8 @@ StatusOr DataConnectionImpl::PrepareQuery( grpc::ClientContext& context, Options const& options, google::bigtable::v2::PrepareQueryRequest const& request) { operation_context->PreCall(context); - auto const& result = stub->PrepareQuery(context, options, request); + auto const& result = + stub->PrepareQuery(context, options, request, *operation_context); operation_context->PostCall(context, result.status()); return result; }, @@ -858,8 +864,9 @@ StatusOr DataConnectionImpl::PrepareQuery( google::cloud::internal::ImmutableOptions options, google::bigtable::v2::PrepareQueryRequest const& request) { operation_context->PreCall(*context); - auto f = stub->AsyncPrepareQuery(cq, context, - std::move(options), request); + auto f = + stub->AsyncPrepareQuery(cq, context, std::move(options), + request, *operation_context); return f.then( [operation_context, context = std::move(context)](auto f) { auto s = f.get(); @@ -909,7 +916,7 @@ future> DataConnectionImpl::AsyncPrepareQuery( google::bigtable::v2::PrepareQueryRequest const& request) { operation_context->PreCall(*context); auto f = stub->AsyncPrepareQuery(cq, context, std::move(options), - request); + request, *operation_context); return f.then( [operation_context, context = std::move(context)](auto f) { auto s = f.get(); @@ -960,7 +967,8 @@ future> DataConnectionImpl::AsyncPrepareQuery( request) { operation_context->PreCall(*context); auto f = stub->AsyncPrepareQuery( - cq, context, std::move(options), request); + cq, context, std::move(options), request, + *operation_context); return f.then([operation_context, context = std::move(context)](auto f) { auto s = f.get(); @@ -1143,7 +1151,8 @@ bigtable::RowStream DataConnectionImpl::ExecuteQuery( auto const& options = internal::CurrentOptions(); internal::ConfigureContext(*context, options); operation_context->PreCall(*context); - auto stream = stub->ExecuteQuery(context, options, request); + auto stream = + stub->ExecuteQuery(context, options, request, *operation_context); std::unique_ptr reader = std::make_unique( std::move(context), operation_context, std::move(stream)); diff --git a/google/cloud/bigtable/internal/data_connection_impl_test.cc b/google/cloud/bigtable/internal/data_connection_impl_test.cc index 7874bb8e2e5f0..3d6559decfb32 100644 --- a/google/cloud/bigtable/internal/data_connection_impl_test.cc +++ b/google/cloud/bigtable/internal/data_connection_impl_test.cc @@ -16,6 +16,7 @@ #include "google/cloud/bigtable/data_connection.h" #include "google/cloud/bigtable/internal/crc32c.h" #include "google/cloud/bigtable/internal/defaults.h" +#include "google/cloud/bigtable/internal/operation_context.h" #include "google/cloud/bigtable/internal/query_plan.h" #ifdef GOOGLE_CLOUD_CPP_BIGTABLE_WITH_OTEL_METRICS #include "google/cloud/bigtable/internal/metrics.h" @@ -358,7 +359,8 @@ class MockMetric : public Metric { ElementDeliveryParams const&), (override)); MOCK_METHOD(std::unique_ptr, clone, - (ResourceLabels resource_labels, DataLabels data_labels), + (TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels), (const, override)); }; @@ -373,7 +375,8 @@ class CloningMetric : public Metric { std::reverse(metrics_.begin(), metrics_.end()); } - std::unique_ptr clone(ResourceLabels, DataLabels) const override { + std::unique_ptr clone(TableResourceLabels const&, + TableDataLabels const&) const override { auto m = std::move(metrics_.back()); metrics_.pop_back(); return m; @@ -385,7 +388,7 @@ class CloningMetric : public Metric { class FakeOperationContextFactory : public OperationContextFactory { public: - FakeOperationContextFactory(ResourceLabels r, DataLabels d, + FakeOperationContextFactory(TableResourceLabels r, TableDataLabels d, std::shared_ptr metric, std::shared_ptr clock) : resource_labels_(std::move(r)), @@ -442,8 +445,8 @@ class FakeOperationContextFactory : public OperationContextFactory { metrics_, clock_); } - ResourceLabels resource_labels_; - DataLabels data_labels_; + TableResourceLabels resource_labels_; + TableDataLabels data_labels_; std::vector> metrics_; std::shared_ptr clock_; }; @@ -499,7 +502,7 @@ TEST_F(DataConnectionTest, ApplySuccess) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -507,7 +510,8 @@ TEST_F(DataConnectionTest, ApplySuccess) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRow) .WillOnce([](grpc::ClientContext&, Options const&, - v2::MutateRowRequest const& request) { + v2::MutateRowRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ("row", request.row_key()); @@ -531,7 +535,7 @@ TEST_F(DataConnectionTest, ApplyPermanentFailure) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -539,7 +543,8 @@ TEST_F(DataConnectionTest, ApplyPermanentFailure) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRow) .WillOnce([](grpc::ClientContext&, Options const&, - v2::MutateRowRequest const& request) { + v2::MutateRowRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ("row", request.row_key()); @@ -563,7 +568,7 @@ TEST_F(DataConnectionTest, ApplyRetryThenSuccess) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -571,14 +576,16 @@ TEST_F(DataConnectionTest, ApplyRetryThenSuccess) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRow) .WillOnce([](grpc::ClientContext&, Options const&, - v2::MutateRowRequest const& request) { + v2::MutateRowRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ("row", request.row_key()); return TransientError(); }) .WillOnce([](grpc::ClientContext&, Options const&, - v2::MutateRowRequest const& request) { + v2::MutateRowRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ("row", request.row_key()); @@ -602,7 +609,7 @@ TEST_F(DataConnectionTest, ApplyRetryExhausted) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -611,7 +618,8 @@ TEST_F(DataConnectionTest, ApplyRetryExhausted) { EXPECT_CALL(*mock, MutateRow) .Times(kNumRetries + 1) .WillRepeatedly([](grpc::ClientContext&, Options const&, - v2::MutateRowRequest const& request) { + v2::MutateRowRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ("row", request.row_key()); @@ -636,7 +644,8 @@ TEST_F(DataConnectionTest, ApplyRetryIdempotency) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRow) .WillOnce([](grpc::ClientContext&, Options const&, - v2::MutateRowRequest const& request) { + v2::MutateRowRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ("row", request.row_key()); @@ -669,7 +678,7 @@ TEST_F(DataConnectionTest, ApplyBigtableCookie) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -677,14 +686,16 @@ TEST_F(DataConnectionTest, ApplyBigtableCookie) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRow) .WillOnce([this](grpc::ClientContext& context, Options const&, - v2::MutateRowRequest const&) { + v2::MutateRowRequest const&, + bigtable_internal::OperationContext&) { // Return a bigtable cookie in the first request. metadata_fixture_.SetServerMetadata( context, {{}, {{"x-goog-cbt-cookie-routing", "routing"}}}); return TransientError(); }) .WillOnce([this](grpc::ClientContext& context, Options const&, - v2::MutateRowRequest const&) { + v2::MutateRowRequest const&, + bigtable_internal::OperationContext&) { // Verify that the next request includes the bigtable cookie from // above. auto headers = metadata_fixture_.GetMetadata(context); @@ -719,7 +730,7 @@ TEST_F(DataConnectionTest, AsyncApplySuccess) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -727,7 +738,8 @@ TEST_F(DataConnectionTest, AsyncApplySuccess) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncMutateRow) .WillOnce([](google::cloud::CompletionQueue&, auto, auto, - v2::MutateRowRequest const& request) { + v2::MutateRowRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ("row", request.row_key()); @@ -744,7 +756,8 @@ TEST_F(DataConnectionTest, AsyncApplyPermanentFailure) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncMutateRow) .WillOnce([](google::cloud::CompletionQueue&, auto, auto, - v2::MutateRowRequest const& request) { + v2::MutateRowRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ("row", request.row_key()); @@ -769,7 +782,7 @@ TEST_F(DataConnectionTest, AsyncApplyRetryExhausted) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -778,7 +791,8 @@ TEST_F(DataConnectionTest, AsyncApplyRetryExhausted) { EXPECT_CALL(*mock, AsyncMutateRow) .Times(kNumRetries + 1) .WillRepeatedly([](google::cloud::CompletionQueue&, auto, auto, - v2::MutateRowRequest const& request) { + v2::MutateRowRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ("row", request.row_key()); @@ -804,7 +818,8 @@ TEST_F(DataConnectionTest, AsyncApplyRetryIdempotency) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncMutateRow) .WillOnce([](google::cloud::CompletionQueue&, auto, auto, - v2::MutateRowRequest const& request) { + v2::MutateRowRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ("row", request.row_key()); @@ -832,7 +847,8 @@ TEST_F(DataConnectionTest, AsyncApplyBigtableCookie) { EXPECT_CALL(*mock, AsyncMutateRow) .WillOnce([this](CompletionQueue&, std::shared_ptr const& context, - auto, v2::MutateRowRequest const&) { + auto, v2::MutateRowRequest const&, + bigtable_internal::OperationContext&) { // Return a bigtable cookie in the first request. metadata_fixture_.SetServerMetadata( *context, {{}, {{"x-goog-cbt-cookie-routing", "routing"}}}); @@ -841,7 +857,8 @@ TEST_F(DataConnectionTest, AsyncApplyBigtableCookie) { }) .WillOnce([this](CompletionQueue&, std::shared_ptr const& context, - auto, v2::MutateRowRequest const&) { + auto, v2::MutateRowRequest const&, + bigtable_internal::OperationContext&) { // Verify that the next request includes the bigtable cookie from above. auto headers = metadata_fixture_.GetMetadata(*context); EXPECT_THAT(headers, @@ -876,7 +893,7 @@ TEST_F(DataConnectionTest, BulkApplyEmpty) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -898,7 +915,7 @@ TEST_F(DataConnectionTest, BulkApplySuccess) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -909,7 +926,8 @@ TEST_F(DataConnectionTest, BulkApplySuccess) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_THAT(request.entries(), ElementsAre(Entry("r0"), Entry("r1"))); @@ -941,7 +959,7 @@ TEST_F(DataConnectionTest, BulkApplyRetryMutationPolicy) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -957,7 +975,8 @@ TEST_F(DataConnectionTest, BulkApplyRetryMutationPolicy) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); auto stream = std::make_unique(); @@ -974,7 +993,8 @@ TEST_F(DataConnectionTest, BulkApplyRetryMutationPolicy) { return stream; }) .WillOnce([](auto, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_THAT(request.entries(), @@ -1006,7 +1026,7 @@ TEST_F(DataConnectionTest, BulkApplyIncompleteStreamRetried) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -1017,7 +1037,8 @@ TEST_F(DataConnectionTest, BulkApplyIncompleteStreamRetried) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); auto stream = std::make_unique(); @@ -1030,7 +1051,8 @@ TEST_F(DataConnectionTest, BulkApplyIncompleteStreamRetried) { return stream; }) .WillOnce([](auto, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_THAT(request.entries(), ElementsAre(Entry("forgotten"))); @@ -1061,7 +1083,7 @@ TEST_F(DataConnectionTest, BulkApplyStreamRetryExhausted) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -1071,15 +1093,15 @@ TEST_F(DataConnectionTest, BulkApplyStreamRetryExhausted) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRows) .Times(kNumRetries + 1) - .WillRepeatedly( - [](auto, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { - EXPECT_EQ(kAppProfile, request.app_profile_id()); - EXPECT_EQ(kTableName, request.table_name()); - auto stream = std::make_unique(); - EXPECT_CALL(*stream, Read).WillOnce(Return(TransientError())); - return stream; - }); + .WillRepeatedly([](auto, auto const&, + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { + EXPECT_EQ(kAppProfile, request.app_profile_id()); + EXPECT_EQ(kTableName, request.table_name()); + auto stream = std::make_unique(); + EXPECT_CALL(*stream, Read).WillOnce(Return(TransientError())); + return stream; + }); auto mock_b = std::make_unique(); EXPECT_CALL(*mock_b, clone).WillOnce([]() { @@ -1106,7 +1128,7 @@ TEST_F(DataConnectionTest, BulkApplyStreamPermanentError) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -1117,7 +1139,8 @@ TEST_F(DataConnectionTest, BulkApplyStreamPermanentError) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); auto stream = std::make_unique(); @@ -1138,7 +1161,8 @@ TEST_F(DataConnectionTest, BulkApplyNoSleepIfNoPendingMutations) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); auto stream = std::make_unique(); @@ -1170,23 +1194,22 @@ TEST_F(DataConnectionTest, BulkApplyRetriesOkStreamWithFailedMutations) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRows) .Times(kNumRetries + 1) - .WillRepeatedly( - [](auto, auto const&, - google::bigtable::v2::MutateRowsRequest const& request) { - EXPECT_EQ(kAppProfile, request.app_profile_id()); - EXPECT_EQ(kTableName, request.table_name()); - auto stream = std::make_unique(); - // The overall stream succeeds, but it contains failed mutations. - // Our retry and backoff policies should take effect. - EXPECT_CALL(*stream, Read) - .WillOnce([](google::bigtable::v2::MutateRowsResponse* r) { - *r = MakeBulkApplyResponse( - {{0, grpc::StatusCode::UNAVAILABLE}}); - return std::nullopt; - }) - .WillOnce(Return(Status())); - return stream; - }); + .WillRepeatedly([](auto, auto const&, + google::bigtable::v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { + EXPECT_EQ(kAppProfile, request.app_profile_id()); + EXPECT_EQ(kTableName, request.table_name()); + auto stream = std::make_unique(); + // The overall stream succeeds, but it contains failed mutations. + // Our retry and backoff policies should take effect. + EXPECT_CALL(*stream, Read) + .WillOnce([](google::bigtable::v2::MutateRowsResponse* r) { + *r = MakeBulkApplyResponse({{0, grpc::StatusCode::UNAVAILABLE}}); + return std::nullopt; + }) + .WillOnce(Return(Status())); + return stream; + }); auto mock_b = std::make_unique(); EXPECT_CALL(*mock_b, clone).WillOnce([]() { @@ -1207,14 +1230,16 @@ TEST_F(DataConnectionTest, BulkApplyRetryInfoHeeded) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRows) - .WillOnce([](auto, auto const&, v2::MutateRowsRequest const&) { + .WillOnce([](auto, auto const&, v2::MutateRowsRequest const&, + bigtable_internal::OperationContext&) { auto status = PermanentError(); internal::SetRetryInfo(status, internal::RetryInfo{ms(0)}); auto stream = std::make_unique(); EXPECT_CALL(*stream, Read).WillOnce(Return(status)); return stream; }) - .WillOnce([](auto, auto const&, v2::MutateRowsRequest const&) { + .WillOnce([](auto, auto const&, v2::MutateRowsRequest const&, + bigtable_internal::OperationContext&) { auto stream = std::make_unique(); EXPECT_CALL(*stream, Read) .WillOnce([](google::bigtable::v2::MutateRowsResponse* r) { @@ -1238,7 +1263,8 @@ TEST_F(DataConnectionTest, BulkApplyRetryInfoIgnored) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRows) - .WillOnce([](auto, auto const&, v2::MutateRowsRequest const&) { + .WillOnce([](auto, auto const&, v2::MutateRowsRequest const&, + bigtable_internal::OperationContext&) { auto status = PermanentError(); internal::SetRetryInfo(status, internal::RetryInfo{ms(0)}); auto stream = std::make_unique(); @@ -1289,7 +1315,8 @@ TEST_F(DataConnectionTest, AsyncBulkApply) { .WillOnce(Return(ByMove(make_ready_future()))); EXPECT_CALL(*mock_stub, AsyncMutateRows) .WillOnce([](CompletionQueue const&, auto, auto, - v2::MutateRowsRequest const& request) { + v2::MutateRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); using ErrorStream = @@ -1310,7 +1337,8 @@ TEST_F(DataConnectionTest, ReadRows) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ(42, request.rows_limit()); @@ -1333,7 +1361,8 @@ TEST_F(DataConnectionTest, ReadRowsReverseScan) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_TRUE(request.reversed()); auto stream = std::make_unique(); @@ -1353,7 +1382,8 @@ TEST_F(DataConnectionTest, ReadRowsFull) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ(42, request.rows_limit()); @@ -1376,20 +1406,22 @@ TEST_F(DataConnectionTest, ReadRowsFull) { TEST_F(DataConnectionTest, ReadRowsRetryInfoHeeded) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) - .WillOnce( - [](auto, auto const&, google::bigtable::v2::ReadRowsRequest const&) { - auto status = PermanentError(); - internal::SetRetryInfo(status, internal::RetryInfo{ms(0)}); - auto stream = std::make_unique(); - EXPECT_CALL(*stream, Read).WillOnce(Return(status)); - return stream; - }) - .WillOnce( - [](auto, auto const&, google::bigtable::v2::ReadRowsRequest const&) { - auto stream = std::make_unique(); - EXPECT_CALL(*stream, Read).WillOnce(Return(Status())); - return stream; - }); + .WillOnce([](auto, auto const&, + google::bigtable::v2::ReadRowsRequest const&, + bigtable_internal::OperationContext&) { + auto status = PermanentError(); + internal::SetRetryInfo(status, internal::RetryInfo{ms(0)}); + auto stream = std::make_unique(); + EXPECT_CALL(*stream, Read).WillOnce(Return(status)); + return stream; + }) + .WillOnce([](auto, auto const&, + google::bigtable::v2::ReadRowsRequest const&, + bigtable_internal::OperationContext&) { + auto stream = std::make_unique(); + EXPECT_CALL(*stream, Read).WillOnce(Return(Status())); + return stream; + }); auto conn = TestConnection(std::move(mock)); internal::OptionsSpan span( @@ -1402,14 +1434,15 @@ TEST_F(DataConnectionTest, ReadRowsRetryInfoHeeded) { TEST_F(DataConnectionTest, ReadRowsRetryInfoIgnored) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) - .WillOnce( - [](auto, auto const&, google::bigtable::v2::ReadRowsRequest const&) { - auto status = PermanentError(); - internal::SetRetryInfo(status, internal::RetryInfo{ms(0)}); - auto stream = std::make_unique(); - EXPECT_CALL(*stream, Read).WillOnce(Return(status)); - return stream; - }); + .WillOnce([](auto, auto const&, + google::bigtable::v2::ReadRowsRequest const&, + bigtable_internal::OperationContext&) { + auto status = PermanentError(); + internal::SetRetryInfo(status, internal::RetryInfo{ms(0)}); + auto stream = std::make_unique(); + EXPECT_CALL(*stream, Read).WillOnce(Return(status)); + return stream; + }); auto conn = TestConnection(std::move(mock)); internal::OptionsSpan span( @@ -1432,7 +1465,7 @@ TEST_F(DataConnectionTest, ReadRowEmpty) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -1440,7 +1473,8 @@ TEST_F(DataConnectionTest, ReadRowEmpty) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ(1, request.rows_limit()); @@ -1470,7 +1504,7 @@ TEST_F(DataConnectionTest, ReadRowSuccess) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -1478,7 +1512,8 @@ TEST_F(DataConnectionTest, ReadRowSuccess) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ(1, request.rows_limit()); @@ -1520,7 +1555,7 @@ TEST_F(DataConnectionTest, ReadRowFailure) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -1528,7 +1563,8 @@ TEST_F(DataConnectionTest, ReadRowFailure) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ(1, request.rows_limit()); @@ -1566,7 +1602,7 @@ TEST_F(DataConnectionTest, CheckAndMutateRowSuccess) { auto fake_metric = std::make_shared(std::move(v)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -1578,7 +1614,8 @@ TEST_F(DataConnectionTest, CheckAndMutateRowSuccess) { auto mock = std::make_shared(); EXPECT_CALL(*mock, CheckAndMutateRow) .WillOnce([&](grpc::ClientContext&, Options const&, - v2::CheckAndMutateRowRequest const& request) { + v2::CheckAndMutateRowRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ("row", request.row_key()); @@ -1593,7 +1630,8 @@ TEST_F(DataConnectionTest, CheckAndMutateRowSuccess) { return resp; }) .WillOnce([&](grpc::ClientContext&, Options const&, - v2::CheckAndMutateRowRequest const& request) { + v2::CheckAndMutateRowRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ("row", request.row_key()); @@ -1630,7 +1668,8 @@ TEST_F(DataConnectionTest, CheckAndMutateRowIdempotency) { auto mock = std::make_shared(); EXPECT_CALL(*mock, CheckAndMutateRow) .WillOnce([&](grpc::ClientContext&, Options const&, - v2::CheckAndMutateRowRequest const& request) { + v2::CheckAndMutateRowRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ("row", request.row_key()); @@ -1670,7 +1709,7 @@ TEST_F(DataConnectionTest, CheckAndMutateRowPermanentError) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -1683,7 +1722,8 @@ TEST_F(DataConnectionTest, CheckAndMutateRowPermanentError) { auto mock = std::make_shared(); EXPECT_CALL(*mock, CheckAndMutateRow) .WillOnce([&](grpc::ClientContext&, Options const&, - v2::CheckAndMutateRowRequest const& request) { + v2::CheckAndMutateRowRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ("row", request.row_key()); @@ -1713,7 +1753,7 @@ TEST_F(DataConnectionTest, CheckAndMutateRowRetryExhausted) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -1727,7 +1767,8 @@ TEST_F(DataConnectionTest, CheckAndMutateRowRetryExhausted) { EXPECT_CALL(*mock, CheckAndMutateRow) .Times(kNumRetries + 1) .WillRepeatedly([&](grpc::ClientContext&, Options const&, - v2::CheckAndMutateRowRequest const& request) { + v2::CheckAndMutateRowRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ("row", request.row_key()); @@ -1766,14 +1807,16 @@ TEST_F(DataConnectionTest, CheckAndMutateRowBigtableCookie) { auto mock = std::make_shared(); EXPECT_CALL(*mock, CheckAndMutateRow) .WillOnce([this](grpc::ClientContext& context, Options const&, - v2::CheckAndMutateRowRequest const&) { + v2::CheckAndMutateRowRequest const&, + bigtable_internal::OperationContext&) { // Return a bigtable cookie in the first request. metadata_fixture_.SetServerMetadata( context, {{}, {{"x-goog-cbt-cookie-routing", "routing"}}}); return TransientError(); }) .WillOnce([this](grpc::ClientContext& context, Options const&, - v2::CheckAndMutateRowRequest const&) { + v2::CheckAndMutateRowRequest const&, + bigtable_internal::OperationContext&) { // Verify that the next request includes the bigtable cookie from above. auto headers = metadata_fixture_.GetMetadata(context); EXPECT_THAT(headers, @@ -1819,7 +1862,7 @@ TEST_F(DataConnectionTest, AsyncCheckAndMutateRowSuccess) { auto fake_metric = std::make_shared(std::move(v)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -1832,7 +1875,8 @@ TEST_F(DataConnectionTest, AsyncCheckAndMutateRowSuccess) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncCheckAndMutateRow) .WillOnce([&](google::cloud::CompletionQueue&, auto, auto, - v2::CheckAndMutateRowRequest const& request) { + v2::CheckAndMutateRowRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ("row", request.row_key()); @@ -1847,7 +1891,8 @@ TEST_F(DataConnectionTest, AsyncCheckAndMutateRowSuccess) { return make_ready_future(make_status_or(resp)); }) .WillOnce([&](google::cloud::CompletionQueue&, auto, auto, - v2::CheckAndMutateRowRequest const& request) { + v2::CheckAndMutateRowRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ("row", request.row_key()); @@ -1886,7 +1931,8 @@ TEST_F(DataConnectionTest, AsyncCheckAndMutateRowIdempotency) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncCheckAndMutateRow) .WillOnce([&](google::cloud::CompletionQueue&, auto, auto, - v2::CheckAndMutateRowRequest const& request) { + v2::CheckAndMutateRowRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ("row", request.row_key()); @@ -1927,7 +1973,7 @@ TEST_F(DataConnectionTest, AsyncCheckAndMutateRowPermanentError) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -1940,7 +1986,8 @@ TEST_F(DataConnectionTest, AsyncCheckAndMutateRowPermanentError) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncCheckAndMutateRow) .WillOnce([&](google::cloud::CompletionQueue&, auto, auto, - v2::CheckAndMutateRowRequest const& request) { + v2::CheckAndMutateRowRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ("row", request.row_key()); @@ -1971,7 +2018,7 @@ TEST_F(DataConnectionTest, AsyncCheckAndMutateRowRetryExhausted) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -1985,7 +2032,8 @@ TEST_F(DataConnectionTest, AsyncCheckAndMutateRowRetryExhausted) { EXPECT_CALL(*mock, AsyncCheckAndMutateRow) .Times(kNumRetries + 1) .WillRepeatedly([&](google::cloud::CompletionQueue&, auto, auto, - v2::CheckAndMutateRowRequest const& request) { + v2::CheckAndMutateRowRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ("row", request.row_key()); @@ -2026,7 +2074,8 @@ TEST_F(DataConnectionTest, AsyncCheckAndMutateRowBigtableCookie) { EXPECT_CALL(*mock, AsyncCheckAndMutateRow) .WillOnce([this](CompletionQueue&, std::shared_ptr const& context, - auto, v2::CheckAndMutateRowRequest const&) { + auto, v2::CheckAndMutateRowRequest const&, + bigtable_internal::OperationContext&) { // Return a bigtable cookie in the first request. metadata_fixture_.SetServerMetadata( *context, {{}, {{"x-goog-cbt-cookie-routing", "routing"}}}); @@ -2035,7 +2084,8 @@ TEST_F(DataConnectionTest, AsyncCheckAndMutateRowBigtableCookie) { }) .WillOnce([this](CompletionQueue&, std::shared_ptr const& context, - auto, v2::CheckAndMutateRowRequest const&) { + auto, v2::CheckAndMutateRowRequest const&, + bigtable_internal::OperationContext&) { // Verify that the next request includes the bigtable cookie from above. auto headers = metadata_fixture_.GetMetadata(*context); EXPECT_THAT(headers, @@ -2073,7 +2123,7 @@ TEST_F(DataConnectionTest, SampleRowsSuccess) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -2081,7 +2131,8 @@ TEST_F(DataConnectionTest, SampleRowsSuccess) { auto mock = std::make_shared(); EXPECT_CALL(*mock, SampleRowKeys) .WillOnce([this](auto client_context, auto const&, - v2::SampleRowKeysRequest const& request) { + v2::SampleRowKeysRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -2123,14 +2174,15 @@ TEST_F(DataConnectionTest, SampleRowsRetryResetsSamples) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif auto mock = std::make_shared(); EXPECT_CALL(*mock, SampleRowKeys) - .WillOnce([](auto, auto const&, v2::SampleRowKeysRequest const& request) { + .WillOnce([](auto, auto const&, v2::SampleRowKeysRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); auto stream = std::make_unique(); @@ -2142,7 +2194,8 @@ TEST_F(DataConnectionTest, SampleRowsRetryResetsSamples) { .WillOnce(Return(TransientError())); return stream; }) - .WillOnce([](auto, auto const&, v2::SampleRowKeysRequest const& request) { + .WillOnce([](auto, auto const&, v2::SampleRowKeysRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); auto stream = std::make_unique(); @@ -2175,7 +2228,7 @@ TEST_F(DataConnectionTest, SampleRowsRetryExhausted) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -2184,7 +2237,8 @@ TEST_F(DataConnectionTest, SampleRowsRetryExhausted) { EXPECT_CALL(*mock, SampleRowKeys) .Times(kNumRetries + 1) .WillRepeatedly([this](auto context, auto const&, - v2::SampleRowKeysRequest const& request) { + v2::SampleRowKeysRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -2222,7 +2276,7 @@ TEST_F(DataConnectionTest, SampleRowsPermanentError) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -2230,7 +2284,8 @@ TEST_F(DataConnectionTest, SampleRowsPermanentError) { auto mock = std::make_shared(); EXPECT_CALL(*mock, SampleRowKeys) .WillOnce([this](auto client_context, auto const&, - v2::SampleRowKeysRequest const& request) { + v2::SampleRowKeysRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -2252,26 +2307,28 @@ TEST_F(DataConnectionTest, SampleRowsPermanentError) { TEST_F(DataConnectionTest, SampleRowsBigtableCookie) { auto mock = std::make_shared(); EXPECT_CALL(*mock, SampleRowKeys) - .WillOnce( - [this](auto context, auto const&, v2::SampleRowKeysRequest const&) { - // Return a bigtable cookie in the first request. - metadata_fixture_.SetServerMetadata( - *context, {{}, {{"x-goog-cbt-cookie-routing", "routing"}}}); - auto stream = std::make_unique(); - EXPECT_CALL(*stream, Read).WillOnce(Return(TransientError())); - return stream; - }) - .WillOnce( - [this](auto context, auto const&, v2::SampleRowKeysRequest const&) { - // Verify that the next request includes the bigtable cookie from - // above. - auto headers = metadata_fixture_.GetMetadata(*context); - EXPECT_THAT(headers, - Contains(Pair("x-goog-cbt-cookie-routing", "routing"))); - auto stream = std::make_unique(); - EXPECT_CALL(*stream, Read).WillOnce(Return(PermanentError())); - return stream; - }); + .WillOnce([this](auto context, auto const&, + v2::SampleRowKeysRequest const&, + bigtable_internal::OperationContext&) { + // Return a bigtable cookie in the first request. + metadata_fixture_.SetServerMetadata( + *context, {{}, {{"x-goog-cbt-cookie-routing", "routing"}}}); + auto stream = std::make_unique(); + EXPECT_CALL(*stream, Read).WillOnce(Return(TransientError())); + return stream; + }) + .WillOnce([this](auto context, auto const&, + v2::SampleRowKeysRequest const&, + bigtable_internal::OperationContext&) { + // Verify that the next request includes the bigtable cookie from + // above. + auto headers = metadata_fixture_.GetMetadata(*context); + EXPECT_THAT(headers, + Contains(Pair("x-goog-cbt-cookie-routing", "routing"))); + auto stream = std::make_unique(); + EXPECT_CALL(*stream, Read).WillOnce(Return(PermanentError())); + return stream; + }); auto mock_b = std::make_unique(); EXPECT_CALL(*mock_b, clone).WillOnce([]() { @@ -2291,14 +2348,16 @@ TEST_F(DataConnectionTest, SampleRowsBigtableCookie) { TEST_F(DataConnectionTest, SampleRowsRetryInfoHeeded) { auto mock = std::make_shared(); EXPECT_CALL(*mock, SampleRowKeys) - .WillOnce([](auto, auto const&, v2::SampleRowKeysRequest const&) { + .WillOnce([](auto, auto const&, v2::SampleRowKeysRequest const&, + bigtable_internal::OperationContext&) { auto status = PermanentError(); internal::SetRetryInfo(status, internal::RetryInfo{ms(0)}); auto stream = std::make_unique(); EXPECT_CALL(*stream, Read).WillOnce(Return(status)); return stream; }) - .WillOnce([](auto, auto const&, v2::SampleRowKeysRequest const&) { + .WillOnce([](auto, auto const&, v2::SampleRowKeysRequest const&, + bigtable_internal::OperationContext&) { auto stream = std::make_unique(); EXPECT_CALL(*stream, Read).WillOnce(Return(Status())); return stream; @@ -2314,7 +2373,8 @@ TEST_F(DataConnectionTest, SampleRowsRetryInfoHeeded) { TEST_F(DataConnectionTest, SampleRowsRetryInfoIgnored) { auto mock = std::make_shared(); EXPECT_CALL(*mock, SampleRowKeys) - .WillOnce([](auto, auto const&, v2::SampleRowKeysRequest const&) { + .WillOnce([](auto, auto const&, v2::SampleRowKeysRequest const&, + bigtable_internal::OperationContext&) { auto status = PermanentError(); internal::SetRetryInfo(status, internal::RetryInfo{ms(0)}); auto stream = std::make_unique(); @@ -2335,7 +2395,8 @@ TEST_F(DataConnectionTest, AsyncSampleRows) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncSampleRowKeys) .WillOnce([](CompletionQueue const&, auto, auto, - v2::SampleRowKeysRequest const& request) { + v2::SampleRowKeysRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); using ErrorStream = @@ -2360,7 +2421,7 @@ TEST_F(DataConnectionTest, ReadModifyWriteRowSuccess) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -2382,7 +2443,8 @@ TEST_F(DataConnectionTest, ReadModifyWriteRowSuccess) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadModifyWriteRow) .WillOnce([&response](grpc::ClientContext&, Options const&, - v2::ReadModifyWriteRowRequest const& request) { + v2::ReadModifyWriteRowRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ("row", request.row_key()); @@ -2415,7 +2477,7 @@ TEST_F(DataConnectionTest, ReadModifyWriteRowPermanentError) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -2423,7 +2485,8 @@ TEST_F(DataConnectionTest, ReadModifyWriteRowPermanentError) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadModifyWriteRow) .WillOnce([](grpc::ClientContext&, Options const&, - v2::ReadModifyWriteRowRequest const& request) { + v2::ReadModifyWriteRowRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ("row", request.row_key()); @@ -2452,7 +2515,7 @@ TEST_F(DataConnectionTest, ReadModifyWriteRowTransientErrorNotRetried) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -2460,7 +2523,8 @@ TEST_F(DataConnectionTest, ReadModifyWriteRowTransientErrorNotRetried) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadModifyWriteRow) .WillOnce([](grpc::ClientContext&, Options const&, - v2::ReadModifyWriteRowRequest const& request) { + v2::ReadModifyWriteRowRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ("row", request.row_key()); @@ -2497,7 +2561,7 @@ TEST_F(DataConnectionTest, AsyncReadModifyWriteRowSuccess) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -2518,9 +2582,10 @@ TEST_F(DataConnectionTest, AsyncReadModifyWriteRowSuccess) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadModifyWriteRow) - .WillOnce([&response, this]( - google::cloud::CompletionQueue&, auto client_context, auto, - v2::ReadModifyWriteRowRequest const& request) { + .WillOnce([&response, this](google::cloud::CompletionQueue&, + auto client_context, auto, + v2::ReadModifyWriteRowRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -2558,7 +2623,7 @@ TEST_F(DataConnectionTest, AsyncReadModifyWriteRowPermanentError) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -2566,7 +2631,8 @@ TEST_F(DataConnectionTest, AsyncReadModifyWriteRowPermanentError) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadModifyWriteRow) .WillOnce([this](google::cloud::CompletionQueue&, auto client_context, - auto, v2::ReadModifyWriteRowRequest const& request) { + auto, v2::ReadModifyWriteRowRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -2594,7 +2660,8 @@ TEST_F(DataConnectionTest, AsyncReadModifyWriteRowTransientErrorNotRetried) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadModifyWriteRow) .WillOnce([this](google::cloud::CompletionQueue&, auto client_context, - auto, v2::ReadModifyWriteRowRequest const& request) { + auto, v2::ReadModifyWriteRowRequest const& request, + bigtable_internal::OperationContext&) { metadata_fixture_.SetServerMetadata(*client_context, {}); EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); @@ -2633,7 +2700,8 @@ TEST_F(DataConnectionTest, AsyncReadRows) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) .WillOnce([](CompletionQueue const&, auto, auto, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ(42, request.rows_limit()); @@ -2663,7 +2731,8 @@ TEST_F(DataConnectionTest, AsyncReadRowsReverseScan) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) .WillOnce([](CompletionQueue const&, auto, auto, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_TRUE(request.reversed()); using ErrorStream = internal::AsyncStreamingReadRpcError; @@ -2696,7 +2765,7 @@ TEST_F(DataConnectionTest, AsyncReadRowEmpty) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -2704,7 +2773,8 @@ TEST_F(DataConnectionTest, AsyncReadRowEmpty) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) .WillOnce([](CompletionQueue const&, auto, auto, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ(1, request.rows_limit()); @@ -2742,7 +2812,7 @@ TEST_F(DataConnectionTest, AsyncReadRowSuccess) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -2750,7 +2820,8 @@ TEST_F(DataConnectionTest, AsyncReadRowSuccess) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) .WillOnce([](CompletionQueue const&, auto, auto, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ(1, request.rows_limit()); @@ -2799,7 +2870,7 @@ TEST_F(DataConnectionTest, AsyncReadRowFailure) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -2807,7 +2878,8 @@ TEST_F(DataConnectionTest, AsyncReadRowFailure) { auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncReadRows) .WillOnce([](CompletionQueue const&, auto, auto, - v2::ReadRowsRequest const& request) { + v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ(kTableName, request.table_name()); EXPECT_EQ(1, request.rows_limit()); @@ -2836,7 +2908,7 @@ TEST_F(DataConnectionTest, PrepareQuerySuccess) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -2861,7 +2933,8 @@ TEST_F(DataConnectionTest, PrepareQuerySuccess) { kResultMetadataText, response.mutable_metadata())); EXPECT_CALL(*mock, PrepareQuery) .WillOnce([&](grpc::ClientContext&, Options const&, - v2::PrepareQueryRequest const& request) { + v2::PrepareQueryRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ("projects/the-project/instances/the-instance", request.instance_name()); @@ -2902,15 +2975,17 @@ TEST_F(DataConnectionTest, PrepareQueryPermanentError) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif auto mock = std::make_shared(); EXPECT_CALL(*mock, PrepareQuery) - .WillOnce( - [](grpc::ClientContext&, Options const&, - v2::PrepareQueryRequest const&) { return PermanentError(); }); + .WillOnce([](grpc::ClientContext&, Options const&, + v2::PrepareQueryRequest const&, + bigtable_internal::OperationContext&) { + return PermanentError(); + }); auto conn = TestConnection(std::move(mock), std::move(factory)); internal::OptionsSpan span(CallOptions()); @@ -2932,7 +3007,7 @@ TEST_F(DataConnectionTest, AsyncPrepareQuerySuccess) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -2954,7 +3029,8 @@ TEST_F(DataConnectionTest, AsyncPrepareQuerySuccess) { kResultMetadataText, response.mutable_metadata())); EXPECT_CALL(*mock, AsyncPrepareQuery) .WillOnce([&](CompletionQueue const&, auto, auto, - v2::PrepareQueryRequest const& request) { + v2::PrepareQueryRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(kAppProfile, request.app_profile_id()); EXPECT_EQ("projects/the-project/instances/the-instance", request.instance_name()); @@ -2994,17 +3070,17 @@ TEST_F(DataConnectionTest, AsyncPrepareQueryPermanentError) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncPrepareQuery) - .WillOnce( - [](CompletionQueue&, auto, auto, v2::PrepareQueryRequest const&) { - return make_ready_future>( - PermanentError()); - }); + .WillOnce([](CompletionQueue&, auto, auto, v2::PrepareQueryRequest const&, + bigtable_internal::OperationContext&) { + return make_ready_future>( + PermanentError()); + }); auto conn = TestConnection(std::move(mock), std::move(factory)); internal::OptionsSpan span(CallOptions()); @@ -3028,7 +3104,7 @@ TEST_F(DataConnectionTest, ExecuteQuerySuccessWithTransientErrors) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -3065,19 +3141,22 @@ TEST_F(DataConnectionTest, ExecuteQuerySuccessWithTransientErrors) { }; EXPECT_CALL(*mock, ExecuteQuery) .WillOnce([&](auto, auto const&, - google::bigtable::v2::ExecuteQueryRequest const&) { + google::bigtable::v2::ExecuteQueryRequest const&, + bigtable_internal::OperationContext&) { auto error_stream = std::make_unique(); EXPECT_CALL(*error_stream, Read).WillOnce(Return(TransientError())); return error_stream; }) .WillOnce([&](auto, auto const&, - google::bigtable::v2::ExecuteQueryRequest const&) { + google::bigtable::v2::ExecuteQueryRequest const&, + bigtable_internal::OperationContext&) { auto error_stream = std::make_unique(); EXPECT_CALL(*error_stream, Read).WillOnce(Return(TransientError())); return error_stream; }) .WillOnce([&](auto, auto const&, - google::bigtable::v2::ExecuteQueryRequest const& request) { + google::bigtable::v2::ExecuteQueryRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(request.app_profile_id(), kAppProfile); EXPECT_EQ(request.instance_name(), "projects/test-project/instances/test-instance"); @@ -3143,7 +3222,7 @@ TEST_F(DataConnectionTest, ExecuteQueryFailure) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -3165,7 +3244,7 @@ TEST_F(DataConnectionTest, ExecuteQueryFailure) { std::move(refresh_fn)); EXPECT_CALL(*mock, ExecuteQuery) - .WillOnce([&](auto, auto const&, auto const&) { + .WillOnce([&](auto, auto const&, auto const&, auto&) { auto stream = std::make_unique(); EXPECT_CALL(*stream, Read).WillOnce(Return(PermanentError())); return stream; @@ -3197,7 +3276,7 @@ TEST_F(DataConnectionTest, ExecuteQueryOperationRetryExhausted) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -3220,7 +3299,7 @@ TEST_F(DataConnectionTest, ExecuteQueryOperationRetryExhausted) { EXPECT_CALL(*mock, ExecuteQuery) .Times(3) - .WillRepeatedly([&](auto, auto const&, auto const&) { + .WillRepeatedly([&](auto, auto const&, auto const&, auto&) { auto stream = std::make_unique(); EXPECT_CALL(*stream, Read).WillOnce(Return(TransientError())); return stream; @@ -3252,7 +3331,7 @@ TEST_F(DataConnectionTest, ExecuteQuerySuccessWithQueryPlanRefresh) { auto fake_metric = std::make_shared(std::move(mock_metric)); auto clock = std::make_shared(); auto factory = std::make_unique( - ResourceLabels{}, DataLabels{}, fake_metric, clock); + TableResourceLabels{}, TableDataLabels{}, fake_metric, clock); #else auto factory = std::make_unique(); #endif @@ -3319,14 +3398,16 @@ TEST_F(DataConnectionTest, ExecuteQuerySuccessWithQueryPlanRefresh) { EXPECT_CALL(*mock, ExecuteQuery) .WillOnce([&](auto, auto const&, - google::bigtable::v2::ExecuteQueryRequest const& request) { + google::bigtable::v2::ExecuteQueryRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(request.prepared_query(), "test-pq-id-initial"); auto error_stream = std::make_unique(); EXPECT_CALL(*error_stream, Read).WillOnce(Return(QueryPlanError())); return error_stream; }) .WillOnce([&](auto, auto const&, - google::bigtable::v2::ExecuteQueryRequest const& request) { + google::bigtable::v2::ExecuteQueryRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(request.app_profile_id(), kAppProfile); EXPECT_EQ(request.instance_name(), "projects/test-project/instances/test-instance"); @@ -3401,15 +3482,17 @@ TEST_F(DataConnectionTest, PrepareAndExecuteQuerySuccessWithQueryPlanRefresh) { std::vector> metrics; metrics.push_back(std::make_shared( std::move(mock_metric_prepared_query_initial))); - return std::make_shared( - ResourceLabels{}, DataLabels{}, std::move(metrics), clock); + return std::make_shared(TableResourceLabels{}, + TableDataLabels{}, + std::move(metrics), clock); }) .WillOnce([&](auto const&, auto const&) { std::vector> metrics; metrics.push_back(std::make_shared( std::move(mock_metric_prepared_query_refresh))); - return std::make_shared( - ResourceLabels{}, DataLabels{}, std::move(metrics), clock); + return std::make_shared(TableResourceLabels{}, + TableDataLabels{}, + std::move(metrics), clock); }); auto mock_metric_execute_query = std::make_unique(); @@ -3423,8 +3506,8 @@ TEST_F(DataConnectionTest, PrepareAndExecuteQuerySuccessWithQueryPlanRefresh) { std::vector> metrics; metrics.push_back( std::make_shared(std::move(mock_metric_execute_query))); - return std::make_shared(ResourceLabels{}, DataLabels{}, - std::move(metrics), clock); + return std::make_shared( + TableResourceLabels{}, TableDataLabels{}, std::move(metrics), clock); }); #else auto factory = std::make_unique(); @@ -3478,22 +3561,27 @@ TEST_F(DataConnectionTest, PrepareAndExecuteQuerySuccessWithQueryPlanRefresh) { auto mock = std::make_shared(); EXPECT_CALL(*mock, PrepareQuery) - .WillOnce( - [&](grpc::ClientContext&, Options const&, - PrepareQueryRequest const&) { return initial_pq_response; }); + .WillOnce([&](grpc::ClientContext&, Options const&, + PrepareQueryRequest const&, + bigtable_internal::OperationContext&) { + return initial_pq_response; + }); EXPECT_CALL(*mock, AsyncPrepareQuery) .WillOnce([&](CompletionQueue const&, auto, auto, - v2::PrepareQueryRequest const&) { + v2::PrepareQueryRequest const&, + bigtable_internal::OperationContext&) { return make_ready_future(make_status_or(refresh_pq_response)); }); EXPECT_CALL(*mock, ExecuteQuery) - .WillOnce([&](auto, auto const&, ExecuteQueryRequest const& request) { + .WillOnce([&](auto, auto const&, ExecuteQueryRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(request.prepared_query(), "test-pq-id-initial"); auto error_stream = std::make_unique(); EXPECT_CALL(*error_stream, Read).WillOnce(Return(QueryPlanError())); return error_stream; }) - .WillOnce([&](auto, auto const&, ExecuteQueryRequest const& request) { + .WillOnce([&](auto, auto const&, ExecuteQueryRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(request.app_profile_id(), kAppProfile); EXPECT_EQ(request.instance_name(), "projects/test-project/instances/test-instance"); @@ -3567,15 +3655,17 @@ TEST_F(DataConnectionTest, std::vector> metrics; metrics.push_back(std::make_shared( std::move(mock_metric_prepared_query_initial))); - return std::make_shared( - ResourceLabels{}, DataLabels{}, std::move(metrics), clock); + return std::make_shared(TableResourceLabels{}, + TableDataLabels{}, + std::move(metrics), clock); }) .WillOnce([&](auto const&, auto const&) { std::vector> metrics; metrics.push_back(std::make_shared( std::move(mock_metric_prepared_query_refresh))); - return std::make_shared( - ResourceLabels{}, DataLabels{}, std::move(metrics), clock); + return std::make_shared(TableResourceLabels{}, + TableDataLabels{}, + std::move(metrics), clock); }); auto mock_metric_execute_query = std::make_unique(); @@ -3589,8 +3679,8 @@ TEST_F(DataConnectionTest, std::vector> metrics; metrics.push_back( std::make_shared(std::move(mock_metric_execute_query))); - return std::make_shared(ResourceLabels{}, DataLabels{}, - std::move(metrics), clock); + return std::make_shared( + TableResourceLabels{}, TableDataLabels{}, std::move(metrics), clock); }); #else auto factory = std::make_unique(); @@ -3644,22 +3734,26 @@ TEST_F(DataConnectionTest, auto mock = std::make_shared(); EXPECT_CALL(*mock, AsyncPrepareQuery) - .WillOnce( - [&](CompletionQueue const&, auto, auto, PrepareQueryRequest const&) { - return make_ready_future(make_status_or(initial_pq_response)); - }) - .WillOnce( - [&](CompletionQueue const&, auto, auto, PrepareQueryRequest const&) { - return make_ready_future(make_status_or(refresh_pq_response)); - }); + .WillOnce([&](CompletionQueue const&, auto, auto, + PrepareQueryRequest const&, + bigtable_internal::OperationContext&) { + return make_ready_future(make_status_or(initial_pq_response)); + }) + .WillOnce([&](CompletionQueue const&, auto, auto, + PrepareQueryRequest const&, + bigtable_internal::OperationContext&) { + return make_ready_future(make_status_or(refresh_pq_response)); + }); EXPECT_CALL(*mock, ExecuteQuery) - .WillOnce([&](auto, auto const&, ExecuteQueryRequest const& request) { + .WillOnce([&](auto, auto const&, ExecuteQueryRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(request.prepared_query(), "test-pq-id-initial"); auto error_stream = std::make_unique(); EXPECT_CALL(*error_stream, Read).WillOnce(Return(QueryPlanError())); return error_stream; }) - .WillOnce([&](auto, auto const&, ExecuteQueryRequest const& request) { + .WillOnce([&](auto, auto const&, ExecuteQueryRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(request.app_profile_id(), kAppProfile); EXPECT_EQ(request.instance_name(), "projects/test-project/instances/test-instance"); @@ -3751,7 +3845,8 @@ TEST_F(DataConnectionTest, ExecuteQueryFailureWithSchemaChange) { }; EXPECT_CALL(*mock, ExecuteQuery) .WillOnce([&](auto, auto const&, - google::bigtable::v2::ExecuteQueryRequest const& request) { + google::bigtable::v2::ExecuteQueryRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_EQ(request.app_profile_id(), kAppProfile); EXPECT_EQ(request.instance_name(), "projects/test-project/instances/test-instance"); @@ -3815,8 +3910,10 @@ TEST_F(DataConnectionTest, PrepareQueryFailsOnInvalidType) { std::chrono::system_clock::now() + std::chrono::seconds(3600)); EXPECT_CALL(*mock, PrepareQuery) - .WillOnce([&](grpc::ClientContext&, Options const&, - v2::PrepareQueryRequest const&) { return pq_response; }); + .WillOnce( + [&](grpc::ClientContext&, Options const&, + v2::PrepareQueryRequest const&, + bigtable_internal::OperationContext&) { return pq_response; }); auto conn = TestConnection(std::move(mock)); internal::OptionsSpan span(CallOptions()); @@ -3860,7 +3957,8 @@ TEST_F(DataConnectionTest, AsyncPrepareQueryFailsOnInvalidType) { EXPECT_CALL(*mock, AsyncPrepareQuery) .WillOnce([&](CompletionQueue const&, auto, auto, - v2::PrepareQueryRequest const&) { + v2::PrepareQueryRequest const&, + bigtable_internal::OperationContext&) { return make_ready_future(make_status_or(pq_response)); }); diff --git a/google/cloud/bigtable/internal/default_row_reader.cc b/google/cloud/bigtable/internal/default_row_reader.cc index 7cb45dad72e67..9fe42d02dfbc0 100644 --- a/google/cloud/bigtable/internal/default_row_reader.cc +++ b/google/cloud/bigtable/internal/default_row_reader.cc @@ -68,7 +68,8 @@ void DefaultRowReader::MakeRequest() { internal::ConfigureContext(*client_context_, options); operation_context_->PreCall(*client_context_); called_post_call_ = false; - stream_ = stub_->ReadRows(client_context_, options, request); + stream_ = + stub_->ReadRows(client_context_, options, request, *operation_context_); stream_is_open_ = true; parser_ = bigtable::internal::ReadRowsParserFactory().Create(reverse_); diff --git a/google/cloud/bigtable/internal/default_row_reader_test.cc b/google/cloud/bigtable/internal/default_row_reader_test.cc index dbe83f709093a..e7b9754b4ccaa 100644 --- a/google/cloud/bigtable/internal/default_row_reader_test.cc +++ b/google/cloud/bigtable/internal/default_row_reader_test.cc @@ -13,6 +13,7 @@ // limitations under the License. #include "google/cloud/bigtable/internal/default_row_reader.h" +#include "google/cloud/bigtable/internal/operation_context.h" #include "google/cloud/bigtable/row_reader.h" #include "google/cloud/bigtable/testing/mock_bigtable_stub.h" #include "google/cloud/bigtable/testing/mock_policies.h" @@ -126,7 +127,8 @@ class MockMetric : public Metric { ElementDeliveryParams const&), (override)); MOCK_METHOD(std::unique_ptr, clone, - (ResourceLabels resource_labels, DataLabels data_labels), + (TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels), (const, override)); }; @@ -135,7 +137,8 @@ class CloningMetric : public Metric { public: explicit CloningMetric(std::unique_ptr metric) : metric_(std::move(metric)) {} - std::unique_ptr clone(ResourceLabels, DataLabels) const override { + std::unique_ptr clone(TableResourceLabels const&, + TableDataLabels const&) const override { return std::move(metric_); } @@ -193,7 +196,8 @@ TEST_F(DefaultRowReaderTest, EmptyReaderHasNoRows) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); EXPECT_CALL(*stream, Read) @@ -237,7 +241,8 @@ TEST_F(DefaultRowReaderTest, ReadOneRow) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); EXPECT_CALL(*stream, Read) @@ -284,7 +289,8 @@ TEST_F(DefaultRowReaderTest, StreamIsDrained) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); ::testing::InSequence s; @@ -348,7 +354,8 @@ TEST_F(DefaultRowReaderTest, RetryThenSuccess) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); EXPECT_CALL(*stream, Read) @@ -356,7 +363,8 @@ TEST_F(DefaultRowReaderTest, RetryThenSuccess) { return stream; }) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); EXPECT_CALL(*stream, Read) @@ -403,7 +411,8 @@ TEST_F(DefaultRowReaderTest, NoRetryOnPermanentError) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); EXPECT_CALL(*stream, Read) @@ -448,7 +457,8 @@ TEST_F(DefaultRowReaderTest, RetryPolicyExhausted) { EXPECT_CALL(*mock, ReadRows) .Times(kNumRetries + 1) .WillRepeatedly([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); EXPECT_CALL(*stream, Read) @@ -501,7 +511,8 @@ TEST_F(DefaultRowReaderTest, RetrySkipsAlreadyReadRows) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); // We should have two rows in the initial request: "r1" and "r2". EXPECT_THAT(request.rows().row_keys(), ElementsAre("r1", "r2")); @@ -515,7 +526,8 @@ TEST_F(DefaultRowReaderTest, RetrySkipsAlreadyReadRows) { return stream; }) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); // We have read "r1". The new request should only contain: "r2". EXPECT_THAT(request.rows().row_keys(), ElementsAre("r2")); @@ -559,7 +571,8 @@ TEST_F(DefaultRowReaderTest, RetrySkipsAlreadyScannedRows) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); // We start our call with 3 rows in the set: "r1", "r2", "r3". EXPECT_THAT(request.rows().row_keys(), ElementsAre("r1", "r2", "r3")); @@ -580,7 +593,8 @@ TEST_F(DefaultRowReaderTest, RetrySkipsAlreadyScannedRows) { return stream; }) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); // We retry the remaining rows. We have "r1" returned, but the service // has also told us that "r2" was scanned. This means there is only one @@ -626,7 +640,8 @@ TEST_F(DefaultRowReaderTest, FailedParseIsRetried) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); EXPECT_CALL(*stream, Read) @@ -637,7 +652,8 @@ TEST_F(DefaultRowReaderTest, FailedParseIsRetried) { return stream; }) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); EXPECT_CALL(*stream, Read) @@ -689,7 +705,8 @@ TEST_F(DefaultRowReaderTest, FailedParseSkipsAlreadyReadRows) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); // We should have two rows in the initial request: "r1" and "r2". EXPECT_THAT(request.rows().row_keys(), ElementsAre("r1", "r2")); @@ -706,7 +723,8 @@ TEST_F(DefaultRowReaderTest, FailedParseSkipsAlreadyReadRows) { return stream; }) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); // We have read "r1". The new request should only contain: "r2". EXPECT_THAT(request.rows().row_keys(), ElementsAre("r2")); @@ -755,7 +773,8 @@ TEST_F(DefaultRowReaderTest, FailedParseSkipsAlreadyScannedRows) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); // We start our call with 3 rows in the set: "r1", "r2", "r3". EXPECT_THAT(request.rows().row_keys(), ElementsAre("r1", "r2", "r3")); @@ -779,7 +798,8 @@ TEST_F(DefaultRowReaderTest, FailedParseSkipsAlreadyScannedRows) { return stream; }) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); // We retry the remaining rows. We have "r1" returned, but the service // has also told us that "r2" was scanned. This means there is only one @@ -830,7 +850,8 @@ TEST_F(DefaultRowReaderTest, FailedParseWithPermanentError) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); ::testing::InSequence s; @@ -880,7 +901,8 @@ TEST_F(DefaultRowReaderTest, NoRetryOnEmptyRowSet) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); EXPECT_CALL(*stream, Read) @@ -928,7 +950,8 @@ TEST_F(DefaultRowReaderTest, RowLimitIsSent) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); EXPECT_THAT(request, RequestWithRowsLimit(42)); auto stream = std::make_unique(); @@ -969,7 +992,8 @@ TEST_F(DefaultRowReaderTest, RowLimitIsDecreasedOnRetry) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); EXPECT_THAT(request, RequestWithRowsLimit(42)); auto stream = std::make_unique(); @@ -982,7 +1006,8 @@ TEST_F(DefaultRowReaderTest, RowLimitIsDecreasedOnRetry) { return stream; }) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); EXPECT_THAT(request, RequestWithRowsLimit(41)); auto stream = std::make_unique(); @@ -1023,7 +1048,8 @@ TEST_F(DefaultRowReaderTest, NoRetryIfRowLimitReached) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); EXPECT_THAT(request, RequestWithRowsLimit(1)); auto stream = std::make_unique(); @@ -1071,7 +1097,8 @@ TEST_F(DefaultRowReaderTest, CancelDrainsStream) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); ::testing::InSequence s; @@ -1220,7 +1247,8 @@ TEST_F(DefaultRowReaderTest, RetryUsesNewContext) { EXPECT_CALL(*mock, ReadRows) .Times(kNumRetries + 1) .WillRepeatedly([](auto context, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { // This is a hack. A new request will have the default compression // algorithm (GRPC_COMPRESS_NONE). We then change the value in this // call. If the context is reused, it will no longer have the default @@ -1271,7 +1299,8 @@ TEST_F(DefaultRowReaderTest, ReverseScanSuccess) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_TRUE(request.reversed()); auto stream = std::make_unique(); ::testing::InSequence s; @@ -1329,7 +1358,8 @@ TEST_F(DefaultRowReaderTest, ReverseScanFailsOnIncreasingRowKeyOrder) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_TRUE(request.reversed()); auto stream = std::make_unique(); ::testing::InSequence s; @@ -1387,7 +1417,8 @@ TEST_F(DefaultRowReaderTest, ReverseScanResumption) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_TRUE(request.reversed()); // We start our call with 3 rows in the set: "r1", "r2", "r3". EXPECT_THAT(request.rows().row_keys(), ElementsAre("r1", "r2", "r3")); @@ -1408,7 +1439,8 @@ TEST_F(DefaultRowReaderTest, ReverseScanResumption) { return stream; }) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_TRUE(request.reversed()); // We retry the remaining rows. We have "r3" returned, but the service // has also told us that "r2" was scanned. This means there is only one @@ -1434,7 +1466,8 @@ TEST_F(DefaultRowReaderTest, BigtableCookies) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([this](auto context, auto const&, - google::bigtable::v2::ReadRowsRequest const&) { + google::bigtable::v2::ReadRowsRequest const&, + bigtable_internal::OperationContext&) { // Return a bigtable cookie in the first request. metadata_fixture_.SetServerMetadata( *context, {{}, {{"x-goog-cbt-cookie-routing", "routing"}}}); @@ -1444,7 +1477,8 @@ TEST_F(DefaultRowReaderTest, BigtableCookies) { return stream; }) .WillOnce([this](auto context, auto const&, - google::bigtable::v2::ReadRowsRequest const&) { + google::bigtable::v2::ReadRowsRequest const&, + bigtable_internal::OperationContext&) { // Verify that the next request includes the bigtable cookie from above. auto headers = metadata_fixture_.GetMetadata(*context); EXPECT_THAT(headers, @@ -1490,7 +1524,8 @@ TEST_F(DefaultRowReaderTest, RetryInfoHeeded) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([delay](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const&) { + google::bigtable::v2::ReadRowsRequest const&, + bigtable_internal::OperationContext&) { auto stream = std::make_unique(); EXPECT_CALL(*stream, Read) .WillOnce([delay](google::bigtable::v2::ReadRowsResponse*) { @@ -1501,7 +1536,8 @@ TEST_F(DefaultRowReaderTest, RetryInfoHeeded) { return stream; }) .WillOnce([](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const& request) { + google::bigtable::v2::ReadRowsRequest const& request, + bigtable_internal::OperationContext&) { EXPECT_THAT(request, HasCorrectResourceNames()); auto stream = std::make_unique(); EXPECT_CALL(*stream, Read) @@ -1552,7 +1588,8 @@ TEST_F(DefaultRowReaderTest, RetryInfoIgnored) { auto mock = std::make_shared(); EXPECT_CALL(*mock, ReadRows) .WillOnce([delay](auto, auto const&, - google::bigtable::v2::ReadRowsRequest const&) { + google::bigtable::v2::ReadRowsRequest const&, + bigtable_internal::OperationContext&) { auto stream = std::make_unique(); EXPECT_CALL(*stream, Read) .WillOnce([delay](google::bigtable::v2::ReadRowsResponse*) { diff --git a/google/cloud/bigtable/internal/dynamic_channel_pool.h b/google/cloud/bigtable/internal/dynamic_channel_pool.h index 7bec7a78bf8a7..9b3994b19ff69 100644 --- a/google/cloud/bigtable/internal/dynamic_channel_pool.h +++ b/google/cloud/bigtable/internal/dynamic_channel_pool.h @@ -18,6 +18,8 @@ #include "google/cloud/bigtable/instance_resource.h" #include "google/cloud/bigtable/internal/channel_usage.h" #include "google/cloud/bigtable/internal/connection_refresh_state.h" +#include "google/cloud/bigtable/internal/defaults.h" +#include "google/cloud/bigtable/internal/metrics.h" #include "google/cloud/bigtable/internal/stub_manager.h" #include "google/cloud/bigtable/options.h" #include "google/cloud/completion_queue.h" @@ -42,6 +44,12 @@ GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN // remove ChannelUsage objects per the configuration present in the // DynamicChannelPoolSizingPolicyOption. // +template +struct SelectedChannel { + std::shared_ptr> channel; + int outstanding_rpcs; +}; + template class DynamicChannelPool : public std::enable_shared_from_this> { @@ -97,6 +105,8 @@ class DynamicChannelPool return sizing_policy_; } + TransportType transport_type() const { return transport_type_; } + // Calls CheckPoolChannelHealth before picking a channel. // // Pick two random channels from channels_ and return the channel with the @@ -111,7 +121,7 @@ class DynamicChannelPool // // If there are no healthy channels in channels_, create a new channel and // use that one. Also call ScheduleAddChannels to replenish channels_. - std::shared_ptr> GetChannelRandomTwoLeastUsed() { + SelectedChannel GetChannelRandomTwoLeastUsed() { std::scoped_lock lk(mu_); CheckPoolChannelHealth(lk); @@ -136,12 +146,13 @@ class DynamicChannelPool // This is the most common case so we try it first. if (d.channel_1_rpcs.ok() && d.channel_2_rpcs.ok()) { - return *d.channel_1_rpcs < *d.channel_2_rpcs ? *d.channel_1_iter - : *d.channel_2_iter; + return *d.channel_1_rpcs < *d.channel_2_rpcs + ? SelectedChannel{*d.channel_1_iter, *d.channel_1_rpcs} + : SelectedChannel{*d.channel_2_iter, *d.channel_2_rpcs}; } if (d.iterators.size() == 1 && d.channel_1_rpcs.ok()) { // Pool contains exactly 1 good channel. - return *d.channel_1_iter; + return SelectedChannel{*d.channel_1_iter, *d.channel_1_rpcs}; } if (d.iterators.empty()) { // Pool is empty, create a channel immediately and return it. While the @@ -150,7 +161,8 @@ class DynamicChannelPool channels_.push_back(stub_factory_fn_(next_channel_id_++, instance_name_, StubManager::Priming::kNoPriming) .value()); - return channels_.front(); + auto sor = channels_.front()->instant_outstanding_rpcs(); + return SelectedChannel{channels_.front(), sor ? *sor : 0}; } return HandleBadChannels(lk, d); } @@ -204,8 +216,8 @@ class DynamicChannelPool // We have one or more bad channels. Spending time finding a good channel // will be cheaper than trying to use a bad channel in the long run. - std::shared_ptr> HandleBadChannels( - std::scoped_lock const& lk, ChannelSelectionData& d) { + SelectedChannel HandleBadChannels(std::scoped_lock const& lk, + ChannelSelectionData& d) { std::vector bad_channel_iters; if (d.shuffle_iter != d.iterators.end()) ++d.shuffle_iter; ChannelSelectionData::FindGoodChannel(d.iterators, d.channel_1_iter, @@ -216,14 +228,22 @@ class DynamicChannelPool bad_channel_iters); std::shared_ptr> channel; + int outstanding_rpcs = 0; if (d.channel_1_rpcs.ok() || d.channel_2_rpcs.ok()) { if (d.channel_1_rpcs.ok() && d.channel_2_rpcs.ok()) { - channel = *d.channel_1_rpcs < *d.channel_2_rpcs ? *d.channel_1_iter - : *d.channel_2_iter; + if (*d.channel_1_rpcs < *d.channel_2_rpcs) { + channel = *d.channel_1_iter; + outstanding_rpcs = *d.channel_1_rpcs; + } else { + channel = *d.channel_2_iter; + outstanding_rpcs = *d.channel_2_rpcs; + } } else if (d.channel_1_rpcs.ok()) { channel = *d.channel_1_iter; + outstanding_rpcs = *d.channel_1_rpcs; } else if (d.channel_2_rpcs.ok()) { channel = *d.channel_2_iter; + outstanding_rpcs = *d.channel_2_rpcs; } // Wait until we no longer need valid iterators to call EvictBadChannels. EvictBadChannels(lk, bad_channel_iters); @@ -240,9 +260,11 @@ class DynamicChannelPool .value()); std::swap(channels_.front(), channels_.back()); channel = channels_.front(); + auto sor = channel->instant_outstanding_rpcs(); + outstanding_rpcs = sor ? *sor : 0; } ScheduleRemoveChannels(lk); - return channel; + return SelectedChannel{std::move(channel), outstanding_rpcs}; } // Determines the number of channels to add and reserves the channel ids to @@ -433,6 +455,9 @@ class DynamicChannelPool future> pool_size_decrease_cooldown_timer_; std::uint32_t next_channel_id_; + TransportType const transport_type_ = bigtable::internal::IsDirectPath() + ? TransportType::kDirectPath + : TransportType::kCloudPath; }; GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END diff --git a/google/cloud/bigtable/internal/dynamic_channel_pool_test.cc b/google/cloud/bigtable/internal/dynamic_channel_pool_test.cc index c2f7d9a2675d9..e3f3be2b56bfe 100644 --- a/google/cloud/bigtable/internal/dynamic_channel_pool_test.cc +++ b/google/cloud/bigtable/internal/dynamic_channel_pool_test.cc @@ -13,6 +13,7 @@ // limitations under the License. #include "google/cloud/bigtable/internal/dynamic_channel_pool.h" +#include "google/cloud/bigtable/internal/operation_context.h" #include "google/cloud/bigtable/testing/mock_bigtable_stub.h" #include "google/cloud/internal/make_status.h" #include "google/cloud/testing_util/fake_completion_queue_impl.h" @@ -35,7 +36,7 @@ class DynamicChannelPoolTestWrapper { using ChannelSelectionData = DynamicChannelPool::ChannelSelectionData; - std::shared_ptr> HandleBadChannels( + SelectedChannel HandleBadChannels( std::scoped_lock const& lk, DynamicChannelPool::ChannelSelectionData& d) { return pool_->HandleBadChannels(lk, d); @@ -151,14 +152,15 @@ TEST_F(DynamicChannelPoolTest, SelectLeastUsedFromTwoChannels) { auto mock_stub_0 = std::make_shared(); EXPECT_CALL(*mock_stub_0, CheckAndMutateRow) .WillOnce([](grpc::ClientContext&, Options const&, - google::bigtable::v2::CheckAndMutateRowRequest const&) { + google::bigtable::v2::CheckAndMutateRowRequest const&, + OperationContext&) { google::bigtable::v2::CheckAndMutateRowResponse response; response.set_predicate_matched(true); return response; }); auto mock_stub_1 = std::make_shared(); EXPECT_CALL(*mock_stub_1, CheckAndMutateRow).Times(0); - int initial_rpc_count = 0; + int initial_rpc_count = 5; channels.push_back(std::make_shared>( std::move(mock_stub_0), initial_rpc_count++)); channels.push_back(std::make_shared>( @@ -179,10 +181,14 @@ TEST_F(DynamicChannelPoolTest, SelectLeastUsedFromTwoChannels) { auto pool = DynamicChannelPool::Create( instance_name, CompletionQueue(mock_cq_impl_), channels, refresh_state, stub_factory_fn, sizing_policy); - auto selected_stub = pool->GetChannelRandomTwoLeastUsed(); + auto selected = pool->GetChannelRandomTwoLeastUsed(); + EXPECT_THAT(selected.outstanding_rpcs, Eq(5)); grpc::ClientContext context; - auto response = - selected_stub->AcquireStub()->CheckAndMutateRow(context, {}, {}); + OperationContext default_oc; + auto response = selected.channel->AcquireStub()->CheckAndMutateRow( + context, {}, {}, default_oc); + EXPECT_THAT(selected.channel->instant_outstanding_rpcs(), + testing_util::IsOkAndHolds(Eq(6))); ASSERT_STATUS_OK(response); EXPECT_TRUE(response->predicate_matched()); @@ -215,7 +221,8 @@ TEST_F(DynamicChannelPoolTest, OneInitialChannel) { auto mock_stub_0 = std::make_shared(); EXPECT_CALL(*mock_stub_0, CheckAndMutateRow) .WillOnce([](grpc::ClientContext&, Options const&, - google::bigtable::v2::CheckAndMutateRowRequest const&) { + google::bigtable::v2::CheckAndMutateRowRequest const&, + OperationContext&) { google::bigtable::v2::CheckAndMutateRowResponse response; response.set_predicate_matched(true); return response; @@ -243,10 +250,12 @@ TEST_F(DynamicChannelPoolTest, OneInitialChannel) { stub_factory_fn.AsStdFunction(), sizing_policy); EXPECT_THAT(pool->size(), Eq(1)); - auto selected_stub = pool->GetChannelRandomTwoLeastUsed(); + auto selected = pool->GetChannelRandomTwoLeastUsed(); + EXPECT_THAT(selected.outstanding_rpcs, Eq(0)); grpc::ClientContext context; - auto response = - selected_stub->AcquireStub()->CheckAndMutateRow(context, {}, {}); + OperationContext default_oc; + auto response = selected.channel->AcquireStub()->CheckAndMutateRow( + context, {}, {}, default_oc); ASSERT_STATUS_OK(response); EXPECT_TRUE(response->predicate_matched()); } @@ -270,7 +279,8 @@ TEST_F(DynamicChannelPoolTest, EmptyInitialPool) { auto mock_stub = std::make_shared(); EXPECT_CALL(*mock_stub, CheckAndMutateRow) .WillOnce([](grpc::ClientContext&, Options const&, - google::bigtable::v2::CheckAndMutateRowRequest const&) { + google::bigtable::v2::CheckAndMutateRowRequest const&, + OperationContext&) { google::bigtable::v2::CheckAndMutateRowResponse response; response.set_predicate_matched(true); return response; @@ -304,10 +314,12 @@ TEST_F(DynamicChannelPoolTest, EmptyInitialPool) { EXPECT_THAT(*pool, ::testing::IsEmpty()); - auto selected_stub = pool->GetChannelRandomTwoLeastUsed(); + auto selected = pool->GetChannelRandomTwoLeastUsed(); + EXPECT_THAT(selected.outstanding_rpcs, Eq(0)); grpc::ClientContext context; - auto response = - selected_stub->AcquireStub()->CheckAndMutateRow(context, {}, {}); + OperationContext default_oc; + auto response = selected.channel->AcquireStub()->CheckAndMutateRow( + context, {}, {}, default_oc); ASSERT_STATUS_OK(response); EXPECT_TRUE(response->predicate_matched()); @@ -743,7 +755,8 @@ TEST_F(DynamicChannelPoolTest, HandleBadChannelsTwoChannelsOneBad) { auto mock_stub = std::make_shared(); EXPECT_CALL(*mock_stub, CheckAndMutateRow) .WillOnce([](grpc::ClientContext&, Options const&, - google::bigtable::v2::CheckAndMutateRowRequest const&) { + google::bigtable::v2::CheckAndMutateRowRequest const&, + OperationContext&) { google::bigtable::v2::CheckAndMutateRowResponse response; response.set_predicate_matched(true); return response; @@ -770,7 +783,7 @@ TEST_F(DynamicChannelPoolTest, HandleBadChannelsTwoChannelsOneBad) { DynamicChannelPoolTestWrapper wrapper(pool); auto draining_channels = wrapper.SetDrainingChannels({}); - std::shared_ptr> selected_stub; + SelectedChannel selected_stub; { auto lock = wrapper.CreateLock(); selected_stub = wrapper.HandleBadChannels(lock, data); @@ -778,8 +791,9 @@ TEST_F(DynamicChannelPoolTest, HandleBadChannelsTwoChannelsOneBad) { EXPECT_THAT(draining_channels, IsEmpty()); grpc::ClientContext context; - auto response = - selected_stub->AcquireStub()->CheckAndMutateRow(context, {}, {}); + OperationContext default_oc; + auto response = selected_stub.channel->AcquireStub()->CheckAndMutateRow( + context, {}, {}, default_oc); ASSERT_STATUS_OK(response); EXPECT_TRUE(response->predicate_matched()); EXPECT_THAT(pool->size(), Eq(1)); @@ -822,7 +836,8 @@ TEST_F(DynamicChannelPoolTest, HandleBadChannelsTwoChannelsOtherOneBad) { auto mock_stub = std::make_shared(); EXPECT_CALL(*mock_stub, CheckAndMutateRow) .WillOnce([](grpc::ClientContext&, Options const&, - google::bigtable::v2::CheckAndMutateRowRequest const&) { + google::bigtable::v2::CheckAndMutateRowRequest const&, + OperationContext&) { google::bigtable::v2::CheckAndMutateRowResponse response; response.set_predicate_matched(true); return response; @@ -851,7 +866,7 @@ TEST_F(DynamicChannelPoolTest, HandleBadChannelsTwoChannelsOtherOneBad) { DynamicChannelPoolTestWrapper wrapper(pool); auto draining_channels = wrapper.SetDrainingChannels({}); - std::shared_ptr> selected_stub; + SelectedChannel selected_stub; { auto lock = wrapper.CreateLock(); selected_stub = wrapper.HandleBadChannels(lock, data); @@ -859,8 +874,9 @@ TEST_F(DynamicChannelPoolTest, HandleBadChannelsTwoChannelsOtherOneBad) { EXPECT_THAT(draining_channels, IsEmpty()); grpc::ClientContext context; - auto response = - selected_stub->AcquireStub()->CheckAndMutateRow(context, {}, {}); + OperationContext default_oc; + auto response = selected_stub.channel->AcquireStub()->CheckAndMutateRow( + context, {}, {}, default_oc); ASSERT_STATUS_OK(response); EXPECT_TRUE(response->predicate_matched()); EXPECT_THAT(pool->size(), Eq(1)); @@ -903,7 +919,8 @@ TEST_F(DynamicChannelPoolTest, HandleBadChannelsThreeChannelsOneBad) { auto mock_stub_0 = std::make_shared(); EXPECT_CALL(*mock_stub_0, CheckAndMutateRow) .WillOnce([](grpc::ClientContext&, Options const&, - google::bigtable::v2::CheckAndMutateRowRequest const&) { + google::bigtable::v2::CheckAndMutateRowRequest const&, + OperationContext&) { google::bigtable::v2::CheckAndMutateRowResponse response; response.set_predicate_matched(true); return response; @@ -938,7 +955,7 @@ TEST_F(DynamicChannelPoolTest, HandleBadChannelsThreeChannelsOneBad) { DynamicChannelPoolTestWrapper wrapper(pool); auto draining_channels = wrapper.SetDrainingChannels({}); - std::shared_ptr> selected_stub; + SelectedChannel selected_stub; { auto lock = wrapper.CreateLock(); selected_stub = wrapper.HandleBadChannels(lock, data); @@ -946,8 +963,9 @@ TEST_F(DynamicChannelPoolTest, HandleBadChannelsThreeChannelsOneBad) { EXPECT_THAT(draining_channels, IsEmpty()); grpc::ClientContext context; - auto response = - selected_stub->AcquireStub()->CheckAndMutateRow(context, {}, {}); + OperationContext default_oc; + auto response = selected_stub.channel->AcquireStub()->CheckAndMutateRow( + context, {}, {}, default_oc); ASSERT_STATUS_OK(response); EXPECT_TRUE(response->predicate_matched()); EXPECT_THAT(pool->size(), Eq(2)); @@ -990,7 +1008,8 @@ TEST_F(DynamicChannelPoolTest, HandleBadChannelsAllChannelsBad) { auto mock_stub = std::make_shared(); EXPECT_CALL(*mock_stub, CheckAndMutateRow) .WillOnce([](grpc::ClientContext&, Options const&, - google::bigtable::v2::CheckAndMutateRowRequest const&) { + google::bigtable::v2::CheckAndMutateRowRequest const&, + OperationContext&) { google::bigtable::v2::CheckAndMutateRowResponse response; response.set_predicate_matched(true); return response; @@ -1035,7 +1054,7 @@ TEST_F(DynamicChannelPoolTest, HandleBadChannelsAllChannelsBad) { DynamicChannelPoolTestWrapper wrapper(pool); auto draining_channels = wrapper.SetDrainingChannels({}); - std::shared_ptr> selected_stub; + SelectedChannel selected_stub; { auto lock = wrapper.CreateLock(); selected_stub = wrapper.HandleBadChannels(lock, data); @@ -1043,8 +1062,9 @@ TEST_F(DynamicChannelPoolTest, HandleBadChannelsAllChannelsBad) { EXPECT_THAT(draining_channels, IsEmpty()); grpc::ClientContext context; - auto response = - selected_stub->AcquireStub()->CheckAndMutateRow(context, {}, {}); + OperationContext default_oc; + auto response = selected_stub.channel->AcquireStub()->CheckAndMutateRow( + context, {}, {}, default_oc); ASSERT_STATUS_OK(response); EXPECT_TRUE(response->predicate_matched()); EXPECT_THAT(pool->size(), Eq(1)); diff --git a/google/cloud/bigtable/internal/grpc_metrics_exporter.cc b/google/cloud/bigtable/internal/grpc_metrics_exporter.cc index 6e86b93ad3dc6..7c8e3fe023304 100644 --- a/google/cloud/bigtable/internal/grpc_metrics_exporter.cc +++ b/google/cloud/bigtable/internal/grpc_metrics_exporter.cc @@ -197,23 +197,11 @@ std::shared_ptr MakeGrpcMeterProvider( std::move(provider)); } -MonitoredResourceResult MakeMonitoredResource( - opentelemetry::sdk::metrics::PointDataAttributes const& pda, - opentelemetry::sdk::resource::Resource const& detected_resource, - Options const& options, std::string const& client_uid) { +ClientResourceLabels MakeClientResourceLabels( + std::string project_id, std::string instance, std::string app_profile, + Options const& options, std::string const& client_uid, + opentelemetry::sdk::resource::Resource const& detected_resource) { namespace sc = ::opentelemetry::semconv; - google::api::MonitoredResource resource; - resource.set_type("bigtable.googleapis.com/Client"); - auto& labels = *resource.mutable_labels(); - auto const& attributes = pda.attributes.GetAttributes(); - auto get_attr = [&](std::string const& key) { - auto it = attributes.find(key); - if (it == attributes.end() || - !opentelemetry::nostd::holds_alternative(it->second)) { - return std::string{}; - } - return opentelemetry::nostd::get(it->second); - }; auto const& detected_attributes = detected_resource.GetAttributes(); auto by_name = [&](std::string const& name, std::string default_value = {}) { auto const l = detected_attributes.find(name); @@ -224,10 +212,6 @@ MonitoredResourceResult MakeMonitoredResource( return opentelemetry::nostd::get(l->second); }; - auto project_id = get_attr("project_id"); - // Fall back to resolving the project ID from `InstanceChannelAffinityOption` - // if the static `ProjectIdOption` is not set on the client (which is common - // in client configurations with instance routing / channel affinity). if (project_id.empty() && options.has()) { auto const& instances = @@ -240,29 +224,65 @@ MonitoredResourceResult MakeMonitoredResource( project_id = by_name(sc::cloud::kCloudAccountId); } - labels["project_id"] = project_id; - labels["instance"] = get_attr("instance"); - labels["app_profile"] = options.get(); - labels["client_name"] = "cpp.Bigtable/" + bigtable::version_string(); - labels["uuid"] = client_uid; - auto client_project = by_name(sc::cloud::kCloudAccountId); if (client_project.empty()) { client_project = project_id; } - if (!client_project.empty()) { - labels["client_project"] = client_project; + + ClientResourceLabels labels; + labels.project_id = std::move(project_id); + labels.instance = std::move(instance); + labels.app_profile = std::move(app_profile); + labels.client_name = "cpp.Bigtable/" + bigtable::version_string(); + labels.client_uid = client_uid; + labels.client_project = std::move(client_project); + labels.location = by_name(sc::cloud::kCloudAvailabilityZone, + by_name(sc::cloud::kCloudRegion, "global")); + labels.cloud_platform = by_name(sc::cloud::kCloudPlatform, "unknown"); + labels.host_id = by_name("faas.id", by_name(sc::host::kHostId, "unknown")); + labels.hostname = by_name(sc::host::kHostName); + return labels; +} + +MonitoredResourceResult MakeMonitoredResource( + opentelemetry::sdk::metrics::PointDataAttributes const& pda, + opentelemetry::sdk::resource::Resource const& detected_resource, + Options const& options, std::string const& client_uid) { + google::api::MonitoredResource resource; + resource.set_type("bigtable.googleapis.com/Client"); + auto const& attributes = pda.attributes.GetAttributes(); + auto get_attr = [&](std::string const& key) { + auto it = attributes.find(key); + if (it == attributes.end() || + !opentelemetry::nostd::holds_alternative(it->second)) { + return std::string{}; + } + return opentelemetry::nostd::get(it->second); + }; + auto project_id = get_attr("project_id"); + auto instance = get_attr("instance"); + auto app_profile = options.get(); + auto labels = MakeClientResourceLabels( + std::move(project_id), std::move(instance), std::move(app_profile), + options, client_uid, detected_resource); + + auto& m = *resource.mutable_labels(); + m["project_id"] = labels.project_id; + m["instance"] = labels.instance; + m["app_profile"] = labels.app_profile; + m["client_name"] = labels.client_name; + m["uuid"] = labels.client_uid; + if (!labels.client_project.empty()) { + m["client_project"] = labels.client_project; } - labels["location"] = by_name(sc::cloud::kCloudAvailabilityZone, - by_name(sc::cloud::kCloudRegion, "global")); - labels["cloud_platform"] = by_name(sc::cloud::kCloudPlatform, "unknown"); - labels["host_id"] = by_name("faas.id", by_name(sc::host::kHostId, "unknown")); - auto hostname = by_name(sc::host::kHostName); - if (!hostname.empty()) { - labels["hostname"] = hostname; + m["location"] = labels.location; + m["cloud_platform"] = labels.cloud_platform; + m["host_id"] = labels.host_id; + if (!labels.hostname.empty()) { + m["hostname"] = labels.hostname; } - return MonitoredResourceResult{std::move(project_id), std::move(resource)}; + return MonitoredResourceResult{labels.project_id, std::move(resource)}; } void EnableGrpcMetrics( diff --git a/google/cloud/bigtable/internal/grpc_metrics_exporter.h b/google/cloud/bigtable/internal/grpc_metrics_exporter.h index 8ab268b7c2fb1..14662cdde3c01 100644 --- a/google/cloud/bigtable/internal/grpc_metrics_exporter.h +++ b/google/cloud/bigtable/internal/grpc_metrics_exporter.h @@ -23,6 +23,7 @@ #include #ifdef GOOGLE_CLOUD_CPP_BIGTABLE_WITH_GRPC_OTEL_METRICS +#include "google/cloud/bigtable/internal/metrics.h" #include "google/cloud/monitoring/v3/metric_connection.h" #include "google/api/monitored_resource.pb.h" #include @@ -75,6 +76,11 @@ struct MonitoredResourceResult { google::api::MonitoredResource resource; }; +ClientResourceLabels MakeClientResourceLabels( + std::string project_id, std::string instance, std::string app_profile, + Options const& options, std::string const& client_uid, + opentelemetry::sdk::resource::Resource const& detected_resource); + MonitoredResourceResult MakeMonitoredResource( opentelemetry::sdk::metrics::PointDataAttributes const& pda, opentelemetry::sdk::resource::Resource const& detected_resource, diff --git a/google/cloud/bigtable/internal/grpc_metrics_exporter_test.cc b/google/cloud/bigtable/internal/grpc_metrics_exporter_test.cc index a9a9ab2c933cd..26c70fd311028 100644 --- a/google/cloud/bigtable/internal/grpc_metrics_exporter_test.cc +++ b/google/cloud/bigtable/internal/grpc_metrics_exporter_test.cc @@ -213,6 +213,28 @@ TEST(GrpcMetricsExporterTest, MakeMonitoredResource) { EXPECT_THAT(labels.at("host_id"), Eq("unknown")); } +TEST(GrpcMetricsExporterTest, MakeClientResourceLabels) { + Options options; + options.set("test-app-profile"); + std::string const client_uid = "test-client-uid"; + + auto labels = MakeClientResourceLabels( + "test-project", "test-instance", "test-app-profile", options, client_uid, + opentelemetry::sdk::resource::Resource::Create({})); + + EXPECT_THAT(labels.project_id, Eq("test-project")); + EXPECT_THAT(labels.instance, Eq("test-instance")); + EXPECT_THAT(labels.app_profile, Eq("test-app-profile")); + EXPECT_THAT(labels.client_name, + Eq("cpp.Bigtable/" + bigtable::version_string())); + EXPECT_THAT(labels.client_uid, Eq("test-client-uid")); + EXPECT_THAT(labels.client_project, Eq("test-project")); + EXPECT_THAT(labels.location, Eq("global")); + EXPECT_THAT(labels.cloud_platform, Eq("unknown")); + EXPECT_THAT(labels.host_id, Eq("unknown")); + EXPECT_THAT(labels.hostname, IsEmpty()); +} + TEST(GrpcMetricsExporterTest, MakeMonitoredResourceMissingAttributes) { Options options; std::string const client_uid = "test-client-uid"; diff --git a/google/cloud/bigtable/internal/metrics.cc b/google/cloud/bigtable/internal/metrics.cc index aa6193ed60b33..73f802d9de64c 100644 --- a/google/cloud/bigtable/internal/metrics.cc +++ b/google/cloud/bigtable/internal/metrics.cc @@ -35,10 +35,10 @@ auto constexpr kMeterInstrumentationScopeVersion = "v1"; } // namespace // TODO(#15329): Refactor how we're handling different data labels for -// the various RPCs. Adding a function to each metric type to add its DataLabels -// to the map should be more performant than performing a set_difference every -// time. -LabelMap IntoLabelMap(ResourceLabels const& r, DataLabels const& d, +// the various RPCs. Adding a function to each metric type to add its +// TableDataLabels to the map should be more performant than performing a +// set_difference every time. +LabelMap IntoLabelMap(TableResourceLabels const& r, TableDataLabels const& d, std::set const& filtered_data_labels, std::optional const& peer_info_labels) { LabelMap labels = { @@ -89,6 +89,76 @@ LabelMap IntoLabelMap(ResourceLabels const& r, DataLabels const& d, return labels; } +std::string ToString(ChannelPoolLbPolicy policy) { + switch (policy) { + case ChannelPoolLbPolicy::kRoundRobin: + return "ROUND_ROBIN"; + case ChannelPoolLbPolicy::kRandomTwoLeastUsed: + return "RANDOM_TWO_LEAST_USED"; + } + return "ROUND_ROBIN"; +} + +std::string ToString(TransportType type) { + switch (type) { + case TransportType::kCloudPath: + return "CloudPath"; + case TransportType::kDirectPath: + return "DirectPath"; + } + return "CloudPath"; +} + +std::string ToString(RpcType streaming) { + switch (streaming) { + case RpcType::kUnary: + return "false"; + case RpcType::kStreaming: + return "true"; + } + return "false"; +} + +LabelMap IntoLabelMap(ClientResourceLabels const& r, + ClientOutstandingRpcLabels const& d, + std::set const& filtered_data_labels) { + LabelMap labels = { + {"project_id", r.project_id}, {"instance", r.instance}, + {"app_profile", r.app_profile}, {"client_name", r.client_name}, + {"client_uid", r.client_uid}, {"client_project", r.client_project}, + {"location", r.location}, {"cloud_platform", r.cloud_platform}, + {"host_id", r.host_id}, {"hostname", r.hostname}}; + + std::map data = {{ + {"transport_type", ToString(d.transport_type)}, + {"channel_pool_lb_policy", ToString(d.channel_pool_lb_policy)}, + {"streaming", ToString(d.streaming)}, + }}; + + if (filtered_data_labels.empty()) { + labels.insert(data.begin(), data.end()); + return labels; + } + + struct Compare { + bool operator()(std::pair const& a, + std::string const& b) { + return a.first < b; + } + + bool operator()(std::string const& a, + std::pair const& b) { + return a < b.first; + } + }; + + std::set_difference(data.begin(), data.end(), filtered_data_labels.begin(), + filtered_data_labels.end(), + std::inserter(labels, labels.begin()), Compare()); + + return labels; +} + bool HasServerTiming(grpc::ClientContext const& client_context) { auto const& initial_metadata = client_context.GetServerInitialMetadata(); auto it = initial_metadata.find("server-timing"); @@ -219,11 +289,12 @@ void OperationLatency::OnDone(opentelemetry::context::Context const& context, context); } -std::unique_ptr OperationLatency::clone(ResourceLabels resource_labels, - DataLabels data_labels) const { +std::unique_ptr OperationLatency::clone( + TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels) const { auto m = std::make_unique(*this); - m->resource_labels_ = std::move(resource_labels); - m->data_labels_ = std::move(data_labels); + m->resource_labels_ = resource_labels; + m->data_labels_ = data_labels; return m; } @@ -256,11 +327,12 @@ void AttemptLatency::PostCall(opentelemetry::context::Context const& context, attempt_latencies_->Record(attempt_elapsed.count(), std::move(m), context); } -std::unique_ptr AttemptLatency::clone(ResourceLabels resource_labels, - DataLabels data_labels) const { +std::unique_ptr AttemptLatency::clone( + TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels) const { auto m = std::make_unique(*this); - m->resource_labels_ = std::move(resource_labels); - m->data_labels_ = std::move(data_labels); + m->resource_labels_ = resource_labels; + m->data_labels_ = data_labels; return m; } @@ -306,11 +378,12 @@ void AttemptLatency2::PostCall(opentelemetry::context::Context const& context, attempt_latencies2_->Record(attempt_elapsed.count(), std::move(m), context); } -std::unique_ptr AttemptLatency2::clone(ResourceLabels resource_labels, - DataLabels data_labels) const { +std::unique_ptr AttemptLatency2::clone( + TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels) const { auto m = std::make_unique(*this); - m->resource_labels_ = std::move(resource_labels); - m->data_labels_ = std::move(data_labels); + m->resource_labels_ = resource_labels; + m->data_labels_ = data_labels; return m; } @@ -350,11 +423,12 @@ void RetryCount::OnDone(opentelemetry::context::Context const& context, context); } -std::unique_ptr RetryCount::clone(ResourceLabels resource_labels, - DataLabels data_labels) const { +std::unique_ptr RetryCount::clone( + TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels) const { auto m = std::make_unique(*this); - m->resource_labels_ = std::move(resource_labels); - m->data_labels_ = std::move(data_labels); + m->resource_labels_ = resource_labels; + m->data_labels_ = data_labels; return m; } @@ -405,10 +479,11 @@ void FirstResponseLatency::OnDone( } std::unique_ptr FirstResponseLatency::clone( - ResourceLabels resource_labels, DataLabels data_labels) const { + TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels) const { auto m = std::make_unique(*this); - m->resource_labels_ = std::move(resource_labels); - m->data_labels_ = std::move(data_labels); + m->resource_labels_ = resource_labels; + m->data_labels_ = data_labels; return m; } @@ -437,11 +512,12 @@ void ServerLatency::PostCall(opentelemetry::context::Context const& context, } } -std::unique_ptr ServerLatency::clone(ResourceLabels resource_labels, - DataLabels data_labels) const { +std::unique_ptr ServerLatency::clone( + TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels) const { auto m = std::make_unique(*this); - m->resource_labels_ = std::move(resource_labels); - m->data_labels_ = std::move(data_labels); + m->resource_labels_ = resource_labels; + m->data_labels_ = data_labels; return m; } @@ -480,10 +556,11 @@ void ConnectivityErrorCount::OnDone( } std::unique_ptr ConnectivityErrorCount::clone( - ResourceLabels resource_labels, DataLabels data_labels) const { + TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels) const { auto m = std::make_unique(*this); - m->resource_labels_ = std::move(resource_labels); - m->data_labels_ = std::move(data_labels); + m->resource_labels_ = resource_labels; + m->data_labels_ = data_labels; return m; } @@ -530,10 +607,43 @@ void ApplicationBlockingLatency::OnDone( } std::unique_ptr ApplicationBlockingLatency::clone( - ResourceLabels resource_labels, DataLabels data_labels) const { + TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels) const { auto m = std::make_unique(*this); - m->resource_labels_ = std::move(resource_labels); - m->data_labels_ = std::move(data_labels); + m->resource_labels_ = resource_labels; + m->data_labels_ = data_labels; + return m; +} + +OutstandingRpcs::OutstandingRpcs( + std::string const& instrumentation_scope, + opentelemetry::nostd::shared_ptr< + opentelemetry::metrics::MeterProvider> const& provider) + : outstanding_rpcs_( + provider + ->GetMeter(instrumentation_scope, + kMeterInstrumentationScopeVersion) + ->CreateDoubleHistogram( + "bigtable.googleapis.com/internal/client/connection_pool/" + "outstanding_rpcs", + "Instantaneous count of outstanding RPCs on the selected " + "channel.", + "1")) {} + +void OutstandingRpcs::StubSelection( + opentelemetry::context::Context const& context, + StubSelectionParams const& p) { + ClientOutstandingRpcLabels data_labels{p.transport_type, + p.channel_pool_lb_policy, p.streaming}; + outstanding_rpcs_->Record(static_cast(p.outstanding_rpcs), + IntoLabelMap(resource_labels_, data_labels), + context); +} + +std::unique_ptr OutstandingRpcs::clone( + ClientResourceLabels const& resource_labels) const { + auto m = std::make_unique(*this); + m->resource_labels_ = resource_labels; return m; } diff --git a/google/cloud/bigtable/internal/metrics.h b/google/cloud/bigtable/internal/metrics.h index 54dcac4733725..3f63a4d12a9a5 100644 --- a/google/cloud/bigtable/internal/metrics.h +++ b/google/cloud/bigtable/internal/metrics.h @@ -36,7 +36,7 @@ namespace cloud { namespace bigtable_internal { GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN -struct ResourceLabels { +struct TableResourceLabels { std::string project_id; std::string instance; std::string table; @@ -44,7 +44,26 @@ struct ResourceLabels { std::string zone; }; -struct DataLabels { +enum class ChannelPoolLbPolicy { + kRoundRobin, + kRandomTwoLeastUsed, +}; + +enum class TransportType { + kCloudPath, + kDirectPath, +}; + +enum class RpcType { + kUnary, + kStreaming, +}; + +std::string ToString(ChannelPoolLbPolicy policy); +std::string ToString(TransportType type); +std::string ToString(RpcType streaming); + +struct TableDataLabels { std::string method; std::string streaming; std::string client_name; @@ -53,6 +72,32 @@ struct DataLabels { std::string status; }; +struct ClientResourceLabels { + std::string project_id; + std::string instance; + std::string app_profile; + std::string client_name; + std::string client_uid; + std::string client_project; + std::string location; + std::string cloud_platform; + std::string host_id; + std::string hostname; +}; + +struct ClientOutstandingRpcLabels { + TransportType transport_type; + ChannelPoolLbPolicy channel_pool_lb_policy; + RpcType streaming; +}; + +struct StubSelectionParams { + std::int64_t outstanding_rpcs; + ChannelPoolLbPolicy channel_pool_lb_policy; + TransportType transport_type; + RpcType streaming; +}; + // Labels populated from the peer info metadata. struct PeerInfoLabels { std::string transport_type; @@ -63,10 +108,14 @@ struct PeerInfoLabels { using LabelMap = std::unordered_map; // `peer_info_labels` is optional because only AttemptLatency2 populates it. LabelMap IntoLabelMap( - ResourceLabels const& r, DataLabels const& d, + TableResourceLabels const& r, TableDataLabels const& d, std::set const& filtered_data_labels = {}, std::optional const& peer_info_labels = std::nullopt); +LabelMap IntoLabelMap(ClientResourceLabels const& r, + ClientOutstandingRpcLabels const& d, + std::set const& filtered_data_labels = {}); + bool HasServerTiming(grpc::ClientContext const& client_context); bool IsConnectivityError(google::cloud::Status const& status, grpc::ClientContext const& client_context); @@ -109,6 +158,8 @@ class Metric { using LatencyDuration = std::chrono::duration; virtual ~Metric() = 0; + virtual void StubSelection(opentelemetry::context::Context const&, + StubSelectionParams const&) {} virtual void PreCall(opentelemetry::context::Context const&, PreCallParams const&) {} virtual void PostCall(opentelemetry::context::Context const&, @@ -119,8 +170,15 @@ class Metric { ElementRequestParams const&) {} virtual void ElementDelivery(opentelemetry::context::Context const&, ElementDeliveryParams const&) {} - virtual std::unique_ptr clone(ResourceLabels resource_labels, - DataLabels data_labels) const = 0; + virtual std::unique_ptr clone( + TableResourceLabels const& /*resource_labels*/, + TableDataLabels const& /*data_labels*/) const { + return nullptr; + } + virtual std::unique_ptr clone( + ClientResourceLabels const& /*resource_labels*/) const { + return nullptr; + } }; class OperationLatency : public Metric { @@ -136,12 +194,13 @@ class OperationLatency : public Metric { PostCallParams const& p) override; void OnDone(opentelemetry::context::Context const& context, OnDoneParams const& p) override; - std::unique_ptr clone(ResourceLabels resource_labels, - DataLabels data_labels) const override; + std::unique_ptr clone( + TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels) const override; private: - ResourceLabels resource_labels_; - DataLabels data_labels_; + TableResourceLabels resource_labels_; + TableDataLabels data_labels_; opentelemetry::nostd::shared_ptr> operation_latencies_; OperationContext::Clock::time_point operation_start_; @@ -157,12 +216,13 @@ class AttemptLatency : public Metric { void PostCall(opentelemetry::context::Context const& context, grpc::ClientContext const& client_context, PostCallParams const& p) override; - std::unique_ptr clone(ResourceLabels resource_labels, - DataLabels data_labels) const override; + std::unique_ptr clone( + TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels) const override; private: - ResourceLabels resource_labels_; - DataLabels data_labels_; + TableResourceLabels resource_labels_; + TableDataLabels data_labels_; opentelemetry::nostd::shared_ptr> attempt_latencies_; OperationContext::Clock::time_point attempt_start_; @@ -179,12 +239,13 @@ class AttemptLatency2 : public Metric { void PostCall(opentelemetry::context::Context const& context, grpc::ClientContext const& client_context, PostCallParams const& p) override; - std::unique_ptr clone(ResourceLabels resource_labels, - DataLabels data_labels) const override; + std::unique_ptr clone( + TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels) const override; private: - ResourceLabels resource_labels_; - DataLabels data_labels_; + TableResourceLabels resource_labels_; + TableDataLabels data_labels_; PeerInfoLabels peer_info_labels_; opentelemetry::nostd::shared_ptr> attempt_latencies2_; @@ -203,12 +264,13 @@ class RetryCount : public Metric { PostCallParams const& p) override; void OnDone(opentelemetry::context::Context const& context, OnDoneParams const& p) override; - std::unique_ptr clone(ResourceLabels resource_labels, - DataLabels data_labels) const override; + std::unique_ptr clone( + TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels) const override; private: - ResourceLabels resource_labels_; - DataLabels data_labels_; + TableResourceLabels resource_labels_; + TableDataLabels data_labels_; std::uint64_t num_retries_ = 0; opentelemetry::nostd::shared_ptr< opentelemetry::metrics::Counter> @@ -231,12 +293,13 @@ class FirstResponseLatency : public Metric { void OnDone(opentelemetry::context::Context const& context, OnDoneParams const& p) override; - std::unique_ptr clone(ResourceLabels resource_labels, - DataLabels data_labels) const override; + std::unique_ptr clone( + TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels) const override; private: - ResourceLabels resource_labels_; - DataLabels data_labels_; + TableResourceLabels resource_labels_; + TableDataLabels data_labels_; opentelemetry::nostd::shared_ptr> first_response_latencies_; OperationContext::Clock::time_point operation_start_; @@ -252,12 +315,13 @@ class ServerLatency : public Metric { grpc::ClientContext const& client_context, PostCallParams const& p) override; - std::unique_ptr clone(ResourceLabels resource_labels, - DataLabels data_labels) const override; + std::unique_ptr clone( + TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels) const override; private: - ResourceLabels resource_labels_; - DataLabels data_labels_; + TableResourceLabels resource_labels_; + TableDataLabels data_labels_; opentelemetry::nostd::shared_ptr> server_latencies_; }; @@ -278,12 +342,13 @@ class ApplicationBlockingLatency : public Metric { void OnDone(opentelemetry::context::Context const& context, OnDoneParams const&) override; - std::unique_ptr clone(ResourceLabels resource_labels, - DataLabels data_labels) const override; + std::unique_ptr clone( + TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels) const override; private: - ResourceLabels resource_labels_; - DataLabels data_labels_; + TableResourceLabels resource_labels_; + TableDataLabels data_labels_; opentelemetry::nostd::shared_ptr> application_blocking_latencies_; OperationContext::Clock::time_point element_delivery_time_; @@ -301,17 +366,34 @@ class ConnectivityErrorCount : public Metric { PostCallParams const& p) override; void OnDone(opentelemetry::context::Context const& context, OnDoneParams const&) override; - std::unique_ptr clone(ResourceLabels resource_labels, - DataLabels data_labels) const override; + std::unique_ptr clone( + TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels) const override; private: - ResourceLabels resource_labels_; - DataLabels data_labels_; + TableResourceLabels resource_labels_; + TableDataLabels data_labels_; std::uint64_t num_errors_ = 0; opentelemetry::nostd::shared_ptr< opentelemetry::metrics::Counter> connectivity_error_count_; }; + +class OutstandingRpcs : public Metric { + public: + OutstandingRpcs(std::string const& instrumentation_scope, + opentelemetry::nostd::shared_ptr< + opentelemetry::metrics::MeterProvider> const& provider); + void StubSelection(opentelemetry::context::Context const&, + StubSelectionParams const& p) override; + std::unique_ptr clone( + ClientResourceLabels const& resource_labels) const override; + + private: + ClientResourceLabels resource_labels_; + opentelemetry::nostd::shared_ptr> + outstanding_rpcs_; +}; GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END } // namespace bigtable_internal } // namespace cloud diff --git a/google/cloud/bigtable/internal/metrics_test.cc b/google/cloud/bigtable/internal/metrics_test.cc index d4d924493d737..1789fe382af57 100644 --- a/google/cloud/bigtable/internal/metrics_test.cc +++ b/google/cloud/bigtable/internal/metrics_test.cc @@ -228,10 +228,10 @@ class MockMeterProvider : public opentelemetry::metrics::MeterProvider { }; TEST(LabelMap, IntoLabelMap) { - ResourceLabels r{"my-project", "my-instance", "my-table", "my-cluster", - "my-zone"}; - DataLabels d{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", "my-status"}; + TableResourceLabels r{"my-project", "my-instance", "my-table", "my-cluster", + "my-zone"}; + TableDataLabels d{"my-method", "my-streaming", "my-client-name", + "my-client-uid", "my-app-profile", "my-status"}; auto label_map = IntoLabelMap(r, d); EXPECT_THAT( label_map, @@ -246,9 +246,9 @@ TEST(LabelMap, IntoLabelMap) { } TEST(LabelMap, IntoLabelMapWithDefaults) { - ResourceLabels r{"my-project", "my-instance", "my-table", "", ""}; - DataLabels d{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", "my-status"}; + TableResourceLabels r{"my-project", "my-instance", "my-table", "", ""}; + TableDataLabels d{"my-method", "my-streaming", "my-client-name", + "my-client-uid", "my-app-profile", "my-status"}; auto label_map = IntoLabelMap(r, d); EXPECT_THAT( label_map, @@ -510,10 +510,11 @@ TEST(OperationLatencyTest, FirstAttemptSuccess) { }); OperationLatency operation_latency("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = operation_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -583,10 +584,11 @@ TEST(OperationLatencyTest, ThirdAttemptSuccess) { }); OperationLatency operation_latency("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = operation_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -671,10 +673,11 @@ TEST(OperationLatencyTest, UsesDefaultClusterAndZone) { }); OperationLatency operation_latency("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = operation_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -743,10 +746,11 @@ TEST(AttemptLatencyTest, NoRetry) { }); AttemptLatency attempt_latency("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = attempt_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -846,10 +850,11 @@ TEST(AttemptLatencyTest, ThreeAttempts) { }); AttemptLatency attempt_latency("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = attempt_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -929,10 +934,11 @@ TEST(AttemptLatencyTest, UsesDefaultClusterAndZone) { }); AttemptLatency attempt_latency("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = attempt_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -1004,10 +1010,11 @@ TEST(AttemptLatency2Test, NoRetry) { }); AttemptLatency2 attempt_latency("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = attempt_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -1114,10 +1121,11 @@ TEST(AttemptLatency2Test, ThreeAttempts) { }); AttemptLatency2 attempt_latency("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = attempt_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -1198,10 +1206,11 @@ TEST(AttemptLatency2Test, UsesDefaultClusterAndZone) { }); AttemptLatency2 attempt_latency("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = attempt_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -1270,10 +1279,11 @@ TEST(RetryCountTest, NoRetry) { }); RetryCount retry_count("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = retry_count.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -1343,10 +1353,11 @@ TEST(RetryCountTest, ThreeAttempts) { }); RetryCount retry_count("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = retry_count.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -1426,10 +1437,11 @@ TEST(RetryCountTest, UsesDefaultClusterAndZone) { }); RetryCount retry_count("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = retry_count.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -1499,10 +1511,11 @@ TEST(FirstResponseLatency, Success) { FirstResponseLatency first_response_latency("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = first_response_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -1564,10 +1577,11 @@ TEST(FirstResponseLatency, NoDataReceived) { FirstResponseLatency first_response_latency("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = first_response_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -1642,10 +1656,11 @@ TEST(FirstResponseLatency, UsesDefaultClusterAndZone) { FirstResponseLatency first_response_latency("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = first_response_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -1756,10 +1771,11 @@ TEST(ServerLatency, SingleSuccess) { }); ServerLatency server_latency("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = server_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -1831,10 +1847,11 @@ TEST(ServerLatency, UsesDefaultClusterAndZone) { }); ServerLatency server_latency("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = server_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -1916,10 +1933,11 @@ TEST(ServerLatency, TwoAttempts) { }); ServerLatency server_latency("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = server_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -1986,10 +2004,11 @@ TEST(ServerLatency, NoServerTiming) { }); ServerLatency server_latency("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = server_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -2053,10 +2072,11 @@ TEST(ConnectivityErrorCount, MissingResourceLabels) { ConnectivityErrorCount connectivity_error_count("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = connectivity_error_count.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -2133,10 +2153,11 @@ TEST(ConnectivityErrorCount, Success) { ConnectivityErrorCount connectivity_error_count("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = connectivity_error_count.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -2209,10 +2230,11 @@ TEST(ConnectivityErrorCount, OkAndMissingServerTiming) { ConnectivityErrorCount connectivity_error_count("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = connectivity_error_count.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -2278,10 +2300,11 @@ TEST(ConnectivityErrorCount, DeadlineExceededAndMissingServerTiming) { ConnectivityErrorCount connectivity_error_count("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = connectivity_error_count.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -2346,10 +2369,11 @@ TEST(ApplicationBlockingLatency, Success) { ApplicationBlockingLatency application_blocking_latency("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = application_blocking_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -2433,10 +2457,11 @@ TEST(ApplicationBlockingLatency, StreamingData) { ApplicationBlockingLatency application_blocking_latency("my-instrument-scope", mock_provider); - ResourceLabels resource_labels{"my-project-id", "my-instance", "my-table", "", - ""}; - DataLabels data_labels{"my-method", "my-streaming", "my-client-name", - "my-client-uid", "my-app-profile", ""}; + TableResourceLabels resource_labels{"my-project-id", "my-instance", + "my-table", "", ""}; + TableDataLabels data_labels{"my-method", "my-streaming", + "my-client-name", "my-client-uid", + "my-app-profile", ""}; auto clone = application_blocking_latency.clone(resource_labels, data_labels); grpc::ClientContext client_context; @@ -2455,6 +2480,96 @@ TEST(ApplicationBlockingLatency, StreamingData) { clone->PostCall(otel_context, client_context, {clock->Now(), Status{}}); clone->OnDone(otel_context, {clock->Now(), Status{}}); } + +TEST(MetricsTest, IntoLabelMapClient) { + ClientResourceLabels resource{"p-1", "i-1", "app-1", "client-1", + "uid-1", "cp-1", "loc-1", "cloud-1", + "host-1", "hostname-1"}; + ClientOutstandingRpcLabels data{TransportType::kDirectPath, + ChannelPoolLbPolicy::kRandomTwoLeastUsed, + RpcType::kStreaming}; + auto actual = IntoLabelMap(resource, data); + EXPECT_THAT(actual, + UnorderedElementsAre( + Pair("project_id", "p-1"), Pair("instance", "i-1"), + Pair("app_profile", "app-1"), Pair("client_name", "client-1"), + Pair("client_uid", "uid-1"), Pair("client_project", "cp-1"), + Pair("location", "loc-1"), Pair("cloud_platform", "cloud-1"), + Pair("host_id", "host-1"), Pair("hostname", "hostname-1"), + Pair("transport_type", "DirectPath"), + Pair("channel_pool_lb_policy", "RANDOM_TWO_LEAST_USED"), + Pair("streaming", "true"))); +} + +TEST(MetricsTest, OutstandingRpcsMetric) { + auto mock_histogram = std::make_unique>(); + EXPECT_CALL( + *mock_histogram, + Record(A(), A(), + A())) + .WillOnce([](double value, + opentelemetry::common::KeyValueIterable const& attrs, + opentelemetry::context::Context const&) { + EXPECT_THAT(value, Eq(42.0)); + EXPECT_THAT( + MakeAttributesMap(attrs), + UnorderedElementsAre( + Pair("project_id", "my-project"), + Pair("instance", "my-instance"), + Pair("app_profile", "my-app-profile"), + Pair("client_name", "my-client-name"), + Pair("client_uid", "my-uid"), + Pair("client_project", "my-client-project"), + Pair("location", "us-east1"), Pair("cloud_platform", "gcp"), + Pair("host_id", "my-host"), Pair("hostname", "my-hostname"), + Pair("transport_type", "DirectPath"), + Pair("channel_pool_lb_policy", "RANDOM_TWO_LEAST_USED"), + Pair("streaming", "false"))); + }); + + opentelemetry::nostd::shared_ptr mock_meter = + std::make_shared(); + EXPECT_CALL(*mock_meter, CreateDoubleHistogram) + .WillOnce([mock = std::move(mock_histogram)]( + opentelemetry::nostd::string_view name, + opentelemetry::nostd::string_view, + opentelemetry::nostd::string_view) mutable { + EXPECT_THAT(name, Eq("bigtable.googleapis.com/internal/client/" + "connection_pool/outstanding_rpcs")); + return std::move(mock); + }); + + opentelemetry::nostd::shared_ptr mock_provider = + std::make_shared(); + EXPECT_CALL(*mock_provider, GetMeter) +#if OPENTELEMETRY_ABI_VERSION_NO >= 2 + .WillOnce([&](opentelemetry::nostd::string_view scope, + opentelemetry::nostd::string_view scope_version, + opentelemetry::nostd::string_view, + opentelemetry::common::KeyValueIterable const*) mutable { +#else + .WillOnce([&](opentelemetry::nostd::string_view scope, + opentelemetry::nostd::string_view scope_version, + opentelemetry::nostd::string_view) mutable { +#endif + EXPECT_THAT(scope, Eq("my-instrument-scope")); + EXPECT_THAT(scope_version, Eq("v1")); + return mock_meter; + }); + + OutstandingRpcs outstanding_rpcs("my-instrument-scope", mock_provider); + ClientResourceLabels resource_labels{ + "my-project", "my-instance", "my-app-profile", "my-client-name", + "my-uid", "my-client-project", "us-east1", "gcp", + "my-host", "my-hostname"}; + auto clone = outstanding_rpcs.clone(resource_labels); + + auto otel_context = opentelemetry::context::RuntimeContext::GetCurrent(); + clone->StubSelection( + otel_context, + StubSelectionParams{42, ChannelPoolLbPolicy::kRandomTwoLeastUsed, + TransportType::kDirectPath, RpcType::kUnary}); +} } // namespace GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END } // namespace bigtable_internal diff --git a/google/cloud/bigtable/internal/operation_context.cc b/google/cloud/bigtable/internal/operation_context.cc index 6a2a0ab8791f8..39c8cdfcb8848 100644 --- a/google/cloud/bigtable/internal/operation_context.cc +++ b/google/cloud/bigtable/internal/operation_context.cc @@ -28,12 +28,20 @@ GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN #ifdef GOOGLE_CLOUD_CPP_BIGTABLE_WITH_OTEL_METRICS namespace { std::vector> CloneMetrics( - ResourceLabels const& resource_labels, DataLabels const& data_labels, + TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels, + ClientResourceLabels const& client_resource_labels, std::vector> const& metrics) { std::vector> v; v.reserve(metrics.size()); for (auto const& m : metrics) { - v.emplace_back(m->clone(resource_labels, data_labels)); + auto clone = m->clone(resource_labels, data_labels); + if (!clone) { + clone = m->clone(client_resource_labels); + } + if (clone) { + v.emplace_back(std::move(clone)); + } } return v; } @@ -54,12 +62,30 @@ void OperationContext::ProcessMetadata( #ifdef GOOGLE_CLOUD_CPP_BIGTABLE_WITH_OTEL_METRICS OperationContext::OperationContext( - ResourceLabels const& resource_labels, DataLabels const& data_labels, + TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels, + std::vector> const& metrics, + std::shared_ptr clock) + : OperationContext(resource_labels, data_labels, ClientResourceLabels{}, + metrics, std::move(clock)) {} + +OperationContext::OperationContext( + TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels, + ClientResourceLabels const& client_resource_labels, std::vector> const& metrics, std::shared_ptr clock) - : cloned_metrics_(CloneMetrics(resource_labels, data_labels, metrics)), + : cloned_metrics_(CloneMetrics(resource_labels, data_labels, + client_resource_labels, metrics)), clock_(std::move(clock)) {} +void OperationContext::StubSelection(StubSelectionParams const& params) { + auto otel_context = opentelemetry::context::RuntimeContext::GetCurrent(); + for (auto& m : cloned_metrics_) { + m->StubSelection(otel_context, params); + } +} + void OperationContext::PreCall(grpc::ClientContext& client_context) { auto otel_context = opentelemetry::context::RuntimeContext::GetCurrent(); auto attempt_start = clock_->Now(); @@ -121,7 +147,7 @@ void OperationContext::ElementDelivery(grpc::ClientContext const&) { #else // GOOGLE_CLOUD_CPP_BIGTABLE_WITH_OTEL_METRICS OperationContext::OperationContext( - ResourceLabels const&, DataLabels const&, + TableResourceLabels const&, TableDataLabels const&, std::vector> const&, std::shared_ptr) { } diff --git a/google/cloud/bigtable/internal/operation_context.h b/google/cloud/bigtable/internal/operation_context.h index 345677a213620..040c91a740d17 100644 --- a/google/cloud/bigtable/internal/operation_context.h +++ b/google/cloud/bigtable/internal/operation_context.h @@ -30,8 +30,10 @@ namespace cloud { namespace bigtable_internal { GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN -struct DataLabels; -struct ResourceLabels; +struct ClientResourceLabels; +struct TableDataLabels; +struct TableResourceLabels; +struct StubSelectionParams; class Metric; /** @@ -64,11 +66,18 @@ class OperationContext { // The default constructor is used when metric support is unavailable or // disabled. OperationContext() = default; - OperationContext(ResourceLabels const& resource_labels, - DataLabels const& data_labels, + OperationContext(TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels, + std::vector> const& metrics, + std::shared_ptr clock); + OperationContext(TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels, + ClientResourceLabels const& client_resource_labels, std::vector> const& metrics, std::shared_ptr clock); + // 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); // Called after receiving RPC attempt response. diff --git a/google/cloud/bigtable/internal/operation_context_factory.cc b/google/cloud/bigtable/internal/operation_context_factory.cc index fe563eb6d7c86..8a06bddadebb8 100644 --- a/google/cloud/bigtable/internal/operation_context_factory.cc +++ b/google/cloud/bigtable/internal/operation_context_factory.cc @@ -37,27 +37,28 @@ GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN #if GOOGLE_CLOUD_CPP_BIGTABLE_WITH_OTEL_METRICS namespace { -ResourceLabels ResourceLabelsFromTableName(std::string const& table_name) { +TableResourceLabels TableResourceLabelsFromTableName( + std::string const& table_name) { // split table_name into component pieces // projects//instances//tables/ std::vector name_parts = absl::StrSplit(table_name, '/'); if (name_parts.size() < 6) return {}; - ResourceLabels resource_labels = { + TableResourceLabels resource_labels = { std::string(name_parts[1]), std::string(name_parts[3]), std::string(name_parts[5]), "" /*=cluster*/, "" /*=zone*/}; return resource_labels; } -ResourceLabels ResourceLabelsFromInstanceName( +TableResourceLabels TableResourceLabelsFromInstanceName( std::string const& instance_name) { // split instance_name into component pieces // projects//instances/ std::vector name_parts = absl::StrSplit(instance_name, '/'); if (name_parts.size() < 4) return {}; - ResourceLabels resource_labels = {std::string(name_parts[1]), - std::string(name_parts[3]), "", - "" /*=cluster*/, "" /*=zone*/}; + TableResourceLabels resource_labels = {std::string(name_parts[1]), + std::string(name_parts[3]), "", + "" /*=cluster*/, "" /*=zone*/}; return resource_labels; } @@ -148,6 +149,17 @@ std::shared_ptr SimpleOperationContextFactory::ExecuteQuery( #ifdef GOOGLE_CLOUD_CPP_BIGTABLE_WITH_OTEL_METRICS +namespace { +ClientResourceLabels MakeClientLabels( + ClientResourceLabels base, TableResourceLabels const& resource_labels, + std::string const& app_profile) { + base.project_id = resource_labels.project_id; + base.instance = resource_labels.instance; + base.app_profile = app_profile; + return base; +} +} // namespace + MetricsOperationContextFactory::MetricsOperationContextFactory( std::string client_uid, std::shared_ptr conn, @@ -288,19 +300,22 @@ std::shared_ptr MetricsOperationContextFactory::ReadRow( std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); + v.emplace_back(std::make_shared(kRpc, provider_)); swap(read_row_metrics_.metrics, v); }); - auto resource_labels = ResourceLabelsFromTableName(table_name); - DataLabels data_labels = {kRpc, - "true", /*=streaming*/ - "cpp.Bigtable/" + version_string(), - client_uid_, - app_profile, - "" /*=status*/}; + auto resource_labels = TableResourceLabelsFromTableName(table_name); + TableDataLabels data_labels = {kRpc, + "true", /*=streaming*/ + "cpp.Bigtable/" + version_string(), + client_uid_, + app_profile, + "" /*=status*/}; - return std::make_shared(resource_labels, data_labels, - read_row_metrics_.metrics, clock_); + return std::make_shared( + resource_labels, data_labels, + MakeClientLabels(client_resource_labels_, resource_labels, app_profile), + read_row_metrics_.metrics, clock_); } std::shared_ptr MetricsOperationContextFactory::ReadRows( @@ -317,19 +332,22 @@ std::shared_ptr MetricsOperationContextFactory::ReadRows( std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); + v.emplace_back(std::make_shared(kRpc, provider_)); swap(read_rows_metrics_.metrics, v); }); - auto resource_labels = ResourceLabelsFromTableName(table_name); - DataLabels data_labels = {kRpc, - "true", /*=streaming*/ - "cpp.Bigtable/" + version_string(), - client_uid_, - app_profile, - "" /*=status*/}; + auto resource_labels = TableResourceLabelsFromTableName(table_name); + TableDataLabels data_labels = {kRpc, + "true", /*=streaming*/ + "cpp.Bigtable/" + version_string(), + client_uid_, + app_profile, + "" /*=status*/}; - return std::make_shared(resource_labels, data_labels, - read_rows_metrics_.metrics, clock_); + return std::make_shared( + resource_labels, data_labels, + MakeClientLabels(client_resource_labels_, resource_labels, app_profile), + read_rows_metrics_.metrics, clock_); } std::shared_ptr MetricsOperationContextFactory::MutateRow( @@ -345,19 +363,22 @@ std::shared_ptr MetricsOperationContextFactory::MutateRow( std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); + v.emplace_back(std::make_shared(kRpc, provider_)); swap(mutate_row_metrics_.metrics, v); }); - auto resource_labels = ResourceLabelsFromTableName(table_name); - DataLabels data_labels = {kRpc, - "false", /*=streaming*/ - "cpp.Bigtable/" + version_string(), - client_uid_, - app_profile, - "" /*=status*/}; + auto resource_labels = TableResourceLabelsFromTableName(table_name); + TableDataLabels data_labels = {kRpc, + "false", /*=streaming*/ + "cpp.Bigtable/" + version_string(), + client_uid_, + app_profile, + "" /*=status*/}; return std::make_shared( - resource_labels, data_labels, mutate_row_metrics_.metrics, clock_); + resource_labels, data_labels, + MakeClientLabels(client_resource_labels_, resource_labels, app_profile), + mutate_row_metrics_.metrics, clock_); } std::shared_ptr MetricsOperationContextFactory::MutateRows( @@ -373,19 +394,22 @@ std::shared_ptr MetricsOperationContextFactory::MutateRows( std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); + v.emplace_back(std::make_shared(kRpc, provider_)); swap(mutate_rows_metrics_.metrics, v); }); - auto resource_labels = ResourceLabelsFromTableName(table_name); - DataLabels data_labels = {kRpc, - "true", /*=streaming*/ - "cpp.Bigtable/" + version_string(), - client_uid_, - app_profile, - "" /*=status*/}; + auto resource_labels = TableResourceLabelsFromTableName(table_name); + TableDataLabels data_labels = {kRpc, + "true", /*=streaming*/ + "cpp.Bigtable/" + version_string(), + client_uid_, + app_profile, + "" /*=status*/}; return std::make_shared( - resource_labels, data_labels, mutate_rows_metrics_.metrics, clock_); + resource_labels, data_labels, + MakeClientLabels(client_resource_labels_, resource_labels, app_profile), + mutate_rows_metrics_.metrics, clock_); } std::shared_ptr @@ -402,20 +426,22 @@ MetricsOperationContextFactory::CheckAndMutateRow( std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); + v.emplace_back(std::make_shared(kRpc, provider_)); swap(check_and_mutate_row_metrics_.metrics, v); }); - auto resource_labels = ResourceLabelsFromTableName(table_name); - DataLabels data_labels = {kRpc, - "false", /*=streaming*/ - "cpp.Bigtable/" + version_string(), - client_uid_, - app_profile, - "" /*=status*/}; + auto resource_labels = TableResourceLabelsFromTableName(table_name); + TableDataLabels data_labels = {kRpc, + "false", /*=streaming*/ + "cpp.Bigtable/" + version_string(), + client_uid_, + app_profile, + "" /*=status*/}; return std::make_shared( - resource_labels, data_labels, check_and_mutate_row_metrics_.metrics, - clock_); + resource_labels, data_labels, + MakeClientLabels(client_resource_labels_, resource_labels, app_profile), + check_and_mutate_row_metrics_.metrics, clock_); } std::shared_ptr MetricsOperationContextFactory::SampleRowKeys( @@ -431,19 +457,22 @@ std::shared_ptr MetricsOperationContextFactory::SampleRowKeys( std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); + v.emplace_back(std::make_shared(kRpc, provider_)); swap(sample_row_keys_metrics_.metrics, v); }); - auto resource_labels = ResourceLabelsFromTableName(table_name); - DataLabels data_labels = {kRpc, - "true", /*=streaming*/ - "cpp.Bigtable/" + version_string(), - client_uid_, - app_profile, - "" /*=status*/}; + auto resource_labels = TableResourceLabelsFromTableName(table_name); + TableDataLabels data_labels = {kRpc, + "true", /*=streaming*/ + "cpp.Bigtable/" + version_string(), + client_uid_, + app_profile, + "" /*=status*/}; return std::make_shared( - resource_labels, data_labels, sample_row_keys_metrics_.metrics, clock_); + resource_labels, data_labels, + MakeClientLabels(client_resource_labels_, resource_labels, app_profile), + sample_row_keys_metrics_.metrics, clock_); } std::shared_ptr @@ -460,20 +489,22 @@ MetricsOperationContextFactory::ReadModifyWriteRow( std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); + v.emplace_back(std::make_shared(kRpc, provider_)); swap(read_modify_write_row_metrics_.metrics, v); }); - auto resource_labels = ResourceLabelsFromTableName(table_name); - DataLabels data_labels = {kRpc, - "false", /*=streaming*/ - "cpp.Bigtable/" + version_string(), - client_uid_, - app_profile, - "" /*=status*/}; + auto resource_labels = TableResourceLabelsFromTableName(table_name); + TableDataLabels data_labels = {kRpc, + "false", /*=streaming*/ + "cpp.Bigtable/" + version_string(), + client_uid_, + app_profile, + "" /*=status*/}; return std::make_shared( - resource_labels, data_labels, read_modify_write_row_metrics_.metrics, - clock_); + resource_labels, data_labels, + MakeClientLabels(client_resource_labels_, resource_labels, app_profile), + read_modify_write_row_metrics_.metrics, clock_); } std::shared_ptr MetricsOperationContextFactory::PrepareQuery( @@ -487,19 +518,22 @@ std::shared_ptr MetricsOperationContextFactory::PrepareQuery( v.emplace_back(std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); + v.emplace_back(std::make_shared(kRpc, provider_)); swap(prepare_query_metrics_.metrics, v); }); - auto resource_labels = ResourceLabelsFromInstanceName(instance_name); - DataLabels data_labels = {kRpc, - "false", /*=streaming*/ - "cpp.Bigtable/" + version_string(), - client_uid_, - app_profile, - "" /*=status*/}; + auto resource_labels = TableResourceLabelsFromInstanceName(instance_name); + TableDataLabels data_labels = {kRpc, + "false", /*=streaming*/ + "cpp.Bigtable/" + version_string(), + client_uid_, + app_profile, + "" /*=status*/}; return std::make_shared( - resource_labels, data_labels, prepare_query_metrics_.metrics, clock_); + resource_labels, data_labels, + MakeClientLabels(client_resource_labels_, resource_labels, app_profile), + prepare_query_metrics_.metrics, clock_); } std::shared_ptr MetricsOperationContextFactory::ExecuteQuery( @@ -516,19 +550,22 @@ std::shared_ptr MetricsOperationContextFactory::ExecuteQuery( std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); v.emplace_back(std::make_shared(kRpc, provider_)); + v.emplace_back(std::make_shared(kRpc, provider_)); swap(execute_query_metrics_.metrics, v); }); - auto resource_labels = ResourceLabelsFromInstanceName(instance_name); - DataLabels data_labels = {kRpc, - "true", /*=streaming*/ - "cpp.Bigtable/" + version_string(), - client_uid_, - app_profile, - "" /*=status*/}; + auto resource_labels = TableResourceLabelsFromInstanceName(instance_name); + TableDataLabels data_labels = {kRpc, + "true", /*=streaming*/ + "cpp.Bigtable/" + version_string(), + client_uid_, + app_profile, + "" /*=status*/}; return std::make_shared( - resource_labels, data_labels, execute_query_metrics_.metrics, clock_); + resource_labels, data_labels, + MakeClientLabels(client_resource_labels_, resource_labels, app_profile), + execute_query_metrics_.metrics, clock_); } #endif // GOOGLE_CLOUD_CPP_BIGTABLE_WITH_OTEL_METRICS diff --git a/google/cloud/bigtable/internal/operation_context_factory.h b/google/cloud/bigtable/internal/operation_context_factory.h index c8d77b4f7c6b0..a57ab04f4c2c4 100644 --- a/google/cloud/bigtable/internal/operation_context_factory.h +++ b/google/cloud/bigtable/internal/operation_context_factory.h @@ -18,6 +18,7 @@ #include "google/cloud/bigtable/internal/operation_context.h" #include "google/cloud/bigtable/version.h" #ifdef GOOGLE_CLOUD_CPP_BIGTABLE_WITH_OTEL_METRICS +#include "google/cloud/bigtable/internal/metrics.h" #include "google/cloud/monitoring/v3/metric_connection.h" #include "absl/base/call_once.h" #include @@ -140,6 +141,7 @@ class MetricsOperationContextFactory : public OperationContextFactory { std::string client_uid_; std::shared_ptr clock_; std::shared_ptr provider_; + ClientResourceLabels client_resource_labels_; // These vectors are initialized exactly once and the initialization is // delayed until the first time the corresponding method is called. diff --git a/google/cloud/bigtable/internal/operation_context_factory_test.cc b/google/cloud/bigtable/internal/operation_context_factory_test.cc index ed731dd862928..31e73b0d997f1 100644 --- a/google/cloud/bigtable/internal/operation_context_factory_test.cc +++ b/google/cloud/bigtable/internal/operation_context_factory_test.cc @@ -26,10 +26,12 @@ namespace { using ::testing::Eq; using ::testing::IsEmpty; +using ::testing::NotNull; class MockMetric : public Metric { public: - MOCK_METHOD(std::unique_ptr, clone, (ResourceLabels, DataLabels), + MOCK_METHOD(std::unique_ptr, clone, + (TableResourceLabels const&, TableDataLabels const&), (const, override)); }; @@ -40,8 +42,8 @@ TEST(MetricsOperationContextFactoryTest, ReadRow) { auto mock_metric = std::make_shared(); EXPECT_CALL(*mock_metric, clone) - .WillOnce([&](ResourceLabels const& resource_labels, - DataLabels const& data_labels) { + .WillOnce([&](TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels) { EXPECT_THAT(resource_labels.project_id, Eq("my-project")); EXPECT_THAT(resource_labels.instance, Eq("my-instance")); EXPECT_THAT(resource_labels.table, Eq("my-table")); @@ -68,11 +70,12 @@ TEST(MetricsOperationContextFactoryTest, ReadRows) { auto mock_metric = std::make_shared(); EXPECT_CALL(*mock_metric, clone) - .WillOnce([&](ResourceLabels const&, DataLabels const& data_labels) { - EXPECT_THAT(data_labels.method, Eq("ReadRows")); - EXPECT_THAT(data_labels.streaming, Eq("true")); - return std::make_unique(); - }); + .WillOnce( + [&](TableResourceLabels const&, TableDataLabels const& data_labels) { + EXPECT_THAT(data_labels.method, Eq("ReadRows")); + EXPECT_THAT(data_labels.streaming, Eq("true")); + return std::make_unique(); + }); MetricsOperationContextFactory factory({}, mock_metric); auto operation_context = factory.ReadRows(table_full_name, app_profile); @@ -85,11 +88,12 @@ TEST(MetricsOperationContextFactoryTest, MutateRow) { auto mock_metric = std::make_shared(); EXPECT_CALL(*mock_metric, clone) - .WillOnce([&](ResourceLabels const&, DataLabels const& data_labels) { - EXPECT_THAT(data_labels.method, Eq("MutateRow")); - EXPECT_THAT(data_labels.streaming, Eq("false")); - return std::make_unique(); - }); + .WillOnce( + [&](TableResourceLabels const&, TableDataLabels const& data_labels) { + EXPECT_THAT(data_labels.method, Eq("MutateRow")); + EXPECT_THAT(data_labels.streaming, Eq("false")); + return std::make_unique(); + }); MetricsOperationContextFactory factory({}, mock_metric); auto operation_context = factory.MutateRow(table_full_name, app_profile); @@ -102,11 +106,12 @@ TEST(MetricsOperationContextFactoryTest, MutateRows) { auto mock_metric = std::make_shared(); EXPECT_CALL(*mock_metric, clone) - .WillOnce([&](ResourceLabels const&, DataLabels const& data_labels) { - EXPECT_THAT(data_labels.method, Eq("MutateRows")); - EXPECT_THAT(data_labels.streaming, Eq("true")); - return std::make_unique(); - }); + .WillOnce( + [&](TableResourceLabels const&, TableDataLabels const& data_labels) { + EXPECT_THAT(data_labels.method, Eq("MutateRows")); + EXPECT_THAT(data_labels.streaming, Eq("true")); + return std::make_unique(); + }); MetricsOperationContextFactory factory({}, mock_metric); auto operation_context = factory.MutateRows(table_full_name, app_profile); @@ -119,11 +124,12 @@ TEST(MetricsOperationContextFactoryTest, CheckAndMutateRow) { auto mock_metric = std::make_shared(); EXPECT_CALL(*mock_metric, clone) - .WillOnce([&](ResourceLabels const&, DataLabels const& data_labels) { - EXPECT_THAT(data_labels.method, Eq("CheckAndMutateRow")); - EXPECT_THAT(data_labels.streaming, Eq("false")); - return std::make_unique(); - }); + .WillOnce( + [&](TableResourceLabels const&, TableDataLabels const& data_labels) { + EXPECT_THAT(data_labels.method, Eq("CheckAndMutateRow")); + EXPECT_THAT(data_labels.streaming, Eq("false")); + return std::make_unique(); + }); MetricsOperationContextFactory factory({}, mock_metric); auto operation_context = @@ -137,11 +143,12 @@ TEST(MetricsOperationContextFactoryTest, SampleRowKeys) { auto mock_metric = std::make_shared(); EXPECT_CALL(*mock_metric, clone) - .WillOnce([&](ResourceLabels const&, DataLabels const& data_labels) { - EXPECT_THAT(data_labels.method, Eq("SampleRowKeys")); - EXPECT_THAT(data_labels.streaming, Eq("true")); - return std::make_unique(); - }); + .WillOnce( + [&](TableResourceLabels const&, TableDataLabels const& data_labels) { + EXPECT_THAT(data_labels.method, Eq("SampleRowKeys")); + EXPECT_THAT(data_labels.streaming, Eq("true")); + return std::make_unique(); + }); MetricsOperationContextFactory factory({}, mock_metric); auto operation_context = factory.SampleRowKeys(table_full_name, app_profile); @@ -154,11 +161,12 @@ TEST(MetricsOperationContextFactoryTest, ReadModifyWriteRow) { auto mock_metric = std::make_shared(); EXPECT_CALL(*mock_metric, clone) - .WillOnce([&](ResourceLabels const&, DataLabels const& data_labels) { - EXPECT_THAT(data_labels.method, Eq("ReadModifyWriteRow")); - EXPECT_THAT(data_labels.streaming, Eq("false")); - return std::make_unique(); - }); + .WillOnce( + [&](TableResourceLabels const&, TableDataLabels const& data_labels) { + EXPECT_THAT(data_labels.method, Eq("ReadModifyWriteRow")); + EXPECT_THAT(data_labels.streaming, Eq("false")); + return std::make_unique(); + }); MetricsOperationContextFactory factory({}, mock_metric); auto operation_context = @@ -171,11 +179,12 @@ TEST(MetricsOperationContextFactoryTest, PrepareQuery) { auto mock_metric = std::make_shared(); EXPECT_CALL(*mock_metric, clone) - .WillOnce([&](ResourceLabels const&, DataLabels const& data_labels) { - EXPECT_THAT(data_labels.method, Eq("PrepareQuery")); - EXPECT_THAT(data_labels.streaming, Eq("false")); - return std::make_unique(); - }); + .WillOnce( + [&](TableResourceLabels const&, TableDataLabels const& data_labels) { + EXPECT_THAT(data_labels.method, Eq("PrepareQuery")); + EXPECT_THAT(data_labels.streaming, Eq("false")); + return std::make_unique(); + }); MetricsOperationContextFactory factory({}, mock_metric); auto operation_context = @@ -188,17 +197,32 @@ TEST(MetricsOperationContextFactoryTest, ExecuteQuery) { auto mock_metric = std::make_shared(); EXPECT_CALL(*mock_metric, clone) - .WillOnce([&](ResourceLabels const&, DataLabels const& data_labels) { - EXPECT_THAT(data_labels.method, Eq("ExecuteQuery")); - EXPECT_THAT(data_labels.streaming, Eq("true")); - return std::make_unique(); - }); + .WillOnce( + [&](TableResourceLabels const&, TableDataLabels const& data_labels) { + EXPECT_THAT(data_labels.method, Eq("ExecuteQuery")); + EXPECT_THAT(data_labels.streaming, Eq("true")); + return std::make_unique(); + }); MetricsOperationContextFactory factory({}, mock_metric); auto operation_context = factory.ExecuteQuery(instance_full_name, app_profile); } +TEST(MetricsOperationContextFactoryTest, IncludesOutstandingRpcs) { + std::string app_profile = "my-app-profile"; + std::string table_full_name = + "projects/my-project/instances/my-instance/tables/my-table"; + MetricsOperationContextFactory factory( + "test-uid", + std::shared_ptr(nullptr)); + auto operation_context = factory.ReadRow(table_full_name, app_profile); + EXPECT_THAT(operation_context, NotNull()); + operation_context->StubSelection( + StubSelectionParams{10, ChannelPoolLbPolicy::kRandomTwoLeastUsed, + TransportType::kDirectPath, RpcType::kUnary}); +} + } // namespace GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END } // namespace bigtable_internal diff --git a/google/cloud/bigtable/internal/operation_context_test.cc b/google/cloud/bigtable/internal/operation_context_test.cc index 7bf7e4fe40b0a..d7f1c13d5c3df 100644 --- a/google/cloud/bigtable/internal/operation_context_test.cc +++ b/google/cloud/bigtable/internal/operation_context_test.cc @@ -141,7 +141,8 @@ class MockMetric : public Metric { ElementDeliveryParams const&), (override)); MOCK_METHOD(std::unique_ptr, clone, - (ResourceLabels resource_labels, DataLabels data_labels), + (TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels), (const, override)); }; @@ -150,7 +151,8 @@ class CloningMetric : public Metric { public: explicit CloningMetric(std::unique_ptr metric) : metric_(std::move(metric)) {} - std::unique_ptr clone(ResourceLabels, DataLabels) const override { + std::unique_ptr clone(TableResourceLabels const&, + TableDataLabels const&) const override { return std::move(metric_); } diff --git a/google/cloud/bigtable/internal/partial_result_set_source_test.cc b/google/cloud/bigtable/internal/partial_result_set_source_test.cc index 7106940c68b0c..394ec54d5d36c 100644 --- a/google/cloud/bigtable/internal/partial_result_set_source_test.cc +++ b/google/cloud/bigtable/internal/partial_result_set_source_test.cc @@ -109,7 +109,8 @@ class MockMetric : public Metric { ElementDeliveryParams const&), (override)); MOCK_METHOD(std::unique_ptr, clone, - (ResourceLabels resource_labels, DataLabels data_labels), + (TableResourceLabels const& resource_labels, + TableDataLabels const& data_labels), (const, override)); }; @@ -118,7 +119,8 @@ class CloningMetric : public Metric { public: explicit CloningMetric(std::unique_ptr metric) : metric_(std::move(metric)) {} - std::unique_ptr clone(ResourceLabels, DataLabels) const override { + std::unique_ptr clone(TableResourceLabels const&, + TableDataLabels const&) const override { return std::move(metric_); } diff --git a/google/cloud/bigtable/internal/stub_manager_test.cc b/google/cloud/bigtable/internal/stub_manager_test.cc index 70c460d736371..299862b1d9f51 100644 --- a/google/cloud/bigtable/internal/stub_manager_test.cc +++ b/google/cloud/bigtable/internal/stub_manager_test.cc @@ -14,6 +14,7 @@ #include "google/cloud/bigtable/internal/stub_manager.h" #include "google/cloud/bigtable/instance_resource.h" +#include "google/cloud/bigtable/internal/operation_context.h" #include "google/cloud/bigtable/table_resource.h" #include "google/cloud/bigtable/testing/mock_bigtable_stub.h" #include "google/cloud/testing_util/scoped_log.h" @@ -41,7 +42,8 @@ TEST(StubManagerTest, NoAffinity) { auto mock = std::make_shared(); EXPECT_CALL(*mock, MutateRow) .WillOnce([&](grpc::ClientContext&, Options const&, - google::bigtable::v2::MutateRowRequest const& request) { + google::bigtable::v2::MutateRowRequest const& request, + OperationContext&) { EXPECT_THAT(request.table_name(), Eq(expected_table_name)); return google::bigtable::v2::MutateRowResponse{}; }); @@ -52,7 +54,8 @@ TEST(StubManagerTest, NoAffinity) { grpc::ClientContext context; google::bigtable::v2::MutateRowRequest request; request.set_table_name(expected_table_name); - auto result = stub->MutateRow(context, {}, request); + OperationContext default_oc; + auto result = stub->MutateRow(context, {}, request, default_oc); EXPECT_THAT(result, IsOk()); } @@ -65,7 +68,8 @@ TEST(StubManagerTest, AffinityToExistingInstance) { EXPECT_CALL(*mock_stub, MutateRow) .WillOnce([instance_name = instance.FullName()]( grpc::ClientContext&, Options const&, - google::bigtable::v2::MutateRowRequest const& request) { + google::bigtable::v2::MutateRowRequest const& request, + OperationContext&) { EXPECT_THAT(request.table_name(), StartsWith(instance_name)); return google::bigtable::v2::MutateRowResponse{}; }); @@ -85,7 +89,8 @@ TEST(StubManagerTest, AffinityToExistingInstance) { grpc::ClientContext context; google::bigtable::v2::MutateRowRequest request; request.set_table_name(expected_table_name); - auto result = stub->MutateRow(context, {}, request); + OperationContext default_oc; + auto result = stub->MutateRow(context, {}, request, default_oc); EXPECT_THAT(result, IsOk()); } @@ -105,7 +110,8 @@ TEST(StubManagerTest, AffinityToMissingInstance) { EXPECT_CALL(*mock_stub, MutateRow) .WillOnce([instance_name = std::string{instance_name}]( grpc::ClientContext&, Options const&, - google::bigtable::v2::MutateRowRequest const& request) { + google::bigtable::v2::MutateRowRequest const& request, + OperationContext&) { EXPECT_THAT(request.table_name(), StartsWith(instance_name)); return google::bigtable::v2::MutateRowResponse{}; }); @@ -121,7 +127,8 @@ TEST(StubManagerTest, AffinityToMissingInstance) { grpc::ClientContext context; google::bigtable::v2::MutateRowRequest request; request.set_table_name(expected_table_name); - auto result = stub->MutateRow(context, {}, request); + OperationContext default_oc; + auto result = stub->MutateRow(context, {}, request, default_oc); EXPECT_THAT(result, IsOk()); EXPECT_THAT(log.ExtractLines(), Contains(HasSubstr( diff --git a/google/cloud/bigtable/testing/mock_bigtable_stub.h b/google/cloud/bigtable/testing/mock_bigtable_stub.h index b2632564d7a47..3170eda24e874 100644 --- a/google/cloud/bigtable/testing/mock_bigtable_stub.h +++ b/google/cloud/bigtable/testing/mock_bigtable_stub.h @@ -30,48 +30,57 @@ class MockBigtableStub : public bigtable_internal::BigtableStub { google::bigtable::v2::ReadRowsResponse>>, ReadRows, (std::shared_ptr, Options const&, - google::bigtable::v2::ReadRowsRequest const&), + google::bigtable::v2::ReadRowsRequest const&, + bigtable_internal::OperationContext&), (override)); MOCK_METHOD(std::unique_ptr>, SampleRowKeys, (std::shared_ptr, Options const&, - google::bigtable::v2::SampleRowKeysRequest const&), + google::bigtable::v2::SampleRowKeysRequest const&, + bigtable_internal::OperationContext&), (override)); MOCK_METHOD(StatusOr, MutateRow, (grpc::ClientContext&, Options const&, - google::bigtable::v2::MutateRowRequest const&), + google::bigtable::v2::MutateRowRequest const&, + bigtable_internal::OperationContext&), (override)); MOCK_METHOD(std::unique_ptr>, MutateRows, (std::shared_ptr, Options const&, - google::bigtable::v2::MutateRowsRequest const&), + google::bigtable::v2::MutateRowsRequest const&, + bigtable_internal::OperationContext&), (override)); MOCK_METHOD(StatusOr, CheckAndMutateRow, (grpc::ClientContext&, Options const&, - google::bigtable::v2::CheckAndMutateRowRequest const&), + google::bigtable::v2::CheckAndMutateRowRequest const&, + bigtable_internal::OperationContext&), (override)); MOCK_METHOD(StatusOr, PingAndWarm, (grpc::ClientContext&, Options const&, - google::bigtable::v2::PingAndWarmRequest const&), + google::bigtable::v2::PingAndWarmRequest const&, + bigtable_internal::OperationContext&), (override)); MOCK_METHOD(StatusOr, ReadModifyWriteRow, (grpc::ClientContext&, Options const&, - google::bigtable::v2::ReadModifyWriteRowRequest const&), + google::bigtable::v2::ReadModifyWriteRowRequest const&, + bigtable_internal::OperationContext&), (override)); MOCK_METHOD(StatusOr, PrepareQuery, (grpc::ClientContext&, Options const&, - google::bigtable::v2::PrepareQueryRequest const&), + google::bigtable::v2::PrepareQueryRequest const&, + bigtable_internal::OperationContext&), (override)); MOCK_METHOD(std::unique_ptr>, ExecuteQuery, (std::shared_ptr, Options const&, - google::bigtable::v2::ExecuteQueryRequest const&), + google::bigtable::v2::ExecuteQueryRequest const&, + bigtable_internal::OperationContext&), (override)); MOCK_METHOD(std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::ReadRowsResponse>>, @@ -79,7 +88,8 @@ class MockBigtableStub : public bigtable_internal::BigtableStub { (google::cloud::CompletionQueue const&, std::shared_ptr, google::cloud::internal::ImmutableOptions, - google::bigtable::v2::ReadRowsRequest const&), + google::bigtable::v2::ReadRowsRequest const&, + bigtable_internal::OperationContext&), (override)); MOCK_METHOD(std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::SampleRowKeysResponse>>, @@ -87,14 +97,16 @@ class MockBigtableStub : public bigtable_internal::BigtableStub { (google::cloud::CompletionQueue const&, std::shared_ptr, google::cloud::internal::ImmutableOptions, - google::bigtable::v2::SampleRowKeysRequest const&), + google::bigtable::v2::SampleRowKeysRequest const&, + bigtable_internal::OperationContext&), (override)); MOCK_METHOD(future>, AsyncMutateRow, (google::cloud::CompletionQueue&, std::shared_ptr, google::cloud::internal::ImmutableOptions, - google::bigtable::v2::MutateRowRequest const&), + google::bigtable::v2::MutateRowRequest const&, + bigtable_internal::OperationContext&), (override)); MOCK_METHOD(std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc< google::bigtable::v2::MutateRowsResponse>>, @@ -102,40 +114,46 @@ class MockBigtableStub : public bigtable_internal::BigtableStub { (google::cloud::CompletionQueue const&, std::shared_ptr, google::cloud::internal::ImmutableOptions, - google::bigtable::v2::MutateRowsRequest const&), + google::bigtable::v2::MutateRowsRequest const&, + bigtable_internal::OperationContext&), (override)); MOCK_METHOD(future>, AsyncCheckAndMutateRow, (google::cloud::CompletionQueue&, std::shared_ptr, google::cloud::internal::ImmutableOptions, - google::bigtable::v2::CheckAndMutateRowRequest const&), + google::bigtable::v2::CheckAndMutateRowRequest const&, + bigtable_internal::OperationContext&), (override)); MOCK_METHOD(future>, AsyncPingAndWarm, (google::cloud::CompletionQueue&, std::shared_ptr, google::cloud::internal::ImmutableOptions, - google::bigtable::v2::PingAndWarmRequest const&), + google::bigtable::v2::PingAndWarmRequest const&, + bigtable_internal::OperationContext&), (override)); MOCK_METHOD( future>, AsyncReadModifyWriteRow, (google::cloud::CompletionQueue&, std::shared_ptr, google::cloud::internal::ImmutableOptions, - google::bigtable::v2::ReadModifyWriteRowRequest const&), + google::bigtable::v2::ReadModifyWriteRowRequest const&, + bigtable_internal::OperationContext&), (override)); MOCK_METHOD(future>, AsyncPrepareQuery, (google::cloud::CompletionQueue&, std::shared_ptr, google::cloud::internal::ImmutableOptions, - google::bigtable::v2::PrepareQueryRequest const&), + google::bigtable::v2::PrepareQueryRequest const&, + bigtable_internal::OperationContext&), (override)); MOCK_METHOD(StatusOr, GetClientConfiguration, (grpc::ClientContext&, Options const&, - google::bigtable::v2::GetClientConfigurationRequest const&), + google::bigtable::v2::GetClientConfigurationRequest const&, + bigtable_internal::OperationContext&), (override)); MOCK_METHOD((std::unique_ptr