From f0d52852faed440869524379df9021985c4ab2a2 Mon Sep 17 00:00:00 2001
From: Gabriel Dos Santos <dss.gabriel@protonmail.com>
Date: Fri, 24 May 2024 13:37:29 +0200
Subject: [PATCH] fix(sessions): only free communicator if it isn't already
 null

---
 src/KokkosComm_communicator.hpp      | 8 +++++++-
 src/impl/KokkosComm_MPI_instance.hpp | 6 +++++-
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/src/KokkosComm_communicator.hpp b/src/KokkosComm_communicator.hpp
index af1b6dc0..c3ecd7c1 100644
--- a/src/KokkosComm_communicator.hpp
+++ b/src/KokkosComm_communicator.hpp
@@ -37,7 +37,13 @@ class Communicator {
   Communicator(MPI_Comm comm) : _comm(comm) {}
   Communicator(const Communicator& other) = delete;
   Communicator(const Communicator&& other) { _comm = std::move(other._comm); }
-  ~Communicator() { /*MPI_Comm_free(&_comm);*/ }
+  ~Communicator() {
+    // Only free the communicator if it hasn't been set to `MPI_COMM_NULL` before. This is to prevent double freeing
+    // when we explicitly call the communicator's dtor in the `Context` dtor.
+    if (MPI_COMM_NULL != _comm) {
+      MPI_Comm_free(&_comm);
+    }
+  }
 
   static auto dup_raw(MPI_Comm raw) -> Communicator {
     MPI_Comm new_comm;
diff --git a/src/impl/KokkosComm_MPI_instance.hpp b/src/impl/KokkosComm_MPI_instance.hpp
index bfde3d34..df86776a 100644
--- a/src/impl/KokkosComm_MPI_instance.hpp
+++ b/src/impl/KokkosComm_MPI_instance.hpp
@@ -33,7 +33,11 @@ class Context {
  public:
   Context(MPI_Session shandle, MPI_Comm comm) : _shandle(shandle), _comm(Communicator<ExecSpace>(comm)) {}
 
-  ~Context() { MPI_Session_finalize(&_shandle); }
+  ~Context() {
+    // Ensure the session-associated communicator is destroyed before the session is finalized.
+    _comm.~Communicator();
+    MPI_Session_finalize(&_shandle);
+  }
 
   auto comm(void) -> const Communicator<ExecSpace>& { return _comm; }
 };