UCT: check completion counter and update status

Signed-off-by: Changcheng Liu <jerrliu@nvidia.com>

src/ucp/core/ucp_worker.c

@@ -2334,15 +2334,24 @@ ucp_worker_discard_uct_ep_pending_cb(uct_pending_req_t *self)
     uct_ep_h uct_ep    = req->send.discard_uct_ep.uct_ep;
     ucs_status_t status;
 
+    ++req->send.state.uct_comp.count;
     status = uct_ep_flush(uct_ep, req->send.discard_uct_ep.ep_flush_flags,
                           &req->send.state.uct_comp);
     if (status == UCS_INPROGRESS) {
         return UCS_OK;
-    } else if (status == UCS_ERR_NO_RESOURCE) {
+    }
+
+    --req->send.state.uct_comp.count;
+    ucs_assert(req->send.state.uct_comp.count == 0);
+
+    if (status == UCS_ERR_NO_RESOURCE) {
         return UCS_ERR_NO_RESOURCE;
     }
 
-    /* UCS_OK, handle flush completion here */
+    if (status != UCS_OK && (req->send.state.uct_comp.status == UCS_OK)) {
+        /* store first failure status */
+        req->send.state.uct_comp.status = status;
+    }
     ucp_worker_discard_uct_ep_flush_comp(&req->send.state.uct_comp, status);
     return UCS_OK;
 }
@@ -2403,7 +2412,8 @@ ucp_worker_discard_tl_uct_ep(ucp_worker_h worker, uct_ep_h uct_ep,
     ucs_assert(!ucp_wireup_ep_test(uct_ep));
     req->send.uct.func                      = ucp_worker_discard_uct_ep_pending_cb;
     req->send.state.uct_comp.func           = ucp_worker_discard_uct_ep_flush_comp;
-    req->send.state.uct_comp.count          = 1;
+    req->send.state.uct_comp.count          = 0;
+    req->send.state.uct_comp.status         = UCS_OK;
     req->send.discard_uct_ep.ucp_worker     = worker;
     req->send.discard_uct_ep.uct_ep         = uct_ep;
     req->send.discard_uct_ep.ep_flush_flags = ep_flush_flags;