From cf5c0424c473530fe743203c435ff9a3f5355ebd Mon Sep 17 00:00:00 2001 From: ctgg Date: Wed, 4 Nov 2026 01:30:37 +0800 Subject: [PATCH 1/2] [Store] Bound concurrent object uploads during batch offload Upload independent object-storage keys with at most eight workers, then report the successful keys and sizes once after every worker has joined. Keep the existing durable upload fencing and partial-failure semantics. Add bounded-overlap and partial-failure regression tests. Related: MindLab-Research/mint#2000 --- .../distributed_storage_backend.cpp | 81 ++++++++++-- .../tests/object_storage_adapter_test.cpp | 119 +++++++++++++++++- 2 files changed, 184 insertions(+), 16 deletions(-) diff --git a/mooncake-store/src/storage/distributed/distributed_storage_backend.cpp b/mooncake-store/src/storage/distributed/distributed_storage_backend.cpp index 9e9376ca..1b70f2d7 100644 --- a/mooncake-store/src/storage/distributed/distributed_storage_backend.cpp +++ b/mooncake-store/src/storage/distributed/distributed_storage_backend.cpp @@ -1,8 +1,11 @@ #include "storage/distributed/distributed_storage_backend.h" #include +#include #include #include +#include +#include #include "storage/distributed/dfs_global_allocator.h" #include "types.h" @@ -12,6 +15,10 @@ namespace mooncake { namespace { +// Upper bound on concurrent object PUTs in one offload batch. Bounded so a +// large batch cannot open an unbounded number of cloud connections at once. +constexpr size_t kMaxConcurrentObjectUploads = 8; + bool IsDfsDescriptorRangeValid(const DistributedFSDescriptor& desc, const DistributedStorageConfig& config) { if (config.alignment == 0 || desc.object_size == 0 || @@ -129,8 +136,16 @@ tl::expected DistributedStorageBackend::BatchOffload( return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); } - std::vector success_keys; - std::vector success_metas; + struct PreparedObject { + const std::string* key; + std::vector iovs; + size_t total_size = 0; + // Written by exactly one upload worker; read only after every worker + // has been joined. + bool uploaded = false; + }; + std::vector prepared; + prepared.reserve(batch_object.size()); for (const auto& [key, slices] : batch_object) { if (slices.size() > static_cast(std::numeric_limits::max())) { @@ -138,7 +153,6 @@ tl::expected DistributedStorageBackend::BatchOffload( << ": slice count exceeds INT_MAX"; continue; } - const int iovcnt = static_cast(slices.size()); std::vector iovs; iovs.reserve(slices.size()); @@ -157,17 +171,60 @@ tl::expected DistributedStorageBackend::BatchOffload( << ": total slice size overflows size_t"; continue; } - - auto result = object_storage_adapter_->PutV(key, iovs.data(), iovcnt); - if (!result) { - LOG(WARNING) << "Failed to offload key " << key << ": " - << static_cast(result.error()); - continue; + prepared.push_back({&key, std::move(iovs), total_size}); + } + + // Each PutV is an independent, self-fenced upload of one logical key + // (its own admission marker and multipart upload id), and the adapter's + // object-storage client is safe for concurrent requests. Uploading the + // batch serially made every object in it wait for all the others: one + // cloud PUT per object, end to end. The caller still holds the durable + // object lifecycle lock for the whole batch, and completion is still + // reported once below, so Master-visible semantics are unchanged. + std::atomic next_object{0}; + auto upload_worker = [&]() { + for (size_t i = next_object.fetch_add(1); i < prepared.size(); + i = next_object.fetch_add(1)) { + auto& object = prepared[i]; + auto result = object_storage_adapter_->PutV( + *object.key, object.iovs.data(), + static_cast(object.iovs.size())); + if (!result) { + LOG(WARNING) << "Failed to offload key " << *object.key << ": " + << static_cast(result.error()); + continue; + } + object.uploaded = true; + } + }; + const size_t worker_count = + std::min(prepared.size(), kMaxConcurrentObjectUploads); + std::vector workers; + if (worker_count > 1) { + workers.reserve(worker_count - 1); + for (size_t i = 1; i < worker_count; ++i) { + try { + workers.emplace_back(upload_worker); + } catch (const std::system_error& e) { + // Fewer workers only means less overlap; the calling thread + // below drains whatever the started workers do not take. + LOG(WARNING) + << "Offload upload worker not started: " << e.what(); + break; + } } + } + upload_worker(); + for (auto& worker : workers) worker.join(); - success_keys.push_back(key); - success_metas.emplace_back(-1, 0, static_cast(key.size()), - static_cast(total_size), ""); + std::vector success_keys; + std::vector success_metas; + for (const auto& object : prepared) { + if (!object.uploaded) continue; + success_keys.push_back(*object.key); + success_metas.emplace_back(-1, 0, + static_cast(object.key->size()), + static_cast(object.total_size), ""); } if (complete_handler && !success_keys.empty()) { diff --git a/mooncake-store/tests/object_storage_adapter_test.cpp b/mooncake-store/tests/object_storage_adapter_test.cpp index 90669e2b..3c75c63c 100644 --- a/mooncake-store/tests/object_storage_adapter_test.cpp +++ b/mooncake-store/tests/object_storage_adapter_test.cpp @@ -5,14 +5,18 @@ #include #include +#include +#include #include #include #include #include #include #include +#include #include #include +#include #include #include #include @@ -71,20 +75,34 @@ class FakeObjectStorageAdapter : public ObjectStorageAdapter { tl::expected PutV(const std::string& logical_key, const iovec* iov, int iovcnt) override { - ++putv_calls; - if (fail_keys.contains(logical_key)) { - return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); + // BatchOffload uploads objects concurrently; the fake's shared state + // is guarded so tests observe a real adapter's thread safety. + const int in_flight = ++putv_in_flight; + int seen = max_putv_in_flight.load(); + while (in_flight > seen && + !max_putv_in_flight.compare_exchange_weak(seen, in_flight)) { } - + if (putv_delay.count() > 0) std::this_thread::sleep_for(putv_delay); std::string buffer; for (int i = 0; i < iovcnt; ++i) { buffer.append(static_cast(iov[i].iov_base), iov[i].iov_len); } + std::lock_guard lock(putv_mutex); + --putv_in_flight; + ++putv_calls; + if (fail_keys.contains(logical_key)) { + return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); + } objects[logical_key] = std::move(buffer); return {}; } + std::mutex putv_mutex; + std::atomic putv_in_flight{0}; + std::atomic max_putv_in_flight{0}; + std::chrono::milliseconds putv_delay{0}; + tl::expected Get(const std::string& logical_key, void* buf, size_t len) override { ++get_calls; @@ -509,6 +527,99 @@ TEST_F(ObjectStorageAdapterTest, BatchOffloadAllowsPartialFailure) { EXPECT_FALSE(adapter->objects.contains("bad")); } +// mint#1981: a batch used to upload one object at a time, so every object in +// it waited for all the others. Uploads now overlap, bounded per batch. +TEST_F(ObjectStorageAdapterTest, BatchOffloadUploadsObjectsConcurrently) { + FakeObjectStorageAdapter* adapter = nullptr; + auto backend = MakeObjectStorageBackend(adapter); + ASSERT_TRUE(backend->Init()); + adapter->putv_delay = std::chrono::milliseconds(50); + + constexpr int kObjects = 16; + std::vector values; + values.reserve(kObjects); + std::unordered_map> batch; + for (int i = 0; i < kObjects; ++i) { + values.push_back("value-" + std::to_string(i)); + batch["key-" + std::to_string(i)] = { + {values.back().data(), values.back().size()}}; + } + int handler_calls = 0; + std::vector completed_keys; + + auto result = + backend->BatchOffload(batch, [&](const std::vector& keys, + std::vector&) { + ++handler_calls; + completed_keys = keys; + return ErrorCode::OK; + }); + + ASSERT_TRUE(result); + EXPECT_EQ(*result, kObjects); + EXPECT_EQ(adapter->putv_calls, kObjects); + EXPECT_GT(adapter->max_putv_in_flight.load(), 1); + EXPECT_LE(adapter->max_putv_in_flight.load(), 8); + // Completion is still reported once, for the whole batch. + EXPECT_EQ(handler_calls, 1); + EXPECT_EQ(completed_keys.size(), static_cast(kObjects)); + for (int i = 0; i < kObjects; ++i) { + EXPECT_EQ(adapter->objects.at("key-" + std::to_string(i)), + "value-" + std::to_string(i)); + } +} + +TEST_F(ObjectStorageAdapterTest, + BatchOffloadReportsOnlyUploadedKeysWhenSomeConcurrentUploadsFail) { + FakeObjectStorageAdapter* adapter = nullptr; + auto backend = MakeObjectStorageBackend(adapter); + ASSERT_TRUE(backend->Init()); + adapter->putv_delay = std::chrono::milliseconds(10); + + constexpr int kObjects = 12; + std::vector values; + values.reserve(kObjects); + std::unordered_map> batch; + std::unordered_set expected_ok; + for (int i = 0; i < kObjects; ++i) { + const auto key = "key-" + std::to_string(i); + values.push_back("value-" + std::to_string(i)); + batch[key] = {{values.back().data(), values.back().size()}}; + if (i % 4 == 0) { + adapter->fail_keys.insert(key); + } else { + expected_ok.insert(key); + } + } + int handler_calls = 0; + std::unordered_set completed_keys; + std::map completed_sizes; + + auto result = backend->BatchOffload( + batch, [&](const std::vector& keys, + std::vector& metadata) { + ++handler_calls; + for (size_t i = 0; i < keys.size(); ++i) { + completed_keys.insert(keys[i]); + completed_sizes[keys[i]] = metadata[i].data_size; + } + return ErrorCode::OK; + }); + + ASSERT_TRUE(result); + EXPECT_EQ(*result, static_cast(expected_ok.size())); + EXPECT_EQ(handler_calls, 1); + EXPECT_EQ(completed_keys, expected_ok); + // Each reported key carries its own size, not a neighbour's. + for (const auto& key : expected_ok) { + EXPECT_EQ(completed_sizes.at(key), + static_cast(adapter->objects.at(key).size())); + } + for (const auto& key : adapter->fail_keys) { + EXPECT_FALSE(adapter->objects.contains(key)); + } +} + TEST_F(ObjectStorageAdapterTest, BatchOffloadAllowsNullCompleteHandler) { FakeObjectStorageAdapter* adapter = nullptr; auto backend = MakeObjectStorageBackend(adapter); From 599255d39791c0e2a8e76c0febe3f527fb74bc77 Mon Sep 17 00:00:00 2001 From: ctgg Date: Wed, 4 Nov 2026 02:01:22 +0800 Subject: [PATCH 2/2] [Store] Isolate upload exceptions within offload batches Treat throwing adapters as failed uploads so one key cannot terminate the provider worker or promote an incomplete batch. Test mixed standard and non-standard exceptions, successful keys and sizes, and all-failed caller/worker batches. Related: MindLab-Research/mint#2000 --- .../distributed_storage_backend.cpp | 21 ++-- .../tests/object_storage_adapter_test.cpp | 96 +++++++++++++++++++ 2 files changed, 110 insertions(+), 7 deletions(-) diff --git a/mooncake-store/src/storage/distributed/distributed_storage_backend.cpp b/mooncake-store/src/storage/distributed/distributed_storage_backend.cpp index 1b70f2d7..c911666f 100644 --- a/mooncake-store/src/storage/distributed/distributed_storage_backend.cpp +++ b/mooncake-store/src/storage/distributed/distributed_storage_backend.cpp @@ -186,15 +186,22 @@ tl::expected DistributedStorageBackend::BatchOffload( for (size_t i = next_object.fetch_add(1); i < prepared.size(); i = next_object.fetch_add(1)) { auto& object = prepared[i]; - auto result = object_storage_adapter_->PutV( - *object.key, object.iovs.data(), - static_cast(object.iovs.size())); - if (!result) { + try { + auto result = object_storage_adapter_->PutV( + *object.key, object.iovs.data(), + static_cast(object.iovs.size())); + if (!result) { + LOG(WARNING) << "Failed to offload key " << *object.key + << ": " << static_cast(result.error()); + continue; + } + object.uploaded = true; + } catch (...) { + // An exception escaping a worker would terminate the provider. + // Leave this key pending for the caller's normal retry path. LOG(WARNING) << "Failed to offload key " << *object.key << ": " - << static_cast(result.error()); - continue; + << "object storage adapter threw an exception"; } - object.uploaded = true; } }; const size_t worker_count = diff --git a/mooncake-store/tests/object_storage_adapter_test.cpp b/mooncake-store/tests/object_storage_adapter_test.cpp index 3c75c63c..9500c91f 100644 --- a/mooncake-store/tests/object_storage_adapter_test.cpp +++ b/mooncake-store/tests/object_storage_adapter_test.cpp @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -33,6 +34,8 @@ class FakeObjectStorageAdapter : public ObjectStorageAdapter { public: std::map objects; std::unordered_set fail_keys; + std::unordered_set throw_keys; + std::unordered_set unknown_throw_keys; int init_calls = 0; int putv_calls = 0; int get_calls = 0; @@ -91,6 +94,12 @@ class FakeObjectStorageAdapter : public ObjectStorageAdapter { std::lock_guard lock(putv_mutex); --putv_in_flight; ++putv_calls; + if (throw_keys.contains(logical_key)) { + throw std::runtime_error("injected upload exception"); + } + if (unknown_throw_keys.contains(logical_key)) { + throw 42; + } if (fail_keys.contains(logical_key)) { return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); } @@ -620,6 +629,93 @@ TEST_F(ObjectStorageAdapterTest, } } +TEST_F(ObjectStorageAdapterTest, + BatchOffloadContainsUploadExceptionsAndCompletesOtherKeys) { + FakeObjectStorageAdapter* adapter = nullptr; + auto backend = MakeObjectStorageBackend(adapter); + ASSERT_TRUE(backend->Init()); + adapter->putv_delay = std::chrono::milliseconds(10); + + constexpr int kObjects = 16; + std::vector values; + values.reserve(kObjects); + std::unordered_map> batch; + std::map expected; + for (int i = 0; i < kObjects; ++i) { + const auto key = "key-" + std::to_string(i); + values.emplace_back(i + 1, 'x'); + batch[key] = {{values.back().data(), values.back().size()}}; + if (i % 4 == 0) { + adapter->throw_keys.insert(key); + } else if (i % 4 == 1) { + adapter->unknown_throw_keys.insert(key); + } else if (i % 4 == 2) { + adapter->fail_keys.insert(key); + } else { + expected[key] = values.back().size(); + } + } + int handler_calls = 0; + std::map completed; + auto result = backend->BatchOffload( + batch, [&](const std::vector& keys, + std::vector& metadata) { + ++handler_calls; + EXPECT_EQ(adapter->putv_in_flight.load(), 0); + EXPECT_EQ(keys.size(), metadata.size()); + for (size_t i = 0; i < std::min(keys.size(), metadata.size()); + ++i) { + completed[keys[i]] = metadata[i].data_size; + } + return ErrorCode::OK; + }); + + ASSERT_TRUE(result); + EXPECT_EQ(*result, static_cast(expected.size())); + EXPECT_EQ(adapter->putv_calls, kObjects); + EXPECT_EQ(handler_calls, 1); + EXPECT_EQ(completed, expected); + EXPECT_EQ(adapter->objects.size(), expected.size()); + for (const auto& [key, size] : expected) { + EXPECT_EQ(adapter->objects.at(key), std::string(size, 'x')); + } +} + +TEST_F(ObjectStorageAdapterTest, + BatchOffloadSkipsCompletionWhenEveryUploadThrows) { + // Cover both the caller-only path and a batch with multiple workers. + for (int count : {1, 16}) { + FakeObjectStorageAdapter* adapter = nullptr; + auto backend = MakeObjectStorageBackend(adapter); + ASSERT_TRUE(backend->Init()); + std::string value = "value"; + std::unordered_map> batch; + for (int i = 0; i < count; ++i) { + const auto key = "key-" + std::to_string(i); + batch[key] = {{value.data(), value.size()}}; + if (i % 2 == 0) { + adapter->throw_keys.insert(key); + } else { + adapter->unknown_throw_keys.insert(key); + } + } + int handler_calls = 0; + auto result = backend->BatchOffload( + batch, [&](const std::vector&, + std::vector&) { + ++handler_calls; + return ErrorCode::OK; + }); + + ASSERT_TRUE(result); + EXPECT_EQ(*result, 0); + EXPECT_EQ(handler_calls, 0); + EXPECT_EQ(adapter->putv_calls, count); + EXPECT_EQ(adapter->putv_in_flight.load(), 0); + EXPECT_TRUE(adapter->objects.empty()); + } +} + TEST_F(ObjectStorageAdapterTest, BatchOffloadAllowsNullCompleteHandler) { FakeObjectStorageAdapter* adapter = nullptr; auto backend = MakeObjectStorageBackend(adapter);