From 6f368da036efb41323a1d7316bb61b9db36130d0 Mon Sep 17 00:00:00 2001 From: Stephen Oost Date: Thu, 29 Sep 2022 15:09:17 -0700 Subject: [PATCH] prov/net: Enhance FI_PEEK to support FI_DISCARD Signed-off-by: Stephen Oost --- prov/net/src/xnet.h | 3 +++ prov/net/src/xnet_srx.c | 29 ++++++++++++++++++++++------- 2 files changed, 25 insertions(+), 7 deletions(-) diff --git a/prov/net/src/xnet.h b/prov/net/src/xnet.h index c8ff734232d..95c9522f10a 100644 --- a/prov/net/src/xnet.h +++ b/prov/net/src/xnet.h @@ -352,6 +352,7 @@ static inline void xnet_signal_progress(struct xnet_progress *progress) #define XNET_NEED_DYN_RBUF BIT(4) #define XNET_ASYNC BIT(5) #define XNET_INJECT_OP BIT(6) +#define XNET_FREE_BUF BIT(7) #define XNET_MULTI_RECV FI_MULTI_RECV /* BIT(16) */ struct xnet_xfer_entry { @@ -544,6 +545,8 @@ static inline void xnet_free_xfer(struct xnet_ep *ep, struct xnet_xfer_entry *xfer) { assert(xnet_progress_locked(xnet_ep2_progress(ep))); + if (xfer->ctrl_flags & XNET_FREE_BUF) + free(xfer->iov[0].iov_base); xfer->hdr.base_hdr.flags = 0; xfer->cq_flags = 0; xfer->cntr_inc = NULL; diff --git a/prov/net/src/xnet_srx.c b/prov/net/src/xnet_srx.c index 674b5c9d285..d22fccf1fbd 100644 --- a/prov/net/src/xnet_srx.c +++ b/prov/net/src/xnet_srx.c @@ -269,6 +269,16 @@ xnet_srx_claim(struct xnet_srx *srx, struct xnet_xfer_entry *recv_entry, if (!ep) return -FI_ENOMSG; + if (flags & FI_DISCARD) { + msg_len = ep->cur_rx.hdr.base_hdr.size - + ep->cur_rx.hdr.base_hdr.hdr_size; + recv_entry->iov[0].iov_base = calloc(1, msg_len); + if (!recv_entry->iov[0].iov_base) + return -FI_ENOMEM; + recv_entry->iov[0].iov_len = msg_len; + recv_entry->ctrl_flags |= XNET_FREE_BUF; + } + ret = xnet_start_recv(ep, recv_entry); if (ret && !OFI_SOCK_TRY_SND_RCV_AGAIN(-ret)) xnet_ep_disable(ep, 0, NULL, 0); @@ -278,12 +288,13 @@ xnet_srx_claim(struct xnet_srx *srx, struct xnet_xfer_entry *recv_entry, return FI_SUCCESS; } -static void +static ssize_t xnet_srx_peek(struct xnet_srx *srx, struct xnet_xfer_entry *recv_entry, uint64_t flags) { struct xnet_ep *ep; struct fi_cq_err_entry err_entry; + ssize_t ret = FI_ENOMSG; assert(xnet_progress_locked(xnet_srx2_progress(srx))); assert(srx->rdm); @@ -296,7 +307,7 @@ xnet_srx_peek(struct xnet_srx *srx, struct xnet_xfer_entry *recv_entry, (size_t) ep->cur_rx.hdr.base_hdr.hdr_size); recv_entry->cq_flags |= xnet_rx_completion_flag(ep); - if (flags & FI_CLAIM) { + if (flags & (FI_CLAIM | FI_DISCARD)) { FI_DBG(&xnet_prov, FI_LOG_EP_DATA, "Marking message for Claim\n"); if (ep->cur_rx.hdr.base_hdr.flags & XNET_REMOTE_CQ_DATA) ep->cur_rx.hdr.tag_data_hdr.tag |= XNET_CLAIM_TAG_BIT; @@ -305,16 +316,20 @@ xnet_srx_peek(struct xnet_srx *srx, struct xnet_xfer_entry *recv_entry, ep->cur_rx.claim_ctx = recv_entry->context; } + if (flags & FI_DISCARD) + return xnet_srx_claim(srx, recv_entry, flags); + xnet_report_success(ep, &srx->cq->util_cq, recv_entry); - return; + return FI_SUCCESS; nomatch: memset(&err_entry, 0, sizeof(err_entry)); err_entry.op_context = recv_entry->context; err_entry.flags = FI_RECV | FI_TAGGED; err_entry.tag = recv_entry->tag; - err_entry.err =FI_ENOMSG; + err_entry.err = ret; ofi_cq_write_error(&srx->cq->util_cq, &err_entry); + return FI_SUCCESS; } /* It's possible that an endpoint may be waiting for the message being @@ -392,9 +407,9 @@ xnet_srx_trecvmsg(struct fid_ep *ep_fid, const struct fi_msg_tagged *msg, recv_entry->context = msg->context; if (flags & FI_PEEK) { - xnet_srx_peek(srx, recv_entry, flags); - ofi_buf_free(recv_entry); - ret = 0; + ret = xnet_srx_peek(srx, recv_entry, flags); + if (ret || !(flags & FI_DISCARD)) + ofi_buf_free(recv_entry); goto unlock; }