SERVER-28311 Make BSONObjBuilder returnable

5 files changed, 57 insertions, 11 deletions

diff --git a/src/mongo/bson/bsonmisc.h b/src/mongo/bson/bsonmisc.h
index 2ca9d39c805..7cf1bb4ac05 100644
--- a/src/mongo/bson/bsonmisc.h
+++ b/src/mongo/bson/bsonmisc.h
@@ -204,8 +204,6 @@ inline BSONObj OR(const BSONObj& a,
 
 // Utility class to implement BSON( key << val ) as described above.
 class BSONObjBuilderValueStream {
-    MONGO_DISALLOW_COPYING(BSONObjBuilderValueStream);
-
 public:
     friend class Labeler;
     BSONObjBuilderValueStream(BSONObjBuilder* builder);
diff --git a/src/mongo/bson/bsonobjbuilder.h b/src/mongo/bson/bsonobjbuilder.h
index 72fd2d901af..835fe4e8312 100644
--- a/src/mongo/bson/bsonobjbuilder.h
+++ b/src/mongo/bson/bsonobjbuilder.h
@@ -124,6 +124,17 @@ public:
         _b.reserveBytes(1);
     }
 
+    // Move constructible, but not assignable due to reference member.
+    BSONObjBuilder(BSONObjBuilder&& other)
+        : _b(&other._b == &other._buf ? _buf : other._b),
+          _buf(std::move(other._buf)),
+          _offset(std::move(other._offset)),
+          _s(std::move(other._s)),
+          _tracker(std::move(other._tracker)),
+          _doneCalled(std::move(other._doneCalled)) {
+        other.abandon();
+    }
+
     ~BSONObjBuilder() {
         // If 'done' has not already been called, and we have a reference to an owning
         // BufBuilder but do not own it ourselves, then we must call _done to write in the
diff --git a/src/mongo/bson/bsonobjbuilder_test.cpp b/src/mongo/bson/bsonobjbuilder_test.cpp
index 003f40c27dd..3cc75b50560 100644
--- a/src/mongo/bson/bsonobjbuilder_test.cpp
+++ b/src/mongo/bson/bsonobjbuilder_test.cpp
@@ -331,5 +331,35 @@ TEST(BSONObjBuilderTest, ResetToEmptyForNestedBuilderOnlyResetsInnerObj) {
     ASSERT_BSONOBJ_EQ(BSON("a" << 3 << "nestedObj" << BSONObj()), bob.obj());
 }
 
+TEST(BSONObjBuilderTest, ReturningAnOwningBSONObjBuilderWorks) {
+    BSONObjBuilder bob = ([] {
+        BSONObjBuilder initial;
+        initial.append("a", 1);
+        return initial;
+    })();
+    ASSERT(bob.owned());
+
+    bob.append("b", 2);
+
+    ASSERT_BSONOBJ_EQ(bob.obj(), BSON("a" << 1 << "b" << 2));
+}
+
+TEST(BSONObjBuilderTest, ReturningANonOwningBSONObjBuilderWorks) {
+    BSONObjBuilder outer;
+    {
+        BSONObjBuilder bob = ([&] {
+            BSONObjBuilder initial(outer.subobjStart("nested"));
+            initial.append("a", 1);
+            return initial;
+        })();
+        ASSERT(!bob.owned());
+        ASSERT_EQ(&bob.bb(), &outer.bb());
+
+        bob.append("b", 2);
+    }
+
+    ASSERT_BSONOBJ_EQ(outer.obj(), BSON("nested" << BSON("a" << 1 << "b" << 2)));
+}
+
 }  // unnamed namespace
 }  // namespace mongo
diff --git a/src/mongo/bson/util/builder.h b/src/mongo/bson/util/builder.h
index b959c59236e..3a0c02d3258 100644
--- a/src/mongo/bson/util/builder.h
+++ b/src/mongo/bson/util/builder.h
@@ -79,6 +79,11 @@ class SharedBufferAllocator {
 public:
     SharedBufferAllocator() = default;
 
+    // Allow moving but not copying. It would be an error for two SharedBufferAllocators to use the
+    // same underlying buffer.
+    SharedBufferAllocator(SharedBufferAllocator&&) = default;
+    SharedBufferAllocator& operator=(SharedBufferAllocator&&) = default;
+
     void malloc(size_t sz) {
         _buf = SharedBuffer::allocate(sz);
     }
@@ -105,6 +110,9 @@ class StackAllocator {
 
 public:
     StackAllocator() = default;
+    ~StackAllocator() {
+        free();
+    }
 
     enum { SZ = 512 };
     void malloc(size_t sz) {
@@ -141,8 +149,6 @@ private:
 
 template <class BufferAllocator>
 class _BufBuilder {
-    MONGO_DISALLOW_COPYING(_BufBuilder);
-
 public:
     _BufBuilder(int initsize = 512) : size(initsize) {
         if (size > 0) {
@@ -151,9 +157,6 @@ public:
         l = 0;
         reservedBytes = 0;
     }
-    ~_BufBuilder() {
-        kill();
-    }
 
     void kill() {
         _buf.free();
@@ -342,6 +345,7 @@ private:
 };
 
 typedef _BufBuilder<SharedBufferAllocator> BufBuilder;
+MONGO_STATIC_ASSERT(std::is_move_constructible<BufBuilder>::value);
 
 /** The StackBufBuilder builds smaller datasets on the stack instead of using malloc.
       this can be significantly faster for small bufs.  However, you can not release() the
@@ -355,6 +359,7 @@ public:
     StackBufBuilder() : _BufBuilder<StackAllocator>(StackAllocator::SZ) {}
     void release() = delete;  // not allowed. not implemented.
 };
+MONGO_STATIC_ASSERT(!std::is_move_constructible<StackBufBuilder>::value);
 
 /** std::stringstream deals with locale so this is a lot faster than std::stringstream for UTF8 */
 template <typename Allocator>
@@ -468,10 +473,6 @@ public:
 private:
     _BufBuilder<Allocator> _buf;
 
-    // non-copyable, non-assignable
-    StringBuilderImpl(const StringBuilderImpl&);
-    StringBuilderImpl& operator=(const StringBuilderImpl&);
-
     template <typename T>
     StringBuilderImpl& SBNUM(T val, int maxSize, const char* macro) {
         int prev = _buf.l;
diff --git a/src/mongo/bson/util/builder_test.cpp b/src/mongo/bson/util/builder_test.cpp
index ad1967b8a27..f56ba5999c8 100644
--- a/src/mongo/bson/util/builder_test.cpp
+++ b/src/mongo/bson/util/builder_test.cpp
@@ -76,4 +76,10 @@ TEST(Builder, BooleanOstreamOperator) {
     sb << "{abc: " << true << ", def: " << false << "}";
     ASSERT_EQUALS("{abc: 1, def: 0}", sb.str());
 }
+
+TEST(Builder, StackAllocatorShouldNotLeak) {
+    StackAllocator stackAlloc;
+    stackAlloc.malloc(StackAllocator::SZ + 1);  // Force heap allocation.
+    // Let the builder go out of scope. If this leaks, it will trip the ASAN leak detector.
+}
 }