From 4fea7d70a4d1dbccb62e04f80106e4980aed60fa Mon Sep 17 00:00:00 2001
From: KunoiSayami <i@leanhe.dev>
Date: Wed, 19 Jan 2022 19:40:07 +0800
Subject: test(skiplist): Fix parallel insert error

Signed-off-by: KunoiSayami <i@leanhe.dev>
---
 db/skiplist.cuh | 80 +++++++++++++++++++++++++++++++++++++++++++++++----------
 1 file changed, 66 insertions(+), 14 deletions(-)

(limited to 'db/skiplist.cuh')

diff --git a/db/skiplist.cuh b/db/skiplist.cuh
index 9f1cdaa..3f7b7e9 100644
--- a/db/skiplist.cuh
+++ b/db/skiplist.cuh
@@ -37,6 +37,43 @@
 
 namespace leveldb {
 
+class CudaSpinLock {
+  static constexpr int UNLOCKED = 0;
+  static constexpr int LOCKED = 1;
+
+  cuda::atomic<int> m_value;
+  bool isFake;
+
+ public:
+
+  __device__ __host__ explicit CudaSpinLock(): m_value(UNLOCKED), isFake(false) {}
+
+  __device__ __host__ explicit CudaSpinLock(bool fake): m_value(UNLOCKED), isFake(fake) {}
+
+  __device__ void lock()
+  {
+    if (!isFake) {
+      while (true) {
+        int expected = UNLOCKED;
+        // this->m_value.wait(LOCKED);
+        if (this->m_value.compare_exchange_weak(expected, LOCKED)) break;
+      }
+    }
+  }
+
+  __device__ void unlock()
+  {
+    if (!isFake) {
+      m_value.store(UNLOCKED);
+    }
+  }
+
+  __device__ bool isLock() {
+    //printf("%d\n", this->m_value.load());
+    return !isFake && this->m_value.load() == LOCKED;
+  }
+};
+
 class Arena;
 
 template <typename Key, class Comparator>
@@ -139,6 +176,9 @@ class SkipList {
   // values are ok.
   cuda::atomic<int> max_height_; // Height of the entire list
 
+  CudaSpinLock arena_lock_;
+  CudaSpinLock find_lock_;
+
   // Read/written only by Insert().
   Random rnd_;
 };
@@ -185,18 +225,21 @@ struct SkipList<Key, Comparator>::Node {
     next_[n].store(x, cuda::memory_order_release);
   }
 
-  __device__ bool SetNextSafe(int n, Node *x, Node * origin) {
+  __device__ bool SetNextSafe(int n, Node *x) {
     assert(n >= 0);
+    Node * origin = next_[n].load(cuda::memory_order_acquire);
+    //assert(origin != x);
     if (origin == nullptr) {
-      next_[n].store(x, cuda::memory_order_release);
+      next_[n].store(x);
+      x->NoBarrier_SetNext(n, origin);
       return true;
     }
-    printf("%d %p %p\n", n, x, origin);
+    if (origin->key < x->key) {
+      return false;
+    }
     bool ret = next_[n].compare_exchange_weak(origin, x, cuda::memory_order_acquire);
-    if (!ret) {
-      Node * rep = next_[n].load(cuda::memory_order_acquire);
-      printf("%p %p\n", rep, origin);
-      assert(rep != origin);
+    if (ret) {
+      x->SetNext(n, origin);
     }
     return ret;
   }
@@ -219,9 +262,12 @@ struct SkipList<Key, Comparator>::Node {
 template <typename Key, class Comparator>
 __device__ typename SkipList<Key, Comparator>::Node* SkipList<Key, Comparator>::NewNode(
     const Key& key, int height) {
+  arena_lock_.lock();
   char* const node_memory = arena_->AllocateAligned(
       sizeof(Node) + sizeof(std::atomic<Node*>) * (height - 1));
-  return new (node_memory) Node(key);
+  Node* ret = new (node_memory) Node(key);
+  arena_lock_.unlock();
+  return ret;
 }
 
 template <typename Key, class Comparator>
@@ -303,6 +349,8 @@ __device__ SkipList<Key, Comparator>::FindGreaterOrEqual(const Key& key,
   int level = GetMaxHeight() - 1;
   while (true) {
     Node* next = x->Next(level);
+    assert(x != next);
+    //printf("%u %lu %x %p\n", threadIdx.x, key, x, next);
     if (KeyIsAfterNode(key, next)) {
       // Keep searching in this list
       x = next;
@@ -365,7 +413,9 @@ __device__ SkipList<Key, Comparator>::SkipList(Comparator cmp, Arena* arena)
       arena_(arena),
       head_(NewNode(0 /* any key will do */, kMaxHeight)),
       max_height_(1),
-      rnd_(0xdeadbeef) {
+      rnd_(0xdeadbeef),
+      arena_lock_(),
+      find_lock_() {
   for (int i = 0; i < kMaxHeight; i++) {
     head_->SetNext(i, nullptr);
   }
@@ -381,6 +431,7 @@ __device__ void SkipList<Key, Comparator>::Insert(const Key& key) {
   // Our data structure does not allow duplicate insertion
   assert(x == nullptr || !Equal(key, x->key));
 
+
   int height = RandomHeight();
   if (height > GetMaxHeight()) {
     for (int i = GetMaxHeight(); i < height; i++) {
@@ -396,19 +447,20 @@ __device__ void SkipList<Key, Comparator>::Insert(const Key& key) {
     max_height_.store(height, cuda::memory_order_relaxed);
   }
 
-  Node * original_next;
   x = NewNode(key, height);
+  //printf("x: %p\n", x);
   for (int i = 0; i < height; i++) {
     // NoBarrier_SetNext() suffices since we will add a barrier when
     // we publish a pointer to "x" in prev[i].
     while (true) {
-      original_next = prev[i]->NoBarrier_Next(i);
-      x->NoBarrier_SetNext(i, original_next);
-      if (prev[i]->SetNextSafe(i, x, original_next))
+      assert(prev[i]->key < key);
+      if (prev[i]->SetNextSafe(i, x))
         break;
-      printf("failed %p\n", original_next);
+      //printf("failed\n");
       FindGreaterOrEqual(key, prev);
     }
+    /*x->NoBarrier_SetNext(i, prev[i]->NoBarrier_Next(i));
+    prev[i]->SetNext(i, x);*/
   }
 }
 
-- 
cgit v1.3.1