PML/UCX: improved error processing in MPI_Recv

- improved error processing in MPI_Recv implementation
  of pml UCX
- added error handling for pml_ucx_mrecv call

Signed-off-by: Sergey Oblomov <sergeyo@nvidia.com>
(cherry picked from commit eb9405d53f)

Conflicts:
	ompi/mca/pml/ucx/pml_ucx.c

ompi/mca/pml/ucx/pml_ucx.c

@@ -625,8 +625,7 @@ int mca_pml_ucx_recv(void *buf, size_t count, ompi_datatype_t *datatype, int src
     MCA_COMMON_UCX_PROGRESS_LOOP(ompi_pml_ucx.ucp_worker) {
         status = ucp_request_test(req, &info);
         if (status != UCS_INPROGRESS) {
-            mca_pml_ucx_set_recv_status_safe(mpi_status, status, &info);
-            return OMPI_SUCCESS;
+            return mca_pml_ucx_set_recv_status_safe(mpi_status, status, &info);
         }
     }
 }
@@ -1070,8 +1069,7 @@ int mca_pml_ucx_mrecv(void *buf, size_t count, ompi_datatype_t *datatype,
 
     PML_UCX_MESSAGE_RELEASE(message);
 
-    ompi_request_wait(&req, status);
-    return OMPI_SUCCESS;
+    return ompi_request_wait(&req, status);
 }
 
 int mca_pml_ucx_start(size_t count, ompi_request_t** requests)

ompi/mca/pml/ucx/pml_ucx_request.h

@@ -165,7 +165,7 @@ static inline void mca_pml_ucx_set_send_status(ompi_status_public_t* mpi_status,
     }
 }
 
-static inline void mca_pml_ucx_set_recv_status(ompi_status_public_t* mpi_status,
+static inline int mca_pml_ucx_set_recv_status(ompi_status_public_t* mpi_status,
                                                ucs_status_t ucp_status,
                                                const ucp_tag_recv_info_t *info)
 {
@@ -186,15 +186,23 @@ static inline void mca_pml_ucx_set_recv_status(ompi_status_public_t* mpi_status,
     } else {
         mpi_status->MPI_ERROR = MPI_ERR_INTERN;
     }
+
+    return mpi_status->MPI_ERROR;
 }
 
-static inline void mca_pml_ucx_set_recv_status_safe(ompi_status_public_t* mpi_status,
-                                                    ucs_status_t ucp_status,
-                                                    const ucp_tag_recv_info_t *info)
+static inline int mca_pml_ucx_set_recv_status_safe(ompi_status_public_t* mpi_status,
+                                                   ucs_status_t ucp_status,
+                                                   const ucp_tag_recv_info_t *info)
 {
     if (mpi_status != MPI_STATUS_IGNORE) {
-        mca_pml_ucx_set_recv_status(mpi_status, ucp_status, info);
+        return mca_pml_ucx_set_recv_status(mpi_status, ucp_status, info);
+    } else if (OPAL_LIKELY(ucp_status == UCS_OK) || (ucp_status == UCS_ERR_CANCELED)) {
+        return UCS_OK;
+    } else if (ucp_status == UCS_ERR_MESSAGE_TRUNCATED) {
+        return MPI_ERR_TRUNCATE;
     }
+
+    return MPI_ERR_INTERN;
 }
 
 OBJ_CLASS_DECLARATION(mca_pml_ucx_persistent_request_t);