From d3f9f9c14f1602c5234fe12bbb3c82467bbb715a Mon Sep 17 00:00:00 2001
From: Peter Rudenko <peterr@mellanox.com>
Date: Tue, 14 Sep 2021 18:58:35 +0300
Subject: [PATCH] UCP: Reregister memh on CM switch.

---
 src/ucp/proto/proto_am.inl          |  1 +
 test/gtest/ucp/test_ucp_sockaddr.cc | 37 ++++++++++++++++++++++++++++-
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/src/ucp/proto/proto_am.inl b/src/ucp/proto/proto_am.inl
index 1fc10f46a11..8169961fbd8 100644
--- a/src/ucp/proto/proto_am.inl
+++ b/src/ucp/proto/proto_am.inl
@@ -306,6 +306,7 @@ ucs_status_t ucp_do_am_zcopy_single(uct_pending_req_t *self, uint8_t am_id,
     ucs_status_t status;
 
     req->send.lane = ucp_ep_get_am_lane(ep);
+    ucp_send_request_add_reg_lane(req, req->send.lane);
 
     status = ucp_am_zcopy_common(req, hdr, hdr_size, user_hdr_desc,
                                  user_hdr_size, 0ul, iov, iovcnt, max_iov,
diff --git a/test/gtest/ucp/test_ucp_sockaddr.cc b/test/gtest/ucp/test_ucp_sockaddr.cc
index a6f208974a2..8801c25075c 100644
--- a/test/gtest/ucp/test_ucp_sockaddr.cc
+++ b/test/gtest/ucp/test_ucp_sockaddr.cc
@@ -571,7 +571,7 @@ class test_ucp_sockaddr : public ucp_test {
                               m_test_addr.get_port());
         EXPECT_EQ(m_test_addr, attr.local_sockaddr);
     }
-    
+
     void one_sided_disconnect(entity &e, enum ucp_ep_close_mode mode) {
         void *req           = e.disconnect_nb(0, 0, mode);
         ucs_time_t deadline = ucs::get_deadline();
@@ -1351,6 +1351,41 @@ UCS_TEST_P(test_ucp_sockaddr_different_tl_rsc, unset_devices_and_communicate)
 
 UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_sockaddr_different_tl_rsc, all, "all")
 
+class test_ucp_sockaddr_cm_switch : public test_ucp_sockaddr {
+protected:
+    ucp_rsc_index_t get_num_cms()
+    {
+        const ucp_worker_h worker    = sender().worker();
+        ucp_rsc_index_t num_cm_cmpts = ucp_worker_num_cm_cmpts(worker);
+        ucp_rsc_index_t num_cms      = 0;
+
+        for (ucp_rsc_index_t cm_idx = 0; cm_idx < num_cm_cmpts; ++cm_idx) {
+            if (worker->cms[cm_idx].cm != NULL) {
+                num_cms++;
+            }
+        }
+
+        return num_cms;
+    }
+
+    void check_cm_fallback()
+    {
+        if (get_num_cms() < 2) {
+            UCS_TEST_SKIP_R("No CM for fallback to");
+        }
+    }
+};
+
+UCS_TEST_P(test_ucp_sockaddr_cm_switch,
+           rereg_memory_on_cm_switch,
+           "ZCOPY_THRESH=0", "TLS=rc",
+           "SOCKADDR_TLS_PRIORITY=rdmacm,tcp")
+{
+    check_cm_fallback();
+    listen_and_communicate(false, SEND_DIRECTION_BIDI);
+}
+
+UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_sockaddr_cm_switch, all, "all")
 
 class test_ucp_sockaddr_cm_private_data : public test_ucp_sockaddr {
 protected: