Fix hanging in transport layer

examples/server.cpp

@@ -1,6 +1,7 @@
 #include <faabric/endpoint/FaabricEndpoint.h>
 #include <faabric/runner/FaabricMain.h>
 #include <faabric/scheduler/ExecutorFactory.h>
+#include <faabric/transport/context.h>
 #include <faabric/util/logging.h>
 
 using namespace faabric::scheduler;
@@ -38,6 +39,7 @@ class ExampleExecutorFactory : public ExecutorFactory
 int main()
 {
     faabric::util::initLogging();
+    faabric::transport::initGlobalMessageContext();
 
     // Start the worker pool
     SPDLOG_INFO("Starting executor pool in the background");
@@ -53,6 +55,7 @@ int main()
 
     SPDLOG_INFO("Shutting down endpoint");
     m.shutdown();
+    faabric::transport::closeGlobalMessageContext();
 
     return EXIT_SUCCESS;
 }

include/faabric/scheduler/FunctionCallClient.h

@@ -2,7 +2,7 @@
 
 #include <faabric/proto/faabric.pb.h>
 #include <faabric/scheduler/FunctionCallApi.h>
-#include <faabric/transport/MessageContext.h>
+#include <faabric/transport/MessageEndpoint.h>
 #include <faabric/transport/MessageEndpointClient.h>
 #include <faabric/util/config.h>
 
@@ -13,13 +13,13 @@ namespace faabric::scheduler {
 // -----------------------------------
 std::vector<std::pair<std::string, faabric::Message>> getFunctionCalls();
 
-std::vector<std::pair<std::string, faabric::ResponseRequest>> getFlushCalls();
+std::vector<std::pair<std::string, faabric::EmptyRequest>> getFlushCalls();
 
 std::vector<
   std::pair<std::string, std::shared_ptr<faabric::BatchExecuteRequest>>>
 getBatchRequests();
 
-std::vector<std::pair<std::string, faabric::ResponseRequest>>
+std::vector<std::pair<std::string, faabric::EmptyRequest>>
 getResourceRequests();
 
 std::vector<std::pair<std::string, faabric::UnregisterRequest>>
@@ -47,7 +47,7 @@ class FunctionCallClient : public faabric::transport::MessageEndpointClient
     void executeFunctions(
       const std::shared_ptr<faabric::BatchExecuteRequest> req);
 
-    void unregister(const faabric::UnregisterRequest& req);
+    void unregister(faabric::UnregisterRequest& req);
 
   private:
     void sendHeader(faabric::scheduler::FunctionCalls call);

include/faabric/scheduler/FunctionCallServer.h

@@ -3,7 +3,6 @@
 #include <faabric/proto/faabric.pb.h>
 #include <faabric/scheduler/FunctionCallApi.h>
 #include <faabric/scheduler/Scheduler.h>
-#include <faabric/transport/MessageEndpointClient.h>
 #include <faabric/transport/MessageEndpointServer.h>
 
 namespace faabric::scheduler {
@@ -13,24 +12,25 @@ class FunctionCallServer final
   public:
     FunctionCallServer();
 
-    void stop() override;
-
   private:
     Scheduler& scheduler;
 
-    void doRecv(faabric::transport::Message& header,
-                faabric::transport::Message& body) override;
-
-    /* Function call server API */
+    void doAsyncRecv(int header,
+                     const uint8_t* buffer,
+                     size_t bufferSize) override;
 
-    void recvFlush(faabric::transport::Message& body);
+    std::unique_ptr<google::protobuf::Message>
+    doSyncRecv(int header, const uint8_t* buffer, size_t bufferSize) override;
 
-    void recvExecuteFunctions(faabric::transport::Message& body);
+    std::unique_ptr<google::protobuf::Message> recvFlush(const uint8_t* buffer,
+                                                         size_t bufferSize);
 
-    void recvGetResources(faabric::transport::Message& body);
+    std::unique_ptr<google::protobuf::Message> recvGetResources(
+      const uint8_t* buffer,
+      size_t bufferSize);
 
-    void recvUnregister(faabric::transport::Message& body);
+    void recvExecuteFunctions(const uint8_t* buffer, size_t bufferSize);
 
-    void recvSetThreadResult(faabric::transport::Message& body);
+    void recvUnregister(const uint8_t* buffer, size_t bufferSize);
 };
 }

include/faabric/scheduler/MpiWorld.h

@@ -10,6 +10,7 @@
 #include <faabric/util/timing.h>
 
 #include <atomic>
+#include <unordered_map>
 
 namespace faabric::scheduler {
 typedef faabric::util::Queue<std::shared_ptr<faabric::MPIMessage>>
@@ -22,12 +23,11 @@ class MpiWorld
 
     void create(const faabric::Message& call, int newId, int newSize);
 
-    void initialiseFromMsg(const faabric::Message& msg,
-                           bool forceLocal = false);
+    void broadcastHostsToRanks();
 
-    std::string getHostForRank(int rank);
+    void initialiseFromMsg(const faabric::Message& msg);
 
-    void setAllRankHostsPorts(const faabric::MpiHostsToRanksMessage& msg);
+    std::string getHostForRank(int rank);
 
     std::string getUser();
 
@@ -181,12 +181,12 @@ class MpiWorld
     double getWTime();
 
   private:
-    int id;
-    int size;
+    int id = -1;
+    int size = -1;
     std::string thisHost;
+    int basePort = DEFAULT_MPI_BASE_PORT;
     faabric::util::TimePoint creationTime;
 
-    std::shared_mutex worldMutex;
     std::atomic_flag isDestroyed = false;
 
     std::string user;
@@ -208,18 +208,29 @@ class MpiWorld
     std::vector<int> basePorts;
     std::vector<int> initLocalBasePorts(
       const std::vector<std::string>& executedAt);
+
     void initRemoteMpiEndpoint(int localRank, int remoteRank);
+
     std::pair<int, int> getPortForRanks(int localRank, int remoteRank);
+
     void sendRemoteMpiMessage(int sendRank,
                               int recvRank,
                               const std::shared_ptr<faabric::MPIMessage>& msg);
+
     std::shared_ptr<faabric::MPIMessage> recvRemoteMpiMessage(int sendRank,
                                                               int recvRank);
+
+    faabric::MpiHostsToRanksMessage recvMpiHostRankMsg();
+
+    void sendMpiHostRankMsg(const std::string& hostIn,
+                            const faabric::MpiHostsToRanksMessage msg);
+
     void closeMpiMessageEndpoints();
 
     // Support for asyncrhonous communications
     std::shared_ptr<MpiMessageBuffer> getUnackedMessageBuffer(int sendRank,
                                                               int recvRank);
+
     std::shared_ptr<faabric::MPIMessage> recvBatchReturnLast(int sendRank,
                                                              int recvRank,
                                                              int batchSize = 0);

include/faabric/scheduler/Scheduler.h

@@ -91,6 +91,8 @@ class Scheduler
 
     void reset();
 
+    void resetThreadLocalCache();
+
     void shutdown();
 
     void broadcastSnapshotDelete(const faabric::Message& msg,
@@ -168,10 +170,6 @@ class Scheduler
 
     ExecGraph getFunctionExecGraph(unsigned int msgId);
 
-    void closeFunctionCallClients();
-
-    void closeSnapshotClients();
-
   private:
     std::string thisHost;
 
@@ -186,15 +184,9 @@ class Scheduler
 
     std::unordered_map<uint32_t, std::promise<int32_t>> threadResults;
 
-    std::shared_mutex functionCallClientsMx;
-    std::unordered_map<std::string, faabric::scheduler::FunctionCallClient>
-      functionCallClients;
     faabric::scheduler::FunctionCallClient& getFunctionCallClient(
       const std::string& otherHost);
 
-    std::shared_mutex snapshotClientsMx;
-    std::unordered_map<std::string, faabric::scheduler::SnapshotClient>
-      snapshotClients;
     faabric::scheduler::SnapshotClient& getSnapshotClient(
       const std::string& otherHost);
 

include/faabric/scheduler/SnapshotClient.h

@@ -2,6 +2,7 @@
 
 #include <faabric/flat/faabric_generated.h>
 #include <faabric/scheduler/SnapshotApi.h>
+#include <faabric/transport/MessageEndpoint.h>
 #include <faabric/transport/MessageEndpointClient.h>
 #include <faabric/util/snapshot.h>
 

include/faabric/scheduler/SnapshotServer.h

@@ -11,21 +11,25 @@ class SnapshotServer final : public faabric::transport::MessageEndpointServer
   public:
     SnapshotServer();
 
-    void stop() override;
-
   protected:
-    void doRecv(faabric::transport::Message& header,
-                faabric::transport::Message& body) override;
+    void doAsyncRecv(int header,
+                     const uint8_t* buffer,
+                     size_t bufferSize) override;
 
-    /* Snapshot server API */
+    std::unique_ptr<google::protobuf::Message>
+    doSyncRecv(int header, const uint8_t* buffer, size_t bufferSize) override;
 
-    void recvPushSnapshot(faabric::transport::Message& msg);
+    std::unique_ptr<google::protobuf::Message> recvPushSnapshot(
+      const uint8_t* buffer,
+      size_t bufferSize);
 
-    void recvDeleteSnapshot(faabric::transport::Message& msg);
+    std::unique_ptr<google::protobuf::Message> recvPushSnapshotDiffs(
+      const uint8_t* buffer,
+      size_t bufferSize);
 
-    void recvPushSnapshotDiffs(faabric::transport::Message& msg);
+    void recvDeleteSnapshot(const uint8_t* buffer, size_t bufferSize);
 
-    void recvThreadResult(faabric::transport::Message& msg);
+    void recvThreadResult(const uint8_t* buffer, size_t bufferSize);
 
   private:
     void applyDiffsToSnapshot(

include/faabric/state/StateClient.h

@@ -3,6 +3,7 @@
 #include <faabric/proto/faabric.pb.h>
 #include <faabric/state/InMemoryStateRegistry.h>
 #include <faabric/state/State.h>
+#include <faabric/transport/MessageEndpoint.h>
 #include <faabric/transport/MessageEndpointClient.h>
 
 namespace faabric::state {
@@ -15,11 +16,6 @@ class StateClient : public faabric::transport::MessageEndpointClient
 
     const std::string user;
     const std::string key;
-    const std::string host;
-
-    InMemoryStateRegistry& reg;
-
-    /* External state client API */
 
     void pushChunks(const std::vector<StateChunk>& chunks);
 
@@ -41,16 +37,8 @@ class StateClient : public faabric::transport::MessageEndpointClient
     void unlock();
 
   private:
-    void sendHeader(faabric::state::StateCalls call);
-
-    // Block, but ignore return value
-    faabric::transport::Message awaitResponse();
-
-    void sendStateRequest(faabric::state::StateCalls header, bool expectReply);
-
     void sendStateRequest(faabric::state::StateCalls header,
-                          const uint8_t* data = nullptr,
-                          int length = 0,
-                          bool expectReply = false);
+                          const uint8_t* data,
+                          int length);
 };
 }

include/faabric/state/StateServer.h

@@ -13,27 +13,42 @@ class StateServer final : public faabric::transport::MessageEndpointServer
   private:
     State& state;
 
-    void doRecv(faabric::transport::Message& header,
-                faabric::transport::Message& body) override;
+    void doAsyncRecv(int header,
+                     const uint8_t* buffer,
+                     size_t bufferSize) override;
 
-    /* State server API */
+    std::unique_ptr<google::protobuf::Message>
+    doSyncRecv(int header, const uint8_t* buffer, size_t bufferSize) override;
 
-    void recvSize(faabric::transport::Message& body);
+    // Sync methods
 
-    void recvPull(faabric::transport::Message& body);
+    std::unique_ptr<google::protobuf::Message> recvSize(const uint8_t* buffer,
+                                                        size_t bufferSize);
 
-    void recvPush(faabric::transport::Message& body);
+    std::unique_ptr<google::protobuf::Message> recvPull(const uint8_t* buffer,
+                                                        size_t bufferSize);
 
-    void recvAppend(faabric::transport::Message& body);
+    std::unique_ptr<google::protobuf::Message> recvPush(const uint8_t* buffer,
+                                                        size_t bufferSize);
 
-    void recvPullAppended(faabric::transport::Message& body);
+    std::unique_ptr<google::protobuf::Message> recvAppend(const uint8_t* buffer,
+                                                          size_t bufferSize);
 
-    void recvClearAppended(faabric::transport::Message& body);
+    std::unique_ptr<google::protobuf::Message> recvPullAppended(
+      const uint8_t* buffer,
+      size_t bufferSize);
 
-    void recvDelete(faabric::transport::Message& body);
+    std::unique_ptr<google::protobuf::Message> recvClearAppended(
+      const uint8_t* buffer,
+      size_t bufferSize);
 
-    void recvLock(faabric::transport::Message& body);
+    std::unique_ptr<google::protobuf::Message> recvDelete(const uint8_t* buffer,
+                                                          size_t bufferSize);
 
-    void recvUnlock(faabric::transport::Message& body);
+    std::unique_ptr<google::protobuf::Message> recvLock(const uint8_t* buffer,
+                                                        size_t bufferSize);
+
+    std::unique_ptr<google::protobuf::Message> recvUnlock(const uint8_t* buffer,
+                                                          size_t bufferSize);
 };
 }