35c4e102bf49d85f01dd08af97a65acdbc864df9
[folly.git] / folly / io / IOBuf.cpp
1 /*
2  * Copyright 2015 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/Conv.h>
22 #include <folly/Likely.h>
23 #include <folly/Malloc.h>
24 #include <folly/Memory.h>
25 #include <folly/ScopeGuard.h>
26 #include <folly/SpookyHashV2.h>
27 #include <folly/io/Cursor.h>
28
29 #include <stdexcept>
30 #include <assert.h>
31 #include <stdint.h>
32 #include <stdlib.h>
33
34 using std::unique_ptr;
35
36 namespace {
37
38 enum : uint16_t {
39   kHeapMagic = 0xa5a5,
40   // This memory segment contains an IOBuf that is still in use
41   kIOBufInUse = 0x01,
42   // This memory segment contains buffer data that is still in use
43   kDataInUse = 0x02,
44 };
45
46 enum : uint64_t {
47   // When create() is called for buffers less than kDefaultCombinedBufSize,
48   // we allocate a single combined memory segment for the IOBuf and the data
49   // together.  See the comments for createCombined()/createSeparate() for more
50   // details.
51   //
52   // (The size of 1k is largely just a guess here.  We could could probably do
53   // benchmarks of real applications to see if adjusting this number makes a
54   // difference.  Callers that know their exact use case can also explicitly
55   // call createCombined() or createSeparate().)
56   kDefaultCombinedBufSize = 1024
57 };
58
59 // Helper function for IOBuf::takeOwnership()
60 void takeOwnershipError(bool freeOnError, void* buf,
61                         folly::IOBuf::FreeFunction freeFn,
62                         void* userData) {
63   if (!freeOnError) {
64     return;
65   }
66   if (!freeFn) {
67     free(buf);
68     return;
69   }
70   try {
71     freeFn(buf, userData);
72   } catch (...) {
73     // The user's free function is not allowed to throw.
74     // (We are already in the middle of throwing an exception, so
75     // we cannot let this exception go unhandled.)
76     abort();
77   }
78 }
79
80 } // unnamed namespace
81
82 namespace folly {
83
84 struct IOBuf::HeapPrefix {
85   HeapPrefix(uint16_t flg)
86     : magic(kHeapMagic),
87       flags(flg) {}
88   ~HeapPrefix() {
89     // Reset magic to 0 on destruction.  This is solely for debugging purposes
90     // to help catch bugs where someone tries to use HeapStorage after it has
91     // been deleted.
92     magic = 0;
93   }
94
95   uint16_t magic;
96   std::atomic<uint16_t> flags;
97 };
98
99 struct IOBuf::HeapStorage {
100   HeapPrefix prefix;
101   // The IOBuf is last in the HeapStorage object.
102   // This way operator new will work even if allocating a subclass of IOBuf
103   // that requires more space.
104   folly::IOBuf buf;
105 };
106
107 struct IOBuf::HeapFullStorage {
108   // Make sure jemalloc allocates from the 64-byte class.  Putting this here
109   // because HeapStorage is private so it can't be at namespace level.
110   static_assert(sizeof(HeapStorage) <= 64,
111                 "IOBuf may not grow over 56 bytes!");
112
113   HeapStorage hs;
114   SharedInfo shared;
115   std::max_align_t align;
116 };
117
118 IOBuf::SharedInfo::SharedInfo()
119   : freeFn(nullptr),
120     userData(nullptr) {
121   // Use relaxed memory ordering here.  Since we are creating a new SharedInfo,
122   // no other threads should be referring to it yet.
123   refcount.store(1, std::memory_order_relaxed);
124 }
125
126 IOBuf::SharedInfo::SharedInfo(FreeFunction fn, void* arg)
127   : freeFn(fn),
128     userData(arg) {
129   // Use relaxed memory ordering here.  Since we are creating a new SharedInfo,
130   // no other threads should be referring to it yet.
131   refcount.store(1, std::memory_order_relaxed);
132 }
133
134 void* IOBuf::operator new(size_t size) {
135   size_t fullSize = offsetof(HeapStorage, buf) + size;
136   auto* storage = static_cast<HeapStorage*>(malloc(fullSize));
137   // operator new is not allowed to return NULL
138   if (UNLIKELY(storage == nullptr)) {
139     throw std::bad_alloc();
140   }
141
142   new (&storage->prefix) HeapPrefix(kIOBufInUse);
143   return &(storage->buf);
144 }
145
146 void* IOBuf::operator new(size_t size, void* ptr) {
147   return ptr;
148 }
149
150 void IOBuf::operator delete(void* ptr) {
151   auto* storageAddr = static_cast<uint8_t*>(ptr) - offsetof(HeapStorage, buf);
152   auto* storage = reinterpret_cast<HeapStorage*>(storageAddr);
153   releaseStorage(storage, kIOBufInUse);
154 }
155
156 void IOBuf::releaseStorage(HeapStorage* storage, uint16_t freeFlags) {
157   CHECK_EQ(storage->prefix.magic, static_cast<uint16_t>(kHeapMagic));
158
159   // Use relaxed memory order here.  If we are unlucky and happen to get
160   // out-of-date data the compare_exchange_weak() call below will catch
161   // it and load new data with memory_order_acq_rel.
162   auto flags = storage->prefix.flags.load(std::memory_order_acquire);
163   DCHECK_EQ((flags & freeFlags), freeFlags);
164
165   while (true) {
166     uint16_t newFlags = (flags & ~freeFlags);
167     if (newFlags == 0) {
168       // The storage space is now unused.  Free it.
169       storage->prefix.HeapPrefix::~HeapPrefix();
170       free(storage);
171       return;
172     }
173
174     // This storage segment still contains portions that are in use.
175     // Just clear the flags specified in freeFlags for now.
176     auto ret = storage->prefix.flags.compare_exchange_weak(
177         flags, newFlags, std::memory_order_acq_rel);
178     if (ret) {
179       // We successfully updated the flags.
180       return;
181     }
182
183     // We failed to update the flags.  Some other thread probably updated them
184     // and cleared some of the other bits.  Continue around the loop to see if
185     // we are the last user now, or if we need to try updating the flags again.
186   }
187 }
188
189 void IOBuf::freeInternalBuf(void* buf, void* userData) {
190   auto* storage = static_cast<HeapStorage*>(userData);
191   releaseStorage(storage, kDataInUse);
192 }
193
194 IOBuf::IOBuf(CreateOp, uint64_t capacity)
195   : next_(this),
196     prev_(this),
197     data_(nullptr),
198     length_(0),
199     flagsAndSharedInfo_(0) {
200   SharedInfo* info;
201   allocExtBuffer(capacity, &buf_, &info, &capacity_);
202   setSharedInfo(info);
203   data_ = buf_;
204 }
205
206 IOBuf::IOBuf(CopyBufferOp op, const void* buf, uint64_t size,
207              uint64_t headroom, uint64_t minTailroom)
208   : IOBuf(CREATE, headroom + size + minTailroom) {
209   advance(headroom);
210   memcpy(writableData(), buf, size);
211   append(size);
212 }
213
214 IOBuf::IOBuf(CopyBufferOp op, ByteRange br,
215              uint64_t headroom, uint64_t minTailroom)
216   : IOBuf(op, br.data(), br.size(), headroom, minTailroom) {
217 }
218
219 unique_ptr<IOBuf> IOBuf::create(uint64_t capacity) {
220   // For smaller-sized buffers, allocate the IOBuf, SharedInfo, and the buffer
221   // all with a single allocation.
222   //
223   // We don't do this for larger buffers since it can be wasteful if the user
224   // needs to reallocate the buffer but keeps using the same IOBuf object.
225   // In this case we can't free the data space until the IOBuf is also
226   // destroyed.  Callers can explicitly call createCombined() or
227   // createSeparate() if they know their use case better, and know if they are
228   // likely to reallocate the buffer later.
229   if (capacity <= kDefaultCombinedBufSize) {
230     return createCombined(capacity);
231   }
232   return createSeparate(capacity);
233 }
234
235 unique_ptr<IOBuf> IOBuf::createCombined(uint64_t capacity) {
236   // To save a memory allocation, allocate space for the IOBuf object, the
237   // SharedInfo struct, and the data itself all with a single call to malloc().
238   size_t requiredStorage = offsetof(HeapFullStorage, align) + capacity;
239   size_t mallocSize = goodMallocSize(requiredStorage);
240   auto* storage = static_cast<HeapFullStorage*>(malloc(mallocSize));
241
242   new (&storage->hs.prefix) HeapPrefix(kIOBufInUse | kDataInUse);
243   new (&storage->shared) SharedInfo(freeInternalBuf, storage);
244
245   uint8_t* bufAddr = reinterpret_cast<uint8_t*>(&storage->align);
246   uint8_t* storageEnd = reinterpret_cast<uint8_t*>(storage) + mallocSize;
247   size_t actualCapacity = storageEnd - bufAddr;
248   unique_ptr<IOBuf> ret(new (&storage->hs.buf) IOBuf(
249         InternalConstructor(), packFlagsAndSharedInfo(0, &storage->shared),
250         bufAddr, actualCapacity, bufAddr, 0));
251   return ret;
252 }
253
254 unique_ptr<IOBuf> IOBuf::createSeparate(uint64_t capacity) {
255   return make_unique<IOBuf>(CREATE, capacity);
256 }
257
258 unique_ptr<IOBuf> IOBuf::createChain(
259     size_t totalCapacity, uint64_t maxBufCapacity) {
260   unique_ptr<IOBuf> out = create(
261       std::min(totalCapacity, size_t(maxBufCapacity)));
262   size_t allocatedCapacity = out->capacity();
263
264   while (allocatedCapacity < totalCapacity) {
265     unique_ptr<IOBuf> newBuf = create(
266         std::min(totalCapacity - allocatedCapacity, size_t(maxBufCapacity)));
267     allocatedCapacity += newBuf->capacity();
268     out->prependChain(std::move(newBuf));
269   }
270
271   return out;
272 }
273
274 IOBuf::IOBuf(TakeOwnershipOp, void* buf, uint64_t capacity, uint64_t length,
275              FreeFunction freeFn, void* userData,
276              bool freeOnError)
277   : next_(this),
278     prev_(this),
279     data_(static_cast<uint8_t*>(buf)),
280     buf_(static_cast<uint8_t*>(buf)),
281     length_(length),
282     capacity_(capacity),
283     flagsAndSharedInfo_(packFlagsAndSharedInfo(kFlagFreeSharedInfo, nullptr)) {
284   try {
285     setSharedInfo(new SharedInfo(freeFn, userData));
286   } catch (...) {
287     takeOwnershipError(freeOnError, buf, freeFn, userData);
288     throw;
289   }
290 }
291
292 unique_ptr<IOBuf> IOBuf::takeOwnership(void* buf, uint64_t capacity,
293                                        uint64_t length,
294                                        FreeFunction freeFn,
295                                        void* userData,
296                                        bool freeOnError) {
297   try {
298     // TODO: We could allocate the IOBuf object and SharedInfo all in a single
299     // memory allocation.  We could use the existing HeapStorage class, and
300     // define a new kSharedInfoInUse flag.  We could change our code to call
301     // releaseStorage(kFlagFreeSharedInfo) when this kFlagFreeSharedInfo,
302     // rather than directly calling delete.
303     //
304     // Note that we always pass freeOnError as false to the constructor.
305     // If the constructor throws we'll handle it below.  (We have to handle
306     // allocation failures from make_unique too.)
307     return make_unique<IOBuf>(TAKE_OWNERSHIP, buf, capacity, length,
308                               freeFn, userData, false);
309   } catch (...) {
310     takeOwnershipError(freeOnError, buf, freeFn, userData);
311     throw;
312   }
313 }
314
315 IOBuf::IOBuf(WrapBufferOp, const void* buf, uint64_t capacity)
316   : IOBuf(InternalConstructor(), 0,
317           // We cast away the const-ness of the buffer here.
318           // This is okay since IOBuf users must use unshare() to create a copy
319           // of this buffer before writing to the buffer.
320           static_cast<uint8_t*>(const_cast<void*>(buf)), capacity,
321           static_cast<uint8_t*>(const_cast<void*>(buf)), capacity) {
322 }
323
324 IOBuf::IOBuf(WrapBufferOp op, ByteRange br)
325   : IOBuf(op, br.data(), br.size()) {
326 }
327
328 unique_ptr<IOBuf> IOBuf::wrapBuffer(const void* buf, uint64_t capacity) {
329   return make_unique<IOBuf>(WRAP_BUFFER, buf, capacity);
330 }
331
332 IOBuf::IOBuf() noexcept {
333 }
334
335 IOBuf::IOBuf(IOBuf&& other) noexcept {
336   *this = std::move(other);
337 }
338
339 IOBuf::IOBuf(const IOBuf& other) {
340   other.cloneInto(*this);
341 }
342
343 IOBuf::IOBuf(InternalConstructor,
344              uintptr_t flagsAndSharedInfo,
345              uint8_t* buf,
346              uint64_t capacity,
347              uint8_t* data,
348              uint64_t length)
349   : next_(this),
350     prev_(this),
351     data_(data),
352     buf_(buf),
353     length_(length),
354     capacity_(capacity),
355     flagsAndSharedInfo_(flagsAndSharedInfo) {
356   assert(data >= buf);
357   assert(data + length <= buf + capacity);
358 }
359
360 IOBuf::~IOBuf() {
361   // Destroying an IOBuf destroys the entire chain.
362   // Users of IOBuf should only explicitly delete the head of any chain.
363   // The other elements in the chain will be automatically destroyed.
364   while (next_ != this) {
365     // Since unlink() returns unique_ptr() and we don't store it,
366     // it will automatically delete the unlinked element.
367     (void)next_->unlink();
368   }
369
370   decrementRefcount();
371 }
372
373 IOBuf& IOBuf::operator=(IOBuf&& other) noexcept {
374   if (this == &other) {
375     return *this;
376   }
377
378   // If we are part of a chain, delete the rest of the chain.
379   while (next_ != this) {
380     // Since unlink() returns unique_ptr() and we don't store it,
381     // it will automatically delete the unlinked element.
382     (void)next_->unlink();
383   }
384
385   // Decrement our refcount on the current buffer
386   decrementRefcount();
387
388   // Take ownership of the other buffer's data
389   data_ = other.data_;
390   buf_ = other.buf_;
391   length_ = other.length_;
392   capacity_ = other.capacity_;
393   flagsAndSharedInfo_ = other.flagsAndSharedInfo_;
394   // Reset other so it is a clean state to be destroyed.
395   other.data_ = nullptr;
396   other.buf_ = nullptr;
397   other.length_ = 0;
398   other.capacity_ = 0;
399   other.flagsAndSharedInfo_ = 0;
400
401   // If other was part of the chain, assume ownership of the rest of its chain.
402   // (It's only valid to perform move assignment on the head of a chain.)
403   if (other.next_ != &other) {
404     next_ = other.next_;
405     next_->prev_ = this;
406     other.next_ = &other;
407
408     prev_ = other.prev_;
409     prev_->next_ = this;
410     other.prev_ = &other;
411   }
412
413   // Sanity check to make sure that other is in a valid state to be destroyed.
414   DCHECK_EQ(other.prev_, &other);
415   DCHECK_EQ(other.next_, &other);
416
417   return *this;
418 }
419
420 IOBuf& IOBuf::operator=(const IOBuf& other) {
421   if (this != &other) {
422     *this = IOBuf(other);
423   }
424   return *this;
425 }
426
427 bool IOBuf::empty() const {
428   const IOBuf* current = this;
429   do {
430     if (current->length() != 0) {
431       return false;
432     }
433     current = current->next_;
434   } while (current != this);
435   return true;
436 }
437
438 size_t IOBuf::countChainElements() const {
439   size_t numElements = 1;
440   for (IOBuf* current = next_; current != this; current = current->next_) {
441     ++numElements;
442   }
443   return numElements;
444 }
445
446 uint64_t IOBuf::computeChainDataLength() const {
447   uint64_t fullLength = length_;
448   for (IOBuf* current = next_; current != this; current = current->next_) {
449     fullLength += current->length_;
450   }
451   return fullLength;
452 }
453
454 void IOBuf::prependChain(unique_ptr<IOBuf>&& iobuf) {
455   // Take ownership of the specified IOBuf
456   IOBuf* other = iobuf.release();
457
458   // Remember the pointer to the tail of the other chain
459   IOBuf* otherTail = other->prev_;
460
461   // Hook up prev_->next_ to point at the start of the other chain,
462   // and other->prev_ to point at prev_
463   prev_->next_ = other;
464   other->prev_ = prev_;
465
466   // Hook up otherTail->next_ to point at us,
467   // and prev_ to point back at otherTail,
468   otherTail->next_ = this;
469   prev_ = otherTail;
470 }
471
472 unique_ptr<IOBuf> IOBuf::clone() const {
473   unique_ptr<IOBuf> ret = make_unique<IOBuf>();
474   cloneInto(*ret);
475   return ret;
476 }
477
478 unique_ptr<IOBuf> IOBuf::cloneOne() const {
479   unique_ptr<IOBuf> ret = make_unique<IOBuf>();
480   cloneOneInto(*ret);
481   return ret;
482 }
483
484 void IOBuf::cloneInto(IOBuf& other) const {
485   IOBuf tmp;
486   cloneOneInto(tmp);
487
488   for (IOBuf* current = next_; current != this; current = current->next_) {
489     tmp.prependChain(current->cloneOne());
490   }
491
492   other = std::move(tmp);
493 }
494
495 void IOBuf::cloneOneInto(IOBuf& other) const {
496   SharedInfo* info = sharedInfo();
497   if (info) {
498     setFlags(kFlagMaybeShared);
499   }
500   other = IOBuf(InternalConstructor(),
501                 flagsAndSharedInfo_, buf_, capacity_,
502                 data_, length_);
503   if (info) {
504     info->refcount.fetch_add(1, std::memory_order_acq_rel);
505   }
506 }
507
508 void IOBuf::unshareOneSlow() {
509   // Allocate a new buffer for the data
510   uint8_t* buf;
511   SharedInfo* sharedInfo;
512   uint64_t actualCapacity;
513   allocExtBuffer(capacity_, &buf, &sharedInfo, &actualCapacity);
514
515   // Copy the data
516   // Maintain the same amount of headroom.  Since we maintained the same
517   // minimum capacity we also maintain at least the same amount of tailroom.
518   uint64_t headlen = headroom();
519   memcpy(buf + headlen, data_, length_);
520
521   // Release our reference on the old buffer
522   decrementRefcount();
523   // Make sure kFlagMaybeShared and kFlagFreeSharedInfo are all cleared.
524   setFlagsAndSharedInfo(0, sharedInfo);
525
526   // Update the buffer pointers to point to the new buffer
527   data_ = buf + headlen;
528   buf_ = buf;
529 }
530
531 void IOBuf::unshareChained() {
532   // unshareChained() should only be called if we are part of a chain of
533   // multiple IOBufs.  The caller should have already verified this.
534   assert(isChained());
535
536   IOBuf* current = this;
537   while (true) {
538     if (current->isSharedOne()) {
539       // we have to unshare
540       break;
541     }
542
543     current = current->next_;
544     if (current == this) {
545       // None of the IOBufs in the chain are shared,
546       // so return without doing anything
547       return;
548     }
549   }
550
551   // We have to unshare.  Let coalesceSlow() do the work.
552   coalesceSlow();
553 }
554
555 void IOBuf::makeManagedChained() {
556   assert(isChained());
557
558   IOBuf* current = this;
559   while (true) {
560     current->makeManagedOne();
561     current = current->next_;
562     if (current == this) {
563       break;
564     }
565   }
566 }
567
568 void IOBuf::coalesceSlow() {
569   // coalesceSlow() should only be called if we are part of a chain of multiple
570   // IOBufs.  The caller should have already verified this.
571   DCHECK(isChained());
572
573   // Compute the length of the entire chain
574   uint64_t newLength = 0;
575   IOBuf* end = this;
576   do {
577     newLength += end->length_;
578     end = end->next_;
579   } while (end != this);
580
581   coalesceAndReallocate(newLength, end);
582   // We should be only element left in the chain now
583   DCHECK(!isChained());
584 }
585
586 void IOBuf::coalesceSlow(size_t maxLength) {
587   // coalesceSlow() should only be called if we are part of a chain of multiple
588   // IOBufs.  The caller should have already verified this.
589   DCHECK(isChained());
590   DCHECK_LT(length_, maxLength);
591
592   // Compute the length of the entire chain
593   uint64_t newLength = 0;
594   IOBuf* end = this;
595   while (true) {
596     newLength += end->length_;
597     end = end->next_;
598     if (newLength >= maxLength) {
599       break;
600     }
601     if (end == this) {
602       throw std::overflow_error("attempted to coalesce more data than "
603                                 "available");
604     }
605   }
606
607   coalesceAndReallocate(newLength, end);
608   // We should have the requested length now
609   DCHECK_GE(length_, maxLength);
610 }
611
612 void IOBuf::coalesceAndReallocate(size_t newHeadroom,
613                                   size_t newLength,
614                                   IOBuf* end,
615                                   size_t newTailroom) {
616   uint64_t newCapacity = newLength + newHeadroom + newTailroom;
617
618   // Allocate space for the coalesced buffer.
619   // We always convert to an external buffer, even if we happened to be an
620   // internal buffer before.
621   uint8_t* newBuf;
622   SharedInfo* newInfo;
623   uint64_t actualCapacity;
624   allocExtBuffer(newCapacity, &newBuf, &newInfo, &actualCapacity);
625
626   // Copy the data into the new buffer
627   uint8_t* newData = newBuf + newHeadroom;
628   uint8_t* p = newData;
629   IOBuf* current = this;
630   size_t remaining = newLength;
631   do {
632     assert(current->length_ <= remaining);
633     remaining -= current->length_;
634     memcpy(p, current->data_, current->length_);
635     p += current->length_;
636     current = current->next_;
637   } while (current != end);
638   assert(remaining == 0);
639
640   // Point at the new buffer
641   decrementRefcount();
642
643   // Make sure kFlagMaybeShared and kFlagFreeSharedInfo are all cleared.
644   setFlagsAndSharedInfo(0, newInfo);
645
646   capacity_ = actualCapacity;
647   buf_ = newBuf;
648   data_ = newData;
649   length_ = newLength;
650
651   // Separate from the rest of our chain.
652   // Since we don't store the unique_ptr returned by separateChain(),
653   // this will immediately delete the returned subchain.
654   if (isChained()) {
655     (void)separateChain(next_, current->prev_);
656   }
657 }
658
659 void IOBuf::decrementRefcount() {
660   // Externally owned buffers don't have a SharedInfo object and aren't managed
661   // by the reference count
662   SharedInfo* info = sharedInfo();
663   if (!info) {
664     return;
665   }
666
667   // Decrement the refcount
668   uint32_t newcnt = info->refcount.fetch_sub(
669       1, std::memory_order_acq_rel);
670   // Note that fetch_sub() returns the value before we decremented.
671   // If it is 1, we were the only remaining user; if it is greater there are
672   // still other users.
673   if (newcnt > 1) {
674     return;
675   }
676
677   // We were the last user.  Free the buffer
678   freeExtBuffer();
679
680   // Free the SharedInfo if it was allocated separately.
681   //
682   // This is only used by takeOwnership().
683   //
684   // To avoid this special case handling in decrementRefcount(), we could have
685   // takeOwnership() set a custom freeFn() that calls the user's free function
686   // then frees the SharedInfo object.  (This would require that
687   // takeOwnership() store the user's free function with its allocated
688   // SharedInfo object.)  However, handling this specially with a flag seems
689   // like it shouldn't be problematic.
690   if (flags() & kFlagFreeSharedInfo) {
691     delete sharedInfo();
692   }
693 }
694
695 void IOBuf::reserveSlow(uint64_t minHeadroom, uint64_t minTailroom) {
696   size_t newCapacity = (size_t)length_ + minHeadroom + minTailroom;
697   DCHECK_LT(newCapacity, UINT32_MAX);
698
699   // reserveSlow() is dangerous if anyone else is sharing the buffer, as we may
700   // reallocate and free the original buffer.  It should only ever be called if
701   // we are the only user of the buffer.
702   DCHECK(!isSharedOne());
703
704   // We'll need to reallocate the buffer.
705   // There are a few options.
706   // - If we have enough total room, move the data around in the buffer
707   //   and adjust the data_ pointer.
708   // - If we're using an internal buffer, we'll switch to an external
709   //   buffer with enough headroom and tailroom.
710   // - If we have enough headroom (headroom() >= minHeadroom) but not too much
711   //   (so we don't waste memory), we can try one of two things, depending on
712   //   whether we use jemalloc or not:
713   //   - If using jemalloc, we can try to expand in place, avoiding a memcpy()
714   //   - If not using jemalloc and we don't have too much to copy,
715   //     we'll use realloc() (note that realloc might have to copy
716   //     headroom + data + tailroom, see smartRealloc in folly/Malloc.h)
717   // - Otherwise, bite the bullet and reallocate.
718   if (headroom() + tailroom() >= minHeadroom + minTailroom) {
719     uint8_t* newData = writableBuffer() + minHeadroom;
720     memmove(newData, data_, length_);
721     data_ = newData;
722     return;
723   }
724
725   size_t newAllocatedCapacity = 0;
726   uint8_t* newBuffer = nullptr;
727   uint64_t newHeadroom = 0;
728   uint64_t oldHeadroom = headroom();
729
730   // If we have a buffer allocated with malloc and we just need more tailroom,
731   // try to use realloc()/xallocx() to grow the buffer in place.
732   SharedInfo* info = sharedInfo();
733   if (info && (info->freeFn == nullptr) && length_ != 0 &&
734       oldHeadroom >= minHeadroom) {
735     size_t headSlack = oldHeadroom - minHeadroom;
736     newAllocatedCapacity = goodExtBufferSize(newCapacity + headSlack);
737     if (usingJEMalloc()) {
738       // We assume that tailroom is more useful and more important than
739       // headroom (not least because realloc / xallocx allow us to grow the
740       // buffer at the tail, but not at the head)  So, if we have more headroom
741       // than we need, we consider that "wasted".  We arbitrarily define "too
742       // much" headroom to be 25% of the capacity.
743       if (headSlack * 4 <= newCapacity) {
744         size_t allocatedCapacity = capacity() + sizeof(SharedInfo);
745         void* p = buf_;
746         if (allocatedCapacity >= jemallocMinInPlaceExpandable) {
747           if (xallocx(p, newAllocatedCapacity, 0, 0) == newAllocatedCapacity) {
748             newBuffer = static_cast<uint8_t*>(p);
749             newHeadroom = oldHeadroom;
750           }
751           // if xallocx failed, do nothing, fall back to malloc/memcpy/free
752         }
753       }
754     } else {  // Not using jemalloc
755       size_t copySlack = capacity() - length_;
756       if (copySlack * 2 <= length_) {
757         void* p = realloc(buf_, newAllocatedCapacity);
758         if (UNLIKELY(p == nullptr)) {
759           throw std::bad_alloc();
760         }
761         newBuffer = static_cast<uint8_t*>(p);
762         newHeadroom = oldHeadroom;
763       }
764     }
765   }
766
767   // None of the previous reallocation strategies worked (or we're using
768   // an internal buffer).  malloc/copy/free.
769   if (newBuffer == nullptr) {
770     newAllocatedCapacity = goodExtBufferSize(newCapacity);
771     void* p = malloc(newAllocatedCapacity);
772     if (UNLIKELY(p == nullptr)) {
773       throw std::bad_alloc();
774     }
775     newBuffer = static_cast<uint8_t*>(p);
776     memcpy(newBuffer + minHeadroom, data_, length_);
777     if (sharedInfo()) {
778       freeExtBuffer();
779     }
780     newHeadroom = minHeadroom;
781   }
782
783   uint64_t cap;
784   initExtBuffer(newBuffer, newAllocatedCapacity, &info, &cap);
785
786   if (flags() & kFlagFreeSharedInfo) {
787     delete sharedInfo();
788   }
789
790   setFlagsAndSharedInfo(0, info);
791   capacity_ = cap;
792   buf_ = newBuffer;
793   data_ = newBuffer + newHeadroom;
794   // length_ is unchanged
795 }
796
797 void IOBuf::freeExtBuffer() {
798   SharedInfo* info = sharedInfo();
799   DCHECK(info);
800
801   if (info->freeFn) {
802     try {
803       info->freeFn(buf_, info->userData);
804     } catch (...) {
805       // The user's free function should never throw.  Otherwise we might
806       // throw from the IOBuf destructor.  Other code paths like coalesce()
807       // also assume that decrementRefcount() cannot throw.
808       abort();
809     }
810   } else {
811     free(buf_);
812   }
813 }
814
815 void IOBuf::allocExtBuffer(uint64_t minCapacity,
816                            uint8_t** bufReturn,
817                            SharedInfo** infoReturn,
818                            uint64_t* capacityReturn) {
819   size_t mallocSize = goodExtBufferSize(minCapacity);
820   uint8_t* buf = static_cast<uint8_t*>(malloc(mallocSize));
821   if (UNLIKELY(buf == nullptr)) {
822     throw std::bad_alloc();
823   }
824   initExtBuffer(buf, mallocSize, infoReturn, capacityReturn);
825   *bufReturn = buf;
826 }
827
828 size_t IOBuf::goodExtBufferSize(uint64_t minCapacity) {
829   // Determine how much space we should allocate.  We'll store the SharedInfo
830   // for the external buffer just after the buffer itself.  (We store it just
831   // after the buffer rather than just before so that the code can still just
832   // use free(buf_) to free the buffer.)
833   size_t minSize = static_cast<size_t>(minCapacity) + sizeof(SharedInfo);
834   // Add room for padding so that the SharedInfo will be aligned on an 8-byte
835   // boundary.
836   minSize = (minSize + 7) & ~7;
837
838   // Use goodMallocSize() to bump up the capacity to a decent size to request
839   // from malloc, so we can use all of the space that malloc will probably give
840   // us anyway.
841   return goodMallocSize(minSize);
842 }
843
844 void IOBuf::initExtBuffer(uint8_t* buf, size_t mallocSize,
845                           SharedInfo** infoReturn,
846                           uint64_t* capacityReturn) {
847   // Find the SharedInfo storage at the end of the buffer
848   // and construct the SharedInfo.
849   uint8_t* infoStart = (buf + mallocSize) - sizeof(SharedInfo);
850   SharedInfo* sharedInfo = new(infoStart) SharedInfo;
851
852   *capacityReturn = infoStart - buf;
853   *infoReturn = sharedInfo;
854 }
855
856 fbstring IOBuf::moveToFbString() {
857   // malloc-allocated buffers are just fine, everything else needs
858   // to be turned into one.
859   if (!sharedInfo() ||         // user owned, not ours to give up
860       sharedInfo()->freeFn ||  // not malloc()-ed
861       headroom() != 0 ||       // malloc()-ed block doesn't start at beginning
862       tailroom() == 0 ||       // no room for NUL terminator
863       isShared() ||            // shared
864       isChained()) {           // chained
865     // We might as well get rid of all head and tailroom if we're going
866     // to reallocate; we need 1 byte for NUL terminator.
867     coalesceAndReallocate(0, computeChainDataLength(), this, 1);
868   }
869
870   // Ensure NUL terminated
871   *writableTail() = 0;
872   fbstring str(reinterpret_cast<char*>(writableData()),
873                length(),  capacity(),
874                AcquireMallocatedString());
875
876   if (flags() & kFlagFreeSharedInfo) {
877     delete sharedInfo();
878   }
879
880   // Reset to a state where we can be deleted cleanly
881   flagsAndSharedInfo_ = 0;
882   buf_ = nullptr;
883   clear();
884   return str;
885 }
886
887 IOBuf::Iterator IOBuf::cbegin() const {
888   return Iterator(this, this);
889 }
890
891 IOBuf::Iterator IOBuf::cend() const {
892   return Iterator(nullptr, nullptr);
893 }
894
895 folly::fbvector<struct iovec> IOBuf::getIov() const {
896   folly::fbvector<struct iovec> iov;
897   iov.reserve(countChainElements());
898   appendToIov(&iov);
899   return iov;
900 }
901
902 void IOBuf::appendToIov(folly::fbvector<struct iovec>* iov) const {
903   IOBuf const* p = this;
904   do {
905     // some code can get confused by empty iovs, so skip them
906     if (p->length() > 0) {
907       iov->push_back({(void*)p->data(), folly::to<size_t>(p->length())});
908     }
909     p = p->next();
910   } while (p != this);
911 }
912
913 size_t IOBuf::fillIov(struct iovec* iov, size_t len) const {
914   IOBuf const* p = this;
915   size_t i = 0;
916   while (i < len) {
917     // some code can get confused by empty iovs, so skip them
918     if (p->length() > 0) {
919       iov[i].iov_base = const_cast<uint8_t*>(p->data());
920       iov[i].iov_len = p->length();
921       i++;
922     }
923     p = p->next();
924     if (p == this) {
925       return i;
926     }
927   }
928   return 0;
929 }
930
931 size_t IOBufHash::operator()(const IOBuf& buf) const {
932   folly::hash::SpookyHashV2 hasher;
933   hasher.Init(0, 0);
934   io::Cursor cursor(&buf);
935   for (;;) {
936     auto p = cursor.peek();
937     if (p.second == 0) {
938       break;
939     }
940     hasher.Update(p.first, p.second);
941     cursor.skip(p.second);
942   }
943   uint64_t h1;
944   uint64_t h2;
945   hasher.Final(&h1, &h2);
946   return h1;
947 }
948
949 bool IOBufEqual::operator()(const IOBuf& a, const IOBuf& b) const {
950   io::Cursor ca(&a);
951   io::Cursor cb(&b);
952   for (;;) {
953     auto pa = ca.peek();
954     auto pb = cb.peek();
955     if (pa.second == 0 && pb.second == 0) {
956       return true;
957     } else if (pa.second == 0 || pb.second == 0) {
958       return false;
959     }
960     size_t n = std::min(pa.second, pb.second);
961     DCHECK_GT(n, 0);
962     if (memcmp(pa.first, pb.first, n)) {
963       return false;
964     }
965     ca.skip(n);
966     cb.skip(n);
967   }
968 }
969
970 } // folly