[K/N] Address AtomicStack tsan reports.

kotlin-native/runtime/src/custom_alloc/cpp/AtomicStack.hpp

@@ -35,14 +35,21 @@ public:
     // freeing pages during STW.
     T* Pop() noexcept {
         T* elm = stack_.load(std::memory_order_acquire);
-        while (elm && !stack_.compare_exchange_weak(elm, elm->next_, std::memory_order_acq_rel)) {}
+        while (true) {
-        return elm;
+            if (!elm) {
+                return nullptr;
+            }
+            auto* elmNext = elm->next_.load(std::memory_order_relaxed);
+            if (stack_.compare_exchange_weak(elm, elmNext, std::memory_order_acq_rel)) {
+                return elm;
+            }
+        }
     }
 
     void Push(T* elm) noexcept {
         T* head = nullptr;
         do {
-            elm->next_ = head;
+            elm->next_.store(head, std::memory_order_relaxed);
         } while (!stack_.compare_exchange_weak(head, elm, std::memory_order_acq_rel));
     }
 
@@ -58,14 +65,14 @@ public:
         }
         // `this` stack is not empty. Find the tail of `other`. If no deletions are performed, this is safe.
         T* otherTail = otherHead;
-        while (otherTail->next_) otherTail = otherTail->next_;
+        while (auto* next = otherTail->next_.load(std::memory_order_relaxed)) otherTail = next;
         // can't be because of the loop above
-        RuntimeAssert(otherTail->next_ == nullptr, "otherTail->next_ must be a tail");
+        RuntimeAssert(otherTail->next_.load(std::memory_order_relaxed) == nullptr, "otherTail->next_ must be a tail");
         // Now make `otherTail->next_` point to the current head of `this` and
         // simultaneously make `otherHead` the new current head.
         do {
-            otherTail->next_ = thisHead;
+            otherTail->next_.store(thisHead, std::memory_order_relaxed);
-        } while (!stack_.compare_exchange_weak(thisHead, otherHead, std::memory_order_acq_rel));
+        } while (!stack_.compare_exchange_weak(thisHead, otherHead, std::memory_order_release, std::memory_order_relaxed));
     }
 
     bool isEmpty() const noexcept { return stack_.load(std::memory_order_relaxed) == nullptr; }

kotlin-native/runtime/src/custom_alloc/cpp/AtomicStackTest.cpp

@@ -0,0 +1,59 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include "AtomicStack.hpp"
+
+#include <array>
+
+#include "TestSupport.hpp"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+#include "ScopedThread.hpp"
+
+using namespace kotlin;
+
+namespace {
+
+struct Element {
+    std::atomic<Element*> next_ = nullptr;
+};
+
+} // namespace
+
+TEST(AtomicStackTest, StressPushPop) {
+    alloc::AtomicStack<Element> ready;
+    alloc::AtomicStack<Element> used;
+    std_support::vector<Element> elements(1000);
+    std_support::vector<Element*> expected;
+    for (auto& element : elements) {
+        ready.Push(&element);
+        expected.push_back(&element);
+    }
+
+    std::atomic<bool> canStart = false;
+    std_support::vector<ScopedThread> mutators;
+    for (int i = 0; i < kDefaultThreadCount; ++i) {
+        mutators.emplace_back([&]() NO_INLINE {
+            while (!canStart.load(std::memory_order_relaxed)) {
+            }
+            while (auto* element = ready.Pop()) {
+                used.Push(element);
+            }
+        });
+    }
+
+    canStart.store(true, std::memory_order_relaxed);
+    mutators.clear();
+
+    std_support::vector<Element*> actual;
+    while (auto* element = ready.Pop()) {
+        actual.push_back(element);
+    }
+    while (auto* element = used.Pop()) {
+        actual.push_back(element);
+    }
+    EXPECT_THAT(actual, testing::UnorderedElementsAreArray(expected));
+}

kotlin-native/runtime/src/custom_alloc/cpp/ExtraObjectPage.cpp

@@ -25,10 +25,10 @@ ExtraObjectPage* ExtraObjectPage::Create(uint32_t ignored) noexcept {
 
 ExtraObjectPage::ExtraObjectPage() noexcept {
     CustomAllocInfo("ExtraObjectPage(%p)::ExtraObjectPage()", this);
-    nextFree_ = cells_;
+    nextFree_.store(cells_, std::memory_order_relaxed);
     ExtraObjectCell* end = cells_ + EXTRA_OBJECT_COUNT;
-    for (ExtraObjectCell* cell = cells_; cell < end; cell = cell->next_) {
+    for (ExtraObjectCell* cell = cells_; cell < end; cell = cell->next_.load(std::memory_order_relaxed)) {
-        cell->next_ = cell + 1;
+        cell->next_.store(cell + 1, std::memory_order_relaxed);
     }
 }
 
@@ -37,11 +37,12 @@ void ExtraObjectPage::Destroy() noexcept {
 }
 
 mm::ExtraObjectData* ExtraObjectPage::TryAllocate() noexcept {
-    if (nextFree_ >= cells_ + EXTRA_OBJECT_COUNT) {
+    auto* next = nextFree_.load(std::memory_order_relaxed);
+    if (next >= cells_ + EXTRA_OBJECT_COUNT) {
         return nullptr;
     }
-    ExtraObjectCell* freeBlock = nextFree_;
+    ExtraObjectCell* freeBlock = next;
-    nextFree_ = freeBlock->next_;
+    nextFree_.store(freeBlock->next_.load(std::memory_order_relaxed), std::memory_order_relaxed);
     CustomAllocDebug("ExtraObjectPage(%p)::TryAllocate() = %p", this, freeBlock->Data());
     return freeBlock->Data();
 }
@@ -52,10 +53,10 @@ bool ExtraObjectPage::Sweep(GCSweepScope& sweepHandle, FinalizerQueue& finalizer
     // necessarily match an actual block starting point.
     ExtraObjectCell* end = cells_ + EXTRA_OBJECT_COUNT;
     bool alive = false;
-    ExtraObjectCell** nextFree = &nextFree_;
+    std::atomic<ExtraObjectCell*>* nextFree = &nextFree_;
     for (ExtraObjectCell* cell = cells_; cell < end; ++cell) {
         // If the current cell is free, move on.
-        if (cell == *nextFree) {
+        if (cell == nextFree->load(std::memory_order_relaxed)) {
             nextFree = &cell->next_;
             continue;
         }
@@ -64,8 +65,8 @@ bool ExtraObjectPage::Sweep(GCSweepScope& sweepHandle, FinalizerQueue& finalizer
             alive = true;
         } else {
             // Free the current block and insert it into the free list.
-            cell->next_ = *nextFree;
+            cell->next_.store(nextFree->load(std::memory_order_relaxed), std::memory_order_relaxed);
-            *nextFree = cell;
+            nextFree->store(cell, std::memory_order_relaxed);
             nextFree = &cell->next_;
         }
     }

kotlin-native/runtime/src/custom_alloc/cpp/ExtraObjectPage.hpp

@@ -21,7 +21,7 @@ struct ExtraObjectCell {
 
     // This is used to simultaneously build two lists: a free list and a finalizers queue.
     // A cell cannot exist in both of them, but can be in neither when it's alive.
-    ExtraObjectCell* next_;
+    std::atomic<ExtraObjectCell*> next_;
     struct alignas(mm::ExtraObjectData) {
         uint8_t data_[sizeof(mm::ExtraObjectData)];
     };
@@ -54,8 +54,8 @@ private:
     ExtraObjectPage() noexcept;
 
     // Used for linking pages together in `pages` queue or in `unswept` queue.
-    ExtraObjectPage* next_;
+    std::atomic<ExtraObjectPage*> next_;
-    ExtraObjectCell* nextFree_;
+    std::atomic<ExtraObjectCell*> nextFree_;
     ExtraObjectCell cells_[];
 };
 

kotlin-native/runtime/src/custom_alloc/cpp/FixedBlockPage.hpp

@@ -53,7 +53,7 @@ private:
     friend class AtomicStack<FixedBlockPage>;
 
     // Used for linking pages together in `pages` queue or in `unswept` queue.
-    FixedBlockPage* next_;
+    std::atomic<FixedBlockPage*> next_;
     FixedCellRange nextFree_;
     uint32_t blockSize_;
     uint32_t end_;

kotlin-native/runtime/src/custom_alloc/cpp/NextFitPage.hpp

@@ -46,7 +46,7 @@ private:
     void UpdateCurBlock(uint32_t cellsNeeded) noexcept;
 
     friend class AtomicStack<NextFitPage>;
-    NextFitPage* next_;
+    std::atomic<NextFitPage*> next_;
 
     Cell* curBlock_;
     Cell cells_[]; // cells_[0] is reserved for an empty block

kotlin-native/runtime/src/custom_alloc/cpp/SingleObjectPage.hpp

@@ -41,7 +41,7 @@ private:
     // Testing method
     std_support::vector<uint8_t*> GetAllocatedBlocks() noexcept;
 
-    SingleObjectPage* next_;
+    std::atomic<SingleObjectPage*> next_;
     bool isAllocated_ = false;
     size_t size_;
     struct alignas(8) {