From 8d33ec3ff11612bb4734d5b3d5ac9aa3cfa26921 Mon Sep 17 00:00:00 2001
From: Yifan Wang <yifan.wang@aei.mpg.de>
Date: Sat, 13 Jul 2024 10:42:25 +0200
Subject: [PATCH] fix the memory leakage bug

---
 bin/bank/pycbc_brute_bank |  1 +
 pycbc/pool.py             | 14 +++++++++++++-
 2 files changed, 14 insertions(+), 1 deletion(-)
 mode change 100644 => 100755 bin/bank/pycbc_brute_bank

diff --git a/bin/bank/pycbc_brute_bank b/bin/bank/pycbc_brute_bank
old mode 100644
new mode 100755
index 1ec4fd45364..f323b08e10c
--- a/bin/bank/pycbc_brute_bank
+++ b/bin/bank/pycbc_brute_bank
@@ -259,6 +259,7 @@ class TriangleBank(object):
                 ({k: params[k][idx] for k in params} for idx in range(total_num))
             ):
             waveform_cache += [return_wf]
+        pool.close_pool()
         del pool
 
         for hp in waveform_cache:
diff --git a/pycbc/pool.py b/pycbc/pool.py
index e3606b151cb..f9640c7acd1 100644
--- a/pycbc/pool.py
+++ b/pycbc/pool.py
@@ -91,7 +91,7 @@ def allmap(self, fcn, args):
         return results
 
     def map(self, func, items, chunksize=None):
-        """ Catch keyboard interuppts to allow the pool to exit cleanly.
+        """ Catch keyboard interupts to allow the pool to exit cleanly.
 
         Parameters
         ----------
@@ -113,6 +113,13 @@ def map(self, func, items, chunksize=None):
                 self.join()
                 raise KeyboardInterrupt
 
+    def close_pool(self):
+        """ Close the pool and remove the reference
+        """
+        self.close()
+        self.join()
+        atexit.unregister(_shutdown_pool)
+
 def _dummy_broadcast(self, f, args):
     self.map(f, [args] * self.size)
 
@@ -130,6 +137,11 @@ def map(self, f, items):
     imap = map
     imap_unordered = map
 
+    def close_pool(self):
+        ''' Dummy function to be consistent with BroadcastPool
+        '''
+        pass
+
 def use_mpi(require_mpi=False, log=True):
     """ Get whether MPI is enabled and if so the current size and rank
     """