IOBuf.cpp 33 KB

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