updating original message to avoid bug in faasm + cleanup

include/faabric/scheduler/MpiContext.h

@@ -9,7 +9,7 @@ class MpiContext
   public:
     MpiContext();
 
-    void createWorld(const faabric::Message& msg);
+    void createWorld(faabric::Message& msg);
 
     void joinWorld(const faabric::Message& msg);
 

src/scheduler/MpiContext.cpp

@@ -12,7 +12,7 @@ MpiContext::MpiContext()
   , worldId(-1)
 {}
 
-void MpiContext::createWorld(const faabric::Message& msg)
+void MpiContext::createWorld(faabric::Message& msg)
 {
     const std::shared_ptr<spdlog::logger>& logger = faabric::util::getLogger();
 
@@ -25,6 +25,9 @@ void MpiContext::createWorld(const faabric::Message& msg)
     worldId = (int)faabric::util::generateGid();
     logger->debug("Initialising world {}", worldId);
 
+    // Update the original message to contain the world ID
+    msg.set_mpiworldid(worldId);
+
     // Create the MPI world
     scheduler::MpiWorldRegistry& reg = scheduler::getMpiWorldRegistry();
     reg.createWorld(msg, worldId);

src/scheduler/MpiWorld.cpp

@@ -113,7 +113,7 @@ void MpiWorld::create(const faabric::Message& call, int newId, int newSize)
 
     // Broadcast the resulting rankHostMap to the other hosts
     std::set<std::string> hosts(executedAt.begin(), executedAt.end());
-    // Erase ourselves if we are in the set (could be that we are not)
+    // Erase ourselves if we are in the set
     hosts.erase(thisHost);
     // This will block until all other hosts have processed the message
     faabric::MpiHostRankMsg hostRankMsg;

src/transport/MpiMessageEndpoint.cpp

@@ -5,10 +5,8 @@ faabric::MpiHostRankMsg recvMpiHostRankMsg()
 {
     faabric::transport::RecvMessageEndpoint endpoint(MPI_PORT);
     endpoint.open(faabric::transport::getGlobalMessageContext());
-    // TODO - preempt data size somehow
     faabric::transport::Message m = endpoint.recv();
     PARSE_MSG(faabric::MpiHostRankMsg, m.data(), m.size());
-    // Note - This may be very slow as we poll until unbound
     endpoint.close();
 
     return msg;