Skip to content

Commit 6afccef

Browse files
authored
imp(bigtable): add round-robin stub to stub factory (#8838)
1 parent 8658400 commit 6afccef

File tree

7 files changed

+272
-27
lines changed

7 files changed

+272
-27
lines changed

generator/generator_config.textproto

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,7 @@ service {
151151
initial_copyright_year: "2022"
152152
omit_client: true
153153
omit_connection: true
154+
omit_stub_factory: true
154155
gen_async_rpcs: [
155156
"CheckAndMutateRow",
156157
"MutateRow",

google/cloud/bigtable/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -318,6 +318,7 @@ if (BUILD_TESTING)
318318
internal/async_bulk_apply_test.cc
319319
internal/async_row_sampler_test.cc
320320
internal/bigtable_round_robin_test.cc
321+
internal/bigtable_stub_factory_test.cc
321322
internal/bulk_mutator_test.cc
322323
internal/connection_refresh_state_test.cc
323324
internal/convert_policies_test.cc

google/cloud/bigtable/bigtable_client_unit_tests.bzl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ bigtable_client_unit_tests = [
4141
"internal/async_bulk_apply_test.cc",
4242
"internal/async_row_sampler_test.cc",
4343
"internal/bigtable_round_robin_test.cc",
44+
"internal/bigtable_stub_factory_test.cc",
4445
"internal/bulk_mutator_test.cc",
4546
"internal/connection_refresh_state_test.cc",
4647
"internal/convert_policies_test.cc",

google/cloud/bigtable/internal/bigtable_stub_factory.cc

Lines changed: 45 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -4,51 +4,69 @@
44
// you may not use this file except in compliance with the License.
55
// You may obtain a copy of the License at
66
//
7-
// https://www.apache.org/licenses/LICENSE-2.0
7+
// https://www.apache.org/licenses/LICENSE-2.0
88
//
99
// Unless required by applicable law or agreed to in writing, software
1010
// distributed under the License is distributed on an "AS IS" BASIS,
1111
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1212
// See the License for the specific language governing permissions and
1313
// limitations under the License.
1414

15-
// Generated by the Codegen C++ plugin.
16-
// If you make any local changes, they will be lost.
17-
// source: google/bigtable/v2/bigtable.proto
18-
1915
#include "google/cloud/bigtable/internal/bigtable_stub_factory.h"
2016
#include "google/cloud/bigtable/internal/bigtable_auth_decorator.h"
2117
#include "google/cloud/bigtable/internal/bigtable_logging_decorator.h"
2218
#include "google/cloud/bigtable/internal/bigtable_metadata_decorator.h"
23-
#include "google/cloud/bigtable/internal/bigtable_stub.h"
19+
#include "google/cloud/bigtable/internal/bigtable_round_robin.h"
2420
#include "google/cloud/common_options.h"
2521
#include "google/cloud/grpc_options.h"
2622
#include "google/cloud/internal/algorithm.h"
23+
#include "google/cloud/internal/unified_grpc_credentials.h"
2724
#include "google/cloud/log.h"
28-
#include "google/cloud/options.h"
29-
#include <google/bigtable/v2/bigtable.grpc.pb.h>
30-
#include <memory>
25+
#include <grpcpp/grpcpp.h>
3126

3227
namespace google {
3328
namespace cloud {
3429
namespace bigtable_internal {
3530
GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
31+
namespace {
3632

37-
std::shared_ptr<BigtableStub> CreateDefaultBigtableStub(
38-
google::cloud::CompletionQueue cq, Options const& options) {
33+
std::shared_ptr<grpc::Channel> CreateGrpcChannel(
34+
google::cloud::internal::GrpcAuthenticationStrategy& auth,
35+
Options const& options, int channel_id) {
36+
auto args = internal::MakeChannelArguments(options);
37+
args.SetInt("grpc.channel_id", channel_id);
38+
return auth.CreateChannel(options.get<EndpointOption>(), std::move(args));
39+
}
40+
41+
} // namespace
42+
43+
std::shared_ptr<BigtableStub> CreateBigtableStubRoundRobin(
44+
Options const& options,
45+
std::function<std::shared_ptr<BigtableStub>(int)> child_factory) {
46+
std::vector<std::shared_ptr<BigtableStub>> children(
47+
(std::max)(1, options.get<GrpcNumChannelsOption>()));
48+
int id = 0;
49+
std::generate(children.begin(), children.end(),
50+
[&id, &child_factory] { return child_factory(id++); });
51+
return std::make_shared<BigtableRoundRobin>(std::move(children));
52+
}
53+
54+
std::shared_ptr<BigtableStub> CreateDecoratedStubs(
55+
google::cloud::CompletionQueue cq, Options const& options,
56+
BaseBigtableStubFactory const& base_factory) {
3957
auto auth = google::cloud::internal::CreateAuthenticationStrategy(
4058
std::move(cq), options);
41-
auto channel = auth->CreateChannel(options.get<EndpointOption>(),
42-
internal::MakeChannelArguments(options));
43-
auto service_grpc_stub = google::bigtable::v2::Bigtable::NewStub(channel);
44-
std::shared_ptr<BigtableStub> stub =
45-
std::make_shared<DefaultBigtableStub>(std::move(service_grpc_stub));
46-
59+
auto child_factory = [base_factory, &auth, options](int id) {
60+
auto channel = CreateGrpcChannel(*auth, options, id);
61+
return base_factory(std::move(channel));
62+
};
63+
auto stub = CreateBigtableStubRoundRobin(options, std::move(child_factory));
4764
if (auth->RequiresConfigureContext()) {
4865
stub = std::make_shared<BigtableAuth>(std::move(auth), std::move(stub));
4966
}
5067
stub = std::make_shared<BigtableMetadata>(std::move(stub));
51-
if (internal::Contains(options.get<TracingComponentsOption>(), "rpc")) {
68+
if (google::cloud::internal::Contains(options.get<TracingComponentsOption>(),
69+
"rpc")) {
5270
GCP_LOG(INFO) << "Enabled logging for gRPC calls";
5371
stub = std::make_shared<BigtableLogging>(
5472
std::move(stub), options.get<GrpcTracingOptionsOption>(),
@@ -57,6 +75,15 @@ std::shared_ptr<BigtableStub> CreateDefaultBigtableStub(
5775
return stub;
5876
}
5977

78+
std::shared_ptr<BigtableStub> CreateBigtableStub(
79+
google::cloud::CompletionQueue cq, Options const& options) {
80+
return CreateDecoratedStubs(
81+
std::move(cq), options, [](std::shared_ptr<grpc::Channel> c) {
82+
return std::make_shared<DefaultBigtableStub>(
83+
google::bigtable::v2::Bigtable::NewStub(std::move(c)));
84+
});
85+
}
86+
6087
GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
6188
} // namespace bigtable_internal
6289
} // namespace cloud

google/cloud/bigtable/internal/bigtable_stub_factory.h

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,34 +4,43 @@
44
// you may not use this file except in compliance with the License.
55
// You may obtain a copy of the License at
66
//
7-
// https://www.apache.org/licenses/LICENSE-2.0
7+
// https://www.apache.org/licenses/LICENSE-2.0
88
//
99
// Unless required by applicable law or agreed to in writing, software
1010
// distributed under the License is distributed on an "AS IS" BASIS,
1111
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1212
// See the License for the specific language governing permissions and
1313
// limitations under the License.
1414

15-
// Generated by the Codegen C++ plugin.
16-
// If you make any local changes, they will be lost.
17-
// source: google/bigtable/v2/bigtable.proto
18-
1915
#ifndef GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_INTERNAL_BIGTABLE_STUB_FACTORY_H
2016
#define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_BIGTABLE_INTERNAL_BIGTABLE_STUB_FACTORY_H
2117

2218
#include "google/cloud/bigtable/internal/bigtable_stub.h"
2319
#include "google/cloud/completion_queue.h"
24-
#include "google/cloud/credentials.h"
25-
#include "google/cloud/internal/unified_grpc_credentials.h"
20+
#include "google/cloud/options.h"
2621
#include "google/cloud/version.h"
22+
#include <functional>
2723
#include <memory>
2824

2925
namespace google {
3026
namespace cloud {
3127
namespace bigtable_internal {
3228
GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
3329

34-
std::shared_ptr<BigtableStub> CreateDefaultBigtableStub(
30+
using BaseBigtableStubFactory = std::function<std::shared_ptr<BigtableStub>(
31+
std::shared_ptr<grpc::Channel>)>;
32+
33+
std::shared_ptr<BigtableStub> CreateBigtableStubRoundRobin(
34+
Options const& options,
35+
std::function<std::shared_ptr<BigtableStub>(int)> child_factory);
36+
37+
/// Used in testing to create decorated mocks.
38+
std::shared_ptr<BigtableStub> CreateDecoratedStubs(
39+
google::cloud::CompletionQueue cq, Options const& options,
40+
BaseBigtableStubFactory const& base_factory);
41+
42+
/// Default function used by `DataConnectionImpl`.
43+
std::shared_ptr<BigtableStub> CreateBigtableStub(
3544
google::cloud::CompletionQueue cq, Options const& options);
3645

3746
GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
Lines changed: 207 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,207 @@
1+
// Copyright 2022 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
#include "google/cloud/bigtable/internal/bigtable_stub_factory.h"
16+
#include "google/cloud/bigtable/options.h"
17+
#include "google/cloud/bigtable/testing/mock_bigtable_stub.h"
18+
#include "google/cloud/common_options.h"
19+
#include "google/cloud/credentials.h"
20+
#include "google/cloud/grpc_options.h"
21+
#include "google/cloud/internal/api_client_header.h"
22+
#include "google/cloud/testing_util/scoped_log.h"
23+
#include "google/cloud/testing_util/status_matchers.h"
24+
#include "google/cloud/testing_util/validate_metadata.h"
25+
#include "absl/memory/memory.h"
26+
#include <gmock/gmock.h>
27+
#include <chrono>
28+
29+
namespace google {
30+
namespace cloud {
31+
namespace bigtable_internal {
32+
GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
33+
namespace {
34+
35+
using ::google::cloud::bigtable::testing::MockBigtableStub;
36+
using ::google::cloud::bigtable::testing::MockReadRowsStream;
37+
using ::google::cloud::testing_util::ScopedLog;
38+
using ::google::cloud::testing_util::StatusIs;
39+
using ::google::cloud::testing_util::ValidateMetadataFixture;
40+
using ::testing::Contains;
41+
using ::testing::HasSubstr;
42+
using ::testing::NotNull;
43+
using ::testing::Return;
44+
45+
// The point of these tests is to verify that the `CreateBigtableStub` factory
46+
// function injects the right decorators. We do this by observing the
47+
// side effects of these decorators. All the tests have nearly identical
48+
// structure. They create a fully decorated stub, configured to round-robin
49+
// over kTestChannel mocks. The first mock expects a call, the remaining mocks
50+
// expect no calls. Some of these side effects can only be verified as part of
51+
// the first mock.
52+
53+
class BigtableStubFactory : public ::testing::Test {
54+
protected:
55+
Status IsContextMDValid(grpc::ClientContext& context,
56+
std::string const& method) {
57+
return validate_metadata_fixture_.IsContextMDValid(
58+
context, method, google::cloud::internal::ApiClientHeader("generator"));
59+
}
60+
61+
private:
62+
ValidateMetadataFixture validate_metadata_fixture_;
63+
};
64+
65+
auto constexpr kTestChannels = 3;
66+
67+
using MockFactory = ::testing::MockFunction<std::shared_ptr<BigtableStub>(
68+
std::shared_ptr<grpc::Channel>)>;
69+
static_assert(std::is_same<decltype(MockFactory{}.AsStdFunction()),
70+
BaseBigtableStubFactory>::value,
71+
"Mismatched mock factory type");
72+
73+
std::shared_ptr<BigtableStub> CreateTestStub(
74+
CompletionQueue cq, BaseBigtableStubFactory const& factory) {
75+
auto credentials = google::cloud::MakeAccessTokenCredentials(
76+
"test-only-invalid",
77+
std::chrono::system_clock::now() + std::chrono::minutes(5));
78+
return CreateDecoratedStubs(std::move(cq),
79+
google::cloud::Options{}
80+
.set<GrpcNumChannelsOption>(kTestChannels)
81+
.set<TracingComponentsOption>({"rpc"})
82+
.set<UnifiedCredentialsOption>(credentials),
83+
factory);
84+
}
85+
86+
TEST_F(BigtableStubFactory, ReadRows) {
87+
::testing::InSequence sequence;
88+
MockFactory factory;
89+
EXPECT_CALL(factory, Call)
90+
.WillOnce([this](std::shared_ptr<grpc::Channel> const&) {
91+
auto mock = std::make_shared<MockBigtableStub>();
92+
EXPECT_CALL(*mock, ReadRows)
93+
.WillOnce([this](std::unique_ptr<grpc::ClientContext> context,
94+
google::bigtable::v2::ReadRowsRequest const&) {
95+
// Verify the Auth decorator is present
96+
EXPECT_THAT(context->credentials(), NotNull());
97+
// Verify the Metadata decorator is present
98+
EXPECT_STATUS_OK(IsContextMDValid(
99+
*context, "google.bigtable.v2.Bigtable.ReadRows"));
100+
auto stream = absl::make_unique<MockReadRowsStream>();
101+
EXPECT_CALL(*stream, Read)
102+
.WillOnce(
103+
Return(Status(StatusCode::kUnavailable, "nothing here")));
104+
return stream;
105+
});
106+
return mock;
107+
});
108+
EXPECT_CALL(factory, Call)
109+
.Times(kTestChannels - 1)
110+
.WillRepeatedly([](std::shared_ptr<grpc::Channel> const&) {
111+
return std::make_shared<MockBigtableStub>();
112+
});
113+
114+
ScopedLog log;
115+
CompletionQueue cq;
116+
google::bigtable::v2::ReadRowsRequest req;
117+
req.set_table_name(
118+
"projects/the-project/instances/the-instance/tables/the-table");
119+
auto stub = CreateTestStub(cq, factory.AsStdFunction());
120+
auto stream = stub->ReadRows(absl::make_unique<grpc::ClientContext>(), req);
121+
auto read = stream->Read();
122+
ASSERT_TRUE(absl::holds_alternative<Status>(read));
123+
EXPECT_THAT(absl::get<Status>(read), StatusIs(StatusCode::kUnavailable));
124+
EXPECT_THAT(log.ExtractLines(), Contains(HasSubstr("ReadRows")));
125+
}
126+
127+
TEST_F(BigtableStubFactory, MutateRow) {
128+
::testing::InSequence sequence;
129+
MockFactory factory;
130+
EXPECT_CALL(factory, Call)
131+
.WillOnce([this](std::shared_ptr<grpc::Channel> const&) {
132+
auto mock = std::make_shared<MockBigtableStub>();
133+
EXPECT_CALL(*mock, MutateRow)
134+
.WillOnce([this](grpc::ClientContext& context,
135+
google::bigtable::v2::MutateRowRequest const&) {
136+
// Verify the Auth decorator is present
137+
EXPECT_THAT(context.credentials(), NotNull());
138+
// Verify the Metadata decorator is present
139+
EXPECT_STATUS_OK(IsContextMDValid(
140+
context, "google.bigtable.v2.Bigtable.MutateRow"));
141+
return StatusOr<google::bigtable::v2::MutateRowResponse>(
142+
Status(StatusCode::kUnavailable, "nothing here"));
143+
});
144+
return mock;
145+
});
146+
EXPECT_CALL(factory, Call)
147+
.Times(kTestChannels - 1)
148+
.WillRepeatedly([](std::shared_ptr<grpc::Channel> const&) {
149+
return std::make_shared<MockBigtableStub>();
150+
});
151+
152+
ScopedLog log;
153+
CompletionQueue cq;
154+
grpc::ClientContext context;
155+
google::bigtable::v2::MutateRowRequest req;
156+
req.set_table_name(
157+
"projects/the-project/instances/the-instance/tables/the-table");
158+
auto stub = CreateTestStub(cq, factory.AsStdFunction());
159+
auto response = stub->MutateRow(context, req);
160+
EXPECT_THAT(response, StatusIs(StatusCode::kUnavailable));
161+
EXPECT_THAT(log.ExtractLines(), Contains(HasSubstr("MutateRow")));
162+
}
163+
164+
TEST_F(BigtableStubFactory, AsyncMutateRow) {
165+
::testing::InSequence sequence;
166+
MockFactory factory;
167+
EXPECT_CALL(factory, Call)
168+
.WillOnce([this](std::shared_ptr<grpc::Channel> const&) {
169+
auto mock = std::make_shared<MockBigtableStub>();
170+
EXPECT_CALL(*mock, AsyncMutateRow)
171+
.WillOnce([this](CompletionQueue&,
172+
std::unique_ptr<grpc::ClientContext> context,
173+
google::bigtable::v2::MutateRowRequest const&) {
174+
// Verify the Auth decorator is present
175+
EXPECT_THAT(context->credentials(), NotNull());
176+
// Verify the Metadata decorator is present
177+
EXPECT_STATUS_OK(IsContextMDValid(
178+
*context, "google.bigtable.v2.Bigtable.MutateRow"));
179+
return make_ready_future<
180+
StatusOr<google::bigtable::v2::MutateRowResponse>>(
181+
Status(StatusCode::kUnavailable, "nothing here"));
182+
});
183+
return mock;
184+
});
185+
EXPECT_CALL(factory, Call)
186+
.Times(kTestChannels - 1)
187+
.WillRepeatedly([](std::shared_ptr<grpc::Channel> const&) {
188+
return std::make_shared<MockBigtableStub>();
189+
});
190+
191+
ScopedLog log;
192+
CompletionQueue cq;
193+
google::bigtable::v2::MutateRowRequest req;
194+
req.set_table_name(
195+
"projects/the-project/instances/the-instance/tables/the-table");
196+
auto stub = CreateTestStub(cq, factory.AsStdFunction());
197+
auto response =
198+
stub->AsyncMutateRow(cq, absl::make_unique<grpc::ClientContext>(), req);
199+
EXPECT_THAT(response.get(), StatusIs(StatusCode::kUnavailable));
200+
EXPECT_THAT(log.ExtractLines(), Contains(HasSubstr("AsyncMutateRow")));
201+
}
202+
203+
} // namespace
204+
GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
205+
} // namespace bigtable_internal
206+
} // namespace cloud
207+
} // namespace google

google/cloud/storage/internal/storage_stub_factory.cc

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,6 @@
2222
#include "google/cloud/internal/algorithm.h"
2323
#include "google/cloud/internal/unified_grpc_credentials.h"
2424
#include "google/cloud/log.h"
25-
#include "absl/strings/str_split.h"
2625
#include <grpcpp/grpcpp.h>
2726

2827
namespace google {

0 commit comments

Comments
 (0)