Code Review Updates #3

speedb-io · Aug 25, 2022 · bc26987 · bc26987
1 parent 5410240
commit bc26987
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 35 deletions.
diff --git a/memory/allocator.h b/memory/allocator.h
@@ -54,7 +54,6 @@ class AllocTracker {
 
  private:
   enum class State {
-    kNone,
     kAllocating,
     kDoneAllocating,
     kFreeMemStarted,
@@ -70,7 +69,7 @@ class AllocTracker {
 
  private:
   WriteBufferManager* write_buffer_manager_ = nullptr;
-  State state_ = State::kNone;
+  State state_ = State::kAllocating;
   std::atomic<size_t> bytes_allocated_ = 0U;
 };
 

diff --git a/memtable/alloc_tracker.cc b/memtable/alloc_tracker.cc
@@ -21,72 +21,79 @@ AllocTracker::~AllocTracker() { FreeMem(); }
 
 void AllocTracker::Allocate(size_t bytes) {
   assert(write_buffer_manager_ != nullptr);
-  assert((state_ == State::kNone) || (state_ == State::kAllocating));
+  assert(state_ == State::kAllocating);
 
-  if (ShouldUpdateWriteBufferManager() &&
-      ((state_ == State::kNone) || (state_ == State::kAllocating))) {
-    bytes_allocated_.fetch_add(bytes, std::memory_order_relaxed);
-    write_buffer_manager_->ReserveMem(bytes);
+  if (state_ == State::kAllocating) {
+    if (ShouldUpdateWriteBufferManager()) {
+      bytes_allocated_.fetch_add(bytes, std::memory_order_relaxed);
+      write_buffer_manager_->ReserveMem(bytes);
+    }
   }
-  state_ = State::kAllocating;
 }
 
 void AllocTracker::DoneAllocating() {
   assert(write_buffer_manager_ != nullptr);
+  assert(state_ == State::kAllocating);
 
-  if (ShouldUpdateWriteBufferManager() && (state_ == State::kAllocating)) {
-    write_buffer_manager_->ScheduleFreeMem(
-        bytes_allocated_.load(std::memory_order_relaxed));
-  } else {
-    assert(bytes_allocated_.load(std::memory_order_relaxed) == 0);
+  if (state_ == State::kAllocating) {
+    if (ShouldUpdateWriteBufferManager()) {
+      write_buffer_manager_->ScheduleFreeMem(
+          bytes_allocated_.load(std::memory_order_relaxed));
+    } else {
+      assert(bytes_allocated_.load(std::memory_order_relaxed) == 0);
+    }
+    state_ = State::kDoneAllocating;
   }
-  state_ = State::kDoneAllocating;
 }
 
 void AllocTracker::FreeMemStarted() {
   assert(write_buffer_manager_ != nullptr);
   assert(state_ == State::kDoneAllocating);
 
-  if (ShouldUpdateWriteBufferManager() && (state_ == State::kDoneAllocating)) {
-    write_buffer_manager_->FreeMemStarted(
-        bytes_allocated_.load(std::memory_order_relaxed));
+  if (state_ == State::kDoneAllocating) {
+    if (ShouldUpdateWriteBufferManager()) {
+      write_buffer_manager_->FreeMemStarted(
+          bytes_allocated_.load(std::memory_order_relaxed));
+    }
+    state_ = State::kFreeMemStarted;
   }
-  state_ = State::kFreeMemStarted;
 }
 
 void AllocTracker::FreeMemAborted() {
   assert(write_buffer_manager_ != nullptr);
   assert(state_ == State::kFreeMemStarted);
 
-  if (ShouldUpdateWriteBufferManager() && (state_ == State::kFreeMemStarted)) {
-    write_buffer_manager_->FreeMemAborted(
-        bytes_allocated_.load(std::memory_order_relaxed));
+  if (state_ == State::kFreeMemStarted) {
+    if (ShouldUpdateWriteBufferManager()) {
+      write_buffer_manager_->FreeMemAborted(
+          bytes_allocated_.load(std::memory_order_relaxed));
+    }
+    state_ = State::kDoneAllocating;
   }
-  state_ = State::kDoneAllocating;
 }
 
 void AllocTracker::FreeMem() {
-  if (state_ != State::kNone) {
-    if (state_ == State::kAllocating) {
-      DoneAllocating();
-    }
+  if (state_ == State::kAllocating) {
+    DoneAllocating();
+  }
 
-    // This is necessary so that the WBM will not decrease the memory being
-    // freed twice in case memory freeing was aborted and then freed via this
-    // call
-    if (state_ == State::kDoneAllocating) {
-      FreeMemStarted();
-    }
+  // This is necessary so that the WBM will not decrease the memory being
+  // freed twice in case memory freeing was aborted and then freed via this
+  // call
+  if (state_ == State::kDoneAllocating) {
+    FreeMemStarted();
+  }
 
-    if (ShouldUpdateWriteBufferManager() &&
-        (state_ == State::kFreeMemStarted)) {
+  if (state_ == State::kFreeMemStarted) {
+    if (ShouldUpdateWriteBufferManager()) {
       write_buffer_manager_->FreeMem(
           bytes_allocated_.load(std::memory_order_relaxed));
     } else {
       assert(bytes_allocated_.load(std::memory_order_relaxed) == 0);
     }
   }
-  bytes_allocated_ = 0U;
+
   state_ = State::kFreed;
 }
+
 }  // namespace ROCKSDB_NAMESPACE