Merge pull request ofiwg#5184 from shefty/master

mr/cache: Ensure that we remove the correct entry from the rbtree
shefty · Jul 29, 2019 · 90c4198 · 90c4198
2 parents 37a511c + 0334bd7
commit 90c4198
Show file tree

Hide file tree

Showing 6 changed files with 25 additions and 16 deletions.
diff --git a/include/ofi_tree.h b/include/ofi_tree.h
@@ -84,7 +84,8 @@ void ofi_rbmap_cleanup(struct ofi_rbmap *map);
 struct ofi_rbnode *ofi_rbmap_find(struct ofi_rbmap *map, void *key);
 struct ofi_rbnode *ofi_rbmap_search(struct ofi_rbmap *map, void *key,
 		int (*compare)(struct ofi_rbmap *map, void *key, void *data));
-int ofi_rbmap_insert(struct ofi_rbmap *map, void *key, void *data);
+int ofi_rbmap_insert(struct ofi_rbmap *map, void *key, void *data,
+		struct ofi_rbnode **node);
 void ofi_rbmap_delete(struct ofi_rbmap *map, struct ofi_rbnode *node);
 int ofi_rbmap_empty(struct ofi_rbmap *map);
 

diff --git a/prov/rxd/src/rxd_av.c b/prov/rxd/src/rxd_av.c
@@ -158,13 +158,14 @@ int rxd_av_insert_dg_addr(struct rxd_av *av, const void *addr,
 
 	*rxd_addr = rxd_set_rxd_addr(av, dg_addr);
 
-	ret = ofi_rbmap_insert(&av->rbmap, (void *) addr, (void *) (*rxd_addr));
-	if (ret && ret != -FI_EALREADY) {
+	ret = ofi_rbmap_insert(&av->rbmap, (void *) addr, (void *) (*rxd_addr),
+			       NULL);
+	if (ret) {
+		assert(ret != -FI_EALREADY);
 		fi_av_remove(av->dg_av, &dg_addr, 1, flags);
-		return ret;
 	}
 
-	return 0;
+	return ret;
 }
 
 static int rxd_av_insert(struct fid_av *av_fid, const void *addr, size_t count,

diff --git a/prov/util/src/util_mr_cache.c b/prov/util/src/util_mr_cache.c
@@ -384,7 +384,8 @@ static int ofi_mr_rbt_insert(struct ofi_mr_storage *storage,
 			     struct ofi_mr_info *key,
 			     struct ofi_mr_entry *entry)
 {
-	return ofi_rbmap_insert(storage->storage, (void *) key, (void *) entry);
+	return ofi_rbmap_insert(storage->storage, (void *) key, (void *) entry,
+				NULL);
 }
 
 static int ofi_mr_rbt_erase(struct ofi_mr_storage *storage,

diff --git a/prov/util/src/util_mr_map.c b/prov/util/src/util_mr_map.c
@@ -59,6 +59,7 @@ int ofi_mr_map_insert(struct ofi_mr_map *map, const struct fi_mr_attr *attr,
 		      uint64_t *key, void *context)
 {
 	struct fi_mr_attr *item;
+	int ret;
 
 	item = dup_mr_attr(attr);
 	if (!item)
@@ -67,20 +68,22 @@ int ofi_mr_map_insert(struct ofi_mr_map *map, const struct fi_mr_attr *attr,
 	if (!(map->mode & FI_MR_VIRT_ADDR))
 		item->offset = (uintptr_t) attr->mr_iov[0].iov_base;
 
-	if (!(map->mode & FI_MR_PROV_KEY)) {
-		if (ofi_rbmap_find(map->rbtree, &item->requested_key)) {
-			free(item);
-			return -FI_ENOKEY;
-		}
-	} else {
+	if (map->mode & FI_MR_PROV_KEY)
 		item->requested_key = map->key++;
-	}
 
-	ofi_rbmap_insert(map->rbtree, &item->requested_key, item);
+	ret = ofi_rbmap_insert(map->rbtree, &item->requested_key, item, NULL);
+	if (ret) {
+		if (ret == -FI_EALREADY)
+			ret = -FI_ENOKEY;
+		goto err;
+	}
 	*key = item->requested_key;
 	item->context = context;
 
 	return 0;
+err:
+	free(item);
+	return ret;
 }
 
 void *ofi_mr_map_get(struct ofi_mr_map *map, uint64_t key)

diff --git a/prov/verbs/src/verbs_domain_xrc.c b/prov/verbs/src/verbs_domain_xrc.c
@@ -147,7 +147,7 @@ int fi_ibv_get_shared_ini_conn(struct fi_ibv_xrc_ep *ep,
 	ofi_atomic_initialize32(&conn->ref_cnt, 1);
 
 	ret = ofi_rbmap_insert(domain->xrc.ini_conn_rbmap,
-			       (void *) &key, (void *) conn);
+			       (void *) &key, (void *) conn, NULL);
 	assert(ret != -FI_EALREADY);
 	if (ret) {
 		VERBS_WARN(FI_LOG_EP_CTRL, "INI QP RBTree insert failed %d\n",

diff --git a/src/tree.c b/src/tree.c
@@ -192,7 +192,8 @@ ofi_insert_rebalance(struct ofi_rbmap *map, struct ofi_rbnode *x)
 	map->root->color = BLACK;
 }
 
-int ofi_rbmap_insert(struct ofi_rbmap *map, void *key, void *data)
+int ofi_rbmap_insert(struct ofi_rbmap *map, void *key, void *data,
+		     struct ofi_rbnode **ret_node)
 {
 	struct ofi_rbnode *current, *parent, *node;
 	int ret;
@@ -229,6 +230,8 @@ int ofi_rbmap_insert(struct ofi_rbmap *map, void *key, void *data)
 	}
 
 	ofi_insert_rebalance(map, node);
+	if (ret_node)
+		*ret_node = node;
 	return 0;
 }