Add FinalizerQueue to ObjectFactory (#4725)

(cherry picked from commit af70866594c31c81dfcc41b2e8b33c7a0cdc38f5)

kotlin-native/runtime/src/mm/cpp/ObjectFactory.hpp

@@ -206,6 +206,71 @@ public:
         Node* node_;
     };
 
+    class Consumer : private MoveOnly {
+    public:
+        class Iterator {
+        public:
+            Node& operator*() noexcept { return *node_; }
+            Node* operator->() noexcept { return node_; }
+
+            Iterator& operator++() noexcept {
+                node_ = node_->next_.get();
+                return *this;
+            }
+
+            bool operator==(const Iterator& rhs) const noexcept { return node_ == rhs.node_; }
+            bool operator!=(const Iterator& rhs) const noexcept { return node_ != rhs.node_; }
+
+        private:
+            friend class Consumer;
+            explicit Iterator(Node* node) noexcept : node_(node) {}
+
+            Node* node_;
+        };
+
+        Consumer() noexcept = default;
+
+        Consumer(Consumer&&) noexcept = default;
+        Consumer& operator=(Consumer&&) noexcept = default;
+
+        ~Consumer() {
+            // Make sure not to blow up the stack by nested `~Node` calls.
+            for (auto node = std::move(root_); node != nullptr; node = std::move(node->next_)) {
+            }
+        }
+
+        Iterator begin() noexcept { return Iterator(root_.get()); }
+        Iterator end() noexcept { return Iterator(nullptr); }
+
+    private:
+        friend class ObjectFactoryStorage;
+
+        void Insert(unique_ptr<Node> node) noexcept {
+            AssertCorrect();
+            auto* nodePtr = node.get();
+            if (!root_) {
+                root_ = std::move(node);
+            } else {
+                last_->next_ = std::move(node);
+            }
+
+            last_ = nodePtr;
+            AssertCorrect();
+        }
+
+        ALWAYS_INLINE void AssertCorrect() const noexcept {
+            if (root_ == nullptr) {
+                RuntimeAssert(last_ == nullptr, "last_ must be null");
+            } else {
+                RuntimeAssert(last_ != nullptr, "last_ must not be null");
+                RuntimeAssert(last_->next_ == nullptr, "last_ must not have next");
+            }
+        }
+
+        unique_ptr<Node> root_;
+        Node* last_ = nullptr;
+    };
+
     class Iterable : private MoveOnly {
     public:
         explicit Iterable(ObjectFactoryStorage& owner) noexcept : owner_(owner), guard_(owner_.mutex_) {}
@@ -213,7 +278,16 @@ public:
         Iterator begin() noexcept { return Iterator(nullptr, owner_.root_.get()); }
         Iterator end() noexcept { return Iterator(owner_.last_, nullptr); }
 
-        void EraseAndAdvance(Iterator& iterator) noexcept { iterator.node_ = owner_.EraseUnsafe(iterator.previousNode_); }
+        void EraseAndAdvance(Iterator& iterator) noexcept {
+            auto result = owner_.ExtractUnsafe(iterator.previousNode_);
+            iterator.node_ = result.second;
+        }
+
+        void MoveAndAdvance(Consumer& consumer, Iterator& iterator) noexcept {
+            auto result = owner_.ExtractUnsafe(iterator.previousNode_);
+            iterator.node_ = result.second;
+            consumer.Insert(std::move(result.first));
+        }
 
     private:
         ObjectFactoryStorage& owner_; // weak
@@ -230,18 +304,19 @@ public:
 
 private:
     // Expects `mutex_` to be held by the current thread.
-    Node* EraseUnsafe(Node* previousNode) noexcept {
+    std::pair<unique_ptr<Node>, Node*> ExtractUnsafe(Node* previousNode) noexcept {
         RuntimeAssert(root_ != nullptr, "Must not be empty");
         AssertCorrectUnsafe();
 
         if (previousNode == nullptr) {
-            // Deleting the root.
+            // Extracting the root.
-            root_ = std::move(root_->next_);
+            auto node = std::move(root_);
+            root_ = std::move(node->next_);
             if (!root_) {
                 last_ = nullptr;
             }
             AssertCorrectUnsafe();
-            return root_.get();
+            return {std::move(node), root_.get()};
         }
 
         auto node = std::move(previousNode->next_);
@@ -251,7 +326,7 @@ private:
         }
 
         AssertCorrectUnsafe();
-        return previousNode->next_.get();
+        return {std::move(node), previousNode->next_.get()};
     }
 
     // Expects `mutex_` to be held by the current thread.
@@ -436,6 +511,38 @@ public:
         typename Storage::Iterator iterator_;
     };
 
+    class FinalizerQueue : private MoveOnly {
+    public:
+        class Iterator {
+        public:
+            NodeRef operator*() noexcept { return NodeRef(*iterator_); }
+            NodeRef operator->() noexcept { return NodeRef(*iterator_); }
+
+            Iterator& operator++() noexcept {
+                ++iterator_;
+                return *this;
+            }
+
+            bool operator==(const Iterator& rhs) const noexcept { return iterator_ == rhs.iterator_; }
+            bool operator!=(const Iterator& rhs) const noexcept { return iterator_ != rhs.iterator_; }
+
+        private:
+            friend class ObjectFactory;
+
+            explicit Iterator(typename Storage::Consumer::Iterator iterator) noexcept : iterator_(std::move(iterator)) {}
+
+            typename Storage::Consumer::Iterator iterator_;
+        };
+
+        Iterator begin() noexcept { return Iterator(consumer_.begin()); }
+        Iterator end() noexcept { return Iterator(consumer_.end()); }
+
+    private:
+        friend class ObjectFactory;
+
+        typename Storage::Consumer consumer_;
+    };
+
     class Iterable {
     public:
         Iterable(ObjectFactory& owner) noexcept : iter_(owner.storage_.Iter()) {}
@@ -445,6 +552,10 @@ public:
 
         void EraseAndAdvance(Iterator& iterator) noexcept { iter_.EraseAndAdvance(iterator.iterator_); }
 
+        void MoveAndAdvance(FinalizerQueue& queue, Iterator& iterator) noexcept {
+            iter_.MoveAndAdvance(queue.consumer_, iterator.iterator_);
+        }
+
     private:
         typename Storage::Iterable iter_;
     };

kotlin-native/runtime/src/mm/cpp/ObjectFactoryTest.cpp

@@ -32,6 +32,9 @@ using ObjectFactoryStorageRegular = ObjectFactoryStorage<alignof(void*)>;
 template <typename Storage>
 using Producer = typename Storage::Producer;
 
+template <typename Storage>
+using Consumer = typename Storage::Consumer;
+
 template <size_t DataAlignment>
 KStdVector<void*> Collect(ObjectFactoryStorage<DataAlignment>& storage) {
     KStdVector<void*> result;
@@ -50,6 +53,15 @@ KStdVector<T> Collect(ObjectFactoryStorage<DataAlignment>& storage) {
     return result;
 }
 
+template <typename T, size_t DataAlignment>
+KStdVector<T> Collect(Consumer<ObjectFactoryStorage<DataAlignment>>& consumer) {
+    KStdVector<T> result;
+    for (auto& node : consumer) {
+        result.push_back(*static_cast<T*>(node.Data()));
+    }
+    return result;
+}
+
 struct MoveOnlyImpl : private MoveOnly {
     MoveOnlyImpl(int value1, int value2) : value1(value1), value2(value2) {}
 
@@ -310,6 +322,7 @@ TEST(ObjectFactoryStorageTest, EraseTheOnlyElement) {
         auto iter = storage.Iter();
         auto it = iter.begin();
         iter.EraseAndAdvance(it);
+        EXPECT_THAT(it, iter.end());
     }
 
     auto actual = Collect<int>(storage);
@@ -317,6 +330,174 @@ TEST(ObjectFactoryStorageTest, EraseTheOnlyElement) {
     EXPECT_THAT(actual, testing::IsEmpty());
 }
 
+TEST(ObjectFactoryStorageTest, MoveFirst) {
+    ObjectFactoryStorageRegular storage;
+    Producer<ObjectFactoryStorageRegular> producer(storage, SimpleAllocator());
+    Consumer<ObjectFactoryStorageRegular> consumer;
+
+    producer.Insert<int>(1);
+    producer.Insert<int>(2);
+    producer.Insert<int>(3);
+
+    producer.Publish();
+
+    {
+        auto iter = storage.Iter();
+        for (auto it = iter.begin(); it != iter.end();) {
+            if (it->Data<int>() == 1) {
+                iter.MoveAndAdvance(consumer, it);
+            } else {
+                ++it;
+            }
+        }
+    }
+
+    auto actual = Collect<int>(storage);
+    auto actualConsumer = Collect<int, alignof(void*)>(consumer);
+
+    EXPECT_THAT(actual, testing::ElementsAre(2, 3));
+    EXPECT_THAT(actualConsumer, testing::ElementsAre(1));
+}
+
+TEST(ObjectFactoryStorageTest, MoveMiddle) {
+    ObjectFactoryStorageRegular storage;
+    Producer<ObjectFactoryStorageRegular> producer(storage, SimpleAllocator());
+    Consumer<ObjectFactoryStorageRegular> consumer;
+
+    producer.Insert<int>(1);
+    producer.Insert<int>(2);
+    producer.Insert<int>(3);
+
+    producer.Publish();
+
+    {
+        auto iter = storage.Iter();
+        for (auto it = iter.begin(); it != iter.end();) {
+            if (it->Data<int>() == 2) {
+                iter.MoveAndAdvance(consumer, it);
+            } else {
+                ++it;
+            }
+        }
+    }
+
+    auto actual = Collect<int>(storage);
+    auto actualConsumer = Collect<int, alignof(void*)>(consumer);
+
+    EXPECT_THAT(actual, testing::ElementsAre(1, 3));
+    EXPECT_THAT(actualConsumer, testing::ElementsAre(2));
+}
+
+TEST(ObjectFactoryStorageTest, MoveLast) {
+    ObjectFactoryStorageRegular storage;
+    Producer<ObjectFactoryStorageRegular> producer(storage, SimpleAllocator());
+    Consumer<ObjectFactoryStorageRegular> consumer;
+
+    producer.Insert<int>(1);
+    producer.Insert<int>(2);
+    producer.Insert<int>(3);
+
+    producer.Publish();
+
+    {
+        auto iter = storage.Iter();
+        for (auto it = iter.begin(); it != iter.end();) {
+            if (it->Data<int>() == 3) {
+                iter.MoveAndAdvance(consumer, it);
+            } else {
+                ++it;
+            }
+        }
+    }
+
+    auto actual = Collect<int>(storage);
+    auto actualConsumer = Collect<int, alignof(void*)>(consumer);
+
+    EXPECT_THAT(actual, testing::ElementsAre(1, 2));
+    EXPECT_THAT(actualConsumer, testing::ElementsAre(3));
+}
+
+TEST(ObjectFactoryStorageTest, MoveAll) {
+    ObjectFactoryStorageRegular storage;
+    Producer<ObjectFactoryStorageRegular> producer(storage, SimpleAllocator());
+    Consumer<ObjectFactoryStorageRegular> consumer;
+
+    producer.Insert<int>(1);
+    producer.Insert<int>(2);
+    producer.Insert<int>(3);
+
+    producer.Publish();
+
+    {
+        auto iter = storage.Iter();
+        for (auto it = iter.begin(); it != iter.end();) {
+            iter.MoveAndAdvance(consumer, it);
+        }
+    }
+
+    auto actual = Collect<int>(storage);
+    auto actualConsumer = Collect<int, alignof(void*)>(consumer);
+
+    EXPECT_THAT(actual, testing::IsEmpty());
+    EXPECT_THAT(actualConsumer, testing::ElementsAre(1, 2, 3));
+}
+
+TEST(ObjectFactoryStorageTest, MoveTheOnlyElement) {
+    ObjectFactoryStorageRegular storage;
+    Producer<ObjectFactoryStorageRegular> producer(storage, SimpleAllocator());
+    Consumer<ObjectFactoryStorageRegular> consumer;
+
+    producer.Insert<int>(1);
+
+    producer.Publish();
+
+    {
+        auto iter = storage.Iter();
+        auto it = iter.begin();
+        iter.MoveAndAdvance(consumer, it);
+        EXPECT_THAT(it, iter.end());
+    }
+
+    auto actual = Collect<int>(storage);
+    auto actualConsumer = Collect<int, alignof(void*)>(consumer);
+
+    EXPECT_THAT(actual, testing::IsEmpty());
+    EXPECT_THAT(actualConsumer, testing::ElementsAre(1));
+}
+
+TEST(ObjectFactoryStorageTest, MoveAndErase) {
+    ObjectFactoryStorageRegular storage;
+    Producer<ObjectFactoryStorageRegular> producer(storage, SimpleAllocator());
+    Consumer<ObjectFactoryStorageRegular> consumer;
+
+    producer.Insert<int>(1);
+    producer.Insert<int>(2);
+    producer.Insert<int>(3);
+    producer.Insert<int>(4);
+    producer.Insert<int>(5);
+    producer.Insert<int>(6);
+    producer.Insert<int>(7);
+    producer.Insert<int>(8);
+    producer.Insert<int>(9);
+
+    producer.Publish();
+
+    {
+        auto iter = storage.Iter();
+        for (auto it = iter.begin(); it != iter.end();) {
+            ++it;
+            iter.EraseAndAdvance(it);
+            iter.MoveAndAdvance(consumer, it);
+        }
+    }
+
+    auto actual = Collect<int>(storage);
+    auto actualConsumer = Collect<int, alignof(void*)>(consumer);
+
+    EXPECT_THAT(actual, testing::ElementsAre(1, 4, 7));
+    EXPECT_THAT(actualConsumer, testing::ElementsAre(3, 6, 9));
+}
+
 TEST(ObjectFactoryStorageTest, ConcurrentPublish) {
     ObjectFactoryStorageRegular storage;
     constexpr int kThreadCount = kDefaultThreadCount;
@@ -657,6 +838,50 @@ TEST(ObjectFactoryTest, Erase) {
     }
 }
 
+TEST(ObjectFactoryTest, Move) {
+    auto objectTypeInfo = MakeObjectTypeInfo(24);
+    auto arrayTypeInfo = MakeArrayTypeInfo(24);
+    GC::ThreadData gc;
+    ObjectFactory objectFactory;
+    ObjectFactory::ThreadQueue threadQueue(objectFactory, gc);
+    ObjectFactory::FinalizerQueue finalizerQueue;
+
+    for (int i = 0; i < 10; ++i) {
+        threadQueue.CreateObject(objectTypeInfo.get());
+        threadQueue.CreateArray(arrayTypeInfo.get(), 3);
+    }
+
+    threadQueue.Publish();
+
+    {
+        auto iter = objectFactory.Iter();
+        for (auto it = iter.begin(); it != iter.end();) {
+            if (it->IsArray()) {
+                iter.MoveAndAdvance(finalizerQueue, it);
+            } else {
+                ++it;
+            }
+        }
+    }
+
+    {
+        auto iter = objectFactory.Iter();
+        int count = 0;
+        for (auto it = iter.begin(); it != iter.end(); ++it, ++count) {
+            EXPECT_FALSE(it->IsArray());
+        }
+        EXPECT_THAT(count, 10);
+    }
+
+    {
+        int count = 0;
+        for (auto it = finalizerQueue.begin(); it != finalizerQueue.end(); ++it, ++count) {
+            EXPECT_TRUE(it->IsArray());
+        }
+        EXPECT_THAT(count, 10);
+    }
+}
+
 TEST(ObjectFactoryTest, ConcurrentPublish) {
     auto typeInfo = MakeObjectTypeInfo(24);
     ObjectFactory objectFactory;