From 706b7f8d43b0aecdc75c5ee49d3e4ef5f27b9faf Mon Sep 17 00:00:00 2001
From: ssid <ssid@google.com>
Date: Wed, 11 Nov 2015 03:54:18 -0800
Subject: [PATCH] Resolve race when getting approximate-memory-usage property

The write operations in the table happens without holding the mutex
lock, but concurrent writes are avoided using "writers_" queue.
The Arena::MemoryUsage could access the blocks when write happens.
So, the memory usage is cached in atomic word and can be loaded
from any thread safely.
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=107573379
---
 db/memtable.h |  5 +----
 util/arena.cc |  6 +++---
 util/arena.h  | 10 +++++-----
 3 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/db/memtable.h b/db/memtable.h
index 92e90bb..9f41567 100644
--- a/db/memtable.h
+++ b/db/memtable.h
@@ -36,10 +36,7 @@ class MemTable {
   }
 
   // Returns an estimate of the number of bytes of data in use by this
-  // data structure.
-  //
-  // REQUIRES: external synchronization to prevent simultaneous
-  // operations on the same MemTable.
+  // data structure. It is safe to call when MemTable is being modified.
   size_t ApproximateMemoryUsage();
 
   // Return an iterator that yields the contents of the memtable.
diff --git a/util/arena.cc b/util/arena.cc
index 9367f71..7407821 100644
--- a/util/arena.cc
+++ b/util/arena.cc
@@ -9,8 +9,7 @@ namespace leveldb {
 
 static const int kBlockSize = 4096;
 
-Arena::Arena() {
-  blocks_memory_ = 0;
+Arena::Arena() : memory_usage_(0) {
   alloc_ptr_ = NULL;  // First allocation will allocate a block
   alloc_bytes_remaining_ = 0;
 }
@@ -60,8 +59,9 @@ char* Arena::AllocateAligned(size_t bytes) {
 
 char* Arena::AllocateNewBlock(size_t block_bytes) {
   char* result = new char[block_bytes];
-  blocks_memory_ += block_bytes;
   blocks_.push_back(result);
+  memory_usage_.NoBarrier_Store(
+      reinterpret_cast<void*>(MemoryUsage() + block_bytes + sizeof(char*)));
   return result;
 }
 
diff --git a/util/arena.h b/util/arena.h
index 73bbf1c..48bab33 100644
--- a/util/arena.h
+++ b/util/arena.h
@@ -9,6 +9,7 @@
 #include <assert.h>
 #include <stddef.h>
 #include <stdint.h>
+#include "port/port.h"
 
 namespace leveldb {
 
@@ -24,10 +25,9 @@ class Arena {
   char* AllocateAligned(size_t bytes);
 
   // Returns an estimate of the total memory usage of data allocated
-  // by the arena (including space allocated but not yet used for user
-  // allocations).
+  // by the arena.
   size_t MemoryUsage() const {
-    return blocks_memory_ + blocks_.capacity() * sizeof(char*);
+    return reinterpret_cast<uintptr_t>(memory_usage_.NoBarrier_Load());
   }
 
  private:
@@ -41,8 +41,8 @@ class Arena {
   // Array of new[] allocated memory blocks
   std::vector<char*> blocks_;
 
-  // Bytes of memory in blocks allocated so far
-  size_t blocks_memory_;
+  // Total memory usage of the arena.
+  port::AtomicPointer memory_usage_;
 
   // No copying allowed
   Arena(const Arena&);