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<RdmaEndPoint>
- 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 <staryxchen@tencent.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* 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 <staryxchen@tencent.com>

* 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 <staryxchen@tencent.com>

---------

Signed-off-by: staryxchen <staryxchen@tencent.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint.h

@@ -15,6 +15,7 @@
 #ifndef TENT_ENDPOINT_H
 #define TENT_ENDPOINT_H
 
+#include <memory>
 #include <queue>
 #include <unordered_set>
 #include <vector>
@@ -23,7 +24,7 @@
 
 namespace mooncake {
 namespace tent {
-class RdmaEndPoint {
+class RdmaEndPoint : public std::enable_shared_from_this<RdmaEndPoint> {
     struct WrDepthBlock {
         volatile int value;
         uint64_t padding[7];

mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h

@@ -20,6 +20,7 @@
 #include <cstdint>
 #include <iostream>
 #include <list>
+#include <memory>
 #include <mutex>
 #include <new>
 #include <thread>
@@ -63,7 +64,7 @@ struct RdmaSlice {
     int source_dev_id = -1;
     int target_dev_id = -1;
 
-    RdmaEndPoint* ep_weak_ptr = nullptr;
+    std::weak_ptr<RdmaEndPoint> ep_weak_ptr;
     TransferStatusEnum word = TransferStatusEnum::INITIAL;
     int qp_index = 0;
     int retry_count = 0;

mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp

@@ -535,6 +535,7 @@ int RdmaEndPoint::submitSlices(std::vector<RdmaSlice*>& 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<RdmaSlice*>& 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;

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;

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

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)

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 <gtest/gtest.h>
+
+#include <memory>
+
+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<FakeEndPoint> {
+   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<FakeEndPoint>();
+    std::weak_ptr<FakeEndPoint> weak = ep;
+
+    auto locked = weak.lock();
+    ASSERT_NE(locked, nullptr);
+    locked->acknowledge();
+    EXPECT_EQ(ep->acknowledge_calls, 1);
+}
+
+TEST(EndpointLifecycleTest, WeakPtrExpiresAfterRelease) {
+    std::weak_ptr<FakeEndPoint> weak;
+    {
+        auto ep = std::make_shared<FakeEndPoint>();
+        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<FakeEndPoint>();
+    // Simulate what submitSlices() does: shared_from_this() assigned to
+    // weak_ptr
+    std::weak_ptr<FakeEndPoint> 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<FakeEndPoint> ep_weak_ptr;
+};
+
+TEST(EndpointLifecycleTest, SliceAccessWhileEndpointAlive) {
+    auto ep = std::make_shared<FakeEndPoint>();
+    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<FakeEndPoint>();
+        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<FakeEndPoint>();
+    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<FakeEndPoint>();
+    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