transport(ptp): move from protobuf to fixed-size c-struct

Author: csegarragonz

include/faabric/transport/PointToPointBroker.h

@@ -2,6 +2,7 @@
 
 #include <faabric/batch-scheduler/SchedulingDecision.h>
 #include <faabric/transport/PointToPointClient.h>
+#include <faabric/transport/PointToPointMessage.h>
 #include <faabric/util/config.h>
 #include <faabric/util/locks.h>
 
@@ -120,19 +121,11 @@ class PointToPointBroker
 
     void updateHostForIdx(int groupId, int groupIdx, std::string newHost);
 
-    void sendMessage(int groupId,
-                     int sendIdx,
-                     int recvIdx,
-                     const uint8_t* buffer,
-                     size_t bufferSize,
+    void sendMessage(const PointToPointMessage& msg,
                      std::string hostHint,
                      bool mustOrderMsg = false);
 
-    void sendMessage(int groupId,
-                     int sendIdx,
-                     int recvIdx,
-                     const uint8_t* buffer,
-                     size_t bufferSize,
+    void sendMessage(const PointToPointMessage& msg,
                      bool mustOrderMsg = false,
                      int sequenceNum = NO_SEQUENCE_NUM,
                      std::string hostHint = "");

include/faabric/transport/PointToPointClient.h

@@ -3,18 +3,19 @@
 #include <faabric/proto/faabric.pb.h>
 #include <faabric/transport/MessageEndpointClient.h>
 #include <faabric/transport/PointToPointCall.h>
+#include <faabric/transport/PointToPointMessage.h>
 
 namespace faabric::transport {
 
 std::vector<std::pair<std::string, faabric::PointToPointMappings>>
 getSentMappings();
 
-std::vector<std::pair<std::string, faabric::PointToPointMessage>>
+std::vector<std::pair<std::string, PointToPointMessage>>
 getSentPointToPointMessages();
 
 std::vector<std::tuple<std::string,
                        faabric::transport::PointToPointCall,
-                       faabric::PointToPointMessage>>
+                       PointToPointMessage>>
 getSentLockMessages();
 
 void clearSentMessages();
@@ -26,7 +27,7 @@ class PointToPointClient : public faabric::transport::MessageEndpointClient
 
     void sendMappings(faabric::PointToPointMappings& mappings);
 
-    void sendMessage(faabric::PointToPointMessage& msg,
+    void sendMessage(const PointToPointMessage& msg,
                      int sequenceNum = NO_SEQUENCE_NUM);
 
     void groupLock(int appId,

include/faabric/transport/PointToPointMessage.h

@@ -0,0 +1,36 @@
+#pragma once
+
+#include <cstdint>
+#include <span>
+
+namespace faabric::transport {
+
+/* Simple fixed-size C-struct to capture the state of a PTP message moving
+ * through Faabric.
+ *
+ * We require fixed-size, and no unique pointers to be able to use
+ * high-throughput ring-buffers to send the messages around. This also means
+ * that we manually malloc/free the data pointer. The message size is:
+ * 4 * int32_t = 4 * 4 bytes = 16 bytes
+ * 1 * size_t = 1 * 8 bytes = 8 bytes
+ * 1 * void* = 1 * 8 bytes = 8 bytes
+ * total = 32 bytes = 4 * 8 so the struct is naturally 8 byte-aligned
+ */
+struct PointToPointMessage
+{
+    int32_t appId;
+    int32_t groupId;
+    int32_t sendIdx;
+    int32_t recvIdx;
+    size_t dataSize;
+    void* dataPtr;
+};
+static_assert((sizeof(PointToPointMessage) % 8) == 0,
+              "PTP message mus be 8-aligned!");
+
+// The wire format for a PTP message is very simple: the fixed-size struct,
+// followed by dataSize bytes containing the payload.
+void serializePtpMsg(std::span<uint8_t> buffer, const PointToPointMessage& msg);
+
+void parsePtpMsg(std::span<const uint8_t> buffer, PointToPointMessage* msg);
+}

src/mpi/MpiWorld.cpp

@@ -2,6 +2,7 @@
 #include <faabric/mpi/MpiWorld.h>
 #include <faabric/mpi/mpi.pb.h>
 #include <faabric/planner/PlannerClient.h>
+#include <faabric/transport/PointToPointMessage.h>
 #include <faabric/transport/macros.h>
 #include <faabric/util/ExecGraph.h>
 #include <faabric/util/batch.h>
@@ -60,14 +61,16 @@ void MpiWorld::sendRemoteMpiMessage(std::string dstHost,
         throw std::runtime_error("Error serialising message");
     }
     try {
-        broker.sendMessage(
-          thisRankMsg->groupid(),
-          sendRank,
-          recvRank,
-          reinterpret_cast<const uint8_t*>(serialisedBuffer.data()),
-          serialisedBuffer.size(),
-          dstHost,
-          true);
+        // It is safe to send a pointer to a stack-allocated object
+        // because the broker will make an additional copy (and so will NNG!)
+        faabric::transport::PointToPointMessage msg(
+          { .groupId = thisRankMsg->groupid(),
+            .sendIdx = sendRank,
+            .recvIdx = recvRank,
+            .dataSize = serialisedBuffer.size(),
+            .dataPtr = (void*)serialisedBuffer.data() });
+
+        broker.sendMessage(msg, dstHost, true);
     } catch (std::runtime_error& e) {
         SPDLOG_ERROR("{}:{}:{} Timed out with: MPI - send {} -> {}",
                      thisRankMsg->appid(),

src/proto/faabric.proto

@@ -199,14 +199,6 @@ message StateAppendedResponse {
 // POINT-TO-POINT
 // ---------------------------------------------
 
-message PointToPointMessage {
-    int32 appId = 1;
-    int32 groupId = 2;
-    int32 sendIdx = 3;
-    int32 recvIdx = 4;
-    bytes data = 5;
-}
-
 message PointToPointMappings {
     int32 appId = 1;
     int32 groupId = 2;

src/scheduler/Scheduler.cpp

@@ -456,8 +456,19 @@ Scheduler::checkForMigrationOpportunities(faabric::Message& msg,
         auto groupIdxs = broker.getIdxsRegisteredForGroup(groupId);
         groupIdxs.erase(0);
         for (const auto& recvIdx : groupIdxs) {
-            broker.sendMessage(
-              groupId, 0, recvIdx, BYTES_CONST(&newGroupId), sizeof(int));
+            // It is safe to send a pointer to the stack, because the
+            // transport layer will perform an additional copy of the PTP
+            // message to put it in the message body
+            // TODO(no-inproc): this may not be true once we move the inproc
+            // sockets to in-memory queues
+            faabric::transport::PointToPointMessage msg(
+              { .groupId = groupId,
+                .sendIdx = 0,
+                .recvIdx = recvIdx,
+                .dataSize = sizeof(int),
+                .dataPtr = &newGroupId });
+
+            broker.sendMessage(msg);
         }
     } else if (overwriteNewGroupId == 0) {
         std::vector<uint8_t> bytes = broker.recvMessage(groupId, 0, groupIdx);

src/transport/CMakeLists.txt

@@ -9,6 +9,7 @@ faabric_lib(transport
     MessageEndpointServer.cpp
     PointToPointBroker.cpp
     PointToPointClient.cpp
+    PointToPointMessage.cpp
     PointToPointServer.cpp
 )
 

src/transport/MessageEndpointClient.cpp

@@ -36,6 +36,7 @@ void MessageEndpointClient::asyncSend(int header,
               sequenceNum);
 }
 
+// TODO: consider making an iovec-style scatter/gather alternative signature
 void MessageEndpointClient::asyncSend(int header,
                                       const uint8_t* buffer,
                                       size_t bufferSize,