Simplified the Mailbox/MailboxDir API

... by combining Mailbox::close(), Mailbox::deliver() and
MailboxDir::open() into MailboxDir::deliver().

Author: yilongli

include/Homa/Homa.h

@@ -305,27 +305,6 @@ class Mailbox {
      */
     virtual ~Mailbox() = default;
 
-    /**
-     * Signal that the caller will not access the mailbox after this call.
-     * A mailbox will only be destroyed if it's removed from the directory
-     * and closed by all openers.
-     *
-     * Not meant to be called by users.
-     *
-     * @sa MailboxDir::open()
-     */
-    virtual void close() = 0;
-
-    /**
-     * Used by a transport to deliver an ingress message to this mailbox.
-     *
-     * Not meant to be called by users.
-     *
-     * @param message
-     *      An ingress message just completed by the transport.
-     */
-    virtual void deliver(InMessage* message) = 0;
-
     /**
      * Retrieve a message currently stored in the mailbox.
      *
@@ -382,17 +361,19 @@ class MailboxDir {
     virtual Mailbox* alloc(uint16_t port) = 0;
 
     /**
-     * Find and open the mailbox that matches the given port number.  Once a
-     * mailbox is opened, it's guaranteed to remain usable even if someone else
-     * removes it from the directory.
+     * Used by a transport to deliver an ingress message to a mailbox.
+     *
+     * Not meant to be called by users.
      *
      * @param port
      *      Port number which identifies the mailbox.
+     * @param message
+     *      An ingress message just completed by the transport.
      * @return
-     *      Pointer to the opened mailbox on success; nullptr, if the desired
+     *      True if the message is delivered successfully; false, if the target
      *      mailbox doesn't exist.
      */
-    virtual Mailbox* open(uint16_t port) = 0;
+    virtual bool deliver(uint16_t port, InMessage* message) = 0;
 
     /**
      * Remove the mailbox that matches the given port number.

include/Homa/Utils/SimpleMailboxDir.h

@@ -47,14 +47,14 @@ class SimpleMailboxDir final : public MailboxDir {
     explicit SimpleMailboxDir();
     ~SimpleMailboxDir() override;
     Mailbox* alloc(uint16_t port) override;
-    Mailbox* open(uint16_t port) override;
+    bool deliver(uint16_t port, Homa::InMessage* message) override;
     bool remove(uint16_t port) override;
 
   private:
     /// Monitor-style lock.
     std::unique_ptr<SpinLock> mutex;
 
-    /// Hash table that maps port numbers to mailboxes.
+    /// Hash table that maps port numbers to mailboxes. Protected by mutex.
     std::unordered_map<uint16_t, MailboxImpl*> map;
 };
 

src/Receiver.cc

@@ -165,11 +165,8 @@ Receiver::handleDataPacket(Driver::Packet* packet, IpAddress sourceIp)
             message->setState(Message::State::COMPLETED);
             bucket->resendTimeouts.cancelTimeout(&message->resendTimeout);
             uint16_t dport = be16toh(header->common.prefix.dport);
-            Mailbox* mailbox = mailboxDir->open(dport);
-            if (mailbox) {
-                mailbox->deliver(message);
-                mailbox->close();
-            } else {
+            bool success = mailboxDir->deliver(dport, message);
+            if (!success) {
                 lock_bucket.destroy();
                 ERROR("Unable to deliver the message; message dropped");
                 dropMessage(message);

src/Shenango.cc

@@ -26,12 +26,6 @@ using namespace Homa;
 #define DECLARE_SHENANGO_FUNC(ReturnType, MethodName, ...) \
     extern ReturnType (*shenango_##MethodName)(__VA_ARGS__);
 
-/**
- * Fast thread-local slab-based memory allocation.
- */
-DECLARE_SHENANGO_FUNC(void*, smalloc, size_t)
-DECLARE_SHENANGO_FUNC(void, sfree, void*)
-
 /**
  * Protect RCU read-side critical sections.
  */
@@ -165,30 +159,13 @@ homa_driver_free(homa_driver drv)
 }
 
 /**
- * An almost trivial implementation of Mailbox.  This class is essentially
- * a wrapper around a socket table entry in Shenango (i.e., struct trans_entry).
- *
+ * A trivial implementation of Mailbox for catching errors.
  */
 class ShenangoMailbox final : public Mailbox {
   public:
-    explicit ShenangoMailbox(void* trans_entry)
-        : trans_entry(trans_entry)
-    {}
-
+    explicit ShenangoMailbox() = default;
     ~ShenangoMailbox() override = default;
 
-    void close() override
-    {
-        this->~ShenangoMailbox();
-        shenango_sfree(this);
-        shenango_rcu_read_unlock();
-    }
-
-    void deliver(InMessage* message) override
-    {
-        shenango_homa_mb_deliver(trans_entry, homa_inmsg{message});
-    }
-
     InMessage* retrieve(bool blocking) override
     {
         (void)blocking;
@@ -199,10 +176,6 @@ class ShenangoMailbox final : public Mailbox {
     {
         PANIC("Shenango should never call Homa::Socket::shutdown");
     }
-
-  private:
-    /// An opaque pointer to "struct trans_entry" in Shenango.
-    void* const trans_entry;
 };
 
 /**
@@ -226,21 +199,22 @@ class ShenangoMailboxDir final : MailboxDir {
     {
         // Shenango doesn't rely on Homa::Socket to receive messages,
         // so there is no need to assign a real mailbox to SocketImpl.
-        static ShenangoMailbox dummyMailbox(nullptr);
+        static ShenangoMailbox dummyMailbox;
         (void)port;
         return &dummyMailbox;
     }
 
-    Mailbox* open(uint16_t port) override
+    bool deliver(uint16_t port, InMessage* message) override
     {
-        SocketAddress laddr = {local_ip, port};
+        // The socket table in Shenango is protected by an RCU.
         shenango_rcu_read_lock();
+        SocketAddress laddr = {local_ip, port};
         void* trans_entry = shenango_trans_table_lookup(proto, laddr, {});
-        if (!trans_entry) {
-            return nullptr;
+        if (trans_entry) {
+            shenango_homa_mb_deliver(trans_entry, homa_inmsg{message});
         }
-        void* backing = shenango_smalloc(sizeof(ShenangoMailbox));
-        return new (backing) ShenangoMailbox(trans_entry);
+        shenango_rcu_read_unlock();
+        return trans_entry != nullptr;
     }
 
     bool remove(uint16_t port) override

src/SimpleMailboxDir.cc

@@ -27,8 +27,9 @@ class MailboxImpl : public Mailbox {
   public:
     explicit MailboxImpl();
     ~MailboxImpl() override;
-    void close() override;
-    void deliver(InMessage* message) override;
+    void open();
+    void close();
+    void deliver(InMessage* message);
     InMessage* retrieve(bool blocking) override;
     void socketShutdown() override;
 
@@ -63,7 +64,24 @@ MailboxImpl::~MailboxImpl()
     }
 }
 
-/// See Homa::Mailbox::close()
+/**
+ * Signal that the caller will be accessing the mailbox until close() is called.
+ * Once a mailbox is opened, it's guaranteed to remain usable even if someone
+ * else removes it from the directory.
+ */
+void
+MailboxImpl::open()
+{
+    // Increment the reference count of the mailbox, so this mailbox won't be
+    // deleted even if it's removed from the hash table.
+    openers.fetch_add(1, std::memory_order_relaxed);
+}
+
+/**
+ * Signal that the caller will not access the mailbox after this call.
+ * A mailbox will only be destroyed if it's removed from the directory
+ * and closed by all openers.
+ */
 void
 MailboxImpl::close()
 {
@@ -75,7 +93,12 @@ MailboxImpl::close()
     }
 }
 
-/// See Homa::Mailbox::deliver()
+/**
+ * Deliver an ingress message to this mailbox.
+ *
+ * @param message
+ *      An ingress message just completed by the transport.
+ */
 void
 MailboxImpl::deliver(InMessage* message)
 {
@@ -131,24 +154,27 @@ SimpleMailboxDir::alloc(uint16_t port)
     return mailbox;
 }
 
-Mailbox*
-SimpleMailboxDir::open(uint16_t port)
+bool
+SimpleMailboxDir::deliver(uint16_t port, Homa::InMessage* message)
 {
+    // Find the mailbox.
     MailboxImpl* mailbox = nullptr;
     {
-        // Look up the mailbox
         SpinLock::Lock _(*mutex);
         auto it = map.find(port);
         if (it != map.end()) {
             mailbox = it->second;
         }
+        if (mailbox == nullptr) {
+            return false;
+        }
     }
 
-    // Increment the reference count of the mailbox.
-    if (mailbox) {
-        mailbox->openers.fetch_add(1, std::memory_order_relaxed);
-    }
-    return mailbox;
+    // Deliver the message.
+    mailbox->open();
+    mailbox->deliver(message);
+    mailbox->close();
+    return true;
 }
 
 bool