Fix hanging in transport layer

include/faabric/scheduler/FunctionCallClient.h

@@ -2,7 +2,6 @@
 
 #include <faabric/proto/faabric.pb.h>
 #include <faabric/scheduler/FunctionCallApi.h>
-#include <faabric/transport/MessageContext.h>
 #include <faabric/transport/MessageEndpointClient.h>
 #include <faabric/util/config.h>
 

include/faabric/transport/Message.h

@@ -18,22 +18,19 @@ class Message
 
     Message();
 
-    ~Message();
-
     char* data();
 
     uint8_t* udata();
 
+    std::vector<uint8_t> dataCopy();
+
     int size();
 
     bool more();
 
-    void persist();
-
   private:
-    uint8_t* msg;
-    int _size;
+    std::vector<uint8_t> bytes;
+
     bool _more;
-    bool _persist;
 };
 }

include/faabric/transport/MessageContext.h

@@ -1,40 +1,15 @@
 #pragma once
 
+#include <shared_mutex>
 #include <zmq.hpp>
 
-namespace faabric::transport {
-/* Wrapper around zmq::context_t
- *
+/*
  * The context object is thread safe, and the constructor parameter indicates
  * the number of hardware IO threads to be used. As a rule of thumb, use one
  * IO thread per Gbps of data.
  */
-class MessageContext
-{
-  public:
-    MessageContext();
-
-    // Message context should not be copied as there must only be one ZMQ
-    // context
-    MessageContext(const MessageContext& ctx) = delete;
-
-    MessageContext(int overrideIoThreads);
-
-    ~MessageContext();
-
-    zmq::context_t ctx;
+#define ZMQ_CONTEXT_IO_THREADS 1
 
-    zmq::context_t& get();
-
-    /* Close the message context
-     *
-     * In 0MQ terms, this method calls close() on the context, which in turn
-     * first shuts down (i.e. stop blocking operations) and then closes.
-     */
-    void close();
-
-    bool isContextShutDown;
-};
-
-faabric::transport::MessageContext& getGlobalMessageContext();
+namespace faabric::transport {
+std::shared_ptr<zmq::context_t> getGlobalMessageContext();
 }

include/faabric/transport/MessageEndpoint.h

@@ -3,7 +3,6 @@
 #include <google/protobuf/message.h>
 
 #include <faabric/transport/Message.h>
-#include <faabric/transport/MessageContext.h>
 #include <faabric/util/exception.h>
 
 #include <thread>
@@ -19,6 +18,9 @@
 #define DEFAULT_RECV_TIMEOUT_MS 20000
 #define DEFAULT_SEND_TIMEOUT_MS 20000
 
+// The monitor is checking an asynchronous event has completed, so can be short
+#define MONITOR_TIMEOUT_MS 2000
+
 namespace faabric::transport {
 enum class SocketType
 {
@@ -36,7 +38,9 @@ enum class SocketType
 class MessageEndpoint
 {
   public:
-    MessageEndpoint(const std::string& hostIn, int portIn);
+    MessageEndpoint(SocketType socketTypeIn,
+                    const std::string& hostIn,
+                    int portIn);
 
     // Message endpoints shouldn't be assigned as ZeroMQ sockets are not thread
     // safe
@@ -47,11 +51,9 @@ class MessageEndpoint
 
     ~MessageEndpoint();
 
-    void open(faabric::transport::MessageContext& context,
-              faabric::transport::SocketType sockTypeIn,
-              bool bind);
+    void open();
 
-    void close(bool bind);
+    void close();
 
     void send(uint8_t* serialisedMsg, size_t msgSize, bool more = false);
 
@@ -70,15 +72,21 @@ class MessageEndpoint
     void setSendTimeoutMs(int value);
 
   protected:
+    const SocketType socketType;
     const std::string host;
     const int port;
+    const std::string address;
     std::thread::id tid;
     int id;
 
     int recvTimeoutMs = DEFAULT_RECV_TIMEOUT_MS;
     int sendTimeoutMs = DEFAULT_SEND_TIMEOUT_MS;
 
     void validateTimeout(int value);
+
+    Message recvBuffer(int size);
+
+    Message recvNoBuffer();
 };
 
 /* Send and Recv Message Endpoints */
@@ -88,7 +96,7 @@ class SendMessageEndpoint : public MessageEndpoint
   public:
     SendMessageEndpoint(const std::string& hostIn, int portIn);
 
-    void open(MessageContext& context);
+    void open();
 
     void close();
 };
@@ -98,7 +106,7 @@ class RecvMessageEndpoint : public MessageEndpoint
   public:
     RecvMessageEndpoint(int portIn);
 
-    void open(MessageContext& context);
+    void open();
 
     void close();
 };

include/faabric/transport/MessageEndpointServer.h

@@ -1,7 +1,6 @@
 #pragma once
 
 #include <faabric/transport/Message.h>
-#include <faabric/transport/MessageContext.h>
 #include <faabric/transport/MessageEndpoint.h>
 #include <faabric/transport/MessageEndpointClient.h>
 
@@ -22,27 +21,12 @@ class MessageEndpointServer
   public:
     MessageEndpointServer(int portIn);
 
-    /* Start and stop the server
-     *
-     * Generic methods to start and stop a message endpoint server. They take
-     * a, thread-safe, 0MQ context as an argument. The stop method will block
-     * until _all_ sockets within the context have been closed. Sockets blocking
-     * on a `recv` will be interrupted with ETERM upon context closure.
-     */
-    void start(faabric::transport::MessageContext& context);
-
-    void stop(faabric::transport::MessageContext& context);
-
-    /* Common start and stop entrypoint
-     *
-     * Call the generic methods with the default global message context.
-     */
     void start();
 
     virtual void stop();
 
   protected:
-    int recv(faabric::transport::RecvMessageEndpoint& endpoint);
+    bool recv();
 
     /* Template function to handle message reception
      *
@@ -67,6 +51,8 @@ class MessageEndpointServer
   private:
     const int port;
 
+    std::unique_ptr<RecvMessageEndpoint> recvEndpoint = nullptr;
+
     std::thread servingThread;
 };
 }

src/scheduler/FunctionCallClient.cpp

@@ -82,7 +82,7 @@ void clearMockRequests()
 FunctionCallClient::FunctionCallClient(const std::string& hostIn)
   : faabric::transport::MessageEndpointClient(hostIn, FUNCTION_CALL_PORT)
 {
-    this->open(faabric::transport::getGlobalMessageContext());
+    this->open();
 }
 
 void FunctionCallClient::sendHeader(faabric::scheduler::FunctionCalls call)

src/scheduler/FunctionCallServer.cpp

@@ -18,7 +18,7 @@ void FunctionCallServer::stop()
     faabric::scheduler::getScheduler().closeFunctionCallClients();
 
     // Call the parent stop
-    MessageEndpointServer::stop(faabric::transport::getGlobalMessageContext());
+    MessageEndpointServer::stop();
 }
 
 void FunctionCallServer::doRecv(faabric::transport::Message& header,

src/scheduler/SnapshotClient.cpp

@@ -78,7 +78,7 @@ void clearMockSnapshotRequests()
 SnapshotClient::SnapshotClient(const std::string& hostIn)
   : faabric::transport::MessageEndpointClient(hostIn, SNAPSHOT_PORT)
 {
-    this->open(faabric::transport::getGlobalMessageContext());
+    this->open();
 }
 
 void SnapshotClient::sendHeader(faabric::scheduler::SnapshotCalls call)

src/scheduler/SnapshotServer.cpp

@@ -7,6 +7,8 @@
 #include <faabric/util/func.h>
 #include <faabric/util/logging.h>
 
+#include <sys/mman.h>
+
 namespace faabric::scheduler {
 SnapshotServer::SnapshotServer()
   : faabric::transport::MessageEndpointServer(SNAPSHOT_PORT)
@@ -18,7 +20,7 @@ void SnapshotServer::stop()
     faabric::scheduler::getScheduler().closeSnapshotClients();
 
     // Call the parent stop
-    MessageEndpointServer::stop(faabric::transport::getGlobalMessageContext());
+    MessageEndpointServer::stop();
 }
 
 void SnapshotServer::doRecv(faabric::transport::Message& header,
@@ -60,12 +62,16 @@ void SnapshotServer::recvPushSnapshot(faabric::transport::Message& msg)
     // Set up the snapshot
     faabric::util::SnapshotData data;
     data.size = r->contents()->size();
-    data.data = r->mutable_contents()->Data();
-    reg.takeSnapshot(r->key()->str(), data, true);
 
-    // Note that now the snapshot data is owned by Faabric and will be deleted
-    // later, so we don't want the message to delete it
-    msg.persist();
+    // TODO - avoid this copy by changing server superclass to allow subclasses
+    // to provide a buffer to receive data.
+    // TODO - work out snapshot ownership here, how do we know when to delete
+    // this data?
+    data.data = (uint8_t*)mmap(
+      nullptr, data.size, PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+    std::memcpy(data.data, r->mutable_contents()->Data(), data.size);
+
+    reg.takeSnapshot(r->key()->str(), data, true);
 
     // Send response
     faabric::EmptyResponse response;

src/state/StateClient.cpp

@@ -13,7 +13,7 @@ StateClient::StateClient(const std::string& userIn,
   , host(hostIn)
   , reg(state::getInMemoryStateRegistry())
 {
-    this->open(faabric::transport::getGlobalMessageContext());
+    this->open();
 }
 
 void StateClient::sendHeader(faabric::state::StateCalls call)

src/transport/Message.cpp

@@ -2,56 +2,42 @@
 
 namespace faabric::transport {
 Message::Message(const zmq::message_t& msgIn)
-  : _size(msgIn.size())
+  : bytes(msgIn.size())
   , _more(msgIn.more())
-  , _persist(false)
 {
-    msg = reinterpret_cast<uint8_t*>(malloc(_size * sizeof(uint8_t)));
-    memcpy(msg, msgIn.data(), _size);
+    memcpy(bytes.data(), msgIn.data(), msgIn.size());
 }
 
 Message::Message(int sizeIn)
-  : _size(sizeIn)
+  : bytes(sizeIn)
   , _more(false)
-  , _persist(false)
-{
-    msg = reinterpret_cast<uint8_t*>(malloc(_size * sizeof(uint8_t)));
-}
+{}
 
 // Empty message signals shutdown
-Message::Message()
-  : msg(nullptr)
-{}
+Message::Message() {}
 
-Message::~Message()
+char* Message::data()
 {
-    if (!_persist) {
-        free(reinterpret_cast<void*>(msg));
-    }
+    return reinterpret_cast<char*>(bytes.data());
 }
 
-char* Message::data()
+uint8_t* Message::udata()
 {
-    return reinterpret_cast<char*>(msg);
+    return bytes.data();
 }
 
-uint8_t* Message::udata()
+std::vector<uint8_t> Message::dataCopy()
 {
-    return msg;
+    return bytes;
 }
 
 int Message::size()
 {
-    return _size;
+    return bytes.size();
 }
 
 bool Message::more()
 {
     return _more;
 }
-
-void Message::persist()
-{
-    _persist = true;
-}
 }