add finalize hooks to cleanup datatypes

src/datatypes.jl

@@ -66,7 +66,9 @@ end
 
 # datatype attribute to store Julia type
 const JULIA_TYPE_PTR_ATTR = Ref(Cint(0))
-add_init_hook!(() -> JULIA_TYPE_PTR_ATTR[] = create_keyval(Datatype))
+add_init_hook!() do
+    JULIA_TYPE_PTR_ATTR[] = create_keyval(Datatype)
+end
 
 """
     to_type(datatype::Datatype)
@@ -134,6 +136,11 @@ _defined_datatype_methods = nothing
 # constructed in MPI.Get. Without the cache, each Get would commit the
 # same datatype over and over again.
 const created_datatypes = IdDict{Type, Datatype}()
+add_finalize_hook!() do
+    for datatype in values(created_datatypes)
+        free(datatype)
+    end
+end
 
 function Datatype(::Type{T}) where {T}
     global created_datatypes

src/environment.jl

@@ -44,13 +44,22 @@ end
 const mpi_init_hooks = Any[]
 add_init_hook!(f) = push!(mpi_init_hooks, f)
 function run_init_hooks()
-    for f in mpi_init_hooks
+    while !isempty(mpi_init_hooks)
+        f = popfirst!(mpi_init_hooks) # FIFO
         f()
     end
-    empty!(mpi_init_hooks)
     return nothing
 end
 
+const mpi_finalize_hooks = Any[]
+add_finalize_hook!(f) = push!(mpi_finalize_hooks, f)
+function run_finalize_hooks()
+    while !isempty(mpi_finalize_hooks)
+        f = pop!(mpi_finalize_hooks) # LIFO
+        f()
+    end
+    return nothing
+end
 
 """
     Init(;threadlevel=:serialized, finalize_atexit=true, errors_return=true)
@@ -225,7 +234,10 @@ Julia exits, if it hasn't already been called.
 $(_doc_external("MPI_Finalize"))
 """
 function Finalize()
-    MPI.Finalized() || API.MPI_Finalize()
+    if !MPI.Finalized()
+        run_finalize_hooks()
+        API.MPI_Finalize()
+    end
     return nothing
 end