Skip to content

Commit 3852865

Browse files
elinorlielinorli
authored
feat(bigtable): support DirectPath in Cloud Bigtable C++ client (#16088)
1 parent d48786e commit 3852865

5 files changed

Lines changed: 185 additions & 31 deletions

File tree

‎google/cloud/bigtable/internal/bigtable_stub_factory.cc‎

Lines changed: 23 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
#include "google/cloud/bigtable/internal/bigtable_round_robin_decorator.h"
2222
#include "google/cloud/bigtable/internal/bigtable_tracing_stub.h"
2323
#include "google/cloud/bigtable/internal/connection_refresh_state.h"
24+
#include "google/cloud/bigtable/internal/defaults.h"
2425
#include "google/cloud/bigtable/options.h"
2526
#include "google/cloud/common_options.h"
2627
#include "google/cloud/grpc_options.h"
@@ -47,17 +48,29 @@ std::shared_ptr<grpc::Channel> CreateGrpcChannel(
4748
return auth.CreateChannel(options.get<EndpointOption>(), std::move(args));
4849
}
4950

51+
std::string CreateFeaturesMetadata(bool is_direct_path) {
52+
google::bigtable::v2::FeatureFlags proto;
53+
proto.set_reverse_scans(true);
54+
proto.set_last_scanned_row_responses(true);
55+
proto.set_mutate_rows_rate_limit(true);
56+
proto.set_mutate_rows_rate_limit2(true);
57+
proto.set_routing_cookie(true);
58+
proto.set_retry_info(true);
59+
if (is_direct_path) {
60+
proto.set_traffic_director_enabled(true);
61+
proto.set_direct_access_requested(true);
62+
}
63+
return internal::UrlsafeBase64Encode(proto.SerializeAsString());
64+
}
65+
5066
std::string FeaturesMetadata() {
51-
static auto const* const kFeatures = new auto([] {
52-
google::bigtable::v2::FeatureFlags proto;
53-
proto.set_reverse_scans(true);
54-
proto.set_last_scanned_row_responses(true);
55-
proto.set_mutate_rows_rate_limit(true);
56-
proto.set_mutate_rows_rate_limit2(true);
57-
proto.set_routing_cookie(true);
58-
proto.set_retry_info(true);
59-
return internal::UrlsafeBase64Encode(proto.SerializeAsString());
60-
}());
67+
if (bigtable::internal::IsDirectPath()) {
68+
static auto const* const kDirectPathFeatures =
69+
new std::string(CreateFeaturesMetadata(true));
70+
return *kDirectPathFeatures;
71+
}
72+
static auto const* const kFeatures =
73+
new std::string(CreateFeaturesMetadata(false));
6174
return *kFeatures;
6275
}
6376

‎google/cloud/bigtable/internal/bigtable_stub_factory_test.cc‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,14 +20,17 @@
2020
#include "google/cloud/grpc_options.h"
2121
#include "google/cloud/internal/api_client_header.h"
2222
#include "google/cloud/internal/async_streaming_read_rpc_impl.h"
23+
#include "google/cloud/internal/base64_transforms.h"
2324
#include "google/cloud/internal/make_status.h"
2425
#include "google/cloud/testing_util/fake_completion_queue_impl.h"
2526
#include "google/cloud/testing_util/mock_grpc_authentication_strategy.h"
2627
#include "google/cloud/testing_util/opentelemetry_matchers.h"
2728
#include "google/cloud/testing_util/scoped_log.h"
29+
#include "google/cloud/testing_util/setenv.h"
2830
#include "google/cloud/testing_util/status_matchers.h"
2931
#include "google/cloud/testing_util/validate_metadata.h"
3032
#include "google/cloud/testing_util/validate_propagator.h"
33+
#include "google/bigtable/v2/feature_flags.pb.h"
3134
#include <gmock/gmock.h>
3235
#include <chrono>
3336
#include <regex>
@@ -285,6 +288,86 @@ TEST(BigtableStubFactory, LoggingDisabled) {
285288
EXPECT_THAT(log.ExtractLines(), Not(Contains(HasSubstr("MutateRow"))));
286289
}
287290

291+
TEST(BigtableStubFactory, FeaturesFlagsCloudDirectPath) {
292+
MockFactory factory;
293+
EXPECT_CALL(factory, Call)
294+
.WillOnce([](std::shared_ptr<grpc::Channel> const&) {
295+
auto mock = std::make_shared<MockBigtableStub>();
296+
EXPECT_CALL(*mock, MutateRow)
297+
.WillOnce([](grpc::ClientContext& context, Options const&,
298+
google::bigtable::v2::MutateRowRequest const&) {
299+
ValidateMetadataFixture fixture;
300+
auto headers = fixture.GetMetadata(context);
301+
auto it = headers.find("bigtable-features");
302+
EXPECT_NE(it, headers.end());
303+
auto decoded = internal::UrlsafeBase64Decode(it->second);
304+
EXPECT_STATUS_OK(decoded);
305+
if (!decoded) return internal::AbortedError("fail to decode");
306+
google::bigtable::v2::FeatureFlags proto;
307+
EXPECT_TRUE(proto.ParseFromArray(
308+
decoded->data(), static_cast<int>(decoded->size())));
309+
EXPECT_TRUE(proto.traffic_director_enabled());
310+
EXPECT_TRUE(proto.direct_access_requested());
311+
return internal::AbortedError("fail");
312+
});
313+
return mock;
314+
});
315+
316+
testing_util::SetEnv("GOOGLE_CLOUD_ENABLE_DIRECT_PATH", "bigtable");
317+
auto auth = MakeStubFactoryMockAuth();
318+
CompletionQueue cq;
319+
auto stub = CreateDecoratedStubs(
320+
std::move(auth), std::move(cq),
321+
Options{}
322+
.set<EndpointOption>("localhost:1")
323+
.set<GrpcNumChannelsOption>(1)
324+
.set<UnifiedCredentialsOption>(MakeInsecureCredentials()),
325+
factory.AsStdFunction());
326+
grpc::ClientContext context;
327+
(void)stub->MutateRow(context, Options{}, {});
328+
testing_util::UnsetEnv("GOOGLE_CLOUD_ENABLE_DIRECT_PATH");
329+
}
330+
331+
TEST(BigtableStubFactory, FeaturesFlagsBigtableDirectPath) {
332+
MockFactory factory;
333+
EXPECT_CALL(factory, Call)
334+
.WillOnce([](std::shared_ptr<grpc::Channel> const&) {
335+
auto mock = std::make_shared<MockBigtableStub>();
336+
EXPECT_CALL(*mock, MutateRow)
337+
.WillOnce([](grpc::ClientContext& context, Options const&,
338+
google::bigtable::v2::MutateRowRequest const&) {
339+
ValidateMetadataFixture fixture;
340+
auto headers = fixture.GetMetadata(context);
341+
auto it = headers.find("bigtable-features");
342+
EXPECT_NE(it, headers.end());
343+
auto decoded = internal::UrlsafeBase64Decode(it->second);
344+
EXPECT_STATUS_OK(decoded);
345+
if (!decoded) return internal::AbortedError("fail to decode");
346+
google::bigtable::v2::FeatureFlags proto;
347+
EXPECT_TRUE(proto.ParseFromArray(
348+
decoded->data(), static_cast<int>(decoded->size())));
349+
EXPECT_TRUE(proto.traffic_director_enabled());
350+
EXPECT_TRUE(proto.direct_access_requested());
351+
return internal::AbortedError("fail");
352+
});
353+
return mock;
354+
});
355+
356+
testing_util::SetEnv("CBT_ENABLE_DIRECTPATH", "true");
357+
auto auth = MakeStubFactoryMockAuth();
358+
CompletionQueue cq;
359+
auto stub = CreateDecoratedStubs(
360+
std::move(auth), std::move(cq),
361+
Options{}
362+
.set<EndpointOption>("localhost:1")
363+
.set<GrpcNumChannelsOption>(1)
364+
.set<UnifiedCredentialsOption>(MakeInsecureCredentials()),
365+
factory.AsStdFunction());
366+
grpc::ClientContext context;
367+
(void)stub->MutateRow(context, Options{}, {});
368+
testing_util::UnsetEnv("CBT_ENABLE_DIRECTPATH");
369+
}
370+
288371
TEST(BigtableStubFactory, FeaturesFlags) {
289372
MockFactory factory;
290373
EXPECT_CALL(factory, Call)

‎google/cloud/bigtable/internal/defaults.cc‎

Lines changed: 16 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,18 @@ int DefaultConnectionPoolSize() {
126126
cpu_count * BIGTABLE_CLIENT_DEFAULT_CHANNELS_PER_CPU);
127127
}
128128

129+
bool IsDirectPath() {
130+
auto const direct_path =
131+
google::cloud::internal::GetEnv("GOOGLE_CLOUD_ENABLE_DIRECT_PATH")
132+
.value_or("");
133+
// Bigtable specific env var for Direct Path support used by all clients.
134+
auto const cbt_direct_path =
135+
google::cloud::internal::GetEnv("CBT_ENABLE_DIRECTPATH").value_or("");
136+
return absl::c_any_of(absl::StrSplit(direct_path, ','),
137+
[](absl::string_view v) { return v == "bigtable"; }) ||
138+
cbt_direct_path == "true";
139+
}
140+
129141
Options HandleUniverseDomain(Options opts) {
130142
if (!opts.has<::google::cloud::bigtable_internal::DataEndpointOption>()) {
131143
auto ep = google::cloud::internal::UniverseDomainEndpoint(
@@ -171,18 +183,11 @@ Options DefaultOptions(Options opts) {
171183
}
172184
}
173185

174-
auto const direct_path =
175-
GetEnv("GOOGLE_CLOUD_ENABLE_DIRECT_PATH").value_or("");
176-
if (absl::c_any_of(absl::StrSplit(direct_path, ','),
177-
[](absl::string_view v) { return v == "bigtable"; })) {
186+
// Set the specific data endpoints if Direct Path is enabled.
187+
if (IsDirectPath()) {
178188
opts.set<::google::cloud::bigtable_internal::DataEndpointOption>(
179-
"google-c2p:///directpath-bigtable.googleapis.com")
180-
.set<AuthorityOption>("directpath-bigtable.googleapis.com");
181-
182-
// When using DirectPath the gRPC library already does load balancing across
183-
// multiple sockets, it makes little sense to perform additional load
184-
// balancing in the client library.
185-
if (!opts.has<GrpcNumChannelsOption>()) opts.set<GrpcNumChannelsOption>(1);
189+
"c2p:///bigtable.googleapis.com")
190+
.set<AuthorityOption>("bigtable.googleapis.com");
186191
}
187192

188193
auto emulator = GetEnv("BIGTABLE_EMULATOR_HOST");

‎google/cloud/bigtable/internal/defaults.h‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,11 @@ namespace internal {
2626

2727
int DefaultConnectionPoolSize();
2828

29+
/**
30+
* Returns true if Direct Path is enabled for Bigtable.
31+
*/
32+
bool IsDirectPath();
33+
2934
/**
3035
* Returns an `Options` with the appropriate defaults for Bigtable.
3136
*

‎google/cloud/bigtable/internal/defaults_test.cc‎

Lines changed: 58 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -497,15 +497,16 @@ TEST(EndpointEnvTest, UserCredentialsOverrideEmulatorEnv) {
497497
typeid(opts.get<GrpcCredentialOption>()));
498498
}
499499

500-
TEST(EndpointEnvTest, DirectPathEnabled) {
500+
TEST(EndpointEnvTest, CloudDirectPathEnabled) {
501501
ScopedEnvironment emulator("BIGTABLE_EMULATOR_HOST", absl::nullopt);
502502
ScopedEnvironment direct_path("GOOGLE_CLOUD_ENABLE_DIRECT_PATH",
503503
"storage,bigtable");
504+
ScopedEnvironment cbt_direct_path("CBT_ENABLE_DIRECTPATH", absl::nullopt);
504505

505506
auto opts = DefaultOptions();
506-
EXPECT_EQ("google-c2p:///directpath-bigtable.googleapis.com",
507+
EXPECT_EQ("c2p:///bigtable.googleapis.com",
507508
opts.get<::google::cloud::bigtable_internal::DataEndpointOption>());
508-
EXPECT_EQ("directpath-bigtable.googleapis.com", opts.get<AuthorityOption>());
509+
EXPECT_EQ("bigtable.googleapis.com", opts.get<AuthorityOption>());
509510
// Admin endpoints are not affected.
510511
EXPECT_EQ(
511512
"bigtableadmin.googleapis.com",
@@ -514,10 +515,31 @@ TEST(EndpointEnvTest, DirectPathEnabled) {
514515
"bigtableadmin.googleapis.com",
515516
opts.get<
516517
::google::cloud::bigtable_internal::InstanceAdminEndpointOption>());
517-
EXPECT_EQ(1, opts.get<GrpcNumChannelsOption>());
518+
EXPECT_EQ(DefaultConnectionPoolSize(), opts.get<GrpcNumChannelsOption>());
518519
}
519520

520-
TEST(EndpointEnvTest, DirectPathNoMatch) {
521+
TEST(EndpointEnvTest, BigtableDirectPathEnabled) {
522+
ScopedEnvironment emulator("BIGTABLE_EMULATOR_HOST", absl::nullopt);
523+
ScopedEnvironment direct_path("GOOGLE_CLOUD_ENABLE_DIRECT_PATH",
524+
absl::nullopt);
525+
ScopedEnvironment cbt_direct_path("CBT_ENABLE_DIRECTPATH", "true");
526+
527+
auto opts = DefaultOptions();
528+
EXPECT_EQ("c2p:///bigtable.googleapis.com",
529+
opts.get<::google::cloud::bigtable_internal::DataEndpointOption>());
530+
EXPECT_EQ("bigtable.googleapis.com", opts.get<AuthorityOption>());
531+
// Admin endpoints are not affected.
532+
EXPECT_EQ(
533+
"bigtableadmin.googleapis.com",
534+
opts.get<::google::cloud::bigtable_internal::AdminEndpointOption>());
535+
EXPECT_EQ(
536+
"bigtableadmin.googleapis.com",
537+
opts.get<
538+
::google::cloud::bigtable_internal::InstanceAdminEndpointOption>());
539+
EXPECT_EQ(DefaultConnectionPoolSize(), opts.get<GrpcNumChannelsOption>());
540+
}
541+
542+
TEST(EndpointEnvTest, CloudDirectPathNoMatch) {
521543
ScopedEnvironment emulator("BIGTABLE_EMULATOR_HOST", absl::nullopt);
522544
ScopedEnvironment direct_path("GOOGLE_CLOUD_ENABLE_DIRECT_PATH",
523545
"bigtable-not,almost-bigtable");
@@ -527,25 +549,51 @@ TEST(EndpointEnvTest, DirectPathNoMatch) {
527549
EXPECT_EQ("bigtable.googleapis.com", opts.get<AuthorityOption>());
528550
}
529551

530-
TEST(EndpointEnvTest, DirectPathOverridesUserEndpoints) {
552+
TEST(EndpointEnvTest, BigtableDirectPathFalse) {
553+
ScopedEnvironment emulator("BIGTABLE_EMULATOR_HOST", absl::nullopt);
554+
ScopedEnvironment cbt_direct_path("CBT_ENABLE_DIRECTPATH", "false");
555+
556+
auto opts = DefaultDataOptions(Options{});
557+
EXPECT_EQ("bigtable.googleapis.com", opts.get<EndpointOption>());
558+
EXPECT_EQ("bigtable.googleapis.com", opts.get<AuthorityOption>());
559+
}
560+
561+
TEST(EndpointEnvTest, CloudDirectPathOverridesUserEndpoints) {
531562
ScopedEnvironment emulator("BIGTABLE_EMULATOR_HOST", absl::nullopt);
532563
ScopedEnvironment direct_path("GOOGLE_CLOUD_ENABLE_DIRECT_PATH", "bigtable");
533564

534565
auto opts = DefaultDataOptions(
535566
Options{}.set<EndpointOption>("ignored").set<AuthorityOption>("ignored"));
536-
EXPECT_EQ("google-c2p:///directpath-bigtable.googleapis.com",
537-
opts.get<EndpointOption>());
538-
EXPECT_EQ("directpath-bigtable.googleapis.com", opts.get<AuthorityOption>());
567+
EXPECT_EQ("c2p:///bigtable.googleapis.com", opts.get<EndpointOption>());
568+
EXPECT_EQ("bigtable.googleapis.com", opts.get<AuthorityOption>());
539569
}
540570

541-
TEST(EndpointEnvTest, EmulatorOverridesDirectPath) {
571+
TEST(EndpointEnvTest, BigtableDirectPathOverridesUserEndpoints) {
572+
ScopedEnvironment emulator("BIGTABLE_EMULATOR_HOST", absl::nullopt);
573+
ScopedEnvironment cbt_direct_path("CBT_ENABLE_DIRECTPATH", "true");
574+
575+
auto opts = DefaultDataOptions(
576+
Options{}.set<EndpointOption>("ignored").set<AuthorityOption>("ignored"));
577+
EXPECT_EQ("c2p:///bigtable.googleapis.com", opts.get<EndpointOption>());
578+
EXPECT_EQ("bigtable.googleapis.com", opts.get<AuthorityOption>());
579+
}
580+
581+
TEST(EndpointEnvTest, EmulatorOverridesCloudDirectPath) {
542582
ScopedEnvironment emulator("BIGTABLE_EMULATOR_HOST", "emulator-host:8000");
543583
ScopedEnvironment direct_path("GOOGLE_CLOUD_ENABLE_DIRECT_PATH", "bigtable");
544584

545585
auto opts = DefaultDataOptions(Options{});
546586
EXPECT_EQ("emulator-host:8000", opts.get<EndpointOption>());
547587
}
548588

589+
TEST(EndpointEnvTest, EmulatorOverridesBigtableDirectPath) {
590+
ScopedEnvironment emulator("BIGTABLE_EMULATOR_HOST", "emulator-host:8000");
591+
ScopedEnvironment cbt_direct_path("CBT_ENABLE_DIRECTPATH", "true");
592+
593+
auto opts = DefaultDataOptions(Options{});
594+
EXPECT_EQ("emulator-host:8000", opts.get<EndpointOption>());
595+
}
596+
549597
TEST(ConnectionRefreshRange, BothUnset) {
550598
ScopedEnvironment emulator("BIGTABLE_EMULATOR_HOST", absl::nullopt);
551599
auto opts = DefaultOptions();

0 commit comments

Comments
 (0)