From: Amir Shalem <amirshalem@fb.com>
Date: Wed, 21 Dec 2016 06:25:34 +0000 (-0800)
Subject: Fix bug where capacity is not updated correctly after reserveLarge()
X-Git-Tag: v2017.03.06.00~175
X-Git-Url: http://demsky.eecs.uci.edu/git/?a=commitdiff_plain;h=6d04232339c9420554eb203a227da8c31aa5477c;p=folly.git

Fix bug where capacity is not updated correctly after reserveLarge()

Summary:
When reallocating new large string, using reserveLarge(), make sure to update the new capacity for the RefCounted string data
to reflect the amount of data allocated by jemalloc for our block (using goodMallocSize())

Reviewed By: Gownta

Differential Revision: D4355440

fbshipit-source-id: f2d58e8888e973418781220d57ff46f674e20556
---

diff --git a/folly/FBString.h b/folly/FBString.h
index 43f4e33e..351b0de0 100644
--- a/folly/FBString.h
+++ b/folly/FBString.h
@@ -575,8 +575,10 @@ private:
     static RefCounted * reallocate(Char *const data,
                                    const size_t currentSize,
                                    const size_t currentCapacity,
-                                   const size_t newCapacity) {
-      FBSTRING_ASSERT(newCapacity > 0 && newCapacity > currentSize);
+                                   size_t * newCapacity) {
+      FBSTRING_ASSERT(*newCapacity > 0 && *newCapacity > currentSize);
+      const size_t allocNewCapacity = goodMallocSize(
+        sizeof(RefCounted) + *newCapacity * sizeof(Char));
       auto const dis = fromData(data);
       FBSTRING_ASSERT(dis->refCount_.load(std::memory_order_acquire) == 1);
       // Don't forget to allocate one extra Char for the terminating
@@ -586,8 +588,9 @@ private:
              smartRealloc(dis,
                           sizeof(RefCounted) + currentSize * sizeof(Char),
                           sizeof(RefCounted) + currentCapacity * sizeof(Char),
-                          sizeof(RefCounted) + newCapacity * sizeof(Char)));
+                          allocNewCapacity));
       FBSTRING_ASSERT(result->refCount_.load(std::memory_order_acquire) == 1);
+      *newCapacity = (allocNewCapacity - sizeof(RefCounted)) / sizeof(Char);
       return result;
     }
   };
@@ -836,7 +839,7 @@ FOLLY_MALLOC_NOINLINE inline void fbstring_core<Char>::reserveLarge(
     if (minCapacity > ml_.capacity()) {
       // Asking for more memory
       auto const newRC = RefCounted::reallocate(
-          ml_.data_, ml_.size_, ml_.capacity(), minCapacity);
+          ml_.data_, ml_.size_, ml_.capacity(), &minCapacity);
       ml_.data_ = newRC->data_;
       ml_.setCapacity(minCapacity, Category::isLarge);
     }
diff --git a/folly/test/FBStringTest.cpp b/folly/test/FBStringTest.cpp
index fb4a7991..6c6f099c 100644
--- a/folly/test/FBStringTest.cpp
+++ b/folly/test/FBStringTest.cpp
@@ -1270,6 +1270,17 @@ TEST(FBString, testFixedBugs) {
   { // D3698862
     EXPECT_EQ(fbstring().find(fbstring(), 4), fbstring::npos);
   }
+  if (usingJEMalloc()) { // D4355440
+    fbstring str(1337, 'f');
+    str.reserve(3840);
+    EXPECT_NE(str.capacity(), 3840);
+
+    struct {
+      std::atomic<size_t> refCount_;
+      char data_[1];
+    } dummyRefCounted;
+    EXPECT_EQ(str.capacity(), goodMallocSize(3840) - sizeof(dummyRefCounted));
+  }
 }
 
 TEST(FBString, findWithNpos) {