TLS registry (#4554)

kotlin-native/runtime/src/legacymm/cpp/Memory.cpp

@@ -600,10 +600,10 @@ public:
         RuntimeAssert(storage_ == nullptr, "Storage must not be committed");
         auto it = map_->find(key);
         if (it != map_->end()) {
-            RuntimeAssert(it->second.second == size, "Attempt to add TLS record with the same key and different size");
+            RuntimeAssert(it->second.size == size, "Attempt to add TLS record with the same key and different size");
             return;
         }
-        map_->emplace(key, std::make_pair(size_, size));
+        map_->emplace(key, Entry{size_, size});
         size_ += size;
     }
 
@@ -629,15 +629,20 @@ public:
         }
         auto it = map_->find(key);
         RuntimeAssert(it != map_->end(), "Must be there");
-        int offset = it->second.first;
-        RuntimeAssert(offset + index < size_, "Out of bound in TLS access");
+        auto entry = it->second;
+        RuntimeAssert(index < entry.size, "Out of bounds in TLS access");
         lastKey_ = key;
-        lastOffset_ = offset;
-        return storage_ + offset + index;
+        lastOffset_ = entry.offset;
+        return storage_ + entry.offset + index;
     }
 
 private:
-    using Map = KStdUnorderedMap<Key, std::pair<int, int>>;
+    struct Entry {
+        int offset;
+        int size;
+    };
+
+    using Map = KStdUnorderedMap<Key, Entry>;
 
     Map* map_ = nullptr;
     KRef* storage_ = nullptr;

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

@@ -31,6 +31,10 @@ ALWAYS_INLINE mm::ThreadRegistry::Node* FromMemoryState(MemoryState* state) {
     return reinterpret_cast<mm::ThreadRegistry::Node*>(state);
 }
 
+ALWAYS_INLINE mm::ThreadData* GetThreadData(MemoryState* state) {
+    return FromMemoryState(state)->Get();
+}
+
 } // namespace
 
 extern "C" MemoryState* InitMemory(bool firstRuntime) {
@@ -55,3 +59,19 @@ extern "C" RUNTIME_NOTHROW void InitAndRegisterGlobal(ObjHeader** location, cons
     mm::GlobalsRegistry::Instance().RegisterStorageForGlobal(threadData, location);
     RuntimeCheck(false, "Unimplemented");
 }
+
+extern "C" RUNTIME_NOTHROW void AddTLSRecord(MemoryState* memory, void** key, int size) {
+    GetThreadData(memory)->tls().AddRecord(key, size);
+}
+
+extern "C" RUNTIME_NOTHROW void CommitTLSStorage(MemoryState* memory) {
+    GetThreadData(memory)->tls().Commit();
+}
+
+extern "C" RUNTIME_NOTHROW void ClearTLS(MemoryState* memory) {
+    GetThreadData(memory)->tls().Clear();
+}
+
+extern "C" RUNTIME_NOTHROW ObjHeader** LookupTLS(void** key, int index) {
+    return mm::ThreadRegistry::Instance().CurrentThreadData()->tls().Lookup(key, index);
+}

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

@@ -158,22 +158,6 @@ void EnsureNeverFrozen(ObjHeader* obj) {
     RuntimeCheck(false, "Unimplemented");
 }
 
-RUNTIME_NOTHROW void AddTLSRecord(MemoryState* memory, void** key, int size) {
-    RuntimeCheck(false, "Unimplemented");
-}
-
-RUNTIME_NOTHROW void CommitTLSStorage(MemoryState* memory) {
-    RuntimeCheck(false, "Unimplemented");
-}
-
-RUNTIME_NOTHROW void ClearTLS(MemoryState* memory) {
-    RuntimeCheck(false, "Unimplemented");
-}
-
-RUNTIME_NOTHROW ObjHeader** LookupTLS(void** key, int index) {
-    RuntimeCheck(false, "Unimplemented");
-}
-
 RUNTIME_NOTHROW void GC_RegisterWorker(void* worker) {
     RuntimeCheck(false, "Unimplemented");
 }

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

@@ -9,6 +9,7 @@
 #include <pthread.h>
 
 #include "GlobalsRegistry.hpp"
+#include "ThreadLocalStorage.hpp"
 #include "Utils.hpp"
 
 namespace kotlin {
@@ -31,9 +32,12 @@ public:
 
     GlobalsRegistry::ThreadQueue* globalsThreadQueue() noexcept { return &globalsThreadQueue_; }
 
+    ThreadLocalStorage& tls() noexcept { return tls_; }
+
 private:
     const pthread_t threadId_;
     GlobalsRegistry::ThreadQueue globalsThreadQueue_;
+    ThreadLocalStorage tls_;
 };
 
 } // namespace mm

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

@@ -0,0 +1,49 @@
+/*
+ * Copyright 2010-2020 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 "ThreadLocalStorage.hpp"
+
+using namespace kotlin;
+
+void mm::ThreadLocalStorage::AddRecord(Key key, int size) noexcept {
+    RuntimeAssert(state_ == State::kBuilding, "Storage must be in the building state");
+    RuntimeAssert(size >= 0, "Size cannot be negative");
+    auto it = map_.find(key);
+    if (it != map_.end()) {
+        RuntimeAssert(it->second.size == size, "Attempt to add TLS record with the same key, but different size");
+        return;
+    }
+    map_.emplace(key, Entry{size_, size});
+    size_ += size;
+}
+
+void mm::ThreadLocalStorage::Commit() noexcept {
+    RuntimeAssert(state_ == State::kBuilding, "Storage must be in the building state");
+    storage_.resize(size_);
+    state_ = State::kCommitted;
+}
+
+void mm::ThreadLocalStorage::Clear() noexcept {
+    RuntimeAssert(state_ == State::kCommitted, "Storage must be in the committed state");
+    // Just free the storage.
+    storage_.clear();
+    state_ = State::kCleared;
+}
+
+ObjHeader** mm::ThreadLocalStorage::Lookup(Key key, int index) noexcept {
+    RuntimeAssert(state_ == State::kCommitted, "Storage must be in the committed state");
+    if (lastKeyAndEntry_.first == key) {
+        return Lookup(lastKeyAndEntry_.second, index);
+    }
+    auto it = map_.find(key);
+    RuntimeAssert(it != map_.end(), "Unknown TLS key");
+    lastKeyAndEntry_ = *it;
+    return Lookup(it->second, index);
+}
+
+ObjHeader** mm::ThreadLocalStorage::Lookup(Entry entry, int index) noexcept {
+    RuntimeAssert(index < entry.size, "Out of bounds TLS access");
+    return &storage_[entry.offset + index];
+}

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

@@ -0,0 +1,78 @@
+/*
+ * Copyright 2010-2020 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.
+ */
+
+#ifndef RUNTIME_MM_THREAD_LOCAL_STORAGE_H
+#define RUNTIME_MM_THREAD_LOCAL_STORAGE_H
+
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "Memory.h"
+#include "Utils.hpp"
+
+namespace kotlin {
+namespace mm {
+
+class ThreadLocalStorage : Pinned {
+public:
+    using Key = void*;
+
+    class Iterator {
+    public:
+        explicit Iterator(std::vector<ObjHeader*>::iterator iterator) : iterator_(iterator) {}
+
+        ObjHeader** operator*() noexcept { return &*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:
+        std::vector<ObjHeader*>::iterator iterator_;
+    };
+
+    // Add TLS record. Can only be called before `Commit`.
+    void AddRecord(Key key, int size) noexcept;
+    // Prepare storage for records added by `AddRecord`.
+    void Commit() noexcept;
+    // Clear storage. Can only be called after `Commit`.
+    void Clear() noexcept;
+    // Lookup value in storage. Can only be called after `Commit`.
+    ObjHeader** Lookup(Key key, int index) noexcept;
+
+    Iterator begin() noexcept { return Iterator(storage_.begin()); }
+    Iterator end() noexcept { return Iterator(storage_.end()); }
+
+private:
+    enum class State {
+        kBuilding,
+        kCommitted,
+        kCleared,
+    };
+
+    struct Entry {
+        int offset;
+        int size;
+    };
+
+    ObjHeader** Lookup(Entry entry, int index) noexcept;
+
+    std::vector<ObjHeader*> storage_;
+    // TODO: `std::unordered_map` is probably the wrong container here.
+    std::unordered_map<Key, Entry> map_;
+    State state_ = State::kBuilding;
+    int size_ = 0; // Only used in `State::kBuilding`
+    std::pair<Key, Entry> lastKeyAndEntry_;
+};
+
+} // namespace mm
+} // namespace kotlin
+
+#endif // RUNTIME_MM_THREAD_LOCAL_STORAGE_H

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

@@ -0,0 +1,175 @@
+/*
+ * Copyright 2010-2020 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 "ThreadLocalStorage.hpp"
+
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+using namespace kotlin;
+
+namespace {
+
+struct Key {};
+
+} // namespace
+
+TEST(ThreadLocalStorageTest, Lookup) {
+    Key key1;
+    Key key2;
+    mm::ThreadLocalStorage tls;
+
+    tls.AddRecord(&key1, 1);
+    tls.AddRecord(&key2, 2);
+    tls.Commit();
+
+    ObjHeader** location1 = tls.Lookup(&key1, 0);
+    ObjHeader** location2 = tls.Lookup(&key2, 0);
+    ObjHeader** location3 = tls.Lookup(&key2, 1);
+
+    // Locations are not nulls.
+    EXPECT_NE(location1, nullptr);
+    EXPECT_NE(location2, nullptr);
+    EXPECT_NE(location3, nullptr);
+
+    // All three are different.
+    EXPECT_NE(location1, location2);
+    EXPECT_NE(location1, location3);
+    EXPECT_NE(location2, location3);
+
+    // All three can be written into.
+    *location1 = nullptr;
+    *location2 = nullptr;
+    *location3 = nullptr;
+}
+
+TEST(ThreadLocalStorageTest, Iterate) {
+    Key key1;
+    Key key2;
+    mm::ThreadLocalStorage tls;
+
+    tls.AddRecord(&key1, 1);
+    tls.AddRecord(&key2, 2);
+    tls.Commit();
+
+    std::vector<ObjHeader**> expected;
+    expected.push_back(tls.Lookup(&key1, 0));
+    expected.push_back(tls.Lookup(&key2, 0));
+    expected.push_back(tls.Lookup(&key2, 1));
+
+    std::vector<ObjHeader**> actual;
+    for (auto item : tls) {
+        actual.push_back(item);
+    }
+
+    EXPECT_THAT(actual, testing::ElementsAreArray(expected));
+}
+
+TEST(ThreadLocalStorageTest, AddRecordEmpty) {
+    Key key1;
+    Key key2;
+    Key key3;
+    mm::ThreadLocalStorage tls;
+
+    tls.AddRecord(&key1, 1);
+    tls.AddRecord(&key2, 0);
+    tls.AddRecord(&key3, 2);
+    tls.Commit();
+
+    std::vector<ObjHeader**> expected;
+    expected.push_back(tls.Lookup(&key1, 0));
+    expected.push_back(tls.Lookup(&key3, 0));
+    expected.push_back(tls.Lookup(&key3, 1));
+
+    std::vector<ObjHeader**> actual;
+    for (auto item : tls) {
+        actual.push_back(item);
+    }
+
+    EXPECT_THAT(actual, testing::ElementsAreArray(expected));
+}
+
+TEST(ThreadLocalStorageTest, AddRecordSameSize) {
+    Key key1;
+    mm::ThreadLocalStorage tls;
+
+    tls.AddRecord(&key1, 1);
+    tls.AddRecord(&key1, 1);
+    tls.Commit();
+
+    std::vector<ObjHeader**> expected;
+    expected.push_back(tls.Lookup(&key1, 0));
+
+    std::vector<ObjHeader**> actual;
+    for (auto item : tls) {
+        actual.push_back(item);
+    }
+
+    EXPECT_THAT(actual, testing::ElementsAreArray(expected));
+}
+
+TEST(ThreadLocalStorageTest, NoRecords) {
+    mm::ThreadLocalStorage tls;
+
+    tls.Commit();
+
+    std::vector<ObjHeader**> actual;
+    for (auto item : tls) {
+        actual.push_back(item);
+    }
+
+    EXPECT_THAT(actual, testing::IsEmpty());
+}
+
+TEST(ThreadLocalStorageTest, ClearEmpty) {
+    mm::ThreadLocalStorage tls;
+
+    tls.Commit();
+
+    tls.Clear();
+
+    std::vector<ObjHeader**> actual;
+    for (auto item : tls) {
+        actual.push_back(item);
+    }
+
+    EXPECT_THAT(actual, testing::IsEmpty());
+}
+
+TEST(ThreadLocalStorageTest, ClearNonEmpty) {
+    Key key1;
+    mm::ThreadLocalStorage tls;
+
+    tls.AddRecord(&key1, 1);
+    tls.Commit();
+
+    tls.Clear();
+
+    std::vector<ObjHeader**> actual;
+    for (auto item : tls) {
+        actual.push_back(item);
+    }
+
+    EXPECT_THAT(actual, testing::IsEmpty());
+}
+
+TEST(ThreadLocalStorageTest, LookupCaching) {
+    Key key1;
+    Key key2;
+    mm::ThreadLocalStorage tls;
+
+    tls.AddRecord(&key1, 1);
+    tls.AddRecord(&key2, 1);
+    tls.Commit();
+
+    ObjHeader** location1 = tls.Lookup(&key1, 0);
+    ObjHeader** location2 = tls.Lookup(&key2, 0);
+
+    // Lookup same stuff again in different order.
+    EXPECT_EQ(location1, tls.Lookup(&key1, 0));
+    EXPECT_EQ(location2, tls.Lookup(&key2, 0));
+    EXPECT_EQ(location2, tls.Lookup(&key2, 0));
+    EXPECT_EQ(location1, tls.Lookup(&key1, 0));
+}