From f83373a5486bf7cad48fbd662c7f7f9484e79a65 Mon Sep 17 00:00:00 2001
From: KunoiSayami <i@leanhe.dev>
Date: Thu, 17 Feb 2022 19:25:39 +0800
Subject: fix(arena): Fix arena alloc

Signed-off-by: KunoiSayami <i@leanhe.dev>
---
 CMakeLists.txt |  3 ++-
 util/arena.cu  | 30 +++++++++++++++++-------------
 2 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 0890b9f..b6d0d01 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -347,7 +347,6 @@ if(LEVELDB_BUILD_TESTS)
         "db/filename_test.cc"
         "db/log_test.cc"
         "db/recovery_test.cc"
-        "db/skiplist_test.cc"
         "db/version_edit_test.cc"
         "db/version_set_test.cc"
         "db/write_batch_test.cc"
@@ -362,6 +361,7 @@ if(LEVELDB_BUILD_TESTS)
         "util/hash_test.cc"
         "util/logging_test.cc"
     )
+    set_target_properties(leveldb_tests PROPERTIES CUDA_ARCHITECTURES "75")
   endif(NOT BUILD_SHARED_LIBS)
   target_link_libraries(leveldb_tests leveldb gmock gtest gtest_main)
   target_compile_definitions(leveldb_tests
@@ -407,6 +407,7 @@ if(LEVELDB_BUILD_TESTS)
   endfunction(leveldb_test)
 
   leveldb_test("db/c_test.c")
+  leveldb_test("db/skiplist_test.cu")
 
   if(NOT BUILD_SHARED_LIBS)
     # TODO(costan): This test also uses
diff --git a/util/arena.cu b/util/arena.cu
index 6d1b76b..f954577 100644
--- a/util/arena.cu
+++ b/util/arena.cu
@@ -61,30 +61,34 @@ __device__ char* Arena::AllocateAligned(size_t bytes) {
 }
 
 __device__ char* Arena::AllocateNewBlock(size_t block_bytes) {
+  // Allocate new target
   char* result = nullptr;
-  cuda::atomic<ArenaNode*>* block_alloc = nullptr;
+  // Allocate storage target
+  ArenaNode* node = nullptr;
+  // Malloc memory
   cudaMalloc((void **)&result, sizeof(char) * block_bytes);
-  cudaMalloc((void **)&block_alloc, sizeof(cuda::atomic<ArenaNode*>));
+  cudaMalloc((void **)&node, sizeof(ArenaNode));
+
+  node->next.store(nullptr);
+
   while (true) {
       ArenaNode * current_ = this->blocks_.load(cuda::memory_order_acquire);
       if (current_ == nullptr) {
-      // cudaMalloc((void**)&this->blocks_, sizeof(ArenaNode));
-          ArenaNode * current_end = this->blocks_.load();
-          if (!this->blocks_.compare_exchange_weak(current_end, reinterpret_cast<ArenaNode*>(block_alloc)))
+          if (!this->blocks_.compare_exchange_strong(current_, node))
+            continue;
+          if (!this->head_.compare_exchange_strong(current_, node))
             continue;
-          if (!this->head_.compare_exchange_weak(current_end, reinterpret_cast<ArenaNode*>(block_alloc)))
-             continue;
+          node->block = result;
           break;
       }
-      ArenaNode * except_next = current_->next;
-      if (except_next != nullptr)
-          continue;
-      if (!this->blocks_.compare_exchange_weak(current_, reinterpret_cast<ArenaNode*>(block_alloc)))
-          continue;
+      ArenaNode * expect_next = nullptr;
+      if (current_->next.compare_exchange_strong(expect_next, node)) {
+        continue;
+      }
       current_->block = result;
-      current_->next = nullptr;
       break;
   }
+
   memory_usage_.fetch_add(block_bytes + sizeof(char*),
                           cuda::memory_order_relaxed);
   return result;
-- 
cgit v1.3.1