StableRef registry (#4567)

kotlin-native/runtime/src/main/cpp/MultiSourceQueue.hpp

@@ -6,6 +6,7 @@
 #ifndef RUNTIME_MULTI_SOURCE_QUEUE_H
 #define RUNTIME_MULTI_SOURCE_QUEUE_H
 
+#include <atomic>
 #include <list>
 #include <mutex>
 
@@ -17,51 +18,131 @@ namespace kotlin {
 template <typename T>
 class MultiSourceQueue {
 public:
+    class Producer;
+
+    // TODO: Consider switching from `std::list` to `SingleLockList` to hide the constructor
+    // and to not store the iterator.
+    class Node : private Pinned {
+    public:
+        Node(const T& value, Producer* owner) noexcept : value_(value), owner_(owner) {}
+
+        T& operator*() noexcept { return value_; }
+
+    private:
+        friend class MultiSourceQueue;
+
+        T value_;
+        std::atomic<Producer*> owner_; // `nullptr` signifies that `MultiSourceQueue` owns it.
+        typename std::list<Node>::iterator position_;
+    };
+
     class Producer {
     public:
         explicit Producer(MultiSourceQueue& owner) noexcept : owner_(owner) {}
 
         ~Producer() { Publish(); }
 
-        void Insert(const T& value) noexcept { queue_.push_back(value); }
+        Node* Insert(const T& value) noexcept {
+            queue_.emplace_back(value, this);
+            auto& node = queue_.back();
+            node.position_ = std::prev(queue_.end());
+            return &node;
+        }
+
+        void Erase(Node* node) noexcept {
+            if (node->owner_ == this) {
+                // If we own it, delete it immediately.
+                queue_.erase(node->position_);
+                return;
+            }
+            // If it's owned by the global queue or some other `Producer`, queue it.
+            deletionQueue_.push_back(node);
+        }
 
         // Merge `this` queue with owning `MultiSourceQueue`. `this` will have empty queue after the call.
         // This call is performed without heap allocations. TODO: Test that no allocations are happening.
-        void Publish() noexcept { owner_.Collect(*this); }
+        void Publish() noexcept {
+            for (auto& node : queue_) {
+                node.owner_ = nullptr;
+            }
+            std::lock_guard<SimpleMutex> guard(owner_.mutex_);
+            owner_.queue_.splice(owner_.queue_.end(), queue_);
+            owner_.deletionQueue_.splice(owner_.deletionQueue_.end(), deletionQueue_);
+        }
+
+    private:
+        MultiSourceQueue& owner_; // weak
+        std::list<Node> queue_;
+        std::list<Node*> deletionQueue_;
+    };
+
+    class Iterator {
+    public:
+        T& operator*() noexcept { return **position_; }
+
+        Iterator& operator++() noexcept {
+            ++position_;
+            return *this;
+        }
+
+        bool operator==(const Iterator& rhs) const noexcept { return position_ == rhs.position_; }
+
+        bool operator!=(const Iterator& rhs) const noexcept { return position_ != rhs.position_; }
 
     private:
         friend class MultiSourceQueue;
 
-        MultiSourceQueue& owner_; // weak
-        std::list<T> queue_;
-    };
+        explicit Iterator(const typename std::list<Node>::iterator& position) noexcept : position_(position) {}
 
-    using Iterator = typename std::list<T>::iterator;
+        typename std::list<Node>::iterator position_;
+    };
 
     class Iterable : MoveOnly {
     public:
-        explicit Iterable(MultiSourceQueue& owner) noexcept : owner_(owner), guard_(owner_.mutex_) {}
-
-        Iterator begin() noexcept { return owner_.commonQueue_.begin(); }
-        Iterator end() noexcept { return owner_.commonQueue_.end(); }
+        Iterator begin() noexcept { return Iterator(owner_.queue_.begin()); }
+        Iterator end() noexcept { return Iterator(owner_.queue_.end()); }
 
     private:
+        friend class MultiSourceQueue;
+
+        explicit Iterable(MultiSourceQueue& owner) noexcept : owner_(owner), guard_(owner_.mutex_) {}
+
         MultiSourceQueue& owner_; // weak
         std::unique_lock<SimpleMutex> guard_;
     };
 
-    // Lock MultiSourceQueue for safe iteration.
+    // Lock `MultiSourceQueue` for safe iteration. If element was scheduled for deletion,
+    // it'll still be iterated. Use `ApplyDeletions` to remove those elements.
     Iterable Iter() noexcept { return Iterable(*this); }
 
-private:
-    void Collect(Producer& producer) noexcept {
+    // Lock `MultiSourceQueue` and apply deletions. Only deletes elements that were published.
+    void ApplyDeletions() noexcept {
         std::lock_guard<SimpleMutex> guard(mutex_);
-        commonQueue_.splice(commonQueue_.end(), producer.queue_);
+        std::list<Node*> remainingDeletions;
+
+        auto it = deletionQueue_.begin();
+        while (it != deletionQueue_.end()) {
+            auto next = std::next(it);
+            Node* node = *it;
+
+            if (node->owner_ != nullptr) {
+                // If the `Node` is still owned by some `Producer`, skip it.
+                remainingDeletions.splice(remainingDeletions.end(), deletionQueue_, it);
+            } else {
+                queue_.erase(node->position_);
+                // `node` is invalid after this
+            }
+
+            it = next;
+        }
+        deletionQueue_ = std::move(remainingDeletions);
     }
 
+private:
     // Using `std::list` as it allows to implement `Collect` without memory allocations,
     // which is important for GC mark phase.
-    std::list<T> commonQueue_;
+    std::list<Node> queue_;
+    std::list<Node*> deletionQueue_;
     SimpleMutex mutex_;
 };
 

kotlin-native/runtime/src/main/cpp/MultiSourceQueueTest.cpp

@@ -15,16 +15,109 @@
 
 using namespace kotlin;
 
+namespace {
+
+template <typename T>
+std::vector<T> Collect(MultiSourceQueue<T>& queue) {
+    std::vector<T> result;
+    for (const auto& element : queue.Iter()) {
+        result.push_back(element);
+    }
+    return result;
+}
+
+} // namespace
+
 using IntQueue = MultiSourceQueue<int>;
 
+TEST(MultiSourceQueueTest, Insert) {
+    IntQueue queue;
+    IntQueue::Producer producer(queue);
+
+    constexpr int kFirst = 1;
+    constexpr int kSecond = 2;
+
+    auto* node1 = producer.Insert(kFirst);
+    auto* node2 = producer.Insert(kSecond);
+
+    EXPECT_THAT(**node1, kFirst);
+    EXPECT_THAT(**node2, kSecond);
+}
+
+TEST(MultiSourceQueueTest, EraseFromTheSameProducer) {
+    IntQueue queue;
+    IntQueue::Producer producer(queue);
+
+    constexpr int kFirst = 1;
+    constexpr int kSecond = 2;
+
+    producer.Insert(kFirst);
+    auto* node2 = producer.Insert(kSecond);
+    producer.Erase(node2);
+    producer.Publish();
+
+    auto actual = Collect(queue);
+    EXPECT_THAT(actual, testing::ElementsAre(kFirst));
+}
+
+TEST(MultiSourceQueueTest, EraseFromGlobal) {
+    IntQueue queue;
+    IntQueue::Producer producer(queue);
+
+    constexpr int kFirst = 1;
+    constexpr int kSecond = 2;
+
+    producer.Insert(kFirst);
+    auto* node2 = producer.Insert(kSecond);
+    producer.Publish();
+    producer.Erase(node2);
+    producer.Publish();
+
+    auto actual1 = Collect(queue);
+    EXPECT_THAT(actual1, testing::ElementsAre(kFirst, kSecond));
+
+    queue.ApplyDeletions();
+
+    auto actual2 = Collect(queue);
+    EXPECT_THAT(actual2, testing::ElementsAre(kFirst));
+}
+
+TEST(MultiSourceQueueTest, EraseFromOtherProducer) {
+    IntQueue queue;
+    IntQueue::Producer producer1(queue);
+    IntQueue::Producer producer2(queue);
+
+    constexpr int kFirst = 1;
+    constexpr int kSecond = 2;
+
+    producer1.Insert(kFirst);
+    auto* node2 = producer1.Insert(kSecond);
+    producer2.Erase(node2);
+    producer1.Publish();
+
+    auto actual1 = Collect(queue);
+    EXPECT_THAT(actual1, testing::ElementsAre(kFirst, kSecond));
+
+    queue.ApplyDeletions();
+
+    auto actual2 = Collect(queue);
+    EXPECT_THAT(actual2, testing::ElementsAre(kFirst, kSecond));
+
+    producer2.Publish();
+
+    auto actual3 = Collect(queue);
+    EXPECT_THAT(actual3, testing::ElementsAre(kFirst, kSecond));
+
+    queue.ApplyDeletions();
+
+    auto actual4 = Collect(queue);
+    EXPECT_THAT(actual4, testing::ElementsAre(kFirst));
+}
+
 TEST(MultiSourceQueueTest, Empty) {
     IntQueue queue;
 
-    std::vector<int> actual;
-    for (int element : queue.Iter()) {
-        actual.push_back(element);
-    }
-
+    auto actual = Collect(queue);
     EXPECT_THAT(actual, testing::IsEmpty());
 }
 
@@ -35,11 +128,7 @@ TEST(MultiSourceQueueTest, DoNotPublish) {
     producer.Insert(1);
     producer.Insert(2);
 
-    std::vector<int> actual;
-    for (int element : queue.Iter()) {
-        actual.push_back(element);
-    }
-
+    auto actual = Collect(queue);
     EXPECT_THAT(actual, testing::IsEmpty());
 }
 
@@ -56,11 +145,7 @@ TEST(MultiSourceQueueTest, Publish) {
     producer1.Publish();
     producer2.Publish();
 
-    std::vector<int> actual;
-    for (int element : queue.Iter()) {
-        actual.push_back(element);
-    }
-
+    auto actual = Collect(queue);
     EXPECT_THAT(actual, testing::ElementsAre(1, 2, 10, 20));
 }
 
@@ -85,11 +170,7 @@ TEST(MultiSourceQueueTest, PublishSeveralTimes) {
     producer.Insert(5);
     producer.Publish();
 
-    std::vector<int> actual;
-    for (int element : queue.Iter()) {
-        actual.push_back(element);
-    }
-
+    auto actual = Collect(queue);
     EXPECT_THAT(actual, testing::ElementsAre(1, 2, 3, 4, 5));
 }
 
@@ -102,11 +183,7 @@ TEST(MultiSourceQueueTest, PublishInDestructor) {
         producer.Insert(2);
     }
 
-    std::vector<int> actual;
-    for (int element : queue.Iter()) {
-        actual.push_back(element);
-    }
-
+    auto actual = Collect(queue);
     EXPECT_THAT(actual, testing::ElementsAre(1, 2));
 }
 
@@ -137,11 +214,7 @@ TEST(MultiSourceQueueTest, ConcurrentPublish) {
         t.join();
     }
 
-    std::vector<int> actual;
-    for (int element : queue.Iter()) {
-        actual.push_back(element);
-    }
-
+    auto actual = Collect(queue);
     EXPECT_THAT(actual, testing::UnorderedElementsAreArray(expected));
 }
 
@@ -198,10 +271,49 @@ TEST(MultiSourceQueueTest, IterWhileConcurrentPublish) {
 
     EXPECT_THAT(actualBefore, testing::ElementsAreArray(expectedBefore));
 
-    std::vector<int> actualAfter;
-    for (int element : queue.Iter()) {
-        actualAfter.push_back(element);
-    }
-
+    auto actualAfter = Collect(queue);
     EXPECT_THAT(actualAfter, testing::UnorderedElementsAreArray(expectedAfter));
 }
+
+TEST(MultiSourceQueueTest, ConcurrentPublishAndApplyDeletions) {
+    IntQueue queue;
+    constexpr int kThreadCount = kDefaultThreadCount;
+
+    std::atomic<bool> canStart(false);
+    std::atomic<int> readyCount(0);
+    std::atomic<int> startedCount(0);
+    std::vector<std::thread> threads;
+    for (int i = 0; i < kThreadCount; ++i) {
+        threads.emplace_back([&queue, i, &canStart, &readyCount, &startedCount]() {
+            IntQueue::Producer producer(queue);
+            auto* node = producer.Insert(i);
+            producer.Publish();
+            producer.Erase(node);
+            ++readyCount;
+            while (!canStart) {
+            }
+            ++startedCount;
+            producer.Publish();
+        });
+    }
+
+    while (readyCount < kThreadCount) {
+    }
+    canStart = true;
+    while (startedCount < kThreadCount) {
+    }
+
+    queue.ApplyDeletions();
+
+    for (auto& t : threads) {
+        t.join();
+    }
+
+    // We do not know which elements were deleted at this point. Expecting not to crash by this point.
+
+    // This must make the queue empty.
+    queue.ApplyDeletions();
+
+    auto actual = Collect(queue);
+    EXPECT_THAT(actual, testing::IsEmpty());
+}