From 6a581ad7f9cad4bbb06f08e748ebf63c052a99bc Mon Sep 17 00:00:00 2001
From: George Bosilca <bosilca@icl.utk.edu>
Date: Mon, 13 Nov 2023 13:28:41 -0700
Subject: [PATCH] Correctly handle attributes on MPI_COMM_WORLD.

During MPI_Finalize delete all predefined attributes from
MPI_COMM_WORLD, before deleting the attributes placeholder. Note that
user-defined attributes that were still attached to MPI_COMM_WORLD will
leak.

Fixes #12035.

Signed-off-by: George Bosilca <bosilca@icl.utk.edu>
(cherry picked from commit b79004e77f1b5e247a1bd7a6cfc34d1a80aebfec)
---
 ompi/attribute/attribute.h            |  4 +++-
 ompi/attribute/attribute_predefined.c | 20 ++++++++++++++++++++
 ompi/communicator/comm_init.c         | 16 +++++-----------
 3 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/ompi/attribute/attribute.h b/ompi/attribute/attribute.h
index 585a077395f..6292f1b0c55 100644
--- a/ompi/attribute/attribute.h
+++ b/ompi/attribute/attribute.h
@@ -17,6 +17,7 @@
  *                         reserved.
  * Copyright (c) 2022      Amazon.com, Inc. or its affiliates.
  *                         All Rights reserved.
+ * Copyright (c) 2023      NVIDIA Corporation.  All rights reserved.
  * $COPYRIGHT$
  *
  * Additional copyrights may follow
@@ -526,11 +527,12 @@ int ompi_attr_create_predefined_keyvals(void);
 /**
  * \internal
  *
- * Cache predefined attribute keys used in the World Model
+ * Cache and release the predefined attribute keys used in the World Model
  *
  * @returns OMPI_SUCCESS
  */
 int ompi_attr_set_predefined_keyvals_for_wm(void);
+void ompi_attr_delete_predefined_keyvals_for_wm(void);
 
 /**
  * \internal
diff --git a/ompi/attribute/attribute_predefined.c b/ompi/attribute/attribute_predefined.c
index cd16b6aa338..3bc1849dc52 100644
--- a/ompi/attribute/attribute_predefined.c
+++ b/ompi/attribute/attribute_predefined.c
@@ -18,6 +18,7 @@
  *                         All Rights reserved.
  * Copyright (c) 2022      Triad National Security, LLC. All rights
  *                         reserved.
+ * Copyright (c) 2023      NVIDIA Corporation.  All rights reserved.
  * $COPYRIGHT$
  *
  * Additional copyrights may follow
@@ -109,6 +110,7 @@ static int create_win(int target_keyval);
 static int free_win(int keyval);
 
 static int set_f(int keyval, MPI_Fint value);
+static int unset_f(int keyval);
 
 /*
  * We do not need a lock here as this function is invoked when the 
@@ -188,6 +190,17 @@ int ompi_attr_set_predefined_keyvals_for_wm(void)
     return ret;
 }
 
+void ompi_attr_delete_predefined_keyvals_for_wm(void)
+{
+    unset_f(MPI_TAG_UB);
+    unset_f(MPI_HOST);
+    unset_f(MPI_IO);
+    unset_f(MPI_WTIME_IS_GLOBAL);
+    unset_f(MPI_FT);
+    unset_f(MPI_LASTUSEDCODE);
+    unset_f(MPI_UNIVERSE_SIZE);
+    unset_f(MPI_APPNUM);
+}
 
 /*
  * We do not need a lock here as this function is invoked when the 
@@ -292,3 +305,10 @@ static int set_f(int keyval, MPI_Fint value)
                               keyval, value,
                               true);
 }
+
+static int unset_f(int keyval)
+{
+    return ompi_attr_delete(COMM_ATTR, MPI_COMM_WORLD,
+                            MPI_COMM_WORLD->c_keyhash,
+                            keyval, true);
+}
diff --git a/ompi/communicator/comm_init.c b/ompi/communicator/comm_init.c
index 9e90e8461cc..20e1fdde876 100644
--- a/ompi/communicator/comm_init.c
+++ b/ompi/communicator/comm_init.c
@@ -26,6 +26,7 @@
  * Copyright (c) 2018-2022 Triad National Security, LLC. All rights
  *                         reserved.
  * Copyright (c) 2023      Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2023      NVIDIA Corporation.  All rights reserved.
  * $COPYRIGHT$
  *
  * Additional copyrights may follow
@@ -311,18 +312,11 @@ static int ompi_comm_finalize (void)
 
     if (ompi_comm_intrinsic_init) {
         /* tear down MPI-3 predefined communicators (not initialized unless using MPI_Init) */
-        /* Free the attributes on comm world. This is not done in the
-         * destructor as we delete attributes in ompi_comm_free (which
-         * is not called for comm world) */
-        if (NULL != ompi_mpi_comm_world.comm.c_keyhash) {
-            /* Ignore errors when deleting attributes on comm_world */
-            (void) ompi_attr_delete_all(COMM_ATTR, &ompi_mpi_comm_world.comm, ompi_mpi_comm_world.comm.c_keyhash);
-            OBJ_RELEASE(ompi_mpi_comm_world.comm.c_keyhash);
-        }
-
-        /* Shut down MPI_COMM_SELF */
         OBJ_DESTRUCT( &ompi_mpi_comm_self );
-        /* Shut down MPI_COMM_WORLD */
+        ompi_attr_delete_predefined_keyvals_for_wm();
+        /* Destroy the keyhash even is user defined attributes are still attached. */
+        OBJ_DESTRUCT(ompi_mpi_comm_world.comm.c_keyhash);
+        ompi_mpi_comm_world.comm.c_keyhash = NULL;
         OBJ_DESTRUCT( &ompi_mpi_comm_world );
 
         ompi_comm_intrinsic_init = false;