UCP/CORE: Increment completion count before calling uct_ep_flush

openucx · Sep 2, 2021 · df16fe6 · df16fe6
1 parent 24bd8b1
commit df16fe6
Showing 1 changed file with 10 additions and 5 deletions.
diff --git a/src/ucp/core/ucp_worker.c b/src/ucp/core/ucp_worker.c
@@ -2286,22 +2286,27 @@ 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_OK) {
-        ucs_assert(req->send.state.uct_comp.count == 0);
         ucp_worker_discard_uct_ep_flush_comp(&req->send.state.uct_comp);
-        return UCS_OK;
+        goto out_comp_count_dec;
     } else if (status == UCS_INPROGRESS) {
-        req->send.state.uct_comp.count++;
         return UCS_OK;
     } else if (status == UCS_ERR_NO_RESOURCE) {
-        return UCS_ERR_NO_RESOURCE;
+        goto out_comp_count_dec;
+    } else {
+        status = UCS_OK;
     }
 
     uct_completion_update_status(&req->send.state.uct_comp, status);
     ucp_worker_discard_uct_ep_flush_comp(&req->send.state.uct_comp);
-    return UCS_OK;
+
+out_comp_count_dec:
+    --req->send.state.uct_comp.count;
+    ucs_assert(req->send.state.uct_comp.count == 0);
+    return status;
 }
 
 unsigned ucp_worker_discard_uct_ep_progress(void *arg)