prov/efa: Clean up rxe map during rxe release

Signed-off-by: Shi Jin <sjina@amazon.com>
shijin-aws · Feb 19, 2025 · e406f7c · e406f7c
1 parent 1592196
commit e406f7c
Show file tree

Hide file tree

Showing 6 changed files with 20 additions and 16 deletions.
diff --git a/prov/efa/src/rdm/efa_rdm_msg.c b/prov/efa/src/rdm/efa_rdm_msg.c
@@ -795,7 +795,7 @@ struct efa_rdm_ope *efa_rdm_msg_alloc_rxe_for_msgrtm(struct efa_rdm_ep *ep,
 
 	pkt_type = efa_rdm_pke_get_base_hdr(*pkt_entry_ptr)->type;
 	if (efa_rdm_pkt_type_is_mulreq(pkt_type))
-		efa_rdm_rxe_map_insert(&ep->rxe_map, *pkt_entry_ptr, rxe);
+		efa_rdm_rxe_map_insert(&ep->rxe_map, efa_rdm_pke_get_rtm_msg_id(*pkt_entry_ptr), (*pkt_entry_ptr)->addr, rxe);
 
 	return rxe;
 }
@@ -871,7 +871,7 @@ struct efa_rdm_ope *efa_rdm_msg_alloc_rxe_for_tagrtm(struct efa_rdm_ep *ep,
 
 	pkt_type = efa_rdm_pke_get_base_hdr(*pkt_entry_ptr)->type;
 	if (efa_rdm_pkt_type_is_mulreq(pkt_type))
-		efa_rdm_rxe_map_insert(&ep->rxe_map, *pkt_entry_ptr, rxe);
+		efa_rdm_rxe_map_insert(&ep->rxe_map, efa_rdm_pke_get_rtm_msg_id(*pkt_entry_ptr), (*pkt_entry_ptr)->addr, rxe);
 
 	return rxe;
 }

diff --git a/prov/efa/src/rdm/efa_rdm_ope.c b/prov/efa/src/rdm/efa_rdm_ope.c
@@ -157,6 +157,9 @@ void efa_rdm_rxe_release_internal(struct efa_rdm_ope *rxe)
 
 	dlist_remove(&rxe->ep_entry);
 
+	if (efa_rdm_rxe_map_lookup(&rxe->ep->rxe_map, rxe->msg_id, rxe->addr))
+		efa_rdm_rxe_map_remove(&rxe->ep->rxe_map, rxe->msg_id, rxe->peer->efa_fiaddr, rxe);
+
 	for (i = 0; i < rxe->iov_count; i++) {
 		if (rxe->mr[i]) {
 			err = fi_close((struct fid *)rxe->mr[i]);

diff --git a/prov/efa/src/rdm/efa_rdm_pke_rtm.c b/prov/efa/src/rdm/efa_rdm_pke_rtm.c
@@ -281,7 +281,7 @@ ssize_t efa_rdm_pke_proc_msgrtm(struct efa_rdm_pke *pkt_entry)
 
 	rtm_hdr = (struct efa_rdm_rtm_base_hdr *)pkt_entry->wiredata;
 	if (rtm_hdr->flags & EFA_RDM_REQ_READ_NACK) {
-		rxe = efa_rdm_rxe_map_lookup(&ep->rxe_map, pkt_entry);
+		rxe = efa_rdm_rxe_map_lookup(&ep->rxe_map, efa_rdm_pke_get_rtm_msg_id(pkt_entry), pkt_entry->addr);
 		rxe->internal_flags |= EFA_RDM_OPE_READ_NACK;
 	} else {
 		rxe = efa_rdm_msg_alloc_rxe_for_msgrtm(ep, &pkt_entry);
@@ -329,7 +329,7 @@ ssize_t efa_rdm_pke_proc_tagrtm(struct efa_rdm_pke *pkt_entry)
 
 	rtm_hdr = (struct efa_rdm_rtm_base_hdr *) pkt_entry->wiredata;
 	if (rtm_hdr->flags & EFA_RDM_REQ_READ_NACK) {
-		rxe = efa_rdm_rxe_map_lookup(&ep->rxe_map, pkt_entry);
+		rxe = efa_rdm_rxe_map_lookup(&ep->rxe_map, efa_rdm_pke_get_rtm_msg_id(pkt_entry), pkt_entry->addr);
 		rxe->internal_flags |= EFA_RDM_OPE_READ_NACK;
 	} else {
 		rxe = efa_rdm_msg_alloc_rxe_for_tagrtm(ep, &pkt_entry);
@@ -445,7 +445,7 @@ void efa_rdm_pke_handle_rtm_rta_recv(struct efa_rdm_pke *pkt_entry)
 		struct efa_rdm_ope *rxe;
 		struct efa_rdm_pke *unexp_pkt_entry;
 
-		rxe = efa_rdm_rxe_map_lookup(&pkt_entry->ep->rxe_map, pkt_entry);
+		rxe = efa_rdm_rxe_map_lookup(&ep->rxe_map, efa_rdm_pke_get_rtm_msg_id(pkt_entry), pkt_entry->addr);
 		if (rxe) {
 			if (rxe->state == EFA_RDM_RXE_MATCHED) {
 				pkt_entry->ope = rxe;
@@ -1205,7 +1205,7 @@ ssize_t efa_rdm_pke_proc_matched_longread_rtm(struct efa_rdm_pke *pkt_entry)
 						 "NACK packet because memory "
 						 "registration limit was "
 						 "reached on the receiver\n");
-			efa_rdm_rxe_map_insert(&ep->rxe_map, pkt_entry, rxe);
+			efa_rdm_rxe_map_insert(&ep->rxe_map, efa_rdm_pke_get_rtm_msg_id(pkt_entry), pkt_entry->addr, rxe);
 			rxe->internal_flags |= EFA_RDM_OPE_READ_NACK;
 			err = efa_rdm_ope_post_send_or_queue(
 				rxe, EFA_RDM_READ_NACK_PKT);

diff --git a/prov/efa/src/rdm/efa_rdm_pke_utils.h b/prov/efa/src/rdm/efa_rdm_pke_utils.h
@@ -7,6 +7,7 @@
 #include "efa_rdm_pke.h"
 #include "efa_rdm_protocol.h"
 #include "efa_rdm_pkt_type.h"
+#include "efa_rdm_pke_rtm.h"
 #include "efa_mr.h"
 
 /**

diff --git a/prov/efa/src/rdm/efa_rdm_rxe_map.c b/prov/efa/src/rdm/efa_rdm_rxe_map.c
@@ -15,14 +15,14 @@
  * pointer to an RX entry. If such RX entry does not exist, return NULL
 */
 struct efa_rdm_ope *efa_rdm_rxe_map_lookup(struct efa_rdm_rxe_map *rxe_map,
-					   struct efa_rdm_pke *pkt_entry)
+					   uint64_t msg_id, fi_addr_t addr)
 {
 	struct efa_rdm_rxe_map_entry *entry = NULL;
 	struct efa_rdm_rxe_map_key key;
 
 	memset(&key, 0, sizeof(key));
-	key.msg_id = efa_rdm_pke_get_rtm_msg_id(pkt_entry);
-	key.addr = pkt_entry->addr;
+	key.msg_id = msg_id;
+	key.addr = addr;
 	HASH_FIND(hh, rxe_map->head, &key, sizeof(struct efa_rdm_rxe_map_key), entry);
 	return entry ? entry->rxe : NULL;
 }
@@ -39,22 +39,22 @@ struct efa_rdm_ope *efa_rdm_rxe_map_lookup(struct efa_rdm_rxe_map *rxe_map,
  * @param[in]		rxe		RX entry
 */
 void efa_rdm_rxe_map_insert(struct efa_rdm_rxe_map *rxe_map,
-			    struct efa_rdm_pke *pkt_entry,
+			    uint64_t msg_id, fi_addr_t addr,
 			    struct efa_rdm_ope *rxe)
 {
 	struct efa_rdm_rxe_map_entry *entry;
 
-	entry = ofi_buf_alloc(pkt_entry->ep->map_entry_pool);
+	entry = ofi_buf_alloc(rxe->ep->map_entry_pool);
 	if (OFI_UNLIKELY(!entry)) {
 		EFA_WARN(FI_LOG_CQ,
 			"Map entries for medium size message exhausted.\n");
-		efa_base_ep_write_eq_error(&pkt_entry->ep->base_ep, FI_ENOBUFS, FI_EFA_ERR_RXE_POOL_EXHAUSTED);
+		efa_base_ep_write_eq_error(&rxe->ep->base_ep, FI_ENOBUFS, FI_EFA_ERR_RXE_POOL_EXHAUSTED);
 		return;
 	}
 
 	memset(&entry->key, 0, sizeof(entry->key));
-	entry->key.msg_id = efa_rdm_pke_get_rtm_msg_id(pkt_entry);
-	entry->key.addr = pkt_entry->addr;
+	entry->key.msg_id = msg_id;
+	entry->key.addr = addr;
 
 #if ENABLE_DEBUG
 	{

diff --git a/prov/efa/src/rdm/efa_rdm_rxe_map.h b/prov/efa/src/rdm/efa_rdm_rxe_map.h
@@ -49,10 +49,10 @@ void efa_rdm_rxe_map_construct(struct efa_rdm_rxe_map *rxe_map)
 struct efa_rdm_pke;
 
 struct efa_rdm_ope *efa_rdm_rxe_map_lookup(struct efa_rdm_rxe_map *rxe_map,
-					   struct efa_rdm_pke *pkt_entry);
+					   uint64_t msg_id, fi_addr_t addr);
 
 void efa_rdm_rxe_map_insert(struct efa_rdm_rxe_map *rxe_map,
-			    struct efa_rdm_pke *pkt_entry,
+			    uint64_t msg_id, fi_addr_t addr,
 			    struct efa_rdm_ope *rxe);
 
 void efa_rdm_rxe_map_remove(struct efa_rdm_rxe_map *rxe_map, uint64_t msg_id,