re-factor to be used depending on message flag, change the message la…

…yout and rename to ExecGraphDetail

include/faabric/util/exec_graph.h

@@ -0,0 +1,56 @@
+#pragma once
+
+#include <faabric/proto/faabric.pb.h>
+
+#include <functional>
+#include <list>
+#include <map>
+
+namespace faabric::util {
+class ExecGraphDetail
+{
+  public:
+    void startRecording(const faabric::Message& msg);
+
+    void stopRecording(faabric::Message& msg);
+
+    void addDetail(const int msgId,
+                   const std::string& key,
+                   const std::string& value);
+
+    void incrementCounter(const int msgId,
+                          const std::string& key,
+                          const int valueToIncrement = 1);
+
+    static inline std::string const mpiMsgCountPrefix = "mpi-msgcount-torank-";
+
+  private:
+    std::shared_ptr<faabric::Message> linkedMsg = nullptr;
+
+    std::map<std::string, std::string> detailsMap;
+
+    std::map<std::string, int> intDetailsMap;
+
+    void checkMessageLinked(const int msgId);
+
+    void checkMessageNotLinked();
+
+    // ----- Wrappers to no-op the functions if not recording -----
+
+    std::function<void(const int, const std::string&, const std::string&)>
+      doAddDetail;
+
+    void addDetailInternal(const int msgId,
+                           const std::string& key,
+                           const std::string& value);
+
+    std::function<void(const int, const std::string&, const int)>
+      doIncrementCounter;
+
+    void incrementCounterInternal(const int msgId,
+                                  const std::string& key,
+                                  const int valueToIncrement);
+};
+
+ExecGraphDetail& getExecGraphDetail();
+}

src/scheduler/Executor.cpp

@@ -5,14 +5,14 @@
 #include <faabric/util/clock.h>
 #include <faabric/util/config.h>
 #include <faabric/util/environment.h>
+#include <faabric/util/exec_graph.h>
 #include <faabric/util/func.h>
 #include <faabric/util/gids.h>
 #include <faabric/util/logging.h>
 #include <faabric/util/macros.h>
 #include <faabric/util/memory.h>
 #include <faabric/util/queue.h>
 #include <faabric/util/timing.h>
-#include <faabric/util/tracing.h>
 
 #define POOL_SHUTDOWN -1
 
@@ -228,7 +228,7 @@ void Executor::threadPoolThread(int threadPoolIdx)
                      isThreads);
 
         // Start recording calls in non-release builds
-        faabric::util::tracing::getCallRecords().startRecording(msg);
+        faabric::util::getExecGraphDetail().startRecording(msg);
 
         int32_t returnValue;
         try {
@@ -247,7 +247,7 @@ void Executor::threadPoolThread(int threadPoolIdx)
         msg.set_returnvalue(returnValue);
 
         // Stop recording calls
-        faabric::util::tracing::getCallRecords().stopRecording(msg);
+        faabric::util::getExecGraphDetail().stopRecording(msg);
 
         // Decrement the task count
         int oldTaskCount = task.batchCounter->fetch_sub(1);

src/scheduler/MpiWorld.cpp

@@ -1,12 +1,12 @@
 #include <faabric/scheduler/MpiWorld.h>
 #include <faabric/scheduler/Scheduler.h>
 #include <faabric/util/environment.h>
+#include <faabric/util/exec_graph.h>
 #include <faabric/util/func.h>
 #include <faabric/util/gids.h>
 #include <faabric/util/macros.h>
 #include <faabric/util/scheduling.h>
 #include <faabric/util/testing.h>
-#include <faabric/util/tracing.h>
 
 // Each MPI rank runs in a separate thread, thus we use TLS to maintain the
 // per-rank data structures
@@ -581,10 +581,10 @@ void MpiWorld::send(int sendRank,
     }
 
     // In non-release builds, track that we have sent this message
-    faabric::util::tracing::getCallRecords().addRecord(
+    faabric::util::getExecGraphDetail().incrementCounter(
       thisMsgId,
-      faabric::util::tracing::RecordType::MpiPerRankMessageCount,
-      recvRank);
+      faabric::util::ExecGraphDetail::mpiMsgCountPrefix +
+        std::to_string(recvRank));
 }
 
 void MpiWorld::recv(int sendRank,

src/util/CMakeLists.txt

@@ -7,6 +7,7 @@ faabric_lib(util
     crash.cpp
     delta.cpp
     environment.cpp
+    exec_graph.cpp
     files.cpp
     func.cpp
     gids.cpp
@@ -21,7 +22,6 @@ faabric_lib(util
     snapshot.cpp
     state.cpp
     string_tools.cpp
-    tracing.cpp
     timing.cpp
     testing.cpp
 )

src/util/exec_graph.cpp

@@ -0,0 +1,150 @@
+#include <faabric/util/exec_graph.h>
+#include <faabric/util/logging.h>
+#include <faabric/util/testing.h>
+
+namespace faabric::util {
+void ExecGraphDetail::startRecording(const faabric::Message& msg)
+{
+    // In the tests there's not a thread to message mapping as we sometimes
+    // spawn extra threads to mock work. Thus, we skip this check here.
+    if (faabric::util::isTestMode()) {
+        return;
+    }
+
+    checkMessageNotLinked();
+
+    linkedMsg = std::make_shared<faabric::Message>(msg);
+
+    if (!msg.recordexecgraph()) {
+        doAddDetail =
+          [](const int, const std::string&, const std::string&) -> void { ; };
+        doIncrementCounter =
+          [](const int, const std::string&, const int) -> void { ; };
+    } else {
+        doAddDetail = [this](const int msgId,
+                             const std::string& key,
+                             const std::string& value) -> void {
+            this->addDetailInternal(msgId, key, value);
+        };
+        doIncrementCounter = [this](const int msgId,
+                                    const std::string& key,
+                                    const int valueToIncrement) -> void {
+            this->incrementCounterInternal(msgId, key, valueToIncrement);
+        };
+    }
+}
+
+void ExecGraphDetail::stopRecording(faabric::Message& msg)
+{
+    // In the tests there's not a thread to message mapping as we sometimes
+    // spawn extra threads to mock work. Thus, we skip this check here.
+    if (faabric::util::isTestMode()) {
+        return;
+    }
+
+    checkMessageLinked(msg.id());
+
+    for (const auto& it : detailsMap) {
+        faabric::ExecGraphDetail detail;
+        detail.set_key(it.first);
+        detail.set_value(it.second);
+        *msg.add_execgraphdetails() = detail;
+        SPDLOG_TRACE("Adding exec. graph detail to message. id: {} ; {}->{}",
+                     linkedMsg->id(),
+                     it.first,
+                     it.second);
+    }
+
+    for (const auto& it : intDetailsMap) {
+        if (detailsMap.find(it.first) != detailsMap.end()) {
+            SPDLOG_WARN(
+              "Replicated key in the exec graph details: {}->{} and {}->{}",
+              it.first,
+              detailsMap.at(it.first),
+              it.first,
+              it.second);
+        }
+
+        faabric::ExecGraphDetail detail;
+        detail.set_key(it.first);
+        detail.set_value(std::to_string(it.second));
+        *msg.add_execgraphdetails() = detail;
+        SPDLOG_TRACE("Adding exec. graph detail to message. id: {} ; {}->{}",
+                     linkedMsg->id(),
+                     it.first,
+                     it.second);
+    }
+
+    linkedMsg = nullptr;
+    detailsMap.clear();
+    intDetailsMap.clear();
+}
+
+void ExecGraphDetail::checkMessageLinked(const int msgId)
+{
+    if (linkedMsg == nullptr || linkedMsg->id() != msgId) {
+        SPDLOG_ERROR("Error during recording, records not linked to the right"
+                     " message: (linked: {} != provided: {})",
+                     linkedMsg == nullptr ? "nullptr"
+                                          : std::to_string(linkedMsg->id()),
+                     msgId);
+        throw std::runtime_error("CallRecords linked to a different message");
+    }
+}
+
+void ExecGraphDetail::checkMessageNotLinked()
+{
+    if (linkedMsg != nullptr) {
+        SPDLOG_ERROR("Error starting recording, record already linked to"
+                     "another message: {}",
+                     linkedMsg->id());
+        throw std::runtime_error("CallRecords linked to a different message");
+    }
+}
+
+void ExecGraphDetail::addDetail(const int msgId,
+                                const std::string& key,
+                                const std::string& value)
+{
+    doAddDetail(msgId, key, value);
+}
+
+void ExecGraphDetail::addDetailInternal(const int msgId,
+                                        const std::string& key,
+                                        const std::string& value)
+{
+    if (faabric::util::isTestMode()) {
+        return;
+    }
+
+    checkMessageLinked(msgId);
+
+    detailsMap[key] = value;
+}
+
+void ExecGraphDetail::incrementCounter(const int msgId,
+                                       const std::string& key,
+                                       const int valueToIncrement)
+{
+    doIncrementCounter(msgId, key, valueToIncrement);
+}
+
+void ExecGraphDetail::incrementCounterInternal(const int msgId,
+                                               const std::string& key,
+                                               const int valueToIncrement)
+{
+    if (faabric::util::isTestMode()) {
+        return;
+    }
+
+    checkMessageLinked(msgId);
+
+    intDetailsMap[key] += valueToIncrement;
+}
+
+ExecGraphDetail& getExecGraphDetail()
+{
+    static thread_local ExecGraphDetail graphDetail;
+    return graphDetail;
+}
+}

src/util/func.cpp

@@ -103,6 +103,8 @@ faabric::Message messageFactory(const std::string& user,
     std::string thisHost = faabric::util::getSystemConfig().endpointHost;
     msg.set_masterhost(thisHost);
 
+    msg.set_recordexecgraph(false);
+
     return msg;
 }