prov/efa: flush MR cache when fork() is called and fork support is ON

When fork support is on, register memory regions was not
copied to child processes's memory space.

With MR cache turned on, the registered memory regions remain
registered after EFA finished send data from it.

As a result, when both fork support and MR cache are turned
on, child processes cannot access some memory they need.

This patch addresses the issue by flush the MR cache when
fork is called. Flushing MR cache will cause any memory
region that is not actively being used (use_cnt == 0) to
be de-registered, therefore they will be copied to child
process's memory space

Signed-off-by: Wei Zhang <[email protected]>

prov/efa/src/efa_domain.c

@@ -39,6 +39,8 @@
 #include "rxr_cntr.h"
 #include "rxr_atomic.h"
 
+struct dlist_entry g_efa_domain_list;
+
 static int efa_domain_close(fid_t fid);
 
 static struct fi_ops efa_ops_domain_fid = {
@@ -161,6 +163,7 @@ int efa_domain_open(struct fid_fabric *fabric_fid, struct fi_info *info,
 	if (!efa_domain)
 		return -FI_ENOMEM;
 
+	dlist_init(&efa_domain->list_entry);
 	efa_domain->fabric = container_of(fabric_fid, struct efa_fabric,
 					  util_fabric.fabric_fid);
 
@@ -263,6 +266,7 @@ int efa_domain_open(struct fid_fabric *fabric_fid, struct fi_info *info,
 		goto err_free;
 	}
 
+	dlist_insert_tail(&efa_domain->list_entry, &g_efa_domain_list);
 	return 0;
 
 err_free:
@@ -288,6 +292,8 @@ static int efa_domain_close(fid_t fid)
 	efa_domain = container_of(fid, struct efa_domain,
 				  util_domain.domain_fid.fid);
 
+	dlist_remove(&efa_domain->list_entry);
+
 	if (efa_domain->cache) {
 		ofi_mr_cache_cleanup(efa_domain->cache);
 		free(efa_domain->cache);

prov/efa/src/efa_domain.h

@@ -52,8 +52,11 @@ struct efa_domain {
 	bool 			mr_local;
 	uint64_t		rdm_mode;
 	size_t			rdm_cq_size;
+	struct dlist_entry	list_entry; /* linked to g_efa_domain_list */
 };
 
+extern struct dlist_entry g_efa_domain_list;
+
 /*
  * efa_is_cache_available() is a check to see whether a memory registration
  * cache is available to be used by this domain.

prov/efa/src/efa_fabric.c

@@ -285,6 +285,7 @@ int efa_prov_initialize(void)
 		goto err_free;
 	}
 
+	dlist_init(&g_efa_domain_list);
 	return 0;
 
 err_free:

prov/efa/src/efa_fork_support.c

@@ -140,7 +140,7 @@ static int efa_fork_support_is_enabled(struct fid_domain *domain_fid)
  * running on an EC2 instance.
  */
 static
-void efa_atfork_callback()
+void efa_atfork_callback_warn_and_abort()
 {
 	static int visited = 0;
 
@@ -176,6 +176,29 @@ void efa_atfork_callback()
 	abort();
 }
 
+/**
+ * @brief flush all MR caches
+ *
+ * Going through all domains, and flush the MR caches in them
+ * until all inactive MRs are de-registered.
+ * This makes all memory regions that are not actively used
+ * in data transfer visible to the child process.
+ */
+void efa_atfork_callback_flush_mr_cache()
+{
+	struct dlist_entry *tmp;
+	struct efa_domain *efa_domain;
+	bool flush_lru = true;
+
+	dlist_foreach_container_safe(&g_efa_domain_list,
+				     struct efa_domain,
+				     efa_domain, list_entry, tmp) {
+		if (efa_domain->cache) {
+			while(ofi_mr_cache_flush(efa_domain->cache, flush_lru));
+		}
+	}
+}
+
 #ifndef _WIN32
 
 /* @brief Check when fork is requested and abort in unsupported cases
@@ -230,14 +253,21 @@ int efa_fork_support_enable_if_requested(struct fid_domain* domain_fid)
 	 * fork check above. This can move to the provider init once that check
 	 * is gone.
 	 */
-	if (!fork_handler_installed && g_efa_fork_status == EFA_FORK_SUPPORT_OFF) {
-		ret = pthread_atfork(efa_atfork_callback, NULL, NULL);
+	if (!fork_handler_installed && g_efa_fork_status != EFA_FORK_SUPPORT_UNNEEDED) {
+		if (g_efa_fork_status == EFA_FORK_SUPPORT_OFF) {
+			ret = pthread_atfork(efa_atfork_callback_warn_and_abort, NULL, NULL);
+		} else {
+			assert(g_efa_fork_status == EFA_FORK_SUPPORT_ON);
+			ret = pthread_atfork(efa_atfork_callback_flush_mr_cache, NULL, NULL);
+		}
+
 		if (ret) {
 			EFA_WARN(FI_LOG_DOMAIN,
 				 "Unable to register atfork callback: %s\n",
 				 strerror(-ret));
 			return ret;
 		}
+
 		fork_handler_installed = 1;
 	}