From cb69bcc18839ce5b069f038cee7eb8012e2b8d71 Mon Sep 17 00:00:00 2001
From: Mikhail Brinskii <mikhailb@nvidia.com>
Date: Wed, 24 Feb 2021 21:52:13 +0200
Subject: [PATCH] UCP/AM: Use correct ep for short send with reply

---
 src/ucp/core/ucp_am.c         |  4 ++--
 test/gtest/ucp/test_ucp_am.cc | 22 ++++++++++++++++++++--
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/src/ucp/core/ucp_am.c b/src/ucp/core/ucp_am.c
index 899015485ed..2e6b7659e58 100644
--- a/src/ucp/core/ucp_am.c
+++ b/src/ucp/core/ucp_am.c
@@ -518,7 +518,6 @@ ucp_am_send_short_reply(ucp_ep_h ep, uint16_t id, uint16_t flags,
                         const void *header, size_t header_length,
                         const void *payload, size_t length)
 {
-    uct_ep_h am_ep = ucp_ep_get_am_uct_ep(ep);
     size_t tx_length;
     ucp_am_hdr_t hdr;
     const void *data;
@@ -557,7 +556,8 @@ ucp_am_send_short_reply(ucp_ep_h ep, uint16_t id, uint16_t flags,
     memcpy(UCS_PTR_BYTE_OFFSET(tx_buffer, sizeof(ucs_ptr_map_key_t)),
            data, tx_length);
 
-    return uct_ep_am_short(am_ep, UCP_AM_ID_SINGLE_REPLY, hdr.u64, tx_buffer,
+    return uct_ep_am_short(ucp_ep_get_am_uct_ep(ep), UCP_AM_ID_SINGLE_REPLY,
+                           hdr.u64, tx_buffer,
                            tx_length + sizeof(ucs_ptr_map_key_t));
 }
 
diff --git a/test/gtest/ucp/test_ucp_am.cc b/test/gtest/ucp/test_ucp_am.cc
index 55030974849..07c6ffee9ad 100644
--- a/test/gtest/ucp/test_ucp_am.cc
+++ b/test/gtest/ucp/test_ucp_am.cc
@@ -415,8 +415,8 @@ class test_ucp_am_nbx : public test_ucp_am_base {
     {
         size_t small_hdr_size = 8;
 
-        test_am_send_recv(size, small_hdr_size, flags);
         test_am_send_recv(size, 0, flags);
+        test_am_send_recv(size, small_hdr_size, flags);
 
         if (max_am_hdr() > small_hdr_size) {
             test_am_send_recv(size, max_am_hdr(), flags);
@@ -772,6 +772,16 @@ class test_ucp_am_nbx_dts : public test_ucp_am_nbx {
                                        UCS_BIT(UCP_DATATYPE_IOV) |
                                        UCS_BIT(UCP_DATATYPE_GENERIC);
 
+    virtual ucp_ep_params_t get_ep_params()
+    {
+        ucp_ep_params_t ep_params = test_ucp_am_nbx::get_ep_params();
+
+        ep_params.field_mask |= UCP_EP_PARAM_FIELD_ERR_HANDLING_MODE;
+        ep_params.err_mode    = static_cast<ucp_err_handling_mode_t>(
+                                                          get_variant_value(2));
+        return ep_params;
+    }
+
     static void get_test_dts(std::vector<ucp_test_variant>& variants)
     {
         /* coverity[overrun-buffer-val] */
@@ -779,12 +789,20 @@ class test_ucp_am_nbx_dts : public test_ucp_am_nbx {
                            dts_bitmap, ucp_datatype_class_names);
     }
 
-    static void get_test_variants(std::vector<ucp_test_variant>& variants)
+    static void get_test_dts_reply(std::vector<ucp_test_variant>& variants)
     {
         add_variant_values(variants, get_test_dts, 0);
         add_variant_values(variants, get_test_dts, UCP_AM_SEND_REPLY, "reply");
     }
 
+    static void get_test_variants(std::vector<ucp_test_variant>& variants)
+    {
+        add_variant_values(variants, get_test_dts_reply,
+                           UCP_ERR_HANDLING_MODE_NONE);
+        add_variant_values(variants, get_test_dts_reply,
+                           UCP_ERR_HANDLING_MODE_PEER, "errh");
+    }
+
     void init()
     {
         test_ucp_am_nbx::init();