feat: Add user key back

Signed-off-by: KunoiSayami <[email protected]>

1 files changed, 13 insertions, 14 deletions

diff --git a/db/memtable.cu b/db/memtable.cu
index 8b1fab5..e021605 100644
--- a/db/memtable.cu
+++ b/db/memtable.cu
@@ -87,7 +87,7 @@ void MemTable::Add(SequenceNumber s, ValueType type, const Slice& key,
                    const Slice& value) {
   // Format of an entry is concatenation of:
   //  key_size     : varint32 of internal_key.size()
-  //  key point    : point to host memory
+  //  key bytes    : char[internal_key.size()]
   //  tag          : uint64((sequence << 8) | type)
   //  value_size   : varint32 of value.size()
   //  value point  : point to host memory
@@ -97,7 +97,7 @@ void MemTable::Add(SequenceNumber s, ValueType type, const Slice& key,
   //const size_t encoded_len = VarintLength(internal_key_size) +
   //                           internal_key_size + VarintLength(val_size) +
   //                           val_size;
-  const size_t encoded_len = VarintLength(internal_key_size) + 8 +
+  const size_t encoded_len = VarintLength(internal_key_size) + internal_key_size +
                              VarintLength(val_size) + 8;
 
 
@@ -113,8 +113,9 @@ void MemTable::Add(SequenceNumber s, ValueType type, const Slice& key,
   char * insert_val = new char[encoded_len], *cuda_insert = nullptr;
   //  EncodeVarint32(insert_val, encoded_len);
   char * p = EncodeVarint32(insert_val + 8, internal_key_size);
-  EncodeFixed64(p, reinterpret_cast<uint64_t>(key_mem));
-  p += 8;
+  //EncodeFixed64(p, reinterpret_cast<uint64_t>(key_mem));
+  memcpy(p, key_mem, key_size);
+  p += key_size;
   EncodeFixed64(p, (s << 8) | type );
   p += 8;
   //EncodeFixed64(p, reinterpret_cast<uint64_t>(val_size));
@@ -147,11 +148,12 @@ __global__ void Get_(MemTable * met, char * memkey, char ** data, size_t* malloc
     // Check that it belongs to same user key.  We do not check the
     // sequence number since the Seek() call above should have skipped
     // all entries with overly large sequence numbers.
+    size_t key_size;
     const char *entry = iter.key();
     const char * p = GetVarint32PtrCuda(entry, entry + 5,
-                       reinterpret_cast<uint32_t*>(malloc_size));
-    p = GetVarint32PtrCuda(p + 8, p + 13, nullptr);
-    *malloc_size = p - entry;
+                       reinterpret_cast<uint32_t*>(&key_size));
+    p = GetVarint32PtrCuda(p + key_size, p + key_size + 5, nullptr);
+    *malloc_size = (p - entry + 8);
     cudaMalloc((void**)*data, *malloc_size);
     memcpy(*data, entry, *malloc_size);
   }
@@ -192,9 +194,9 @@ bool MemTable::Get(const LookupKey& key, std::string* value, Status* s) {
 
   Get_<<<1,1>>>(this, cuda_mem_key, cuda_skiplist_key, cuda_malloc_size);
   cudaDeviceSynchronize();
+  cudaFree(cuda_mem_key);
 
   if (*cuda_skiplist_key == nullptr) {
-    cudaFree(cuda_mem_key);
     cudaFree(cuda_skiplist_key);
     cudaFree(cuda_malloc_size);
     return false;
@@ -205,23 +207,20 @@ bool MemTable::Get(const LookupKey& key, std::string* value, Status* s) {
   char * entry = this->host_arena_.Allocate(*malloc_size);
   cudaMemcpy(entry, cuda_skiplist_key, *malloc_size, cudaMemcpyDeviceToHost);
 
-  cudaFree(cuda_mem_key);
   cudaFree(cuda_skiplist_key);
   cudaFree(cuda_malloc_size);
   delete malloc_size;
 
     uint32_t key_length;
-    const char* key_ptr_ptr = GetVarint32Ptr(entry, entry + 5, &key_length);
-    const char * key_ptr =
-        reinterpret_cast<const char*>(DecodeFixed64(key_ptr_ptr));
+    const char* key_ptr = GetVarint32Ptr(entry, entry + 5, &key_length);
     if (comparator_.comparator.user_comparator()->Compare(
             Slice(key_ptr, key_length - 8), key.user_key()) == 0) {
       // Correct user key
-      const uint64_t tag = DecodeFixed64(key_ptr_ptr + key_length - 8);
+      const uint64_t tag = DecodeFixed64(key_ptr + key_length - 8);
       switch (static_cast<ValueType>(tag & 0xff)) {
         case kTypeValue: {
           // HOW CAN WE GET IT?
-          Slice v = GetLengthPrefixedSlice(key_ptr_ptr + key_length);
+          Slice v = GetLengthPrefixedSlice(key_ptr + key_length);
           value->assign(v.data(), v.size());
           return true;
         }