RemoteAllocator: dequeue as destructive iterator

This avoids repeated double-tapping domestication of the same pointer in
!QueueHeadsAreTame builds, by keeping the current "front" pointer to the queue
in trusted locations (stack, register) rather than storing it back to possibly
client-accessible memory.

src/mem/corealloc.h

@@ -418,37 +418,36 @@ namespace snmalloc
         [local_state](freelist::QueuePtr p) SNMALLOC_FAST_PATH_LAMBDA {
           return capptr_domesticate<SharedStateHandle>(local_state, p);
         };
-      for (size_t i = 0; i < REMOTE_BATCH; i++)
-      {
-        auto p = message_queue().peek();
-        auto& entry = SharedStateHandle::Pagemap::get_metaentry(
-          local_state, snmalloc::address_cast(p));
-
-        std::pair<freelist::HeadPtr, bool> r;
-        if constexpr (SharedStateHandle::Options.QueueHeadsAreTame)
-        {
-          /*
-           * The front of the queue has already been validated; just change the
-           * annotating type.
-           */
-          auto domesticate_first = [](freelist::QueuePtr p)
-                                     SNMALLOC_FAST_PATH_LAMBDA {
-                                       return freelist::HeadPtr(p.unsafe_ptr());
-                                     };
-          r =
-            message_queue().dequeue(key_global, domesticate_first, domesticate);
-        }
-        else
-        {
-          r = message_queue().dequeue(key_global, domesticate, domesticate);
-        }
-
-        if (unlikely(!r.second))
-          break;
+      size_t i = 0;
+      auto cb = [this, local_state, &need_post, &i](freelist::HeadPtr msg)
+                  SNMALLOC_FAST_PATH_LAMBDA {
 #ifdef SNMALLOC_TRACING
-        std::cout << "Handling remote" << std::endl;
+                    std::cout << "Handling remote" << std::endl;
 #endif
-        handle_dealloc_remote(entry, r.first.as_void(), need_post);
+
+                    auto& entry = SharedStateHandle::Pagemap::get_metaentry(
+                      local_state, snmalloc::address_cast(msg));
+
+                    handle_dealloc_remote(entry, msg.as_void(), need_post);
+
+                    return (i++ < REMOTE_BATCH);
+                  };
+
+      if constexpr (SharedStateHandle::Options.QueueHeadsAreTame)
+      {
+        /*
+         * The front of the queue has already been validated; just change the
+         * annotating type.
+         */
+        auto domesticate_first = [](freelist::QueuePtr p)
+                                   SNMALLOC_FAST_PATH_LAMBDA {
+                                     return freelist::HeadPtr(p.unsafe_ptr());
+                                   };
+        message_queue().dequeue(key_global, domesticate_first, domesticate, cb);
+      }
+      else
+      {
+        message_queue().dequeue(key_global, domesticate, domesticate, cb);
       }
 
       if (need_post)

src/mem/remoteallocator.h

@@ -128,37 +128,56 @@ namespace snmalloc
     }
 
     /**
-     * Returns the front message, or null if not possible to return a message.
+     * Destructively iterate the queue.  Each queue element is removed and fed
+     * to the callback in turn.  The callback may return false to stop iteration
+     * early (but must have processed the element it was given!).
      *
      * Takes a domestication callback for each of "pointers read from head" and
      * "pointers read from queue".  See the commentary on the class.
      */
-    template<typename Domesticator_head, typename Domesticator_queue>
-    std::pair<freelist::HeadPtr, bool> dequeue(
+    template<
+      typename Domesticator_head,
+      typename Domesticator_queue,
+      typename Cb>
+    void dequeue(
       const FreeListKey& key,
       Domesticator_head domesticate_head,
-      Domesticator_queue domesticate_queue)
+      Domesticator_queue domesticate_queue,
+      Cb cb)
     {
       invariant();
-      freelist::HeadPtr first = domesticate_head(front);
-      freelist::HeadPtr next = first->atomic_read_next(key, domesticate_queue);
+      freelist::HeadPtr curr = domesticate_head(front);
+      freelist::HeadPtr next = curr->atomic_read_next(key, domesticate_queue);
 
-      if (next != nullptr)
+      while (next != nullptr)
       {
-        /*
-         * We've domesticate_queue-d next so that we can read through it, but
-         * we're storing it back into client-accessible memory in
-         * !QueueHeadsAreTame builds, so go ahead and consider it Wild again.
-         * On QueueHeadsAreTame builds, the subsequent domesticate_head call
-         * above will also be a type-level sleight of hand, but we can still
-         * justify it by the domesticate_queue that happened in this dequeue().
-         */
-        front = capptr_rewild(next);
-        invariant();
-        return {first, true};
+        if (!cb(curr))
+        {
+          /*
+           * We've domesticate_queue-d next so that we can read through it, but
+           * we're storing it back into client-accessible memory in
+           * !QueueHeadsAreTame builds, so go ahead and consider it Wild again.
+           * On QueueHeadsAreTame builds, the subsequent domesticate_head call
+           * above will also be a type-level sleight of hand, but we can still
+           * justify it by the domesticate_queue that happened in this
+           * dequeue().
+           */
+          front = capptr_rewild(next);
+          invariant();
+          return;
+        }
+
+        curr = next;
+        next = next->atomic_read_next(key, domesticate_queue);
       }
 
-      return {nullptr, false};
+      /*
+       * Here, we've hit the end of the queue: next is nullptr and curr has not
+       * been handed to the callback.  The same considerations about Wildness
+       * above hold here.
+       */
+      front = capptr_rewild(curr);
+      invariant();
     }
 
     alloc_id_t trunc_id()

src/test/func/domestication/domestication.cc

@@ -150,9 +150,6 @@ int main()
    *
    *   - RemoteAllocator::dequeue domesticating the stub's next pointer (p)
    *
-   *   - On !QueueHeadsAreTame builds only, RemoteAllocator::dequeue
-   *     domesticating the front pointer (to p, this time)
-   *
    *   - RemoteAllocator::dequeue domesticating nullptr (p is the last message)
    *
    *   - Metaslab::alloc_free_list, domesticating the successor object in the
@@ -160,7 +157,7 @@ int main()
    *     after q).
    */
   static constexpr size_t expected_count =
-    snmalloc::CustomGlobals::Options.QueueHeadsAreTame ? 3 : 5;
+    snmalloc::CustomGlobals::Options.QueueHeadsAreTame ? 3 : 4;
   SNMALLOC_CHECK(snmalloc::CustomGlobals::domesticate_count == expected_count);
 
   // Prevent the allocators from going out of scope during the above test