Skip to content

Commit 45f7899

Browse files
laramielcopybara-github
laramiel
authored and
copybara-github
committed
Update AWS example, preferring AWS C libraries to AWS CRT CPP wrapped libraries.
Hook Abseil logging into AWS. PiperOrigin-RevId: 722066318 Change-Id: I9998929b4637554154c3b59910e062b82fb4dacd
1 parent 0ee12fe commit 45f7899

File tree

4 files changed

+293
-15
lines changed

4 files changed

+293
-15
lines changed

‎tensorstore/kvstore/s3/BUILD‎

Lines changed: 32 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -368,16 +368,6 @@ tensorstore_cc_test(
368368
],
369369
)
370370

371-
tensorstore_cc_test(
372-
name = "aws_api_test",
373-
size = "small",
374-
srcs = ["aws_api_test.cc"],
375-
deps = [
376-
"@aws_crt_cpp",
377-
"@com_google_googletest//:gtest_main",
378-
],
379-
)
380-
381371
tensorstore_cc_library(
382372
name = "use_conditional_write",
383373
srcs = ["use_conditional_write.cc"],
@@ -390,11 +380,42 @@ tensorstore_cc_library(
390380
],
391381
)
392382

393-
cc_test(
383+
tensorstore_cc_test(
394384
name = "use_conditional_write_test",
395385
srcs = ["use_conditional_write_test.cc"],
396386
deps = [
397387
":use_conditional_write",
398388
"@com_google_googletest//:gtest_main",
399389
],
400390
)
391+
392+
tensorstore_cc_library(
393+
name = "aws_api",
394+
srcs = [
395+
"aws_api.cc",
396+
],
397+
hdrs = ["aws_api.h"],
398+
deps = [
399+
"//tensorstore/internal/log:verbose_flag",
400+
"@aws_c_auth",
401+
"@aws_c_common",
402+
"@aws_c_io",
403+
"@com_google_absl//absl/base",
404+
"@com_google_absl//absl/base:core_headers",
405+
"@com_google_absl//absl/base:log_severity",
406+
"@com_google_absl//absl/debugging:leak_check",
407+
"@com_google_absl//absl/log:absl_log",
408+
],
409+
)
410+
411+
tensorstore_cc_test(
412+
name = "aws_api_test",
413+
size = "small",
414+
srcs = ["aws_api_test.cc"],
415+
# args = ["--tensorstore_verbose_logging=aws"],
416+
deps = [
417+
":aws_api",
418+
"@aws_c_common",
419+
"@com_google_googletest//:gtest_main",
420+
],
421+
)

‎tensorstore/kvstore/s3/aws_api.cc‎

Lines changed: 203 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,203 @@
1+
// Copyright 2025 The TensorStore Authors
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+
// http://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 "tensorstore/kvstore/s3/aws_api.h"
16+
17+
#include <stdint.h>
18+
19+
#include <cstdarg>
20+
#include <cstdio>
21+
#include <cstring>
22+
#include <memory>
23+
#include <string_view>
24+
25+
#include "absl/base/attributes.h"
26+
#include "absl/base/call_once.h"
27+
#include "absl/base/log_severity.h"
28+
#include "absl/debugging/leak_check.h"
29+
#include "absl/log/absl_log.h"
30+
#include <aws/auth/auth.h>
31+
#include <aws/common/allocator.h>
32+
#include <aws/common/byte_buf.h>
33+
#include <aws/common/common.h>
34+
#include <aws/common/error.h>
35+
#include <aws/common/logging.h>
36+
#include <aws/common/zero.h>
37+
#include <aws/io/channel_bootstrap.h>
38+
#include <aws/io/event_loop.h>
39+
#include <aws/io/host_resolver.h>
40+
#include <aws/io/io.h>
41+
#include <aws/io/tls_channel_handler.h>
42+
#include "tensorstore/internal/log/verbose_flag.h"
43+
44+
namespace tensorstore {
45+
namespace internal_kvstore_s3 {
46+
namespace {
47+
48+
constexpr int kLogBufSize = 2000;
49+
50+
// Hook AWS logging into absl logging.
51+
ABSL_CONST_INIT internal_log::VerboseFlag aws_logging("aws");
52+
53+
int absl_log(struct aws_logger *logger, enum aws_log_level log_level,
54+
aws_log_subject_t subject, const char *format, ...) {
55+
absl::LogSeverity severity = absl::LogSeverity::kInfo;
56+
if (log_level <= AWS_LL_FATAL) {
57+
severity = absl::LogSeverity::kFatal;
58+
} else if (log_level <= AWS_LL_ERROR) {
59+
severity = absl::LogSeverity::kError;
60+
} else if (log_level <= AWS_LL_WARN) {
61+
severity = absl::LogSeverity::kWarning;
62+
}
63+
#ifdef ABSL_MIN_LOG_LEVEL
64+
if (severity < static_cast<absl::LogSeverity>(ABSL_MIN_LOG_LEVEL) &&
65+
severity < absl::LogSeverity::kFatal) {
66+
enabled = false;
67+
}
68+
#endif
69+
70+
// AWS Logging doesn't provide a way to get the filename or line number,
71+
// instead use the aws subject name as the filename and the subject itself as
72+
// the line number.
73+
const char *subject_name = aws_log_subject_name(subject);
74+
bool is_valid_subject =
75+
(subject_name != nullptr && strcmp(subject_name, "Unknown") != 0);
76+
77+
char buffer[kLogBufSize];
78+
char *buf = buffer;
79+
size_t size = sizeof(buffer);
80+
81+
va_list argptr;
82+
va_start(argptr, format);
83+
int n = vsnprintf(buf, size, format, argptr);
84+
va_end(argptr);
85+
if (n > 0 && n < size) {
86+
ABSL_LOG(LEVEL(severity))
87+
.AtLocation(is_valid_subject ? subject_name : "aws_api.cc", subject)
88+
<< std::string_view(buf, n);
89+
}
90+
return AWS_OP_SUCCESS;
91+
};
92+
93+
enum aws_log_level absl_get_log_level(struct aws_logger *logger,
94+
aws_log_subject_t subject) {
95+
uintptr_t lvl = reinterpret_cast<uintptr_t>(logger->p_impl);
96+
if (lvl != 0) {
97+
return static_cast<enum aws_log_level>(lvl - 1);
98+
}
99+
if (!aws_logging) {
100+
return AWS_LL_WARN;
101+
}
102+
// NOTE: AWS logging is quite verbose even at AWS_LL_INFO.
103+
if (aws_logging.Level(1)) {
104+
return aws_logging.Level(2) ? AWS_LL_TRACE : AWS_LL_DEBUG;
105+
}
106+
return AWS_LL_INFO;
107+
}
108+
109+
int absl_set_log_level(struct aws_logger *logger, enum aws_log_level lvl) {
110+
if (lvl == AWS_LL_NONE) {
111+
reinterpret_cast<uintptr_t &>(logger->p_impl) = 0;
112+
} else {
113+
reinterpret_cast<uintptr_t &>(logger->p_impl) =
114+
1 + static_cast<uintptr_t>(lvl);
115+
}
116+
return AWS_OP_SUCCESS;
117+
}
118+
119+
void absl_clean_up(struct aws_logger *logger) { (void)logger; }
120+
121+
static aws_logger_vtable s_absl_vtable = {
122+
/*log =*/absl_log,
123+
/*get_log_level =*/absl_get_log_level,
124+
/*clean_up =*/absl_clean_up,
125+
/*set_log_level =*/absl_set_log_level,
126+
};
127+
128+
static aws_logger s_absl_logger = {
129+
/*vtable =*/&s_absl_vtable,
130+
/*.allocator =*/nullptr,
131+
/*.p_impl =*/0,
132+
};
133+
134+
// AWS apis rely on global initialization; do that here.
135+
absl::once_flag g_init_once;
136+
aws_allocator *g_allocator = nullptr;
137+
struct aws_event_loop_group *g_event_loop_group = nullptr;
138+
139+
void InitAwsGlobals() {
140+
absl::LeakCheckDisabler disabler;
141+
aws_logger_set(&s_absl_logger);
142+
g_allocator = aws_default_allocator();
143+
144+
// Initialize the AWS APIs used.
145+
aws_common_library_init(g_allocator);
146+
aws_io_library_init(g_allocator);
147+
aws_auth_library_init(g_allocator);
148+
149+
/* event loop */
150+
g_event_loop_group =
151+
aws_event_loop_group_new_default(g_allocator, 0, nullptr);
152+
}
153+
154+
} // namespace
155+
156+
AwsApiContext::~AwsApiContext() {
157+
aws_tls_ctx_release(tls_ctx);
158+
aws_client_bootstrap_release(client_bootstrap);
159+
aws_host_resolver_release(resolver);
160+
}
161+
162+
std::shared_ptr<AwsApiContext> GetAwsApiContext() {
163+
absl::call_once(g_init_once, InitAwsGlobals);
164+
165+
/* resolver */
166+
aws_host_resolver_default_options resolver_options;
167+
AWS_ZERO_STRUCT(resolver_options);
168+
resolver_options.el_group = g_event_loop_group;
169+
resolver_options.max_entries = 32; // defaults to 8?
170+
aws_host_resolver *resolver =
171+
aws_host_resolver_new_default(g_allocator, &resolver_options);
172+
173+
/* client bootstrap */
174+
aws_client_bootstrap_options bootstrap_options;
175+
AWS_ZERO_STRUCT(bootstrap_options);
176+
bootstrap_options.event_loop_group = g_event_loop_group;
177+
bootstrap_options.host_resolver = resolver;
178+
auto *client_bootstrap =
179+
aws_client_bootstrap_new(g_allocator, &bootstrap_options);
180+
if (client_bootstrap == nullptr) {
181+
ABSL_LOG(FATAL) << "ERROR initializing client bootstrap: "
182+
<< aws_error_debug_str(aws_last_error());
183+
}
184+
185+
aws_tls_ctx_options tls_options;
186+
aws_tls_ctx_options_init_default_client(&tls_options, g_allocator);
187+
aws_tls_ctx *tls_ctx = aws_tls_client_ctx_new(g_allocator, &tls_options);
188+
aws_tls_ctx_options_clean_up(&tls_options);
189+
if (tls_ctx == nullptr) {
190+
ABSL_LOG(FATAL) << "ERROR initializing TLS context: "
191+
<< aws_error_debug_str(aws_last_error());
192+
}
193+
194+
auto result = std::make_shared<AwsApiContext>();
195+
result->allocator = g_allocator;
196+
result->resolver = resolver;
197+
result->client_bootstrap = client_bootstrap;
198+
result->tls_ctx = tls_ctx;
199+
return result;
200+
}
201+
202+
} // namespace internal_kvstore_s3
203+
} // namespace tensorstore

‎tensorstore/kvstore/s3/aws_api.h‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
// Copyright 2025 The TensorStore Authors
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+
// http://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+
#ifndef TENSORSTORE_KVSTORE_S3_AWS_API_H_
16+
#define TENSORSTORE_KVSTORE_S3_AWS_API_H_
17+
18+
#include <memory>
19+
20+
#include <aws/common/allocator.h>
21+
#include <aws/io/channel_bootstrap.h>
22+
#include <aws/io/host_resolver.h>
23+
#include <aws/io/tls_channel_handler.h>
24+
25+
namespace tensorstore {
26+
namespace internal_kvstore_s3 {
27+
28+
/// Context class for using the AWS apis.
29+
class AwsApiContext {
30+
public:
31+
~AwsApiContext();
32+
33+
aws_allocator *allocator; // global
34+
aws_host_resolver *resolver;
35+
aws_client_bootstrap *client_bootstrap;
36+
aws_tls_ctx *tls_ctx;
37+
};
38+
39+
std::shared_ptr<AwsApiContext> GetAwsApiContext();
40+
41+
} // namespace internal_kvstore_s3
42+
} // namespace tensorstore
43+
44+
#endif // TENSORSTORE_KVSTORE_S3_AWS_API_H_

‎tensorstore/kvstore/s3/aws_api_test.cc‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,11 +12,21 @@
1212
// See the License for the specific language governing permissions and
1313
// limitations under the License.
1414

15-
#include <gmock/gmock.h>
15+
#include "tensorstore/kvstore/s3/aws_api.h"
16+
1617
#include <gtest/gtest.h>
17-
#include <aws/crt/Api.h>
18+
#include <aws/common/logging.h>
19+
20+
using ::tensorstore::internal_kvstore_s3::GetAwsApiContext;
21+
22+
namespace {
1823

1924
TEST(AwsApiTest, Basic) {
20-
::Aws::Crt::ApiHandle api_handle;
21-
EXPECT_THAT(api_handle.GetCrtVersion().major, ::testing::Eq(0));
25+
// This does not validate anything; it just initializes the AWS API and
26+
// verifies that it doesn't crash, and then leaks the library setup.
27+
auto context = GetAwsApiContext();
28+
AWS_LOGF_INFO(AWS_LS_COMMON_GENERAL, "info log call");
29+
AWS_LOGF_WARN(AWS_LS_COMMON_GENERAL, "warn log call");
2230
}
31+
32+
} // namespace

0 commit comments

Comments
 (0)