Clean up registration cache code, use std::optional, increase byte si…

…ze in tests

cpp/mrc/src/internal/ucx/registration_cache.cpp

@@ -148,16 +148,17 @@ std::shared_ptr<ucxx::MemoryHandle> RegistrationCache3::add_block(uintptr_t addr
     return this->add_block(reinterpret_cast<void*>(addr), bytes);
 }
 
-std::shared_ptr<ucxx::MemoryHandle> RegistrationCache3::lookup(const void* addr) const noexcept
+std::optional<std::shared_ptr<ucxx::MemoryHandle>> RegistrationCache3::lookup(const void* addr) const noexcept
 {
     std::lock_guard<decltype(m_mutex)> lock(m_mutex);
     if (m_memory_handle_by_address.find(addr) != m_memory_handle_by_address.end())
     {
         return m_memory_handle_by_address.at(addr);
     }
+    return std::nullopt;
 }
 
-std::shared_ptr<ucxx::MemoryHandle> RegistrationCache3::lookup(uintptr_t addr) const noexcept
+std::optional<std::shared_ptr<ucxx::MemoryHandle>> RegistrationCache3::lookup(uintptr_t addr) const noexcept
 {
     return this->lookup(reinterpret_cast<const void*>(addr));
 }

cpp/mrc/src/internal/ucx/registration_cache.hpp

@@ -178,9 +178,8 @@ class RegistrationCache2 final
 /**
  * @brief UCX Registration Cache
  *
- * UCX memory registration object that will both register/deregister memory as well as cache the set of local and remote
- * keys for each registration. The cache can be queried for the original memory block by providing any valid address
- * contained in the contiguous block.
+ * UCX memory registration object that will both register/deregister memory. The cache can be queried for the original
+ * memory block by providing the starting address of the contiguous block.
  */
 class RegistrationCache3 final
 {
@@ -200,39 +199,20 @@ class RegistrationCache3 final
 
     std::shared_ptr<ucxx::MemoryHandle> add_block(uintptr_t addr, std::size_t bytes);
 
-    /**
-     * @brief Deregister a contiguous block of memory from the ucx context and remove the cache entry
-     *
-     * @param addr
-     * @param bytes
-     * @return std::size_t
-     */
-    std::size_t drop_block(const void* addr, std::size_t bytes);
-
-    std::size_t drop_block(uintptr_t addr, std::size_t bytes);
-
     /**
      * @brief Look up the memory registration details for a given address.
      *
-     * This method queries the registration cache to find the UcxMemoryBlock containing the original address and size as
-     * well as the local and remote keys associated with the memory block.
-     *
-     * Any address contained within a registered block can be used to query the UcxMemoryBlock
+     * This method queries the registration cache to find the MemoryHanlde containing the original address and size as
+     * well as the serialized remote keys associated with the memory block.
      *
      * @param addr
-     * @return const MemoryBlock&
+     * @return std::shared_ptr<ucxx::MemoryHandle>
      */
-    std::shared_ptr<ucxx::MemoryHandle> lookup(const void* addr) const noexcept;
+    std::optional<std::shared_ptr<ucxx::MemoryHandle>> lookup(const void* addr) const noexcept;
 
-    std::shared_ptr<ucxx::MemoryHandle> lookup(uintptr_t addr) const noexcept;
+    std::optional<std::shared_ptr<ucxx::MemoryHandle>> lookup(uintptr_t addr) const noexcept;
 
   private:
-    ucp_mem_h register_memory(const void* address, std::size_t bytes);
-
-    std::tuple<ucp_mem_h, void*, std::size_t> register_memory_with_rkey(const void* address, std::size_t bytes);
-
-    void unregister_memory(ucp_mem_h handle, void* rbuffer = nullptr);
-
     mutable std::mutex m_mutex;
     const std::shared_ptr<ucxx::Context> m_context;
     std::map<const void*, std::shared_ptr<ucxx::MemoryHandle>> m_memory_handle_by_address;

cpp/mrc/src/public/runtime/remote_descriptor.cpp

@@ -417,12 +417,17 @@ void Descriptor2::setup_remote_payloads()
 
         auto* deferred_msg = payload.mutable_deferred_msg();
 
-        // Need to register the memory
-        auto ucx_block = m_data_plane_resources.registration_cache3().add_block(deferred_msg->address(),
-                                                                                deferred_msg->bytes());
+        auto ucx_block = m_data_plane_resources.registration_cache3().lookup(deferred_msg->address());
 
-        auto serializedRemoteKey = ucx_block->createRemoteKey()->serialize();
-        deferred_msg->set_remote_key(serializedRemoteKey);
+        if (!ucx_block.has_value())
+        {
+            // Need to register the memory
+            ucx_block = m_data_plane_resources.registration_cache3().add_block(deferred_msg->address(),
+                                                                               deferred_msg->bytes());
+        }
+
+        auto remoteKey = ucx_block.value()->createRemoteKey();
+        deferred_msg->set_remote_key(remoteKey->serialize());
     }
 }
 

cpp/mrc/src/tests/test_network.cpp

@@ -628,7 +628,7 @@ TEST_F(TestNetwork, TransferFullDescriptors)
     static_assert(codable::member_decodable<ComplexObject>);
     static_assert(codable::member_decodable<TransferObject>);
 
-    ComplexObject send_data = {"test", 42, {"test", 42, std::vector<u_int8_t>(8_KiB)}, std::vector<u_int8_t>(8_KiB)};
+    ComplexObject send_data = {"test", 42, {"test", 42, std::vector<u_int8_t>(64_KiB)}, std::vector<u_int8_t>(8_KiB)};
 
     auto send_data_copy = send_data;
 
@@ -706,7 +706,7 @@ TEST_F(TestNetwork, TransferFullDescriptorsBroadcast)
     // Create initial data
     static_assert(codable::decodable<TransferObject>);
 
-    TransferObject send_data = {"test", 42, std::vector<u_int8_t>(8_KiB)};
+    TransferObject send_data = {"test", 42, std::vector<u_int8_t>(64_KiB)};
 
     auto send_data_copy = send_data;