1 files changed, 22 insertions, 9 deletions
diff --git a/util/arena.cu b/util/arena.cu
index 5075318..db5827d 100644
--- a/util/arena.cu
+++ b/util/arena.cu
@@ -9,12 +9,16 @@ namespace leveldb {
 static const int kBlockSize = 4096;
 
 Arena::Arena()
-    : alloc_ptr_(nullptr), alloc_bytes_remaining_(0), memory_usage_(0) {}
+    : alloc_ptr_(nullptr), alloc_bytes_remaining_(0), memory_usage_(0),
+      head_(nullptr), blocks_(nullptr) {}
 
 Arena::~Arena() {
-  for (size_t i = 0; i < blocks_.size(); i++) {
-    //cudaFree(blocks_[i]);
-    delete [] blocks_[i];
+  ArenaNode * current = this->head_;
+  while (current != nullptr) {
+    ArenaNode * next = current->next;
+    cudaFree(current->block);
+    cudaFree(current);
+    current = next;
   }
 }
 
@@ -56,12 +60,21 @@ __device__ char* Arena::AllocateAligned(size_t bytes) {
   return result;
 }
 
-char* Arena::AllocateNewBlock(size_t block_bytes) {
-  char* result = new char[block_bytes];
-  //cudaMallocManaged((void **)&result, sizeof(char) * block_bytes);
-  blocks_.push_back(result);
+__device__ char* Arena::AllocateNewBlock(size_t block_bytes) {
+  char* result = nullptr;
+  cudaMalloc((void **)&result, sizeof(char) * block_bytes);
+  if (this->blocks_ == nullptr) {
+    cudaMalloc((void**)&this->blocks_, sizeof(ArenaNode));
+    // First alloc
+    this->head_ = this->blocks_;
+  } else {
+    cudaMalloc((void**)&this->blocks_->next, sizeof(ArenaNode));
+    this->blocks_ = this->blocks_->next;
+  }
+  this->blocks_->block = result;
+  this->blocks_->next = nullptr;
   memory_usage_.fetch_add(block_bytes + sizeof(char*),
-                          std::memory_order_relaxed);
+                          cuda::memory_order_relaxed);
   return result;
 }