folly/io/IOBuf.cpp

   1 /*
   2  * Copyright 2013 Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #define __STDC_LIMIT_MACROS
  18
  19 #include "folly/io/IOBuf.h"
  20
  21 #include "folly/Malloc.h"
  22 #include "folly/Likely.h"
  23
  24 #include <stdexcept>
  25 #include <assert.h>
  26 #include <stdint.h>
  27 #include <stdlib.h>
  28
  29 using std::unique_ptr;
  30
  31 namespace {
  32
  33 enum : uint16_t {
  34   kHeapMagic = 0xa5a5,
  35   // This memory segment contains an IOBuf that is still in use
  36   kIOBufInUse = 0x01,
  37   // This memory segment contains buffer data that is still in use
  38   kDataInUse = 0x02,
  39 };
  40
  41 enum : uint32_t {
  42   // When create() is called for buffers less than kDefaultCombinedBufSize,
  43   // we allocate a single combined memory segment for the IOBuf and the data
  44   // together.  See the comments for createCombined()/createSeparate() for more
  45   // details.
  46   //
  47   // (The size of 1k is largely just a guess here.  We could could probably do
  48   // benchmarks of real applications to see if adjusting this number makes a
  49   // difference.  Callers that know their exact use case can also explicitly
  50   // call createCombined() or createSeparate().)
  51   kDefaultCombinedBufSize = 1024
  52 };
  53
  54 }
  55
  56 namespace folly {
  57
  58 struct IOBuf::HeapPrefix {
  59   HeapPrefix(uint16_t flg)
  60     : magic(kHeapMagic),
  61       flags(flg) {}
  62   ~HeapPrefix() {
  63     // Reset magic to 0 on destruction.  This is solely for debugging purposes
  64     // to help catch bugs where someone tries to use HeapStorage after it has
  65     // been deleted.
  66     magic = 0;
  67   }
  68
  69   uint16_t magic;
  70   std::atomic<uint16_t> flags;
  71 };
  72
  73 struct IOBuf::HeapStorage {
  74   HeapPrefix prefix;
  75   // The IOBuf is last in the HeapStorage object.
  76   // This way operator new will work even if allocating a subclass of IOBuf
  77   // that requires more space.
  78   folly::IOBuf buf;
  79 };
  80
  81 struct IOBuf::HeapFullStorage {
  82   HeapStorage hs;
  83   SharedInfo shared;
  84   MaxAlign align;
  85 };
  86
  87 IOBuf::SharedInfo::SharedInfo()
  88   : freeFn(NULL),
  89     userData(NULL) {
  90   // Use relaxed memory ordering here.  Since we are creating a new SharedInfo,
  91   // no other threads should be referring to it yet.
  92   refcount.store(1, std::memory_order_relaxed);
  93 }
  94
  95 IOBuf::SharedInfo::SharedInfo(FreeFunction fn, void* arg)
  96   : freeFn(fn),
  97     userData(arg) {
  98   // Use relaxed memory ordering here.  Since we are creating a new SharedInfo,
  99   // no other threads should be referring to it yet.
 100   refcount.store(1, std::memory_order_relaxed);
 101 }
 102
 103 void* IOBuf::operator new(size_t size) {
 104   size_t fullSize = offsetof(HeapStorage, buf) + size;
 105   auto* storage = static_cast<HeapStorage*>(malloc(fullSize));
 106   // operator new is not allowed to return NULL
 107   if (UNLIKELY(storage == nullptr)) {
 108     throw std::bad_alloc();
 109   }
 110
 111   new (&storage->prefix) HeapPrefix(kIOBufInUse);
 112   return &(storage->buf);
 113 }
 114
 115 void* IOBuf::operator new(size_t size, void* ptr) {
 116   return ptr;
 117 }
 118
 119 void IOBuf::operator delete(void* ptr) {
 120   auto* storageAddr = static_cast<uint8_t*>(ptr) - offsetof(HeapStorage, buf);
 121   auto* storage = reinterpret_cast<HeapStorage*>(storageAddr);
 122   releaseStorage(storage, kIOBufInUse);
 123 }
 124
 125 void IOBuf::releaseStorage(HeapStorage* storage, uint16_t freeFlags) {
 126   CHECK_EQ(storage->prefix.magic, kHeapMagic);
 127
 128   // Use relaxed memory order here.  If we are unlucky and happen to get
 129   // out-of-date data the compare_exchange_weak() call below will catch
 130   // it and load new data with memory_order_acq_rel.
 131   auto flags = storage->prefix.flags.load(std::memory_order_acquire);
 132   DCHECK_EQ((flags & freeFlags), freeFlags);
 133
 134   while (true) {
 135     uint16_t newFlags = (flags & ~freeFlags);
 136     if (newFlags == 0) {
 137       // The storage space is now unused.  Free it.
 138       storage->prefix.HeapPrefix::~HeapPrefix();
 139       free(storage);
 140       return;
 141     }
 142
 143     // This storage segment still contains portions that are in use.
 144     // Just clear the flags specified in freeFlags for now.
 145     auto ret = storage->prefix.flags.compare_exchange_weak(
 146         flags, newFlags, std::memory_order_acq_rel);
 147     if (ret) {
 148       // We successfully updated the flags.
 149       return;
 150     }
 151
 152     // We failed to update the flags.  Some other thread probably updated them
 153     // and cleared some of the other bits.  Continue around the loop to see if
 154     // we are the last user now, or if we need to try updating the flags again.
 155   }
 156 }
 157
 158 void IOBuf::freeInternalBuf(void* buf, void* userData) {
 159   auto* storage = static_cast<HeapStorage*>(userData);
 160   releaseStorage(storage, kDataInUse);
 161 }
 162
 163 unique_ptr<IOBuf> IOBuf::create(uint32_t capacity) {
 164   // For smaller-sized buffers, allocate the IOBuf, SharedInfo, and the buffer
 165   // all with a single allocation.
 166   //
 167   // We don't do this for larger buffers since it can be wasteful if the user
 168   // needs to reallocate the buffer but keeps using the same IOBuf object.
 169   // In this case we can't free the data space until the IOBuf is also
 170   // destroyed.  Callers can explicitly call createCombined() or
 171   // createSeparate() if they know their use case better, and know if they are
 172   // likely to reallocate the buffer later.
 173   if (capacity <= kDefaultCombinedBufSize) {
 174     return createCombined(capacity);
 175   }
 176   return createSeparate(capacity);
 177 }
 178
 179 unique_ptr<IOBuf> IOBuf::createCombined(uint32_t capacity) {
 180   // To save a memory allocation, allocate space for the IOBuf object, the
 181   // SharedInfo struct, and the data itself all with a single call to malloc().
 182   size_t requiredStorage = offsetof(HeapFullStorage, align) + capacity;
 183   size_t mallocSize = goodMallocSize(requiredStorage);
 184   auto* storage = static_cast<HeapFullStorage*>(malloc(mallocSize));
 185
 186   new (&storage->hs.prefix) HeapPrefix(kIOBufInUse | kDataInUse);
 187   new (&storage->shared) SharedInfo(freeInternalBuf, storage);
 188
 189   uint8_t* bufAddr = reinterpret_cast<uint8_t*>(&storage->align);
 190   uint8_t* storageEnd = reinterpret_cast<uint8_t*>(storage) + mallocSize;
 191   size_t actualCapacity = storageEnd - bufAddr;
 192   unique_ptr<IOBuf> ret(new (&storage->hs.buf) IOBuf(
 193         kCombinedAlloc, 0, bufAddr, actualCapacity,
 194         bufAddr, 0, &storage->shared));
 195   return ret;
 196 }
 197
 198 unique_ptr<IOBuf> IOBuf::createSeparate(uint32_t capacity) {
 199   // Allocate an external buffer
 200   uint8_t* buf;
 201   SharedInfo* sharedInfo;
 202   uint32_t actualCapacity;
 203   allocExtBuffer(capacity, &buf, &sharedInfo, &actualCapacity);
 204
 205   // Allocate the IOBuf header
 206   try {
 207     return unique_ptr<IOBuf>(new IOBuf(kExtAllocated, 0,
 208                                        buf, actualCapacity,
 209                                        buf, 0,
 210                                        sharedInfo));
 211   } catch (...) {
 212     free(buf);
 213     throw;
 214   }
 215 }
 216
 217 unique_ptr<IOBuf> IOBuf::createChain(
 218     size_t totalCapacity, uint32_t maxBufCapacity) {
 219   unique_ptr<IOBuf> out = create(
 220       std::min(totalCapacity, size_t(maxBufCapacity)));
 221   size_t allocatedCapacity = out->capacity();
 222
 223   while (allocatedCapacity < totalCapacity) {
 224     unique_ptr<IOBuf> newBuf = create(
 225         std::min(totalCapacity - allocatedCapacity, size_t(maxBufCapacity)));
 226     allocatedCapacity += newBuf->capacity();
 227     out->prependChain(std::move(newBuf));
 228   }
 229
 230   return out;
 231 }
 232
 233 unique_ptr<IOBuf> IOBuf::takeOwnership(void* buf, uint32_t capacity,
 234                                        uint32_t length,
 235                                        FreeFunction freeFn,
 236                                        void* userData,
 237                                        bool freeOnError) {
 238   SharedInfo* sharedInfo = NULL;
 239   try {
 240     // TODO: We could allocate the IOBuf object and SharedInfo all in a single
 241     // memory allocation.  We could use the existing HeapStorage class, and
 242     // define a new kSharedInfoInUse flag.  We could change our code to call
 243     // releaseStorage(kFlagFreeSharedInfo) when this kFlagFreeSharedInfo,
 244     // rather than directly calling delete.
 245     sharedInfo = new SharedInfo(freeFn, userData);
 246
 247     uint8_t* bufPtr = static_cast<uint8_t*>(buf);
 248     return unique_ptr<IOBuf>(new IOBuf(kExtUserSupplied, kFlagFreeSharedInfo,
 249                                        bufPtr, capacity,
 250                                        bufPtr, length,
 251                                        sharedInfo));
 252   } catch (...) {
 253     delete sharedInfo;
 254     if (freeOnError) {
 255       if (freeFn) {
 256         try {
 257           freeFn(buf, userData);
 258         } catch (...) {
 259           // The user's free function is not allowed to throw.
 260           abort();
 261         }
 262       } else {
 263         free(buf);
 264       }
 265     }
 266     throw;
 267   }
 268 }
 269
 270 unique_ptr<IOBuf> IOBuf::wrapBuffer(const void* buf, uint32_t capacity) {
 271   // We cast away the const-ness of the buffer here.
 272   // This is okay since IOBuf users must use unshare() to create a copy of
 273   // this buffer before writing to the buffer.
 274   uint8_t* bufPtr = static_cast<uint8_t*>(const_cast<void*>(buf));
 275   return unique_ptr<IOBuf>(new IOBuf(kExtUserSupplied, kFlagUserOwned,
 276                                      bufPtr, capacity,
 277                                      bufPtr, capacity,
 278                                      NULL));
 279 }
 280
 281 IOBuf::IOBuf(ExtBufTypeEnum type,
 282              uint32_t flags,
 283              uint8_t* buf,
 284              uint32_t capacity,
 285              uint8_t* data,
 286              uint32_t length,
 287              SharedInfo* sharedInfo)
 288   : next_(this),
 289     prev_(this),
 290     data_(data),
 291     buf_(buf),
 292     length_(length),
 293     capacity_(capacity),
 294     flags_(flags),
 295     type_(type),
 296     sharedInfo_(sharedInfo) {
 297   assert(data >= buf);
 298   assert(data + length <= buf + capacity);
 299   assert(static_cast<bool>(flags & kFlagUserOwned) ==
 300          (sharedInfo == NULL));
 301 }
 302
 303 IOBuf::~IOBuf() {
 304   // Destroying an IOBuf destroys the entire chain.
 305   // Users of IOBuf should only explicitly delete the head of any chain.
 306   // The other elements in the chain will be automatically destroyed.
 307   while (next_ != this) {
 308     // Since unlink() returns unique_ptr() and we don't store it,
 309     // it will automatically delete the unlinked element.
 310     (void)next_->unlink();
 311   }
 312
 313   decrementRefcount();
 314 }
 315
 316 bool IOBuf::empty() const {
 317   const IOBuf* current = this;
 318   do {
 319     if (current->length() != 0) {
 320       return false;
 321     }
 322     current = current->next_;
 323   } while (current != this);
 324   return true;
 325 }
 326
 327 uint32_t IOBuf::countChainElements() const {
 328   uint32_t numElements = 1;
 329   for (IOBuf* current = next_; current != this; current = current->next_) {
 330     ++numElements;
 331   }
 332   return numElements;
 333 }
 334
 335 uint64_t IOBuf::computeChainDataLength() const {
 336   uint64_t fullLength = length_;
 337   for (IOBuf* current = next_; current != this; current = current->next_) {
 338     fullLength += current->length_;
 339   }
 340   return fullLength;
 341 }
 342
 343 void IOBuf::prependChain(unique_ptr<IOBuf>&& iobuf) {
 344   // Take ownership of the specified IOBuf
 345   IOBuf* other = iobuf.release();
 346
 347   // Remember the pointer to the tail of the other chain
 348   IOBuf* otherTail = other->prev_;
 349
 350   // Hook up prev_->next_ to point at the start of the other chain,
 351   // and other->prev_ to point at prev_
 352   prev_->next_ = other;
 353   other->prev_ = prev_;
 354
 355   // Hook up otherTail->next_ to point at us,
 356   // and prev_ to point back at otherTail,
 357   otherTail->next_ = this;
 358   prev_ = otherTail;
 359 }
 360
 361 unique_ptr<IOBuf> IOBuf::clone() const {
 362   unique_ptr<IOBuf> newHead(cloneOne());
 363
 364   for (IOBuf* current = next_; current != this; current = current->next_) {
 365     newHead->prependChain(current->cloneOne());
 366   }
 367
 368   return newHead;
 369 }
 370
 371 unique_ptr<IOBuf> IOBuf::cloneOne() const {
 372   if (sharedInfo_) {
 373     flags_ |= kFlagMaybeShared;
 374   }
 375   unique_ptr<IOBuf> iobuf(new IOBuf(static_cast<ExtBufTypeEnum>(type_),
 376                                     flags_, buf_, capacity_,
 377                                     data_, length_,
 378                                     sharedInfo_));
 379   if (sharedInfo_) {
 380     sharedInfo_->refcount.fetch_add(1, std::memory_order_acq_rel);
 381   }
 382   return iobuf;
 383 }
 384
 385 void IOBuf::unshareOneSlow() {
 386   // Allocate a new buffer for the data
 387   uint8_t* buf;
 388   SharedInfo* sharedInfo;
 389   uint32_t actualCapacity;
 390   allocExtBuffer(capacity_, &buf, &sharedInfo, &actualCapacity);
 391
 392   // Copy the data
 393   // Maintain the same amount of headroom.  Since we maintained the same
 394   // minimum capacity we also maintain at least the same amount of tailroom.
 395   uint32_t headlen = headroom();
 396   memcpy(buf + headlen, data_, length_);
 397
 398   // Release our reference on the old buffer
 399   decrementRefcount();
 400   // Make sure kFlagUserOwned, kFlagMaybeShared, and kFlagFreeSharedInfo
 401   // are all cleared.
 402   flags_ = 0;
 403
 404   // Update the buffer pointers to point to the new buffer
 405   data_ = buf + headlen;
 406   buf_ = buf;
 407   sharedInfo_ = sharedInfo;
 408 }
 409
 410 void IOBuf::unshareChained() {
 411   // unshareChained() should only be called if we are part of a chain of
 412   // multiple IOBufs.  The caller should have already verified this.
 413   assert(isChained());
 414
 415   IOBuf* current = this;
 416   while (true) {
 417     if (current->isSharedOne()) {
 418       // we have to unshare
 419       break;
 420     }
 421
 422     current = current->next_;
 423     if (current == this) {
 424       // None of the IOBufs in the chain are shared,
 425       // so return without doing anything
 426       return;
 427     }
 428   }
 429
 430   // We have to unshare.  Let coalesceSlow() do the work.
 431   coalesceSlow();
 432 }
 433
 434 void IOBuf::coalesceSlow(size_t maxLength) {
 435   // coalesceSlow() should only be called if we are part of a chain of multiple
 436   // IOBufs.  The caller should have already verified this.
 437   assert(isChained());
 438   assert(length_ < maxLength);
 439
 440   // Compute the length of the entire chain
 441   uint64_t newLength = 0;
 442   IOBuf* end = this;
 443   do {
 444     newLength += end->length_;
 445     end = end->next_;
 446   } while (newLength < maxLength && end != this);
 447
 448   uint64_t newHeadroom = headroom();
 449   uint64_t newTailroom = end->prev_->tailroom();
 450   coalesceAndReallocate(newHeadroom, newLength, end, newTailroom);
 451   // We should be only element left in the chain now
 452   assert(length_ >= maxLength || !isChained());
 453 }
 454
 455 void IOBuf::coalesceAndReallocate(size_t newHeadroom,
 456                                   size_t newLength,
 457                                   IOBuf* end,
 458                                   size_t newTailroom) {
 459   uint64_t newCapacity = newLength + newHeadroom + newTailroom;
 460   if (newCapacity > UINT32_MAX) {
 461     throw std::overflow_error("IOBuf chain too large to coalesce");
 462   }
 463
 464   // Allocate space for the coalesced buffer.
 465   // We always convert to an external buffer, even if we happened to be an
 466   // internal buffer before.
 467   uint8_t* newBuf;
 468   SharedInfo* newInfo;
 469   uint32_t actualCapacity;
 470   allocExtBuffer(newCapacity, &newBuf, &newInfo, &actualCapacity);
 471
 472   // Copy the data into the new buffer
 473   uint8_t* newData = newBuf + newHeadroom;
 474   uint8_t* p = newData;
 475   IOBuf* current = this;
 476   size_t remaining = newLength;
 477   do {
 478     assert(current->length_ <= remaining);
 479     remaining -= current->length_;
 480     memcpy(p, current->data_, current->length_);
 481     p += current->length_;
 482     current = current->next_;
 483   } while (current != end);
 484   assert(remaining == 0);
 485
 486   // Point at the new buffer
 487   decrementRefcount();
 488
 489   // Make sure kFlagUserOwned, kFlagMaybeShared, and kFlagFreeSharedInfo
 490   // are all cleared.
 491   flags_ = 0;
 492
 493   capacity_ = actualCapacity;
 494   type_ = kExtAllocated;
 495   buf_ = newBuf;
 496   sharedInfo_ = newInfo;
 497   data_ = newData;
 498   length_ = newLength;
 499
 500   // Separate from the rest of our chain.
 501   // Since we don't store the unique_ptr returned by separateChain(),
 502   // this will immediately delete the returned subchain.
 503   if (isChained()) {
 504     (void)separateChain(next_, current->prev_);
 505   }
 506 }
 507
 508 void IOBuf::decrementRefcount() {
 509   // Externally owned buffers don't have a SharedInfo object and aren't managed
 510   // by the reference count
 511   if (flags_ & kFlagUserOwned) {
 512     assert(sharedInfo_ == nullptr);
 513     return;
 514   }
 515
 516   // Decrement the refcount
 517   uint32_t newcnt = sharedInfo_->refcount.fetch_sub(
 518       1, std::memory_order_acq_rel);
 519   // Note that fetch_sub() returns the value before we decremented.
 520   // If it is 1, we were the only remaining user; if it is greater there are
 521   // still other users.
 522   if (newcnt > 1) {
 523     return;
 524   }
 525
 526   // We were the last user.  Free the buffer
 527   freeExtBuffer();
 528
 529   // Free the SharedInfo if it was allocated separately.
 530   //
 531   // This is only used by takeOwnership().
 532   //
 533   // To avoid this special case handling in decrementRefcount(), we could have
 534   // takeOwnership() set a custom freeFn() that calls the user's free function
 535   // then frees the SharedInfo object.  (This would require that
 536   // takeOwnership() store the user's free function with its allocated
 537   // SharedInfo object.)  However, handling this specially with a flag seems
 538   // like it shouldn't be problematic.
 539   if (flags_ & kFlagFreeSharedInfo) {
 540     delete sharedInfo_;
 541   }
 542 }
 543
 544 void IOBuf::reserveSlow(uint32_t minHeadroom, uint32_t minTailroom) {
 545   size_t newCapacity = (size_t)length_ + minHeadroom + minTailroom;
 546   DCHECK_LT(newCapacity, UINT32_MAX);
 547
 548   // reserveSlow() is dangerous if anyone else is sharing the buffer, as we may
 549   // reallocate and free the original buffer.  It should only ever be called if
 550   // we are the only user of the buffer.
 551   DCHECK(!isSharedOne());
 552
 553   // We'll need to reallocate the buffer.
 554   // There are a few options.
 555   // - If we have enough total room, move the data around in the buffer
 556   //   and adjust the data_ pointer.
 557   // - If we're using an internal buffer, we'll switch to an external
 558   //   buffer with enough headroom and tailroom.
 559   // - If we have enough headroom (headroom() >= minHeadroom) but not too much
 560   //   (so we don't waste memory), we can try one of two things, depending on
 561   //   whether we use jemalloc or not:
 562   //   - If using jemalloc, we can try to expand in place, avoiding a memcpy()
 563   //   - If not using jemalloc and we don't have too much to copy,
 564   //     we'll use realloc() (note that realloc might have to copy
 565   //     headroom + data + tailroom, see smartRealloc in folly/Malloc.h)
 566   // - Otherwise, bite the bullet and reallocate.
 567   if (headroom() + tailroom() >= minHeadroom + minTailroom) {
 568     uint8_t* newData = writableBuffer() + minHeadroom;
 569     memmove(newData, data_, length_);
 570     data_ = newData;
 571     return;
 572   }
 573
 574   size_t newAllocatedCapacity = goodExtBufferSize(newCapacity);
 575   uint8_t* newBuffer = nullptr;
 576   uint32_t newHeadroom = 0;
 577   uint32_t oldHeadroom = headroom();
 578
 579   // If we have a buffer allocated with malloc and we just need more tailroom,
 580   // try to use realloc()/rallocm() to grow the buffer in place.
 581   if ((flags_ & kFlagUserOwned) == 0 && (sharedInfo_->freeFn == nullptr) &&
 582       length_ != 0 && oldHeadroom >= minHeadroom) {
 583     if (usingJEMalloc()) {
 584       size_t headSlack = oldHeadroom - minHeadroom;
 585       // We assume that tailroom is more useful and more important than
 586       // headroom (not least because realloc / rallocm allow us to grow the
 587       // buffer at the tail, but not at the head)  So, if we have more headroom
 588       // than we need, we consider that "wasted".  We arbitrarily define "too
 589       // much" headroom to be 25% of the capacity.
 590       if (headSlack * 4 <= newCapacity) {
 591         size_t allocatedCapacity = capacity() + sizeof(SharedInfo);
 592         void* p = buf_;
 593         if (allocatedCapacity >= jemallocMinInPlaceExpandable) {
 594           int r = rallocm(&p, &newAllocatedCapacity, newAllocatedCapacity,
 595                           0, ALLOCM_NO_MOVE);
 596           if (r == ALLOCM_SUCCESS) {
 597             newBuffer = static_cast<uint8_t*>(p);
 598             newHeadroom = oldHeadroom;
 599           } else if (r == ALLOCM_ERR_OOM) {
 600             // shouldn't happen as we don't actually allocate new memory
 601             // (due to ALLOCM_NO_MOVE)
 602             throw std::bad_alloc();
 603           }
 604           // if ALLOCM_ERR_NOT_MOVED, do nothing, fall back to
 605           // malloc/memcpy/free
 606         }
 607       }
 608     } else {  // Not using jemalloc
 609       size_t copySlack = capacity() - length_;
 610       if (copySlack * 2 <= length_) {
 611         void* p = realloc(buf_, newAllocatedCapacity);
 612         if (UNLIKELY(p == nullptr)) {
 613           throw std::bad_alloc();
 614         }
 615         newBuffer = static_cast<uint8_t*>(p);
 616         newHeadroom = oldHeadroom;
 617       }
 618     }
 619   }
 620
 621   // None of the previous reallocation strategies worked (or we're using
 622   // an internal buffer).  malloc/copy/free.
 623   if (newBuffer == nullptr) {
 624     void* p = malloc(newAllocatedCapacity);
 625     if (UNLIKELY(p == nullptr)) {
 626       throw std::bad_alloc();
 627     }
 628     newBuffer = static_cast<uint8_t*>(p);
 629     memcpy(newBuffer + minHeadroom, data_, length_);
 630     if ((flags_ & kFlagUserOwned) == 0) {
 631       freeExtBuffer();
 632     }
 633     newHeadroom = minHeadroom;
 634   }
 635
 636   SharedInfo* info;
 637   uint32_t cap;
 638   initExtBuffer(newBuffer, newAllocatedCapacity, &info, &cap);
 639
 640   if (flags_ & kFlagFreeSharedInfo) {
 641     delete sharedInfo_;
 642   }
 643
 644   flags_ = 0;
 645   capacity_ = cap;
 646   type_ = kExtAllocated;
 647   buf_ = newBuffer;
 648   sharedInfo_ = info;
 649   data_ = newBuffer + newHeadroom;
 650   // length_ is unchanged
 651 }
 652
 653 void IOBuf::freeExtBuffer() {
 654   DCHECK((flags_ & kFlagUserOwned) == 0);
 655
 656   if (sharedInfo_->freeFn) {
 657     try {
 658       sharedInfo_->freeFn(buf_, sharedInfo_->userData);
 659     } catch (...) {
 660       // The user's free function should never throw.  Otherwise we might
 661       // throw from the IOBuf destructor.  Other code paths like coalesce()
 662       // also assume that decrementRefcount() cannot throw.
 663       abort();
 664     }
 665   } else {
 666     free(buf_);
 667   }
 668 }
 669
 670 void IOBuf::allocExtBuffer(uint32_t minCapacity,
 671                            uint8_t** bufReturn,
 672                            SharedInfo** infoReturn,
 673                            uint32_t* capacityReturn) {
 674   size_t mallocSize = goodExtBufferSize(minCapacity);
 675   uint8_t* buf = static_cast<uint8_t*>(malloc(mallocSize));
 676   if (UNLIKELY(buf == NULL)) {
 677     throw std::bad_alloc();
 678   }
 679   initExtBuffer(buf, mallocSize, infoReturn, capacityReturn);
 680   *bufReturn = buf;
 681 }
 682
 683 size_t IOBuf::goodExtBufferSize(uint32_t minCapacity) {
 684   // Determine how much space we should allocate.  We'll store the SharedInfo
 685   // for the external buffer just after the buffer itself.  (We store it just
 686   // after the buffer rather than just before so that the code can still just
 687   // use free(buf_) to free the buffer.)
 688   size_t minSize = static_cast<size_t>(minCapacity) + sizeof(SharedInfo);
 689   // Add room for padding so that the SharedInfo will be aligned on an 8-byte
 690   // boundary.
 691   minSize = (minSize + 7) & ~7;
 692
 693   // Use goodMallocSize() to bump up the capacity to a decent size to request
 694   // from malloc, so we can use all of the space that malloc will probably give
 695   // us anyway.
 696   return goodMallocSize(minSize);
 697 }
 698
 699 void IOBuf::initExtBuffer(uint8_t* buf, size_t mallocSize,
 700                           SharedInfo** infoReturn,
 701                           uint32_t* capacityReturn) {
 702   // Find the SharedInfo storage at the end of the buffer
 703   // and construct the SharedInfo.
 704   uint8_t* infoStart = (buf + mallocSize) - sizeof(SharedInfo);
 705   SharedInfo* sharedInfo = new(infoStart) SharedInfo;
 706
 707   size_t actualCapacity = infoStart - buf;
 708   // On the unlikely possibility that the actual capacity is larger than can
 709   // fit in a uint32_t after adding room for the refcount and calling
 710   // goodMallocSize(), truncate downwards if necessary.
 711   if (actualCapacity >= UINT32_MAX) {
 712     *capacityReturn = UINT32_MAX;
 713   } else {
 714     *capacityReturn = actualCapacity;
 715   }
 716
 717   *infoReturn = sharedInfo;
 718 }
 719
 720 fbstring IOBuf::moveToFbString() {
 721   // malloc-allocated buffers are just fine, everything else needs
 722   // to be turned into one.
 723   if ((flags_ & kFlagUserOwned) ||  // user owned, not ours to give up
 724       sharedInfo_->freeFn != nullptr || // not malloc()-ed
 725       headroom() != 0 ||     // malloc()-ed block doesn't start at beginning
 726       tailroom() == 0 ||     // no room for NUL terminator
 727       isShared() ||          // shared
 728       isChained()) {         // chained
 729     // We might as well get rid of all head and tailroom if we're going
 730     // to reallocate; we need 1 byte for NUL terminator.
 731     coalesceAndReallocate(0, computeChainDataLength(), this, 1);
 732   }
 733
 734   // Ensure NUL terminated
 735   *writableTail() = 0;
 736   fbstring str(reinterpret_cast<char*>(writableData()),
 737                length(),  capacity(),
 738                AcquireMallocatedString());
 739
 740   if (flags_ & kFlagFreeSharedInfo) {
 741     delete sharedInfo_;
 742   }
 743
 744   // Reset to a state where we can be deleted cleanly
 745   flags_ = kFlagUserOwned;
 746   sharedInfo_ = nullptr;
 747   buf_ = nullptr;
 748   clear();
 749   return str;
 750 }
 751
 752 IOBuf::Iterator IOBuf::cbegin() const {
 753   return Iterator(this, this);
 754 }
 755
 756 IOBuf::Iterator IOBuf::cend() const {
 757   return Iterator(nullptr, nullptr);
 758 }
 759
 760 folly::fbvector<struct iovec> IOBuf::getIov() const {
 761   folly::fbvector<struct iovec> iov;
 762   iov.reserve(countChainElements());
 763   IOBuf const* p = this;
 764   do {
 765     // some code can get confused by empty iovs, so skip them
 766     if (p->length() > 0) {
 767       iov.push_back({(void*)p->data(), p->length()});
 768     }
 769     p = p->next();
 770   } while (p != this);
 771   return iov;
 772 }
 773
 774 } // folly