UCT/IB/UD: update stale dest ep_id and stale acks from before connection reset

AUTHORS

@@ -119,6 +119,7 @@ Zhongkai Zhang <zhzhang@habana.ai>
 Zhu Yanjun <yanjunz@mellanox.com>
 Zihao Zhao <zizhao@nvidia.com>
 lzhang2 <cherry.zhang@intel.com>
+michaelzli <michaelzli@tencent.com>
 
 In addition we would like to acknowledge the following members of UCX community
 for their participation in annual face-to-face meeting, design discussions, and

src/uct/ib/ud/base/ud_ep.c

@@ -728,6 +728,15 @@ uct_ud_ep_process_ack(uct_ud_iface_t *iface, uct_ud_ep_t *ep,
         return;
     }
 
+    /* Ignore stale ACKs for unsent packets (e.g., after endpoint reset).
+     * Valid ACK PSN must be in (acked_psn, psn). */
+    if (ucs_unlikely(UCT_UD_PSN_COMPARE(ep->tx.psn, <=, ack_psn))) {
+        ucs_debug("ep %p: ignoring invalid ack_psn=%u (tx.psn=%u acked_psn=%u)"
+                  " - possibly stale from previous connection",
+                  ep, ack_psn, ep->tx.psn, ep->tx.acked_psn);
+        return;
+    }
+
     ep->tx.acked_psn = ack_psn;
     ucs_assertv(UCT_UD_PSN_COMPARE(ep->tx.acked_psn, <, ep->tx.psn),
                 "ep %p: flags=0x%x acked_psn=%u must be smaller than"
@@ -829,21 +838,30 @@ static void uct_ud_ep_rx_creq(uct_ud_iface_t *iface, uct_ud_neth_t *neth)
         ep->rx.ooo_pkts.head_sn = neth->psn;
         uct_ud_peer_copy(&ep->peer, ucs_unaligned_ptr(&ctl->peer));
         uct_ud_ep_ctl_op_add(iface, ep, UCT_UD_EP_OP_CREP);
-    } else if (ep->dest_ep_id == UCT_UD_EP_NULL_ID) {
-        /* simultaneous CREQ */
+    } else if (uct_ib_unpack_uint24(ctl->conn_req.ep_addr.ep_id) != ep->dest_ep_id) {
+        if (ep->dest_ep_id == UCT_UD_EP_NULL_ID) {
+            /* simultaneous CREQ */
+            ucs_debug("simultaneous CREQ ep=%p"
+                      "(iface=%p conn_sn=%d ep_id=%d, dest_ep_id=%d rx_psn=%u)",
+                      ep, iface, ep->conn_sn, ep->ep_id,
+                      ep->dest_ep_id, ep->rx.ooo_pkts.head_sn);
+            if (UCT_UD_PSN_COMPARE(ep->tx.psn, >, UCT_UD_INITIAL_PSN)) {
+                /* our own creq was sent, treat incoming creq as ack and remove our
+                 * own from tx window
+                 */
+                uct_ud_ep_process_ack(iface, ep, UCT_UD_INITIAL_PSN, 0);
+            }
+        } else {
+            /* stale EP reuse */
+            ucs_debug("iface=%p: detected stale EP reuse (ep=%p conn_sn=%d "
+                      "old_dest_ep_id=%d new_ep_id=%d), updating dest_ep_id",
+                      iface, ep, ep->conn_sn, ep->dest_ep_id,
+                      uct_ib_unpack_uint24(ctl->conn_req.ep_addr.ep_id));
+        }
+        /* Update dest_ep_id */
         uct_ud_ep_set_dest_ep_id(ep, uct_ib_unpack_uint24(ctl->conn_req.ep_addr.ep_id));
         ep->rx.ooo_pkts.head_sn = neth->psn;
         uct_ud_peer_copy(&ep->peer, ucs_unaligned_ptr(&ctl->peer));
-        ucs_debug("simultaneous CREQ ep=%p"
-                  "(iface=%p conn_sn=%d ep_id=%d, dest_ep_id=%d rx_psn=%u)",
-                  ep, iface, ep->conn_sn, ep->ep_id,
-                  ep->dest_ep_id, ep->rx.ooo_pkts.head_sn);
-        if (UCT_UD_PSN_COMPARE(ep->tx.psn, >, UCT_UD_INITIAL_PSN)) {
-            /* our own creq was sent, treat incoming creq as ack and remove our
-             * own from tx window
-             */
-            uct_ud_ep_process_ack(iface, ep, UCT_UD_INITIAL_PSN, 0);
-        }
         uct_ud_ep_ctl_op_add(iface, ep, UCT_UD_EP_OP_CREP);
     }
 

test/gtest/uct/ib/test_ud.cc

@@ -920,6 +920,96 @@ UCS_TEST_SKIP_COND_P(test_ud, ctls_loss,
 
 UCT_INSTANTIATE_UD_TEST_CASE(test_ud)
 
+#if UCT_UD_EP_DEBUG_HOOKS
+class test_ud_stale_ack : public test_ud {
+public:
+    static ucs_status_t inject_stale_ack_psn(uct_ud_ep_t *ep, uct_ud_neth_t *neth)
+    {
+        if (m_stale_ack_psn_to_inject != 0) {
+            neth->ack_psn = m_stale_ack_psn_to_inject;
+        }
+        return UCS_OK;
+    }
+
+    static ucs_status_t capture_rx_ack_psn(uct_ud_ep_t *ep, uct_ud_neth_t *neth)
+    {
+        m_received_ack_psn = neth->ack_psn;
+        return UCS_OK;
+    }
+
+    static void set_stale_ack_psn(uct_ud_psn_t psn) {
+        m_stale_ack_psn_to_inject = psn;
+    }
+
+    static uct_ud_psn_t get_stale_ack_psn() {
+        return m_stale_ack_psn_to_inject;
+    }
+
+    static uct_ud_psn_t get_received_ack_psn() {
+        return m_received_ack_psn;
+    }
+
+private:
+    /* Stale ACK PSN to inject - simulates ACK from before endpoint reset */
+    static volatile uct_ud_psn_t m_stale_ack_psn_to_inject;
+    /* Captured ACK PSN from received packet */
+    static volatile uct_ud_psn_t m_received_ack_psn;
+};
+
+volatile uct_ud_psn_t test_ud_stale_ack::m_stale_ack_psn_to_inject = 0;
+volatile uct_ud_psn_t test_ud_stale_ack::m_received_ack_psn        = 0;
+
+UCS_TEST_SKIP_COND_P(test_ud_stale_ack, stale_ack_after_reset,
+                     !check_caps(UCT_IFACE_FLAG_AM_SHORT)) {
+    constexpr uct_ud_psn_t STALE_ACK_PSN = 135;
+    constexpr uct_ud_psn_t WINDOW_SIZE = 1024;
+
+    disable_async(m_e1);
+    disable_async(m_e2);
+    connect();
+    set_tx_win(m_e1, WINDOW_SIZE);
+    set_tx_win(m_e2, WINDOW_SIZE);
+
+    uct_ud_ep_t *ud_ep1 = ep(m_e1);
+
+    /* Send some data to advance PSN */
+    for (int i = 0; i < 5; i++) {
+        EXPECT_UCS_OK(tx(m_e1));
+    }
+    flush();
+
+    /* Simulate endpoint reset on m_e1:
+     * Reset PSN to low values as if endpoint was just created. */
+    uct_ud_enter(iface(m_e1));
+    ud_ep1->tx.psn       = 3;   /* Next PSN to send */
+    ud_ep1->tx.acked_psn = 0;   /* Last ACKed PSN */
+    ucs_queue_head_init(&ud_ep1->tx.window);  /* Clear TX window */
+    uct_ud_leave(iface(m_e1));
+
+    /* Set up TX hook on m_e2 to inject stale ack_psn in outgoing packets.
+     * This simulates a delayed/stale packet arriving after reset. */
+    set_stale_ack_psn(STALE_ACK_PSN);
+    ep(m_e2)->tx.tx_hook = test_ud_stale_ack::inject_stale_ack_psn;
+    ep(m_e1)->rx.rx_hook = test_ud_stale_ack::capture_rx_ack_psn;
+
+    /* m_e2 sends a packet. Due to the hook, it will contain ack_psn=135. */
+    EXPECT_UCS_OK(tx(m_e2));
+    short_progress_loop();
+
+    /* Verify the hook actually injected the stale ack_psn */
+    EXPECT_EQ(STALE_ACK_PSN, get_received_ack_psn())
+        << "Invalid value after the hook";
+
+    set_stale_ack_psn(0);
+    ep(m_e2)->tx.tx_hook = uct_ud_ep_null_hook;
+    ep(m_e1)->rx.rx_hook = uct_ud_ep_null_hook;
+
+    /* Verify endpoint state wasn't corrupted by the stale ACK */
+    EXPECT_EQ(0, ud_ep1->tx.acked_psn);  /* Should NOT be STALE_ACK_PSN */
+}
+
+UCT_INSTANTIATE_UD_TEST_CASE(test_ud_stale_ack)
+#endif
 
 class test_ud_peer_failure : public ud_base_test {
 public: