Skip to content

Commit a1cf5ef

Browse files
plusplusjiajiaplusplusjiajia
committed
fix(io): use credential providers for S3 refresh
Adapt the refresh policy to the provider API merged in #899. Install initial credentials and the provider together, and keep the provider across later credential updates. Migrate the refresh tests and cover failed initialization and the REST-to-S3 refresh path with real object storage. AI-Model: gpt-6 AI-Contributed/Feature: 40/40 AI-Contributed/UT: 248/248
1 parent 1fdf718 commit a1cf5ef

3 files changed

Lines changed: 206 additions & 82 deletions

File tree

‎src/iceberg/arrow/s3/arrow_s3_file_io.cc‎

Lines changed: 24 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -271,12 +271,9 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials {
271271
return storage_credentials_;
272272
}
273273

274-
void SetCredentialRefresher(StorageCredentialRefresher refresher) override {
275-
std::unique_lock lock(mutex_);
276-
refresher_ = std::move(refresher);
277-
// A refresh in flight was started for the refresher just replaced.
278-
++credential_generation_;
279-
}
274+
Status InitializeStorageCredentials(
275+
const std::vector<StorageCredential>& storage_credentials,
276+
std::shared_ptr<StorageCredentialProvider> provider) override;
280277

281278
SupportsStorageCredentials* AsSupportsStorageCredentials() override { return this; }
282279

@@ -340,17 +337,29 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials {
340337
mutable std::shared_mutex mutex_;
341338
std::vector<StorageCredential> storage_credentials_;
342339
DelegatesByPrefix file_io_by_prefix_;
343-
StorageCredentialRefresher refresher_;
340+
std::shared_ptr<StorageCredentialProvider> provider_;
344341
std::optional<std::chrono::system_clock::time_point> expires_at_;
345342
std::chrono::steady_clock::time_point retry_refresh_at_;
346-
// Bumped whenever the credentials or the refresher change, so a refresh that
347-
// fetched before one of those happened can tell its result is already stale.
343+
// Bumped on each install so a refresh cannot overwrite newer credentials.
348344
uint64_t credential_generation_ = 0;
349345
// Held across a refresh so concurrent operations skip it. Timed, so waiting
350346
// on it is bounded.
351347
std::timed_mutex refresh_mutex_;
352348
};
353349

350+
Status ArrowS3FileIO::InitializeStorageCredentials(
351+
const std::vector<StorageCredential>& storage_credentials,
352+
std::shared_ptr<StorageCredentialProvider> provider) {
353+
ICEBERG_ASSIGN_OR_RAISE(auto delegates, BuildDelegates(storage_credentials));
354+
auto credentials = storage_credentials;
355+
{
356+
std::unique_lock lock(mutex_);
357+
provider_.swap(provider);
358+
InstallCredentials(credentials, delegates);
359+
}
360+
return {};
361+
}
362+
354363
Status ArrowS3FileIO::SetStorageCredentials(
355364
const std::vector<StorageCredential>& storage_credentials) {
356365
ICEBERG_ASSIGN_OR_RAISE(auto delegates, BuildDelegates(storage_credentials));
@@ -436,7 +445,7 @@ void ArrowS3FileIO::MaybeRefreshCredentials() {
436445
{
437446
// Cheap pre-check, so the common case costs one shared lock and no more.
438447
std::shared_lock lock(mutex_);
439-
if (!refresher_ || !RefreshDue()) {
448+
if (!provider_ || !RefreshDue()) {
440449
return;
441450
}
442451
}
@@ -455,24 +464,23 @@ void ArrowS3FileIO::MaybeRefreshCredentials() {
455464
return;
456465
}
457466
}
458-
// Read together: pairing this refresher with a generation bumped by another
459-
// one installed in between would make its result look current.
460-
StorageCredentialRefresher refresher;
467+
// Snapshot the provider and credentials' generation together.
468+
std::shared_ptr<StorageCredentialProvider> provider;
461469
uint64_t generation = 0;
462470
{
463471
// Whoever held the lock may also have just finished, leaving nothing to do.
464472
std::shared_lock lock(mutex_);
465-
if (!refresher_ || !RefreshDue()) {
473+
if (!provider_ || !RefreshDue()) {
466474
return;
467475
}
468-
refresher = refresher_;
476+
provider = provider_;
469477
generation = credential_generation_;
470478
}
471479

472480
// Outside `mutex_`: both are slow and must not block readers.
473481
Status status;
474482
DelegatesByPrefix delegates;
475-
auto refreshed = refresher();
483+
auto refreshed = provider->Load();
476484
if (refreshed.has_value()) {
477485
auto built = BuildDelegates(*refreshed);
478486
if (!built.has_value()) {

‎src/iceberg/test/arrow_s3_file_io_test.cc‎

Lines changed: 130 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -153,6 +153,11 @@ bool HasWarning(const CapturingLogger& logger, std::string_view substring = {})
153153
});
154154
}
155155

156+
class MockStorageCredentialProvider : public StorageCredentialProvider {
157+
public:
158+
MOCK_METHOD((Result<std::vector<StorageCredential>>), Load, (), (override));
159+
};
160+
156161
constexpr auto kOutlastsAShortenedBackoff = std::chrono::milliseconds(1200);
157162

158163
std::vector<StorageCredential> ExpiringCredentials(std::chrono::milliseconds valid_for,
@@ -267,13 +272,16 @@ TEST_F(ArrowS3FileIOTest, RefreshesCredentialsCloseToExpiry) {
267272

268273
const auto refreshed = ExpiringCredentials(std::chrono::hours(1), "refreshed-key");
269274
int refresh_calls = 0;
270-
credentialed->SetCredentialRefresher([&]() -> Result<std::vector<StorageCredential>> {
271-
++refresh_calls;
272-
return refreshed;
273-
});
274-
ASSERT_THAT(credentialed->SetStorageCredentials(
275-
ExpiringCredentials(std::chrono::minutes(1), "expiring-key")),
275+
auto provider = std::make_shared<MockStorageCredentialProvider>();
276+
EXPECT_CALL(*provider, Load())
277+
.WillRepeatedly([&]() -> Result<std::vector<StorageCredential>> {
278+
++refresh_calls;
279+
return refreshed;
280+
});
281+
ASSERT_THAT(credentialed->InitializeStorageCredentials(
282+
ExpiringCredentials(std::chrono::minutes(1), "expiring-key"), provider),
276283
IsOk());
284+
EXPECT_EQ(refresh_calls, 0);
277285

278286
EXPECT_THAT(result.value()->NewInputFile("s3://bucket/key"), IsOk());
279287
EXPECT_EQ(refresh_calls, 1);
@@ -290,13 +298,15 @@ TEST_F(ArrowS3FileIOTest, DoesNotRefreshCredentialsThatAreNotCloseToExpiry) {
290298
ASSERT_NE(credentialed, nullptr);
291299

292300
int refresh_calls = 0;
293-
credentialed->SetCredentialRefresher([&]() -> Result<std::vector<StorageCredential>> {
294-
++refresh_calls;
295-
return std::vector<StorageCredential>{};
296-
});
297-
298-
ASSERT_THAT(credentialed->SetStorageCredentials(
299-
ExpiringCredentials(std::chrono::hours(1), "access-key")),
301+
auto provider = std::make_shared<MockStorageCredentialProvider>();
302+
EXPECT_CALL(*provider, Load())
303+
.WillRepeatedly([&]() -> Result<std::vector<StorageCredential>> {
304+
++refresh_calls;
305+
return std::vector<StorageCredential>{};
306+
});
307+
308+
ASSERT_THAT(credentialed->InitializeStorageCredentials(
309+
ExpiringCredentials(std::chrono::hours(1), "access-key"), provider),
300310
IsOk());
301311
EXPECT_THAT(result.value()->NewInputFile("s3://bucket/key"), IsOk());
302312
EXPECT_EQ(refresh_calls, 0);
@@ -310,19 +320,58 @@ TEST_F(ArrowS3FileIOTest, DoesNotRefreshCredentialsThatAreNotCloseToExpiry) {
310320
EXPECT_EQ(refresh_calls, 0);
311321
}
312322

323+
TEST_F(ArrowS3FileIOTest, KeepsProviderWhenCredentialsAreUpdated) {
324+
ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO({}));
325+
auto* credentialed = io->AsSupportsStorageCredentials();
326+
ASSERT_NE(credentialed, nullptr);
327+
328+
const auto refreshed = ExpiringCredentials(std::chrono::hours(2), "refreshed-key");
329+
auto provider = std::make_shared<MockStorageCredentialProvider>();
330+
EXPECT_CALL(*provider, Load()).WillOnce(::testing::Return(refreshed));
331+
ASSERT_THAT(credentialed->InitializeStorageCredentials(
332+
ExpiringCredentials(std::chrono::hours(1), "initial-key"), provider),
333+
IsOk());
334+
ASSERT_THAT(credentialed->SetStorageCredentials(
335+
ExpiringCredentials(std::chrono::minutes(1), "expiring-key")),
336+
IsOk());
337+
338+
EXPECT_THAT(io->NewInputFile("s3://bucket/key"), IsOk());
339+
EXPECT_EQ(credentialed->credentials(), refreshed);
340+
}
341+
342+
TEST_F(ArrowS3FileIOTest, FailedInitializationDoesNotInstallProvider) {
343+
ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO({}));
344+
auto* credentialed = io->AsSupportsStorageCredentials();
345+
ASSERT_NE(credentialed, nullptr);
346+
347+
const auto initial = ExpiringCredentials(std::chrono::minutes(1), "initial-key");
348+
ASSERT_THAT(credentialed->InitializeStorageCredentials(initial, nullptr), IsOk());
349+
auto invalid = initial;
350+
invalid.front().prefix.clear();
351+
auto provider = std::make_shared<MockStorageCredentialProvider>();
352+
EXPECT_CALL(*provider, Load()).Times(0);
353+
EXPECT_THAT(credentialed->InitializeStorageCredentials(invalid, provider),
354+
IsError(ErrorKind::kValidationFailed));
355+
356+
EXPECT_THAT(io->NewInputFile("s3://bucket/key"), IsOk());
357+
EXPECT_EQ(credentialed->credentials(), initial);
358+
}
359+
313360
TEST_F(ArrowS3FileIOTest, RefreshesOnceWhenOperationsRaceForIt) {
314361
auto result = MakeS3FileIO({});
315362
ASSERT_THAT(result, IsOk());
316363
auto* credentialed = result.value()->AsSupportsStorageCredentials();
317364
ASSERT_NE(credentialed, nullptr);
318365

319366
std::atomic<int> refresh_calls = 0;
320-
credentialed->SetCredentialRefresher([&]() -> Result<std::vector<StorageCredential>> {
321-
++refresh_calls;
322-
return ExpiringCredentials(std::chrono::hours(1), "refreshed-key");
323-
});
324-
ASSERT_THAT(credentialed->SetStorageCredentials(
325-
ExpiringCredentials(std::chrono::minutes(1), "expiring-key")),
367+
auto provider = std::make_shared<MockStorageCredentialProvider>();
368+
EXPECT_CALL(*provider, Load())
369+
.WillRepeatedly([&]() -> Result<std::vector<StorageCredential>> {
370+
++refresh_calls;
371+
return ExpiringCredentials(std::chrono::hours(1), "refreshed-key");
372+
});
373+
ASSERT_THAT(credentialed->InitializeStorageCredentials(
374+
ExpiringCredentials(std::chrono::minutes(1), "expiring-key"), provider),
326375
IsOk());
327376

328377
constexpr int kThreads = 8;
@@ -358,16 +407,18 @@ TEST_F(ArrowS3FileIOTest, RefreshesOnceWhenCredentialsHaveExpired) {
358407
bool release_refresh = false;
359408
std::atomic<int> refresh_calls = 0;
360409

361-
credentialed->SetCredentialRefresher([&]() -> Result<std::vector<StorageCredential>> {
362-
++refresh_calls;
363-
std::unique_lock lock(mutex);
364-
refresh_started = true;
365-
cv.notify_all();
366-
cv.wait(lock, [&] { return release_refresh; });
367-
return ExpiringCredentials(std::chrono::hours(1), "refreshed-key");
368-
});
369-
ASSERT_THAT(credentialed->SetStorageCredentials(
370-
ExpiringCredentials(-std::chrono::minutes(1), "expired-key")),
410+
auto provider = std::make_shared<MockStorageCredentialProvider>();
411+
EXPECT_CALL(*provider, Load())
412+
.WillRepeatedly([&]() -> Result<std::vector<StorageCredential>> {
413+
++refresh_calls;
414+
std::unique_lock lock(mutex);
415+
refresh_started = true;
416+
cv.notify_all();
417+
cv.wait(lock, [&] { return release_refresh; });
418+
return ExpiringCredentials(std::chrono::hours(1), "refreshed-key");
419+
});
420+
ASSERT_THAT(credentialed->InitializeStorageCredentials(
421+
ExpiringCredentials(-std::chrono::minutes(1), "expired-key"), provider),
371422
IsOk());
372423

373424
std::thread winner(
@@ -410,15 +461,17 @@ TEST_F(ArrowS3FileIOTest, RefreshDoesNotUndoCredentialsInstalledWhileItRan) {
410461
bool refresh_started = false;
411462
bool release_refresh = false;
412463

413-
credentialed->SetCredentialRefresher([&]() -> Result<std::vector<StorageCredential>> {
414-
std::unique_lock lock(mutex);
415-
refresh_started = true;
416-
cv.notify_all();
417-
cv.wait(lock, [&] { return release_refresh; });
418-
return ExpiringCredentials(std::chrono::hours(1), "fetched-by-refresh");
419-
});
420-
ASSERT_THAT(credentialed->SetStorageCredentials(
421-
ExpiringCredentials(std::chrono::minutes(1), "expiring-key")),
464+
auto provider = std::make_shared<MockStorageCredentialProvider>();
465+
EXPECT_CALL(*provider, Load())
466+
.WillRepeatedly([&]() -> Result<std::vector<StorageCredential>> {
467+
std::unique_lock lock(mutex);
468+
refresh_started = true;
469+
cv.notify_all();
470+
cv.wait(lock, [&] { return release_refresh; });
471+
return ExpiringCredentials(std::chrono::hours(1), "fetched-by-refresh");
472+
});
473+
ASSERT_THAT(credentialed->InitializeStorageCredentials(
474+
ExpiringCredentials(std::chrono::minutes(1), "expiring-key"), provider),
422475
IsOk());
423476

424477
std::thread operation(
@@ -430,14 +483,15 @@ TEST_F(ArrowS3FileIOTest, RefreshDoesNotUndoCredentialsInstalledWhileItRan) {
430483
std::unique_lock lock(mutex);
431484
cv.wait(lock, [&] { return refresh_started; });
432485
}
433-
ASSERT_THAT(credentialed->SetStorageCredentials(installed), IsOk());
486+
auto install_status = credentialed->SetStorageCredentials(installed);
434487
{
435488
std::lock_guard lock(mutex);
436489
release_refresh = true;
437490
}
438491
cv.notify_all();
439492
operation.join();
440493

494+
ASSERT_THAT(install_status, IsOk());
441495
EXPECT_EQ(credentialed->credentials(), installed);
442496
}
443497

@@ -450,10 +504,12 @@ TEST_F(ArrowS3FileIOTest, RefreshesSessionCredentialsWithoutAUsableExpiry) {
450504
ASSERT_NE(credentialed, nullptr);
451505

452506
int refresh_calls = 0;
453-
credentialed->SetCredentialRefresher([&]() -> Result<std::vector<StorageCredential>> {
454-
++refresh_calls;
455-
return ExpiringCredentials(std::chrono::hours(1), "refreshed-key");
456-
});
507+
auto provider = std::make_shared<MockStorageCredentialProvider>();
508+
EXPECT_CALL(*provider, Load())
509+
.WillRepeatedly([&]() -> Result<std::vector<StorageCredential>> {
510+
++refresh_calls;
511+
return ExpiringCredentials(std::chrono::hours(1), "refreshed-key");
512+
});
457513

458514
auto logger = std::make_shared<CapturingLogger>();
459515
ScopedDefaultLogger scoped(logger);
@@ -464,8 +520,8 @@ TEST_F(ArrowS3FileIOTest, RefreshesSessionCredentialsWithoutAUsableExpiry) {
464520
if (!expiry.empty()) {
465521
config[std::string(S3Properties::kSessionTokenExpiresAtMs)] = std::string(expiry);
466522
}
467-
ASSERT_THAT(credentialed->SetStorageCredentials(
468-
{{.prefix = "s3", .config = std::move(config)}}),
523+
ASSERT_THAT(credentialed->InitializeStorageCredentials(
524+
{{.prefix = "s3", .config = std::move(config)}}, provider),
469525
IsOk());
470526
EXPECT_TRUE(HasWarning(*logger, "session token"));
471527

@@ -486,11 +542,13 @@ TEST_F(ArrowS3FileIOTest, BacksOffWhenReplacementsAlsoLackAnExpiry) {
486542
{std::string(S3Properties::kSecretAccessKey), "secret"},
487543
{std::string(S3Properties::kSessionToken), "token"}}}};
488544
int refresh_calls = 0;
489-
credentialed->SetCredentialRefresher([&]() -> Result<std::vector<StorageCredential>> {
490-
++refresh_calls;
491-
return undatable;
492-
});
493-
ASSERT_THAT(credentialed->SetStorageCredentials(undatable), IsOk());
545+
auto provider = std::make_shared<MockStorageCredentialProvider>();
546+
EXPECT_CALL(*provider, Load())
547+
.WillRepeatedly([&]() -> Result<std::vector<StorageCredential>> {
548+
++refresh_calls;
549+
return undatable;
550+
});
551+
ASSERT_THAT(credentialed->InitializeStorageCredentials(undatable, provider), IsOk());
494552

495553
EXPECT_THAT(result.value()->NewInputFile("s3://bucket/key"), IsOk());
496554
ASSERT_EQ(refresh_calls, 1);
@@ -507,10 +565,12 @@ TEST_F(ArrowS3FileIOTest, IgnoresUnparseableExpiry) {
507565
ASSERT_NE(credentialed, nullptr);
508566

509567
int refresh_calls = 0;
510-
credentialed->SetCredentialRefresher([&]() -> Result<std::vector<StorageCredential>> {
511-
++refresh_calls;
512-
return std::vector<StorageCredential>{};
513-
});
568+
auto provider = std::make_shared<MockStorageCredentialProvider>();
569+
EXPECT_CALL(*provider, Load())
570+
.WillRepeatedly([&]() -> Result<std::vector<StorageCredential>> {
571+
++refresh_calls;
572+
return std::vector<StorageCredential>{};
573+
});
514574

515575
auto logger = std::make_shared<CapturingLogger>();
516576
ScopedDefaultLogger scoped(logger);
@@ -520,7 +580,7 @@ TEST_F(ArrowS3FileIOTest, IgnoresUnparseableExpiry) {
520580
{std::string(S3Properties::kSessionTokenExpiresAtMs), "not-a-number"}};
521581
const std::vector<StorageCredential> credentials = {
522582
{.prefix = "s3", .config = std::move(config)}};
523-
ASSERT_THAT(credentialed->SetStorageCredentials(credentials), IsOk());
583+
ASSERT_THAT(credentialed->InitializeStorageCredentials(credentials, provider), IsOk());
524584
EXPECT_EQ(credentialed->credentials(), credentials);
525585

526586
EXPECT_THAT(result.value()->NewInputFile("s3://bucket/key"), IsOk());
@@ -534,12 +594,14 @@ TEST_F(ArrowS3FileIOTest, BacksOffWhenTheReplacementIsAlsoCloseToExpiry) {
534594
ASSERT_NE(credentialed, nullptr);
535595

536596
int refresh_calls = 0;
537-
credentialed->SetCredentialRefresher([&]() -> Result<std::vector<StorageCredential>> {
538-
++refresh_calls;
539-
return ExpiringCredentials(std::chrono::minutes(1), "short-lived-key");
540-
});
541-
ASSERT_THAT(credentialed->SetStorageCredentials(
542-
ExpiringCredentials(std::chrono::minutes(1), "expiring-key")),
597+
auto provider = std::make_shared<MockStorageCredentialProvider>();
598+
EXPECT_CALL(*provider, Load())
599+
.WillRepeatedly([&]() -> Result<std::vector<StorageCredential>> {
600+
++refresh_calls;
601+
return ExpiringCredentials(std::chrono::minutes(1), "short-lived-key");
602+
});
603+
ASSERT_THAT(credentialed->InitializeStorageCredentials(
604+
ExpiringCredentials(std::chrono::minutes(1), "expiring-key"), provider),
543605
IsOk());
544606

545607
for (int i = 0; i < 3; ++i) {
@@ -555,12 +617,14 @@ TEST_F(ArrowS3FileIOTest, KeepsCredentialsWhenRefreshFails) {
555617
ASSERT_NE(credentialed, nullptr);
556618

557619
int refresh_calls = 0;
558-
credentialed->SetCredentialRefresher([&]() -> Result<std::vector<StorageCredential>> {
559-
++refresh_calls;
560-
return NotFound("catalog unreachable");
561-
});
620+
auto provider = std::make_shared<MockStorageCredentialProvider>();
621+
EXPECT_CALL(*provider, Load())
622+
.WillRepeatedly([&]() -> Result<std::vector<StorageCredential>> {
623+
++refresh_calls;
624+
return NotFound("catalog unreachable");
625+
});
562626
const auto expiring = ExpiringCredentials(std::chrono::minutes(1), "expiring-key");
563-
ASSERT_THAT(credentialed->SetStorageCredentials(expiring), IsOk());
627+
ASSERT_THAT(credentialed->InitializeStorageCredentials(expiring, provider), IsOk());
564628

565629
auto logger = std::make_shared<CapturingLogger>();
566630
ScopedDefaultLogger scoped(logger);

0 commit comments

Comments
 (0)