transport(tcp): fix race condition in migrations w/ tcp (#416)

transport(ptp): fix race condition in migrations w/ tcp
faasm · Apr 12, 2024 · 9f17c56 · 9f17c56
1 parent 8f9f62a
commit 9f17c56
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 36 deletions.
diff --git a/src/executor/Executor.cpp b/src/executor/Executor.cpp
@@ -362,9 +362,6 @@ void Executor::threadPoolThread(std::stop_token st, int threadPoolIdx)
         // in the same group
         bool isMigration =
           task.req->type() == faabric::BatchExecuteRequest::MIGRATION;
-        if (isMigration) {
-            faabric::transport::getPointToPointBroker().postMigrationHook(msg);
-        }
 
         SPDLOG_TRACE("Thread {}:{} executing task {} ({}, thread={}, group={})",
                      id,
@@ -380,6 +377,14 @@ void Executor::threadPoolThread(std::stop_token st, int threadPoolIdx)
         // Execute the task
         int32_t returnValue;
         try {
+            // Right before executing the task, do any kind of synchronisation
+            // if the task has been migrated. Also benefit from the try/catch
+            // statement in case the migration fails
+            if (isMigration) {
+                faabric::transport::getPointToPointBroker().postMigrationHook(
+                  msg);
+            }
+
             returnValue =
               executeTask(threadPoolIdx, task.messageIndex, task.req);
         } catch (const faabric::util::FunctionMigratedException& ex) {

diff --git a/src/transport/PointToPointBroker.cpp b/src/transport/PointToPointBroker.cpp
@@ -894,42 +894,14 @@ void PointToPointBroker::resetThreadLocalCache()
 
 void PointToPointBroker::postMigrationHook(faabric::Message& msg)
 {
-    int postMigrationOkCode = 1337;
-    int recvCode = 0;
-
-    // TODO: implement this as a broadcast in the PTP broker
-    int mainIdx = 0;
-    if (msg.groupidx() == mainIdx) {
-        auto groupIdxs = getIdxsRegisteredForGroup(msg.groupid());
-        groupIdxs.erase(mainIdx);
-        for (const auto& recvIdx : groupIdxs) {
-            sendMessage(msg.groupid(),
-                        mainIdx,
-                        recvIdx,
-                        BYTES_CONST(&postMigrationOkCode),
-                        sizeof(int));
-        }
-        recvCode = postMigrationOkCode;
-    } else {
-        std::vector<uint8_t> bytes =
-          recvMessage(msg.groupid(), 0, msg.groupidx());
-        recvCode = faabric::util::bytesToInt(bytes);
-    }
-
-    if (recvCode != postMigrationOkCode) {
-        SPDLOG_ERROR("Error in post-migration hook. {}:{}:{} received code {}",
-                     msg.appid(),
-                     msg.groupid(),
-                     msg.groupidx(),
-                     recvCode);
-        throw std::runtime_error("Error in post-migration hook");
-    }
+    // Make sure all threads are here before we move forward
+    PointToPointGroup::getGroup(msg.groupid())->barrier(msg.groupidx());
 
     // Do per-thread MPI initialisation (mostly send/recv TCP sockets)
     if (msg.ismpi()) {
-        auto& mpiWorld =
-          faabric::mpi::getMpiWorldRegistry().getWorld(msg.mpiworldid());
-        mpiWorld.initialiseRankFromMsg(msg);
+        // Get-or-initialise to initialise the world in case we are migrating
+        // to a completely new world
+        faabric::mpi::getMpiWorldRegistry().getOrInitialiseWorld(msg);
     }
 
     SPDLOG_DEBUG("{}:{}:{} exiting post-migration hook",