SERVER-36243 Use sized deallocation.

Added mongoFree to be used when allocating memory with mongoMalloc.
It has an overload taking size utilizing tc_free_sized if built with tcmalloc.

12 files changed, 119 insertions, 33 deletions

diff --git a/src/mongo/SConscript b/src/mongo/SConscript
index 39b12a203b3..b295b98d9d2 100644
--- a/src/mongo/SConscript
+++ b/src/mongo/SConscript
@@ -48,10 +48,29 @@ env.SConscript(
     ],
 )
 
+allocatorEnv = env.Clone()
+if allocatorEnv['MONGO_ALLOCATOR'] == 'tcmalloc':
+    allocatorEnv.Append(
+            CPPDEFINES=[
+                'MONGO_USE_GPERFTOOLS_TCMALLOC',
+            ]
+        )
+
+    if not use_system_version_of_library('tcmalloc'):
+        # Add in the include path for our vendored tcmalloc.
+        allocatorEnv.InjectThirdParty('gperftools')
+
+        # If our changes to tcmalloc are ever upstreamed, this should become set based on a top
+        # level configure check, though its effects should still be scoped just to these files.
+
+    if not use_system_version_of_library('valgrind'):
+        # Include valgrind since tcmalloc disables itself while running under valgrind
+        allocatorEnv.InjectThirdParty('valgrind')
+
 # NOTE: This library does not really belong here. Its presence here is
 # temporary. Do not add to this library, do not remove from it, and do
 # not declare other libraries in this file.
-env.Library(
+allocatorEnv.Library(
     target='base',
     source=[
         'base/data_range.cpp',
diff --git a/src/mongo/base/data_builder.h b/src/mongo/base/data_builder.h
index 00292c92a99..1b70e376f79 100644
--- a/src/mongo/base/data_builder.h
+++ b/src/mongo/base/data_builder.h
@@ -54,10 +54,30 @@ class DataBuilder {
     /**
      * The dtor type used in the unique_ptr which holds the buffer
      */
-    struct FreeBuf {
-        void operator()(char* buf) {
-            std::free(buf);
+    struct BufDeleter {
+        BufDeleter() : _capacity(0) {}
+        explicit BufDeleter(size_t capacity) : _capacity(capacity) {}
+
+        BufDeleter(BufDeleter&& other) : _capacity(other._capacity) {
+            other._capacity = 0;
+        }
+
+        BufDeleter& operator=(BufDeleter&& other) {
+            _capacity = other._capacity;
+            other._capacity = 0;
+            return *this;
+        }
+
+        void operator()(char* buf) const {
+            mongoFree(buf, _capacity);
         }
+
+        size_t capacity() const {
+            return _capacity;
+        }
+
+    private:
+        size_t _capacity;
     };
 
     static const std::size_t kInitialBufferSize = 64;
@@ -79,10 +99,8 @@ public:
 
     DataBuilder& operator=(DataBuilder&& other) {
         _buf = std::move(other._buf);
-        _capacity = other._capacity;
         _unwrittenSpaceCursor = {_buf.get(), _buf.get() + other.size()};
 
-        other._capacity = 0;
         other._unwrittenSpaceCursor = {nullptr, nullptr};
 
         return *this;
@@ -151,14 +169,14 @@ public:
             return 0;
         }
 
-        return _capacity - _unwrittenSpaceCursor.length();
+        return capacity() - _unwrittenSpaceCursor.length();
     }
 
     /**
      * The total size of the buffer, including reserved but not written bytes.
      */
     std::size_t capacity() const {
-        return _capacity;
+        return _buf.get_deleter().capacity();
     }
 
     /**
@@ -166,7 +184,7 @@ public:
      * grow it.
      */
     void resize(std::size_t newSize) {
-        if (newSize == _capacity)
+        if (newSize == capacity())
             return;
 
         if (newSize == 0) {
@@ -178,12 +196,12 @@ public:
 
         auto ptr = _buf.release();
 
-        _buf.reset(static_cast<char*>(mongoRealloc(ptr, newSize)));
-
-        _capacity = newSize;
+        _buf = std::unique_ptr<char, BufDeleter>(static_cast<char*>(mongoRealloc(ptr, newSize)),
+                                                 BufDeleter(newSize));
 
         // If we downsized, truncate. If we upsized keep the old size
-        _unwrittenSpaceCursor = {_buf.get() + std::min(oldSize, _capacity), _buf.get() + _capacity};
+        _unwrittenSpaceCursor = {_buf.get() + std::min(oldSize, capacity()),
+                                 _buf.get() + capacity()};
     }
 
     /**
@@ -194,7 +212,9 @@ public:
     void reserve(std::size_t needed) {
         std::size_t oldSize = size();
 
-        std::size_t newSize = _capacity ? _capacity : kInitialBufferSize;
+        std::size_t newSize = capacity();
+        if (newSize == 0)
+            newSize = kInitialBufferSize;
 
         while ((newSize < oldSize) || (newSize - oldSize < needed)) {
             // growth factor of about 1.5
@@ -212,14 +232,14 @@ public:
      * internal data pointers.
      */
     void clear() {
-        _unwrittenSpaceCursor = {_buf.get(), _buf.get() + _capacity};
+        _unwrittenSpaceCursor = {_buf.get(), _buf.get() + capacity()};
     }
 
     /**
      * Release the buffer. After this the builder is left in the default
      * constructed state.
      */
-    std::unique_ptr<char, FreeBuf> release() {
+    std::unique_ptr<char, BufDeleter> release() {
         auto buf = std::move(_buf);
 
         *this = DataBuilder{};
@@ -252,8 +272,7 @@ private:
         }
     }
 
-    std::unique_ptr<char, FreeBuf> _buf;
-    std::size_t _capacity = 0;
+    std::unique_ptr<char, BufDeleter> _buf;
     DataRangeCursor _unwrittenSpaceCursor = {nullptr, nullptr};
 };
 
diff --git a/src/mongo/base/secure_allocator.h b/src/mongo/base/secure_allocator.h
index 54e7df3b915..00b63b09652 100644
--- a/src/mongo/base/secure_allocator.h
+++ b/src/mongo/base/secure_allocator.h
@@ -61,7 +61,7 @@ inline void deallocateWrapper(void* ptr, std::size_t bytes, bool secure) {
     if (secure) {
         return deallocate(ptr, bytes);
     } else {
-        return free(ptr);
+        return mongoFree(ptr, bytes);
     }
 }
 
diff --git a/src/mongo/bson/util/builder.h b/src/mongo/bson/util/builder.h
index b17b6d4752e..6f19e6e3cd9 100644
--- a/src/mongo/bson/util/builder.h
+++ b/src/mongo/bson/util/builder.h
@@ -136,7 +136,7 @@ public:
     }
     void free() {
         if (_ptr != _buf)
-            ::free(_ptr);
+            mongoFree(_ptr);
         _ptr = _buf;
     }
 
diff --git a/src/mongo/client/cyrus_sasl_client_session.cpp b/src/mongo/client/cyrus_sasl_client_session.cpp
index d2201be2149..67f8d6d2b5d 100644
--- a/src/mongo/client/cyrus_sasl_client_session.cpp
+++ b/src/mongo/client/cyrus_sasl_client_session.cpp
@@ -83,6 +83,10 @@ void* saslOurRealloc(void* ptr, SaslAllocSize sz) {
     return mongoRealloc(ptr, sz);
 }
 
+void saslOurFree(void* ptr) {
+    mongoFree(ptr);
+}
+
 /*
  * Mutex functions to be used by the SASL library, if the client doesn't initialize the library
  * for us.
@@ -111,7 +115,7 @@ void saslMutexFree(void* mutex) {
  * unless the client application has previously initialized the SASL library.
  */
 MONGO_INITIALIZER(CyrusSaslAllocatorsAndMutexes)(InitializerContext*) {
-    sasl_set_alloc(saslOurMalloc, saslOurCalloc, saslOurRealloc, free);
+    sasl_set_alloc(saslOurMalloc, saslOurCalloc, saslOurRealloc, saslOurFree);
 
     sasl_set_mutex(saslMutexAlloc, saslMutexLock, saslMutexUnlock, saslMutexFree);
     return Status::OK();
diff --git a/src/mongo/dbtests/jsobjtests.cpp b/src/mongo/dbtests/jsobjtests.cpp
index 3ef6d5c2eaf..5c3336b8e2a 100644
--- a/src/mongo/dbtests/jsobjtests.cpp
+++ b/src/mongo/dbtests/jsobjtests.cpp
@@ -1853,7 +1853,7 @@ public:
             memcpy(crap, x.objdata(), x.objsize());
             BSONObj y(crap);
             ASSERT_BSONOBJ_EQ(x, y);
-            free(crap);
+            mongoFree(crap);
         }
 
         {
@@ -1869,7 +1869,7 @@ public:
                 state = 2;
                 ASSERT(strstr(e.what(), "_id: 5") != NULL);
             }
-            free(crap);
+            mongoFree(crap);
             ASSERT_EQUALS(2, state);
         }
     }
diff --git a/src/mongo/util/allocator.cpp b/src/mongo/util/allocator.cpp
index 358689a0655..d5beecdca87 100644
--- a/src/mongo/util/allocator.cpp
+++ b/src/mongo/util/allocator.cpp
@@ -33,10 +33,18 @@
 
 #include "mongo/util/signal_handlers_synchronous.h"
 
+#if defined(MONGO_USE_GPERFTOOLS_TCMALLOC)
+#include <gperftools/tcmalloc.h>
+#endif
+
 namespace mongo {
 
 void* mongoMalloc(size_t size) {
+#if defined(MONGO_USE_GPERFTOOLS_TCMALLOC)
+    void* x = tc_malloc(size);
+#else
     void* x = std::malloc(size);
+#endif
     if (x == NULL) {
         reportOutOfMemoryErrorAndExit();
     }
@@ -44,11 +52,31 @@ void* mongoMalloc(size_t size) {
 }
 
 void* mongoRealloc(void* ptr, size_t size) {
+#if defined(MONGO_USE_GPERFTOOLS_TCMALLOC)
+    void* x = tc_realloc(ptr, size);
+#else
     void* x = std::realloc(ptr, size);
+#endif
     if (x == NULL) {
         reportOutOfMemoryErrorAndExit();
     }
     return x;
 }
 
+void mongoFree(void* ptr) {
+#if defined(MONGO_USE_GPERFTOOLS_TCMALLOC)
+    tc_free(ptr);
+#else
+    std::free(ptr);
+#endif
+}
+
+void mongoFree(void* ptr, size_t size) {
+#if defined(MONGO_USE_GPERFTOOLS_TCMALLOC)
+    tc_free_sized(ptr, size);
+#else
+    std::free(ptr);
+#endif
+}
+
 }  // namespace mongo
diff --git a/src/mongo/util/allocator.h b/src/mongo/util/allocator.h
index 14c5f3e459c..785e71cbd98 100644
--- a/src/mongo/util/allocator.h
+++ b/src/mongo/util/allocator.h
@@ -34,15 +34,22 @@
 namespace mongo {
 
 /**
- * Wrapper around std::malloc().
- * If std::malloc() fails, reports error with stack trace and exit.
+ * Wrapper around malloc for allocator in use.
+ * If malloc fails, reports error with stack trace and exit.
  */
 void* mongoMalloc(size_t size);
 
 /**
- * Wrapper around std::realloc().
- * If std::realloc() fails, reports error with stack trace and exit.
+ * Wrapper around realloc for allocator in use.
+ * If realloc fails, reports error with stack trace and exit.
  */
 void* mongoRealloc(void* ptr, size_t size);
 
+/**
+ * Wrapper around free for allocator in use.
+ * If sized free is not available, it decays to regular free(void*)
+ */
+void mongoFree(void* ptr);
+void mongoFree(void* ptr, size_t size);
+
 }  // namespace mongo
diff --git a/src/mongo/util/intrusive_counter.cpp b/src/mongo/util/intrusive_counter.cpp
index 92eeb0b4bb8..44dbecf0991 100644
--- a/src/mongo/util/intrusive_counter.cpp
+++ b/src/mongo/util/intrusive_counter.cpp
@@ -44,8 +44,7 @@ intrusive_ptr<const RCString> RCString::create(StringData s) {
                           << "MB",
             s.size() < static_cast<size_t>(BSONObjMaxUserSize));
 
-    const size_t sizeWithNUL = s.size() + 1;
-    const size_t bytesNeeded = sizeof(RCString) + sizeWithNUL;
+    const size_t bytesNeeded = bytesRequiredForSize(s.size());
 
 #pragma warning(push)
 #pragma warning(disable : 4291)
diff --git a/src/mongo/util/intrusive_counter.h b/src/mongo/util/intrusive_counter.h
index 1ca640e5a4b..771f6db16b5 100644
--- a/src/mongo/util/intrusive_counter.h
+++ b/src/mongo/util/intrusive_counter.h
@@ -29,6 +29,7 @@
 
 #pragma once
 
+#include <atomic>
 #include <boost/intrusive_ptr.hpp>
 #include <stdlib.h>
 
@@ -114,7 +115,11 @@ public:
 #pragma warning(push)
 #pragma warning(disable : 4291)
     void operator delete(void* ptr) {
-        free(ptr);
+        // Accessing ptr here is technically undefined behavior, but the toolchains we use + UBSAN
+        // are okay with it.
+        // TODO When we are on C++20, this should change and use a destroying delete function.
+        // See https://jira.mongodb.org/browse/SERVER-39506
+        mongoFree(ptr, bytesRequiredForSize(reinterpret_cast<RCString*>(ptr)->size()));
     }
 #pragma warning(pop)
 
@@ -125,6 +130,10 @@ private:
         return mongoMalloc(realSize);
     }
 
+    static size_t bytesRequiredForSize(size_t size) {
+        return sizeof(RCString) + size + 1;
+    }
+
     int _size;  // does NOT include trailing NUL byte.
     // char[_size+1] array allocated past end of class
 };
diff --git a/src/mongo/util/shared_buffer.h b/src/mongo/util/shared_buffer.h
index e18db9b7e3c..383439348cc 100644
--- a/src/mongo/util/shared_buffer.h
+++ b/src/mongo/util/shared_buffer.h
@@ -114,8 +114,9 @@ private:
             if (h->_refCount.subtractAndFetch(1) == 0) {
                 // We placement new'ed a Holder in takeOwnership above,
                 // so we must destroy the object here.
+                auto allocation_size = sizeof(Holder) + h->_capacity;
                 h->~Holder();
-                free(h);
+                mongoFree(h, allocation_size);
             }
         }
 
diff --git a/src/mongo/util/text.cpp b/src/mongo/util/text.cpp
index 67340215814..719cba39ec3 100644
--- a/src/mongo/util/text.cpp
+++ b/src/mongo/util/text.cpp
@@ -311,8 +311,8 @@ WindowsCommandLine::WindowsCommandLine(int argc, wchar_t* argvW[], wchar_t* envp
 }
 
 WindowsCommandLine::~WindowsCommandLine() {
-    free(_argv);
-    free(_envp);
+    mongoFree(_argv);
+    mongoFree(_envp);
 }
 
 #endif  // #if defined(_WIN32)