add shutdown method and call it from mpiworld's destroy()

faasm · May 27, 2021 · ceb549a · ceb549a
1 parent d502e61
commit ceb549a
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 10 deletions.
diff --git a/include/faabric/scheduler/MpiThreadPool.h b/include/faabric/scheduler/MpiThreadPool.h
@@ -23,15 +23,15 @@ class MpiAsyncThreadPool
   public:
     explicit MpiAsyncThreadPool(int nThreads);
 
-    ~MpiAsyncThreadPool();
+    void shutdown();
 
     int size;
 
     std::shared_ptr<MpiReqQueue> getMpiReqQueue();
 
   private:
     std::vector<std::thread> threadPool;
-    std::atomic<bool> shutdown;
+    std::atomic<bool> isShutdown;
 
     std::shared_ptr<MpiReqQueue> localReqQueue;
 

diff --git a/include/faabric/scheduler/MpiWorld.h b/include/faabric/scheduler/MpiWorld.h
@@ -47,7 +47,7 @@ class MpiWorld
 
     void destroy();
 
-    void closeFunctionCallClients();
+    void shutdownThreadPool();
 
     void enqueueMessage(faabric::MPIMessage& msg);
 

diff --git a/src/scheduler/MpiThreadPool.cpp b/src/scheduler/MpiThreadPool.cpp
@@ -4,7 +4,7 @@
 namespace faabric::scheduler {
 MpiAsyncThreadPool::MpiAsyncThreadPool(int nThreads)
   : size(nThreads)
-  , shutdown(false)
+  , isShutdown(false)
 {
     faabric::util::getLogger()->debug(
       "Starting an MpiAsyncThreadPool of size {}", nThreads);
@@ -19,7 +19,7 @@ MpiAsyncThreadPool::MpiAsyncThreadPool(int nThreads)
     }
 }
 
-MpiAsyncThreadPool::~MpiAsyncThreadPool()
+void MpiAsyncThreadPool::shutdown()
 {
     faabric::util::getLogger()->debug("Shutting down MpiAsyncThreadPool");
 
@@ -39,7 +39,7 @@ void MpiAsyncThreadPool::entrypoint(int i)
 {
     faabric::scheduler::ReqQueueType req;
 
-    while (!this->shutdown) {
+    while (!this->isShutdown) {
         req = getMpiReqQueue()->dequeue();
 
         int id = std::get<0>(req);
@@ -51,8 +51,8 @@ void MpiAsyncThreadPool::entrypoint(int i)
             // The shutdown tuple includes a TLS cleanup function that we run
             // _once per thread_ and exit
             func();
-            if (!this->shutdown) {
-                this->shutdown = true;
+            if (!this->isShutdown) {
+                this->isShutdown = true;
             }
             break;
         }

diff --git a/src/scheduler/MpiWorld.cpp b/src/scheduler/MpiWorld.cpp
@@ -134,7 +134,7 @@ void MpiWorld::destroy()
 {
     // Destroy once per host
     if (!isDestroyed.test_and_set()) {
-        closeFunctionCallClients();
+        shutdownThreadPool();
 
         // Note - we are deliberately not deleting the KV in the global state
         // TODO - find a way to do this only from the master client
@@ -149,7 +149,7 @@ void MpiWorld::destroy()
     }
 }
 
-void MpiWorld::closeFunctionCallClients()
+void MpiWorld::shutdownThreadPool()
 {
     // When shutting down the thread pool, we also make sure we clean all thread
     // local state by sending a clear message to the queue. Currently, we only
@@ -162,6 +162,8 @@ void MpiWorld::closeFunctionCallClients()
                           std::move(p)));
     }
 
+    threadPool->shutdown();
+
     // Lastly clean the main thread as well
     closeThreadLocalClients();
 }