From 15d99a300246ea3a236745ededa672d64395744e Mon Sep 17 00:00:00 2001 From: Stary Date: Thu, 16 Apr 2026 20:13:33 +0800 Subject: [PATCH] feat(tent): replace raw RdmaEndPoint* with weak_ptr for endpoint lifecycle safety (#1897) * feat(tent): replace raw RdmaEndPoint* with weak_ptr for lifecycle safety RdmaSlice::ep_weak_ptr was a raw pointer despite its name, creating a dangling-pointer risk when endpoints are evicted from the EndpointStore cache while slices are still in-flight. Changes: - RdmaEndPoint now inherits std::enable_shared_from_this (endpoints are already managed as shared_ptr in FIFOEndpointStore/SIEVEEndpointStore) - RdmaSlice::ep_weak_ptr changed from RdmaEndPoint* to std::weak_ptr - submitSlices() assigns via shared_from_this() instead of raw this - All dereference sites in workers.cpp now call .lock() and gracefully handle nullptr (endpoint already destroyed) by marking slices FAILED - Add endpoint_lifecycle_test with 7 test cases verifying weak_ptr semantics: lock-while-alive, expire-after-release, shared_from_this, slice access patterns, multi-slice, and reset Signed-off-by: staryxchen Co-Authored-By: Claude Opus 4.6 (1M context) * perf(tent): hoist shared_from_this() out of submitSlices loop Move the shared_from_this() call before the loop so we create a single shared_ptr and assign it to each slice's weak_ptr, avoiding N redundant atomic refcount increment/decrement pairs per batch submission. Signed-off-by: staryxchen * fix(tent): decrement inflight_slices when endpoint expired on timeout When a slice times out and its endpoint weak_ptr has already expired, the inflight_slices counter was not decremented, causing the worker to never enter suspension and skewing load balancing decisions. Add fetch_sub(1) in the null-endpoint timeout path to keep the counter consistent with the actual number of in-flight slices. Signed-off-by: staryxchen --------- Signed-off-by: staryxchen Co-authored-by: Claude Opus 4.6 (1M context) --- .../include/tent/transport/rdma/endpoint.h | 3 +- .../tent/include/tent/transport/rdma/slice.h | 3 +- .../tent/src/transport/rdma/endpoint.cpp | 3 +- .../src/transport/rdma/rdma_transport.cpp | 2 +- .../tent/src/transport/rdma/workers.cpp | 21 ++- .../tent/tests/CMakeLists.txt | 4 + .../tent/tests/endpoint_lifecycle_test.cpp | 147 ++++++++++++++++++ 7 files changed, 174 insertions(+), 9 deletions(-) create mode 100644 mooncake-transfer-engine/tent/tests/endpoint_lifecycle_test.cpp diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint.h index 9f0e2292..169d4744 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint.h @@ -15,6 +15,7 @@ #ifndef TENT_ENDPOINT_H #define TENT_ENDPOINT_H +#include #include #include #include @@ -23,7 +24,7 @@ namespace mooncake { namespace tent { -class RdmaEndPoint { +class RdmaEndPoint : public std::enable_shared_from_this { struct WrDepthBlock { volatile int value; uint64_t padding[7]; diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h index 8b815380..c4bfe1fb 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -63,7 +64,7 @@ struct RdmaSlice { int source_dev_id = -1; int target_dev_id = -1; - RdmaEndPoint* ep_weak_ptr = nullptr; + std::weak_ptr ep_weak_ptr; TransferStatusEnum word = TransferStatusEnum::INITIAL; int qp_index = 0; int retry_count = 0; diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp index 6088cd8b..a0326de2 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp @@ -535,6 +535,7 @@ int RdmaEndPoint::submitSlices(std::vector& slice_list, ibv_send_wr* bad_wr = nullptr; int sge_idx = 0; + auto self = shared_from_this(); for (int wr_idx = 0; wr_idx < wr_count; ++wr_idx) { auto current = slice_list[wr_idx]; auto& wr = wr_list[wr_idx]; @@ -544,7 +545,7 @@ int RdmaEndPoint::submitSlices(std::vector& slice_list, sge.length = current->length; sge.lkey = current->source_lkey; } - current->ep_weak_ptr = this; + current->ep_weak_ptr = self; current->qp_index = qp_index; current->failed = false; wr.wr_id = (uint64_t)current; diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp index e5543a80..dad739f4 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp @@ -373,7 +373,7 @@ Status RdmaTransport::submitTransferTasks( slice->length = length; slice->task = &task; slice->retry_count = 0; - slice->ep_weak_ptr = nullptr; + slice->ep_weak_ptr.reset(); slice->word = PENDING; slice->next = nullptr; slice->enqueue_ts = enqueue_ts; diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp index 3298a88a..9191b41b 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp @@ -202,9 +202,9 @@ void Workers::disableEndpoint(RdmaSlice* slice) { auto& rail = worker.rails[desc->machine_id]; rail.markFailed(slice->source_dev_id, slice->target_dev_id); } - if (slice->ep_weak_ptr) { - slice->ep_weak_ptr->acknowledge(slice, FAILED); - slice->ep_weak_ptr->reset(); + if (auto ep = slice->ep_weak_ptr.lock()) { + ep->acknowledge(slice, FAILED); + ep->reset(); } } @@ -294,9 +294,15 @@ void Workers::asyncPollCq() { for (auto& slice : worker.inflight_slice_set) { if (slice->word != PENDING) continue; if (current_ts - slice->enqueue_ts > slice_timeout_ns_) { - auto ep = slice->ep_weak_ptr; + auto ep = slice->ep_weak_ptr.lock(); LOG(WARNING) << "Slice " << slice << " failed: transfer timeout (software)"; + if (!ep) { + updateSliceStatus(slice, TIMEOUT); + slice_to_remove.push_back(slice); + worker.inflight_slices.fetch_sub(1); + continue; + } auto num_slices = ep->acknowledge(slice, TIMEOUT); disableEndpoint(slice); worker.inflight_slices.fetch_sub(num_slices); @@ -315,7 +321,7 @@ void Workers::asyncPollCq() { for (int i = 0; i < nr_poll; ++i) { auto slice = (RdmaSlice*)wc[i].wr_id; worker.inflight_slice_set.erase(slice); - auto ep = slice->ep_weak_ptr; + auto ep = slice->ep_weak_ptr.lock(); double enqueue_lat = (slice->submit_ts - slice->enqueue_ts) / 1000.0; double inflight_lat = (poll_ts - slice->submit_ts) / 1000.0; @@ -325,6 +331,11 @@ void Workers::asyncPollCq() { overall_lat_sec); } if (slice->word != PENDING) continue; + if (!ep) { + updateSliceStatus(slice, FAILED); + num_slices++; + continue; + } if (wc[i].status != IBV_WC_SUCCESS) { if (wc[i].status != IBV_WC_WR_FLUSH_ERR) { // TE handles them automatically diff --git a/mooncake-transfer-engine/tent/tests/CMakeLists.txt b/mooncake-transfer-engine/tent/tests/CMakeLists.txt index 85e35ea0..a09a9f85 100644 --- a/mooncake-transfer-engine/tent/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/tests/CMakeLists.txt @@ -43,3 +43,7 @@ target_link_libraries(tent_failover_test PRIVATE tent gtest gtest_main) target_include_directories(tent_failover_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) 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) diff --git a/mooncake-transfer-engine/tent/tests/endpoint_lifecycle_test.cpp b/mooncake-transfer-engine/tent/tests/endpoint_lifecycle_test.cpp new file mode 100644 index 00000000..06c8adea --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/endpoint_lifecycle_test.cpp @@ -0,0 +1,147 @@ +// 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 + +namespace mooncake { +namespace tent { +namespace { + +// --------------------------------------------------------------------------- +// Minimal stub to validate weak_ptr lifecycle without RDMA dependencies. +// The real RdmaEndPoint inherits enable_shared_from_this; we mirror that +// pattern here so the test proves the ownership model works. +// --------------------------------------------------------------------------- + +class FakeEndPoint : public std::enable_shared_from_this { + public: + int acknowledge_calls = 0; + int reset_calls = 0; + + void acknowledge() { ++acknowledge_calls; } + void reset() { ++reset_calls; } +}; + +// --------------------------------------------------------------------------- +// weak_ptr basic lifecycle +// --------------------------------------------------------------------------- + +TEST(EndpointLifecycleTest, WeakPtrLocksWhileAlive) { + auto ep = std::make_shared(); + std::weak_ptr weak = ep; + + auto locked = weak.lock(); + ASSERT_NE(locked, nullptr); + locked->acknowledge(); + EXPECT_EQ(ep->acknowledge_calls, 1); +} + +TEST(EndpointLifecycleTest, WeakPtrExpiresAfterRelease) { + std::weak_ptr weak; + { + auto ep = std::make_shared(); + weak = ep; + EXPECT_FALSE(weak.expired()); + } // ep destroyed here + EXPECT_TRUE(weak.expired()); + EXPECT_EQ(weak.lock(), nullptr); +} + +TEST(EndpointLifecycleTest, SharedFromThisProducesValidWeakPtr) { + auto ep = std::make_shared(); + // Simulate what submitSlices() does: shared_from_this() assigned to + // weak_ptr + std::weak_ptr weak = ep->shared_from_this(); + + auto locked = weak.lock(); + ASSERT_NE(locked, nullptr); + EXPECT_EQ(locked.get(), ep.get()); +} + +// --------------------------------------------------------------------------- +// Simulate the slice → endpoint dereference pattern used in workers.cpp +// --------------------------------------------------------------------------- + +struct FakeSlice { + std::weak_ptr ep_weak_ptr; +}; + +TEST(EndpointLifecycleTest, SliceAccessWhileEndpointAlive) { + auto ep = std::make_shared(); + FakeSlice slice; + slice.ep_weak_ptr = ep; + + // Simulate workers.cpp completion path + if (auto locked = slice.ep_weak_ptr.lock()) { + locked->acknowledge(); + locked->reset(); + } + EXPECT_EQ(ep->acknowledge_calls, 1); + EXPECT_EQ(ep->reset_calls, 1); +} + +TEST(EndpointLifecycleTest, SliceAccessAfterEndpointEvicted) { + FakeSlice slice; + { + auto ep = std::make_shared(); + slice.ep_weak_ptr = ep; + } // endpoint evicted — shared_ptr destroyed + + // Simulate workers.cpp: lock() returns nullptr, gracefully skip + auto locked = slice.ep_weak_ptr.lock(); + EXPECT_EQ(locked, nullptr); + // No crash — the slice safely detected endpoint destruction +} + +TEST(EndpointLifecycleTest, MultipleSlicesSameEndpoint) { + auto ep = std::make_shared(); + FakeSlice slices[3]; + for (auto& s : slices) s.ep_weak_ptr = ep; + + // All slices can lock while endpoint alive + for (auto& s : slices) { + auto locked = s.ep_weak_ptr.lock(); + ASSERT_NE(locked, nullptr); + locked->acknowledge(); + } + EXPECT_EQ(ep->acknowledge_calls, 3); + + // Simulate eviction: drop the owning shared_ptr + ep.reset(); + + // All slices now get nullptr + for (auto& s : slices) { + EXPECT_EQ(s.ep_weak_ptr.lock(), nullptr); + } +} + +TEST(EndpointLifecycleTest, WeakPtrResetClearsReference) { + auto ep = std::make_shared(); + FakeSlice slice; + slice.ep_weak_ptr = ep; + + // Simulate rdma_transport.cpp slice initialization: reset() + slice.ep_weak_ptr.reset(); + EXPECT_TRUE(slice.ep_weak_ptr.expired()); + EXPECT_EQ(slice.ep_weak_ptr.lock(), nullptr); + + // Original endpoint still alive + EXPECT_NE(ep, nullptr); +} + +} // namespace +} // namespace tent +} // namespace mooncake