Retry GET from Redis for messages in the Execution Graph

src/scheduler/Scheduler.cpp

@@ -1012,7 +1012,29 @@ ExecGraphNode Scheduler::getFunctionExecGraphNode(unsigned int messageId)
 
     // Get the result for this message
     std::string statusKey = faabric::util::statusKeyFromMessageId(messageId);
+
+    // We want to make sure the message bytes have been populated by the time
+    // we get them from Redis. For the time being, we retry a number of times
+    // and fail if we don't succeed.
     std::vector<uint8_t> messageBytes = redis.get(statusKey);
+    int maxNumRetries = 3;
+    int numRetries = 0;
+    while (messageBytes.empty() && numRetries < maxNumRetries) {
+        SPDLOG_WARN(
+          "Retry GET message for ExecGraph node with id {} (Retry {}/{})",
+          messageId,
+          numRetries + 1,
+          maxNumRetries);
+        SLEEP_MS(500);
+        messageBytes = redis.get(statusKey);
+    }
+    if (messageBytes.empty()) {
+        SPDLOG_ERROR("Can't GET message from redis (id: {}, key: {})",
+                     messageId,
+                     statusKey);
+        throw std::runtime_error("Message for exec graph not in Redis");
+    }
+
     faabric::Message result;
     result.ParseFromArray(messageBytes.data(), (int)messageBytes.size());
 

tests/dist/scheduler/test_exec_graph.cpp

@@ -0,0 +1,52 @@
+#include <catch2/catch.hpp>
+
+#include "fixtures.h"
+
+#include <faabric/scheduler/Scheduler.h>
+
+namespace tests {
+
+TEST_CASE_METHOD(DistTestsFixture,
+                 "Test generating the execution graph",
+                 "[funcs]")
+{
+    // Set up this host's resources
+    int nLocalSlots = 2;
+    int nFuncs = 4;
+    faabric::HostResources res;
+    res.set_slots(nLocalSlots);
+    sch.setThisHostResources(res);
+
+    // Retry the test a number of times to catch the race-condition where
+    // we get the execution graph before all results have been published
+    int numRetries = 10;
+    for (int r = 0; r < numRetries; r++) {
+        // Set up the messages
+        std::shared_ptr<faabric::BatchExecuteRequest> req =
+          faabric::util::batchExecFactory("funcs", "simple", nFuncs);
+
+        // Add a fictional chaining dependency between functions
+        for (int i = 1; i < nFuncs; i++) {
+            sch.logChainedFunction(req->mutable_messages()->at(0).id(),
+                                   req->mutable_messages()->at(i).id());
+        }
+
+        // Call the functions
+        sch.callFunctions(req);
+
+        faabric::Message& m = req->mutable_messages()->at(0);
+
+        // Wait for the result, and immediately after query for the execution
+        // graph
+        faabric::Message result = sch.getFunctionResult(m.id(), 1000);
+        auto execGraph = sch.getFunctionExecGraph(m.id());
+        REQUIRE(countExecGraphNodes(execGraph) == nFuncs);
+
+        REQUIRE(execGraph.rootNode.msg.id() == m.id());
+        for (int i = 1; i < nFuncs; i++) {
+            auto node = execGraph.rootNode.children.at(i - 1);
+            REQUIRE(node.msg.id() == req->mutable_messages()->at(i).id());
+        }
+    }
+}
+}