2 files changed, 52 insertions, 7 deletions

diff --git a/src/mongo/bson/bsonobjbuilder.h b/src/mongo/bson/bsonobjbuilder.h
index 5cd27a5a198..3c3cd7de363 100644
--- a/src/mongo/bson/bsonobjbuilder.h
+++ b/src/mongo/bson/bsonobjbuilder.h
@@ -68,10 +68,13 @@ namespace mongo {
             , _s(this)
             , _tracker(0)
             , _doneCalled(false) {
-            // Reserve space for a holder object at the beginning of the buffer, followed by
+            // Skip over space for a holder object at the beginning of the buffer, followed by
             // space for the object length. The length is filled in by _done.
             _b.skip(sizeof(BSONObj::Holder));
             _b.skip(sizeof(int));
+
+            // Reserve space for the EOO byte. This means _done() can't fail.
+            _b.reserveBytes(1);
         }
 
         /** @param baseBuilder construct a BSONObjBuilder using an existing BufBuilder
@@ -84,9 +87,13 @@ namespace mongo {
             , _s(this)
             , _tracker(0)
             , _doneCalled(false) {
-            // Reserve space for the object length, which is filled in by _done. We don't need a holder
-            // since we are a sub-builder, and some parent builder has already made the reservation.
+            // Skip over space for the object length, which is filled in by _done. We don't need a
+            // holder since we are a sub-builder, and some parent builder has already made the
+            // reservation.
             _b.skip(sizeof(int));
+
+            // Reserve space for the EOO byte. This means _done() can't fail.
+            _b.reserveBytes(1);
         }
 
         BSONObjBuilder( const BSONSizeTracker & tracker )
@@ -99,6 +106,9 @@ namespace mongo {
             // See the comments in the first constructor for details.
             _b.skip(sizeof(BSONObj::Holder));
             _b.skip(sizeof(int));
+
+            // Reserve space for the EOO byte. This means _done() can't fail.
+            _b.reserveBytes(1);
         }
 
         ~BSONObjBuilder() {
@@ -622,6 +632,7 @@ namespace mongo {
         BSONObj asTempObj() {
             BSONObj temp(_done());
             _b.setlen(_b.len()-1); //next append should overwrite the EOO
+            _b.reserveBytes(1); // Rereserve room for the real EOO
             _doneCalled = false;
             return temp;
         }
@@ -703,8 +714,14 @@ namespace mongo {
                 return _b.buf() + _offset;
 
             _doneCalled = true;
+
+            // TODO remove this or find some way to prevent it from failing. Since this is intended
+            // for use with BSON() literal queries, it is less likely to result in oversized BSON.
             _s.endField();
+
+            _b.claimReservedBytes(1); // Prevents adding EOO from failing.
             _b.appendNum((char) EOO);
+
             char *data = _b.buf() + _offset;
             int size = _b.len() - _offset;
             DataView(data).writeLE(size);
diff --git a/src/mongo/bson/util/builder.h b/src/mongo/bson/util/builder.h
index 1b4979663e4..584b562a609 100644
--- a/src/mongo/bson/util/builder.h
+++ b/src/mongo/bson/util/builder.h
@@ -125,6 +125,7 @@ namespace mongo {
                 data = 0;
             }
             l = 0;
+            reservedBytes = 0;
         }
         ~_BufBuilder() { kill(); }
 
@@ -137,9 +138,11 @@ namespace mongo {
 
         void reset() {
             l = 0;
+            reservedBytes = 0;
         }
         void reset( int maxSize ) {
             l = 0;
+            reservedBytes = 0;
             if ( maxSize && size > maxSize ) {
                 al.Free(data);
                 data = (char*)al.Malloc(maxSize);
@@ -222,13 +225,36 @@ namespace mongo {
         inline char* grow(int by) {
             int oldlen = l;
             int newLen = l + by;
-            if ( newLen > size ) {
-                grow_reallocate(newLen);
+            int minSize = newLen + reservedBytes;
+            if ( minSize > size ) {
+                grow_reallocate(minSize);
             }
             l = newLen;
             return data + oldlen;
         }
 
+        /**
+         * Reserve room for some number of bytes to be claimed at a later time.
+         */
+        void reserveBytes(int bytes) {
+            int minSize = l + reservedBytes + bytes;
+            if (minSize > size)
+                grow_reallocate(minSize);
+
+            // This must happen *after* any attempt to grow.
+            reservedBytes += bytes;
+        }
+
+        /**
+         * Claim an earlier reservation of some number of bytes. These bytes must already have been
+         * reserved. Appends of up to this many bytes immediately following a claim are
+         * guaranteed to succeed without a need to reallocate.
+         */
+        void claimReservedBytes(int bytes) {
+            invariant(reservedBytes >= bytes);
+            reservedBytes -= bytes;
+        }
+
     private:
         template<typename T>
         void appendNumImpl(T t) {
@@ -241,10 +267,11 @@ namespace mongo {
 
 
         /* "slow" portion of 'grow()'  */
-        void NOINLINE_DECL grow_reallocate(int newLen) {
+        void NOINLINE_DECL grow_reallocate(int minSize) {
             int a = 64;
-            while( a < newLen ) 
+            while (a < minSize)
                 a = a * 2;
+
             if ( a > BufferMaxSize ) {
                 std::stringstream ss;
                 ss << "BufBuilder attempted to grow() to " << a << " bytes, past the 64MB limit.";
@@ -259,6 +286,7 @@ namespace mongo {
         char *data;
         int l;
         int size;
+        int reservedBytes; // eagerly grow_reallocate to keep this many bytes of spare room.
 
         friend class StringBuilderImpl<Allocator>;
     };