Add support for at-most-once message delivery

By default, Homa provides at-least-once message delivery.  Add an
option for applications to request at-most-once message delivery.

Author: cstlee

include/Homa/Homa.h

@@ -132,6 +132,16 @@ class OutMessage {
         FAILED,       //< The message failed to be delivered and processed.
     };
 
+    /**
+     * Options with which an OutMessage can be sent.
+     */
+    enum Options {
+        NONE = 0,           //< Default send behavior.
+        NO_RETRY = 1 << 0,  //< Message will not be resent if recoverable send
+                            //< failure occurs; provides at-most-once delivery
+                            //< of messages.
+    };
+
     /**
      * Custom deleter for use with std::unique_ptr.
      */
@@ -206,8 +216,11 @@ class OutMessage {
      *
      * @param destination
      *      Address of the transport to which this message will be sent.
+     * @param options
+     *      Flags to request non-default sending behavior.
      */
-    virtual void send(Driver::Address destination) = 0;
+    virtual void send(Driver::Address destination,
+                      Options options = Options::NONE) = 0;
 
   protected:
     /**

src/Sender.cc

@@ -332,8 +332,32 @@ Sender::handleUnknownPacket(Driver::Packet* packet, Driver* driver)
     }
 
     OutMessage::Status status = message->getStatus();
-    if (status == OutMessage::Status::IN_PROGRESS ||
-        status == OutMessage::Status::SENT) {
+    assert(status != OutMessage::Status::NOT_STARTED);
+    if (status != OutMessage::Status::IN_PROGRESS &&
+        status != OutMessage::Status::SENT) {
+        // The message is already considered "done" so the UNKNOWN packet
+        // must be a stale response to a ping.
+    } else if (message->options & OutMessage::Options::NO_RETRY) {
+        // Option: NO_RETRY
+
+        // Either the Message or the DONE packet was lost; consider the message
+        // failed since the application asked for the message not to be retried.
+
+        // Remove Message from sendQueue.
+        if (message->numPackets > 1) {
+            SpinLock::Lock lock_queue(queueMutex);
+            QueuedMessageInfo* info = &message->queuedMessageInfo;
+            if (message->state == OutMessage::Status::IN_PROGRESS) {
+                assert(sendQueue.contains(&info->sendQueueNode));
+                sendQueue.remove(&info->sendQueueNode);
+            }
+            assert(!sendQueue.contains(&info->sendQueueNode));
+        }
+
+        bucket->messageTimeouts.cancelTimeout(&message->messageTimeout);
+        bucket->pingTimeouts.cancelTimeout(&message->pingTimeout);
+        message->state.store(OutMessage::Status::FAILED);
+    } else {
         // Message isn't done yet so we will restart sending the message.
 
         // Make sure the message is not in the sendQueue before making any
@@ -405,9 +429,6 @@ Sender::handleUnknownPacket(Driver::Packet* packet, Driver* driver)
                 QueuedMessageInfo::ComparePriority());
             sendReady.store(true);
         }
-    } else {
-        // The message is already considered "done" so the UNKNOWN packet
-        // must be a stale response to a ping.
     }
 
     driver->releasePackets(&packet, 1);
@@ -676,9 +697,10 @@ Sender::Message::reserve(size_t count)
  * @copydoc Homa::OutMessage::send()
  */
 void
-Sender::Message::send(Driver::Address destination)
+Sender::Message::send(Driver::Address destination,
+                      Sender::Message::Options options)
 {
-    sender->sendMessage(this, destination);
+    sender->sendMessage(this, destination, options);
 }
 
 /**
@@ -730,11 +752,14 @@ Sender::Message::getOrAllocPacket(size_t index)
  *      Sender::Message to be sent.
  * @param destination
  *      Destination address for this message.
+ * @param options
+ *      Flags indicating requested non-default send behavior.
  *
  * @sa dropMessage()
  */
 void
-Sender::sendMessage(Sender::Message* message, Driver::Address destination)
+Sender::sendMessage(Sender::Message* message, Driver::Address destination,
+                    Sender::Message::Options options)
 {
     // Prepare the message
     assert(message->driver == driver);
@@ -749,6 +774,7 @@ Sender::sendMessage(Sender::Message* message, Driver::Address destination)
 
     message->id = id;
     message->destination = destination;
+    message->options = options;
     message->state.store(OutMessage::Status::IN_PROGRESS);
 
     int actualMessageLen = 0;

src/Sender.h

@@ -139,6 +139,7 @@ class Sender {
                                  TRANSPORT_HEADER_LENGTH)
             , id(0, 0)
             , destination()
+            , options(Options::NONE)
             , start(0)
             , messageLength(0)
             , numPackets(0)
@@ -160,7 +161,8 @@ class Sender {
         virtual void prepend(const void* source, size_t count);
         virtual void release();
         virtual void reserve(size_t count);
-        virtual void send(Driver::Address destination);
+        virtual void send(Driver::Address destination,
+                          Options options = Options::NONE);
 
       private:
         /// Define the maximum number of packets that a message can hold.
@@ -189,6 +191,9 @@ class Sender {
         /// Contains destination address this message.
         Driver::Address destination;
 
+        /// Contains flags for any requested optional send behavior.
+        Options options;
+
         /// First byte where data is or will go if empty.
         int start;
 
@@ -379,7 +384,8 @@ class Sender {
         Protocol::MessageId::Hasher hasher;
     };
 
-    void sendMessage(Sender::Message* message, Driver::Address destination);
+    void sendMessage(Sender::Message* message, Driver::Address destination,
+                     Message::Options options = Message::Options::NONE);
     void cancelMessage(Sender::Message* message);
     void dropMessage(Sender::Message* message);
     uint64_t checkMessageTimeouts();

src/SenderTest.cc

@@ -751,6 +751,50 @@ TEST_F(SenderTest, handleUnknownPacket_singlePacketMessage)
     EXPECT_FALSE(sender->sendReady.load());
 }
 
+TEST_F(SenderTest, handleUnknownPacket_NO_RETRY)
+{
+    Protocol::MessageId id = {42, 1};
+    Driver::Address destination = (Driver::Address)22;
+
+    Sender::MessageBucket* bucket = sender->messageBuckets.getBucket(id);
+    Sender::Message* message =
+        dynamic_cast<Sender::Message*>(sender->allocMessage());
+    message->options = OutMessage::Options::NO_RETRY;
+    std::vector<Homa::Mock::MockDriver::MockPacket*> packets;
+    char payload[5][1028];
+    for (int i = 0; i < 5; ++i) {
+        Homa::Mock::MockDriver::MockPacket* packet =
+            new Homa::Mock::MockDriver::MockPacket(payload[i]);
+        packets.push_back(packet);
+        setMessagePacket(message, i, packet);
+    }
+    message->destination = destination;
+    message->messageLength = 4500;
+    message->state.store(Homa::OutMessage::Status::IN_PROGRESS);
+    SenderTest::addMessage(sender, id, message, true);
+    bucket->messageTimeouts.setTimeout(&message->messageTimeout);
+    bucket->pingTimeouts.setTimeout(&message->pingTimeout);
+    EXPECT_TRUE(
+        sender->sendQueue.contains(&message->queuedMessageInfo.sendQueueNode));
+
+    Protocol::Packet::UnknownHeader* header =
+        static_cast<Protocol::Packet::UnknownHeader*>(mockPacket.payload);
+    header->common.messageId = id;
+
+    EXPECT_CALL(mockDriver, releasePackets(Pointee(&mockPacket), Eq(1)))
+        .Times(1);
+
+    sender->handleUnknownPacket(&mockPacket, &mockDriver);
+
+    EXPECT_FALSE(
+        sender->sendQueue.contains(&message->queuedMessageInfo.sendQueueNode));
+    EXPECT_EQ(nullptr, message->messageTimeout.node.list);
+    EXPECT_EQ(nullptr, message->pingTimeout.node.list);
+    EXPECT_EQ(Homa::OutMessage::Status::FAILED, message->state);
+    EXPECT_EQ(Homa::OutMessage::Status::FAILED, message->state);
+    EXPECT_FALSE(sender->sendReady.load());
+}
+
 TEST_F(SenderTest, handleUnknownPacket_no_message)
 {
     Protocol::MessageId id = {42, 1};
@@ -1303,11 +1347,13 @@ TEST_F(SenderTest, sendMessage_basic)
         .WillOnce(Return(policy));
     EXPECT_CALL(mockDriver, sendPacket(Eq(&mockPacket))).Times(1);
 
-    sender->sendMessage(message, destination);
+    sender->sendMessage(message, destination,
+                        Sender::Message::Options::NO_RETRY);
 
     // Check Message metadata
     EXPECT_EQ(id, message->id);
     EXPECT_EQ(destination, message->destination);
+    EXPECT_EQ(Sender::Message::Options::NO_RETRY, message->options);
 
     // Check packet metadata
     Protocol::Packet::DataHeader* header =