SERVER-29564 BSONObjBuilder can now be seeded with a BSONObj prefix

This will avoid copying whenever it is safe.

1 files changed, 38 insertions, 6 deletions

diff --git a/src/mongo/util/shared_buffer.h b/src/mongo/util/shared_buffer.h
index 09951dc80dc..dc8599f7e6e 100644
--- a/src/mongo/util/shared_buffer.h
+++ b/src/mongo/util/shared_buffer.h
@@ -48,7 +48,7 @@ public:
     }
 
     static SharedBuffer allocate(size_t bytes) {
-        return takeOwnership(mongoMalloc(sizeof(Holder) + bytes));
+        return takeOwnership(mongoMalloc(sizeof(Holder) + bytes), bytes);
     }
 
     /**
@@ -67,7 +67,7 @@ public:
 
         // Get newPtr into _holder with a ref-count of 1 without touching the current pointee of
         // _holder which is now invalid.
-        auto tmp = SharedBuffer::takeOwnership(newPtr);
+        auto tmp = SharedBuffer::takeOwnership(newPtr, size);
         _holder.detach();
         _holder = std::move(tmp._holder);
     }
@@ -80,11 +80,29 @@ public:
         return bool(_holder);
     }
 
+    /**
+     * Returns true if this object has exclusive access to the underlying buffer.
+     * (That is, reference count == 1).
+     */
+    bool isShared() const {
+        return _holder && _holder->isShared();
+    }
+
+    /**
+     * Returns the allocation size of the underlying buffer.
+     * Users of this type must maintain the "used" size separately.
+     */
+    size_t capacity() const {
+        return _holder ? _holder->_capacity : 0;
+    }
+
 private:
     class Holder {
     public:
-        explicit Holder(AtomicUInt32::WordType initial = AtomicUInt32::WordType())
-            : _refCount(initial) {}
+        explicit Holder(AtomicUInt32::WordType initial, size_t capacity)
+            : _refCount(initial), _capacity(capacity) {
+            invariant(capacity == _capacity);
+        }
 
         // these are called automatically by boost::intrusive_ptr
         friend void intrusive_ptr_add_ref(Holder* h) {
@@ -113,6 +131,7 @@ private:
         }
 
         AtomicUInt32 _refCount;
+        uint32_t _capacity;
     };
 
     explicit SharedBuffer(Holder* holder) : _holder(holder, /*add_ref=*/false) {
@@ -127,12 +146,12 @@ private:
      * This class will call free(holderPrefixedData), so it must have been allocated in a way
      * that makes that valid.
      */
-    static SharedBuffer takeOwnership(void* holderPrefixedData) {
+    static SharedBuffer takeOwnership(void* holderPrefixedData, size_t capacity) {
         // Initialize the refcount to 1 so we don't need to increment it in the constructor
         // (see private Holder* constructor above).
         //
         // TODO: Should dassert alignment of holderPrefixedData here if possible.
-        return SharedBuffer(new (holderPrefixedData) Holder(1U));
+        return SharedBuffer(new (holderPrefixedData) Holder(1U, capacity));
     }
 
     boost::intrusive_ptr<Holder> _holder;
@@ -164,6 +183,19 @@ public:
         return bool(_buffer);
     }
 
+    bool isShared() const {
+        return _buffer.isShared();
+    }
+
+    /**
+     * Converts to a mutable SharedBuffer.
+     * This is only legal to call if you have exclusive access to the underlying buffer.
+     */
+    SharedBuffer constCast() && {
+        invariant(!isShared());
+        return std::move(_buffer);
+    }
+
 private:
     SharedBuffer _buffer;
 };