diff --git a/mooncake-transfer-engine/tent/include/tent/transport/fault_proxy/fault_proxy_transport.h b/mooncake-transfer-engine/tent/include/tent/transport/fault_proxy/fault_proxy_transport.h new file mode 100644 index 00000000..0563b468 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/fault_proxy/fault_proxy_transport.h @@ -0,0 +1,213 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef TENT_FAULT_PROXY_TRANSPORT_H +#define TENT_FAULT_PROXY_TRANSPORT_H + +#include +#include +#include +#include +#include +#include + +#include "tent/common/status.h" +#include "tent/runtime/transport.h" + +namespace mooncake { +namespace tent { + +// Configurable fault injection policy. All fields default to "no fault". +struct FaultPolicy { + // Probability [0.0, 1.0] that submitTransferTasks() returns an error. + double submit_fail_rate = 0.0; + + // Probability [0.0, 1.0] that getTransferStatus() flips COMPLETED→FAILED. + double status_corrupt_rate = 0.0; + + // Artificial latency (microseconds) added before each submit. + uint64_t submit_delay_us = 0; + + // Deterministic mode: succeed for the first N submits, then always fail. + // -1 disables this mode (rate-based mode used instead). + int fail_after_n_submits = -1; + + // If true, install() returns an error immediately. + bool fail_install = false; +}; + +// Decorator that wraps any Transport and injects faults according to a +// FaultPolicy. The engine sees a normal Transport; failures trigger the +// real failover / retry machinery without requiring hardware. +class FaultProxyTransport : public Transport { + public: + FaultProxyTransport(std::shared_ptr real, FaultPolicy policy) + : real_(std::move(real)), policy_(policy), submit_count_(0) { + assert(real_ && "FaultProxyTransport: real transport must not be null"); + } + + // -- Lifecycle ----------------------------------------------------------- + + Status install(std::string& local_segment_name, + std::shared_ptr metadata, + std::shared_ptr local_topology, + std::shared_ptr conf = nullptr) override { + if (policy_.fail_install) { + return Status::InternalError( + "fault injected: install failure" LOC_MARK); + } + return real_->install(local_segment_name, metadata, local_topology, + conf); + } + + Status uninstall() override { return real_->uninstall(); } + + // -- Capabilities -------------------------------------------------------- + + const Capabilities capabilities() const override { + return real_->capabilities(); + } + + // -- Batch management ---------------------------------------------------- + + Status allocateSubBatch(SubBatchRef& batch, size_t max_size) override { + return real_->allocateSubBatch(batch, max_size); + } + + Status freeSubBatch(SubBatchRef& batch) override { + return real_->freeSubBatch(batch); + } + + // -- Transfer (primary injection points) --------------------------------- + + Status submitTransferTasks( + SubBatchRef batch, const std::vector& request_list) override { + // Artificial delay + if (policy_.submit_delay_us > 0) { + std::this_thread::sleep_for( + std::chrono::microseconds(policy_.submit_delay_us)); + } + + int count = submit_count_.fetch_add(1, std::memory_order_relaxed); + + // Deterministic count-based failure (takes precedence over rate-based) + if (policy_.fail_after_n_submits >= 0 && + count >= policy_.fail_after_n_submits) { + return Status::InternalError( + "fault injected: submit failure (count)" LOC_MARK); + } + + // Rate-based probabilistic failure + if (policy_.submit_fail_rate > 0.0 && + randomDouble() < policy_.submit_fail_rate) { + return Status::InternalError( + "fault injected: submit failure (rate)" LOC_MARK); + } + + return real_->submitTransferTasks(batch, request_list); + } + + Status getTransferStatus(SubBatchRef batch, int task_id, + TransferStatus& status) override { + auto s = real_->getTransferStatus(batch, task_id, status); + if (!s.ok()) return s; + + // Status corruption: flip COMPLETED → FAILED + if (status.s == TransferStatusEnum::COMPLETED && + policy_.status_corrupt_rate > 0.0 && + randomDouble() < policy_.status_corrupt_rate) { + status.s = TransferStatusEnum::FAILED; + } + return s; + } + + // -- Memory management (pass-through) ------------------------------------ + + Status addMemoryBuffer(BufferDesc& desc, + const MemoryOptions& options) override { + return real_->addMemoryBuffer(desc, options); + } + + Status addMemoryBuffer(std::vector& desc_list, + const MemoryOptions& options) override { + return real_->addMemoryBuffer(desc_list, options); + } + + Status removeMemoryBuffer(BufferDesc& desc) override { + return real_->removeMemoryBuffer(desc); + } + + Status allocateLocalMemory(void** addr, size_t size, + MemoryOptions& options) override { + return real_->allocateLocalMemory(addr, size, options); + } + + Status freeLocalMemory(void* addr, size_t size) override { + return real_->freeLocalMemory(addr, size); + } + + bool warmupMemory(void* addr, size_t length) override { + return real_->warmupMemory(addr, length); + } + + // -- Notifications (pass-through) ---------------------------------------- + + bool supportNotification() const override { + return real_->supportNotification(); + } + + Status sendNotification(SegmentID target_id, + const Notification& notify) override { + return real_->sendNotification(target_id, notify); + } + + Status receiveNotification( + std::vector& notify_list) override { + return real_->receiveNotification(notify_list); + } + + // -- Identity ------------------------------------------------------------ + + const char* getName() const override { return ""; } + + // -- Test helpers -------------------------------------------------------- + // Note: resetPolicy() is intended for single-threaded test scenarios. + // It is NOT safe to call concurrently with submitTransferTasks(). + + void resetPolicy(FaultPolicy new_policy) { + policy_ = new_policy; + submit_count_.store(0, std::memory_order_relaxed); + } + + int submitCount() const { + return submit_count_.load(std::memory_order_relaxed); + } + + private: + // Thread-safe random double in [0.0, 1.0). + static double randomDouble() { + thread_local std::mt19937 rng(std::random_device{}()); + thread_local std::uniform_real_distribution dist(0.0, 1.0); + return dist(rng); + } + + std::shared_ptr real_; + FaultPolicy policy_; + std::atomic submit_count_; +}; + +} // namespace tent +} // namespace mooncake + +#endif // TENT_FAULT_PROXY_TRANSPORT_H diff --git a/mooncake-transfer-engine/tent/tests/CMakeLists.txt b/mooncake-transfer-engine/tent/tests/CMakeLists.txt index 10188a25..bd13e5ef 100644 --- a/mooncake-transfer-engine/tent/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/tests/CMakeLists.txt @@ -57,3 +57,10 @@ add_test(NAME tent_failover_test COMMAND tent_failover_test) add_executable(tent_endpoint_lifecycle_test endpoint_lifecycle_test.cpp) target_link_libraries(tent_endpoint_lifecycle_test PRIVATE gtest gtest_main) add_test(NAME tent_endpoint_lifecycle_test COMMAND tent_endpoint_lifecycle_test) + +add_executable(tent_fault_proxy_test fault_proxy_test.cpp) +target_link_libraries(tent_fault_proxy_test PRIVATE gtest gtest_main + tent_link_group) +target_include_directories(tent_fault_proxy_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME tent_fault_proxy_test COMMAND tent_fault_proxy_test) diff --git a/mooncake-transfer-engine/tent/tests/fault_proxy_test.cpp b/mooncake-transfer-engine/tent/tests/fault_proxy_test.cpp new file mode 100644 index 00000000..147febe4 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/fault_proxy_test.cpp @@ -0,0 +1,365 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include +#include + +#include "tent/common/types.h" +#include "tent/runtime/transfer_engine_impl.h" +#include "tent/transport/fault_proxy/fault_proxy_transport.h" + +namespace mooncake { +namespace tent { +namespace { + +// --------------------------------------------------------------------------- +// Minimal FakeTransport: always succeeds, tracks call counts. +// --------------------------------------------------------------------------- + +class FakeSubBatch : public Transport::SubBatch { + public: + size_t size() const override { return task_count; } + size_t task_count = 0; + // Store status per task for getTransferStatus + std::vector statuses; +}; + +class FakeTransport : public Transport { + public: + int install_calls = 0; + int submit_calls = 0; + int status_calls = 0; + + Status install(std::string& /*local_segment_name*/, + std::shared_ptr /*metadata*/, + std::shared_ptr /*local_topology*/, + std::shared_ptr /*conf*/ = nullptr) override { + ++install_calls; + return Status::OK(); + } + + Status allocateSubBatch(SubBatchRef& batch, size_t /*max_size*/) override { + batch = new FakeSubBatch(); + return Status::OK(); + } + + Status freeSubBatch(SubBatchRef& batch) override { + delete batch; + batch = nullptr; + return Status::OK(); + } + + Status submitTransferTasks( + SubBatchRef batch, const std::vector& request_list) override { + ++submit_calls; + auto* fb = static_cast(batch); + for (size_t i = 0; i < request_list.size(); ++i) { + fb->statuses.push_back( + {TransferStatusEnum::COMPLETED, request_list[i].length}); + fb->task_count++; + } + return Status::OK(); + } + + Status getTransferStatus(SubBatchRef batch, int task_id, + TransferStatus& status) override { + ++status_calls; + auto* fb = static_cast(batch); + if (task_id < 0 || task_id >= (int)fb->statuses.size()) { + return Status::InvalidArgument("bad task_id" LOC_MARK); + } + status = fb->statuses[task_id]; + return Status::OK(); + } + + const char* getName() const override { return ""; } +}; + +// --------------------------------------------------------------------------- +// Helper: create a Request for testing (no real memory needed) +// --------------------------------------------------------------------------- + +static Request makeRequest(size_t length = 4096) { + Request req; + req.opcode = Request::WRITE; + req.source = nullptr; + req.target_id = 1; + req.target_offset = 0; + req.length = length; + return req; +} + +// =========================================================================== +// Group 1: Unit — FaultPolicy and FaultProxyTransport basics +// =========================================================================== + +TEST(FaultPolicyTest, Defaults) { + FaultPolicy policy; + EXPECT_DOUBLE_EQ(policy.submit_fail_rate, 0.0); + EXPECT_DOUBLE_EQ(policy.status_corrupt_rate, 0.0); + EXPECT_EQ(policy.submit_delay_us, 0u); + EXPECT_EQ(policy.fail_after_n_submits, -1); + EXPECT_FALSE(policy.fail_install); +} + +TEST(FaultProxyTest, DelegatesToReal) { + auto fake = std::make_shared(); + FaultPolicy policy; // all defaults — no faults + auto proxy = std::make_shared(fake, policy); + + // install + std::string seg_name = "test"; + ASSERT_TRUE(proxy->install(seg_name, nullptr, nullptr).ok()); + EXPECT_EQ(fake->install_calls, 1); + + // allocateSubBatch + submitTransferTasks + Transport::SubBatchRef batch = nullptr; + ASSERT_TRUE(proxy->allocateSubBatch(batch, 16).ok()); + ASSERT_NE(batch, nullptr); + + auto req = makeRequest(); + ASSERT_TRUE(proxy->submitTransferTasks(batch, {req}).ok()); + EXPECT_EQ(fake->submit_calls, 1); + EXPECT_EQ(proxy->submitCount(), 1); + + // getTransferStatus + TransferStatus ts{}; + ASSERT_TRUE(proxy->getTransferStatus(batch, 0, ts).ok()); + EXPECT_EQ(ts.s, TransferStatusEnum::COMPLETED); + EXPECT_EQ(ts.transferred_bytes, 4096u); + EXPECT_EQ(fake->status_calls, 1); + + proxy->freeSubBatch(batch); +} + +TEST(FaultProxyTest, FailsInstall) { + auto fake = std::make_shared(); + FaultPolicy policy; + policy.fail_install = true; + auto proxy = std::make_shared(fake, policy); + + std::string seg_name = "test"; + auto status = proxy->install(seg_name, nullptr, nullptr); + EXPECT_FALSE(status.ok()); + EXPECT_EQ(fake->install_calls, 0); // never reached the real transport +} + +TEST(FaultProxyTest, FailsSubmitDeterministic) { + auto fake = std::make_shared(); + FaultPolicy policy; + policy.fail_after_n_submits = 2; // succeed twice, then fail + auto proxy = std::make_shared(fake, policy); + + Transport::SubBatchRef batch = nullptr; + ASSERT_TRUE(proxy->allocateSubBatch(batch, 16).ok()); + auto req = makeRequest(); + + // First two submits succeed + EXPECT_TRUE(proxy->submitTransferTasks(batch, {req}).ok()); + EXPECT_TRUE(proxy->submitTransferTasks(batch, {req}).ok()); + EXPECT_EQ(fake->submit_calls, 2); + + // Third submit fails + auto s = proxy->submitTransferTasks(batch, {req}); + EXPECT_FALSE(s.ok()); + EXPECT_EQ(fake->submit_calls, 2); // real transport not called + + // Fourth also fails + EXPECT_FALSE(proxy->submitTransferTasks(batch, {req}).ok()); + + proxy->freeSubBatch(batch); +} + +TEST(FaultProxyTest, CorruptsStatus) { + auto fake = std::make_shared(); + FaultPolicy policy; + policy.status_corrupt_rate = 1.0; // always corrupt + auto proxy = std::make_shared(fake, policy); + + Transport::SubBatchRef batch = nullptr; + ASSERT_TRUE(proxy->allocateSubBatch(batch, 16).ok()); + auto req = makeRequest(); + ASSERT_TRUE(proxy->submitTransferTasks(batch, {req}).ok()); + + TransferStatus ts{}; + ASSERT_TRUE(proxy->getTransferStatus(batch, 0, ts).ok()); + // Real transport returned COMPLETED, but proxy flipped it to FAILED + EXPECT_EQ(ts.s, TransferStatusEnum::FAILED); + + proxy->freeSubBatch(batch); +} + +// =========================================================================== +// Group 2: Integration — Failover state machine driven by proxy +// =========================================================================== + +TEST(FaultProxyFailoverTest, FailoverFromProxyToReal) { + // Simulate: RDMA (proxied, always fails) → TCP (real, always succeeds) + auto fake_rdma = std::make_shared(); + auto fake_tcp = std::make_shared(); + + FaultPolicy rdma_policy; + rdma_policy.submit_fail_rate = 1.0; // RDMA always fails + auto proxied_rdma = + std::make_shared(fake_rdma, rdma_policy); + + // Allocate sub-batches for both transports + Transport::SubBatchRef rdma_batch = nullptr; + Transport::SubBatchRef tcp_batch = nullptr; + ASSERT_TRUE(proxied_rdma->allocateSubBatch(rdma_batch, 16).ok()); + ASSERT_TRUE(fake_tcp->allocateSubBatch(tcp_batch, 16).ok()); + + // Simulate TransferEngineImpl::submitTransfer + resubmitTransferTask + constexpr int kMaxAttempts = 3; + TaskInfo task; + task.type = RDMA; + task.xport_priority = 0; + task.status = TransferStatusEnum::PENDING; + task.failover_count = 0; + + auto req = makeRequest(); + + // Step 1: Submit on "RDMA" — should fail (proxy injects fault) + auto s = proxied_rdma->submitTransferTasks(rdma_batch, {req}); + EXPECT_FALSE(s.ok()); + + // Step 2: Failover logic (mirrors resubmitTransferTask) + task.status = TransferStatusEnum::FAILED; + ++task.failover_count; + EXPECT_LE(task.failover_count, kMaxAttempts); + task.xport_priority++; + task.type = TCP; + task.status = TransferStatusEnum::PENDING; + + // Step 3: Submit on "TCP" — should succeed + s = fake_tcp->submitTransferTasks(tcp_batch, {req}); + EXPECT_TRUE(s.ok()); + + // Step 4: Verify completion + TransferStatus ts{}; + ASSERT_TRUE(fake_tcp->getTransferStatus(tcp_batch, 0, ts).ok()); + EXPECT_EQ(ts.s, TransferStatusEnum::COMPLETED); + task.status = TransferStatusEnum::COMPLETED; + + // Assertions + EXPECT_EQ(task.type, TCP); + EXPECT_EQ(task.failover_count, 1); + EXPECT_EQ(task.status, TransferStatusEnum::COMPLETED); + EXPECT_EQ(fake_rdma->submit_calls, 0); // proxy intercepted, never hit real + EXPECT_EQ(fake_tcp->submit_calls, 1); + + proxied_rdma->freeSubBatch(rdma_batch); + fake_tcp->freeSubBatch(tcp_batch); +} + +TEST(FaultProxyFailoverTest, ExhaustAllTransports) { + // Both RDMA and TCP are proxied with 100% failure + auto fake_rdma = std::make_shared(); + auto fake_tcp = std::make_shared(); + + FaultPolicy always_fail; + always_fail.submit_fail_rate = 1.0; + + auto proxy_rdma = + std::make_shared(fake_rdma, always_fail); + auto proxy_tcp = + std::make_shared(fake_tcp, always_fail); + + Transport::SubBatchRef rdma_batch = nullptr; + Transport::SubBatchRef tcp_batch = nullptr; + ASSERT_TRUE(proxy_rdma->allocateSubBatch(rdma_batch, 16).ok()); + ASSERT_TRUE(proxy_tcp->allocateSubBatch(tcp_batch, 16).ok()); + + constexpr int kMaxAttempts = 3; + TaskInfo task; + task.type = RDMA; + task.xport_priority = 0; + task.status = TransferStatusEnum::PENDING; + task.failover_count = 0; + + auto req = makeRequest(); + + // Attempt on RDMA — fails + EXPECT_FALSE(proxy_rdma->submitTransferTasks(rdma_batch, {req}).ok()); + task.status = TransferStatusEnum::FAILED; + ++task.failover_count; + task.xport_priority++; + task.type = TCP; + task.status = TransferStatusEnum::PENDING; + + // Attempt on TCP — also fails + EXPECT_FALSE(proxy_tcp->submitTransferTasks(tcp_batch, {req}).ok()); + task.status = TransferStatusEnum::FAILED; + ++task.failover_count; + task.xport_priority++; + task.status = TransferStatusEnum::PENDING; + + // Third attempt — fails again + EXPECT_FALSE(proxy_rdma->submitTransferTasks(rdma_batch, {req}).ok()); + task.status = TransferStatusEnum::FAILED; + ++task.failover_count; + EXPECT_LE(task.failover_count, kMaxAttempts); // Still within limit + + task.xport_priority++; + task.status = TransferStatusEnum::PENDING; + + // Fourth attempt — exceeds limit + EXPECT_FALSE(proxy_tcp->submitTransferTasks(tcp_batch, {req}).ok()); + task.status = TransferStatusEnum::FAILED; + ++task.failover_count; + EXPECT_GT(task.failover_count, kMaxAttempts); + + // Task stays FAILED — no more failover + EXPECT_EQ(task.status, TransferStatusEnum::FAILED); + EXPECT_EQ(task.failover_count, kMaxAttempts + 1); + + proxy_rdma->freeSubBatch(rdma_batch); + proxy_tcp->freeSubBatch(tcp_batch); +} + +// =========================================================================== +// Group 3: Policy mutation +// =========================================================================== + +TEST(FaultProxyTest, ResetPolicyMidRun) { + auto fake = std::make_shared(); + FaultPolicy fail_policy; + fail_policy.submit_fail_rate = 1.0; + auto proxy = std::make_shared(fake, fail_policy); + + Transport::SubBatchRef batch = nullptr; + ASSERT_TRUE(proxy->allocateSubBatch(batch, 16).ok()); + auto req = makeRequest(); + + // Should fail + EXPECT_FALSE(proxy->submitTransferTasks(batch, {req}).ok()); + EXPECT_EQ(fake->submit_calls, 0); + + // Reset to clean policy + FaultPolicy clean_policy; + proxy->resetPolicy(clean_policy); + + // Should succeed now + EXPECT_TRUE(proxy->submitTransferTasks(batch, {req}).ok()); + EXPECT_EQ(fake->submit_calls, 1); + EXPECT_EQ(proxy->submitCount(), 1); // counter was reset too + + proxy->freeSubBatch(batch); +} + +} // namespace +} // namespace tent +} // namespace mooncake