folly/io/test/IOBufTest.cpp

   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
  17 #include <folly/io/IOBuf.h>
  18 #include <folly/io/TypedIOBuf.h>
  19
  20 #include <cstddef>
  21
  22 #include <boost/random.hpp>
  23
  24 #include <folly/Range.h>
  25 #include <folly/memory/Malloc.h>
  26 #include <folly/portability/GTest.h>
  27
  28 using folly::fbstring;
  29 using folly::fbvector;
  30 using folly::IOBuf;
  31 using folly::TypedIOBuf;
  32 using folly::StringPiece;
  33 using folly::ByteRange;
  34 using std::unique_ptr;
  35
  36 void append(std::unique_ptr<IOBuf>& buf, StringPiece str) {
  37   EXPECT_LE(str.size(), buf->tailroom());
  38   memcpy(buf->writableData(), str.data(), str.size());
  39   buf->append(str.size());
  40 }
  41
  42 void prepend(std::unique_ptr<IOBuf>& buf, StringPiece str) {
  43   EXPECT_LE(str.size(), buf->headroom());
  44   memcpy(buf->writableData() - str.size(), str.data(), str.size());
  45   buf->prepend(str.size());
  46 }
  47
  48 TEST(IOBuf, Simple) {
  49   unique_ptr<IOBuf> buf(IOBuf::create(100));
  50   uint32_t cap = buf->capacity();
  51   EXPECT_LE(100, cap);
  52   EXPECT_EQ(0, buf->headroom());
  53   EXPECT_EQ(0, buf->length());
  54   EXPECT_EQ(cap, buf->tailroom());
  55
  56   append(buf, "world");
  57   buf->advance(10);
  58   EXPECT_EQ(10, buf->headroom());
  59   EXPECT_EQ(5, buf->length());
  60   EXPECT_EQ(cap - 15, buf->tailroom());
  61
  62   prepend(buf, "hello ");
  63   EXPECT_EQ(4, buf->headroom());
  64   EXPECT_EQ(11, buf->length());
  65   EXPECT_EQ(cap - 15, buf->tailroom());
  66
  67   const char* p = reinterpret_cast<const char*>(buf->data());
  68   EXPECT_EQ("hello world", std::string(p, buf->length()));
  69
  70   buf->clear();
  71   EXPECT_EQ(0, buf->headroom());
  72   EXPECT_EQ(0, buf->length());
  73   EXPECT_EQ(cap, buf->tailroom());
  74 }
  75
  76
  77 void testAllocSize(uint32_t requestedCapacity) {
  78   unique_ptr<IOBuf> iobuf(IOBuf::create(requestedCapacity));
  79   EXPECT_GE(iobuf->capacity(), requestedCapacity);
  80 }
  81
  82 TEST(IOBuf, AllocSizes) {
  83   // Try with a small allocation size that should fit in the internal buffer
  84   testAllocSize(28);
  85
  86   // Try with a large allocation size that will require an external buffer.
  87   testAllocSize(9000);
  88
  89   // 220 bytes is currently the cutoff
  90   // (It would be nice to use the IOBuf::kMaxInternalDataSize constant,
  91   // but it's private and it doesn't seem worth making it public just for this
  92   // test code.)
  93   testAllocSize(220);
  94   testAllocSize(219);
  95   testAllocSize(221);
  96 }
  97
  98 void deleteArrayBuffer(void *buf, void* arg) {
  99   uint32_t* deleteCount = static_cast<uint32_t*>(arg);
 100   ++(*deleteCount);
 101   uint8_t* bufPtr = static_cast<uint8_t*>(buf);
 102   delete[] bufPtr;
 103 }
 104
 105 TEST(IOBuf, TakeOwnership) {
 106   uint32_t size1 = 99;
 107   uint8_t *buf1 = static_cast<uint8_t*>(malloc(size1));
 108   unique_ptr<IOBuf> iobuf1(IOBuf::takeOwnership(buf1, size1));
 109   EXPECT_EQ(buf1, iobuf1->data());
 110   EXPECT_EQ(size1, iobuf1->length());
 111   EXPECT_EQ(buf1, iobuf1->buffer());
 112   EXPECT_EQ(size1, iobuf1->capacity());
 113
 114   uint32_t deleteCount = 0;
 115   uint32_t size2 = 4321;
 116   uint8_t *buf2 = new uint8_t[size2];
 117   unique_ptr<IOBuf> iobuf2(IOBuf::takeOwnership(buf2, size2,
 118                                                 deleteArrayBuffer,
 119                                                 &deleteCount));
 120   EXPECT_EQ(buf2, iobuf2->data());
 121   EXPECT_EQ(size2, iobuf2->length());
 122   EXPECT_EQ(buf2, iobuf2->buffer());
 123   EXPECT_EQ(size2, iobuf2->capacity());
 124   EXPECT_EQ(0, deleteCount);
 125   iobuf2.reset();
 126   EXPECT_EQ(1, deleteCount);
 127
 128   deleteCount = 0;
 129   uint32_t size3 = 3456;
 130   uint8_t *buf3 = new uint8_t[size3];
 131   uint32_t length3 = 48;
 132   unique_ptr<IOBuf> iobuf3(IOBuf::takeOwnership(buf3, size3, length3,
 133                                                 deleteArrayBuffer,
 134                                                 &deleteCount));
 135   EXPECT_EQ(buf3, iobuf3->data());
 136   EXPECT_EQ(length3, iobuf3->length());
 137   EXPECT_EQ(buf3, iobuf3->buffer());
 138   EXPECT_EQ(size3, iobuf3->capacity());
 139   EXPECT_EQ(0, deleteCount);
 140   iobuf3.reset();
 141   EXPECT_EQ(1, deleteCount);
 142
 143   deleteCount = 0;
 144   {
 145     uint32_t size4 = 1234;
 146     uint8_t *buf4 = new uint8_t[size4];
 147     uint32_t length4 = 48;
 148     IOBuf iobuf4(IOBuf::TAKE_OWNERSHIP, buf4, size4, length4,
 149                  deleteArrayBuffer, &deleteCount);
 150     EXPECT_EQ(buf4, iobuf4.data());
 151     EXPECT_EQ(length4, iobuf4.length());
 152     EXPECT_EQ(buf4, iobuf4.buffer());
 153     EXPECT_EQ(size4, iobuf4.capacity());
 154
 155     IOBuf iobuf5 = std::move(iobuf4);
 156     EXPECT_EQ(buf4, iobuf5.data());
 157     EXPECT_EQ(length4, iobuf5.length());
 158     EXPECT_EQ(buf4, iobuf5.buffer());
 159     EXPECT_EQ(size4, iobuf5.capacity());
 160     EXPECT_EQ(0, deleteCount);
 161   }
 162   EXPECT_EQ(1, deleteCount);
 163 }
 164
 165 TEST(IOBuf, WrapBuffer) {
 166   const uint32_t size1 = 1234;
 167   uint8_t buf1[size1];
 168   unique_ptr<IOBuf> iobuf1(IOBuf::wrapBuffer(buf1, size1));
 169   EXPECT_EQ(buf1, iobuf1->data());
 170   EXPECT_EQ(size1, iobuf1->length());
 171   EXPECT_EQ(buf1, iobuf1->buffer());
 172   EXPECT_EQ(size1, iobuf1->capacity());
 173
 174   uint32_t size2 = 0x1234;
 175   unique_ptr<uint8_t[]> buf2(new uint8_t[size2]);
 176   unique_ptr<IOBuf> iobuf2(IOBuf::wrapBuffer(buf2.get(), size2));
 177   EXPECT_EQ(buf2.get(), iobuf2->data());
 178   EXPECT_EQ(size2, iobuf2->length());
 179   EXPECT_EQ(buf2.get(), iobuf2->buffer());
 180   EXPECT_EQ(size2, iobuf2->capacity());
 181
 182   uint32_t size3 = 4321;
 183   unique_ptr<uint8_t[]> buf3(new uint8_t[size3]);
 184   IOBuf iobuf3(IOBuf::WRAP_BUFFER, buf3.get(), size3);
 185   EXPECT_EQ(buf3.get(), iobuf3.data());
 186   EXPECT_EQ(size3, iobuf3.length());
 187   EXPECT_EQ(buf3.get(), iobuf3.buffer());
 188   EXPECT_EQ(size3, iobuf3.capacity());
 189
 190   const uint32_t size4 = 2345;
 191   unique_ptr<uint8_t[]> buf4(new uint8_t[size4]);
 192   IOBuf iobuf4 = IOBuf::wrapBufferAsValue(buf4.get(), size4);
 193   EXPECT_EQ(buf4.get(), iobuf4.data());
 194   EXPECT_EQ(size4, iobuf4.length());
 195   EXPECT_EQ(buf4.get(), iobuf4.buffer());
 196   EXPECT_EQ(size4, iobuf4.capacity());
 197 }
 198
 199 TEST(IOBuf, CreateCombined) {
 200   // Create a combined IOBuf, then destroy it.
 201   // The data buffer and IOBuf both become unused as part of the destruction
 202   {
 203     auto buf = IOBuf::createCombined(256);
 204     EXPECT_FALSE(buf->isShared());
 205   }
 206
 207   // Create a combined IOBuf, clone from it, and then destroy the original
 208   // IOBuf.  The data buffer cannot be deleted until the clone is also
 209   // destroyed.
 210   {
 211     auto bufA = IOBuf::createCombined(256);
 212     EXPECT_FALSE(bufA->isShared());
 213     auto bufB = bufA->clone();
 214     EXPECT_TRUE(bufA->isShared());
 215     EXPECT_TRUE(bufB->isShared());
 216     bufA.reset();
 217     EXPECT_FALSE(bufB->isShared());
 218   }
 219
 220   // Create a combined IOBuf, then call reserve() to get a larger buffer.
 221   // The IOBuf no longer points to the combined data buffer, but the
 222   // overall memory segment cannot be deleted until the IOBuf is also
 223   // destroyed.
 224   {
 225     auto buf = IOBuf::createCombined(256);
 226     buf->reserve(0, buf->capacity() + 100);
 227   }
 228
 229   // Create a combined IOBuf, clone from it, then call unshare() on the original
 230   // buffer.  This creates a situation where bufB is pointing at the combined
 231   // buffer associated with bufA, but bufA is now using a different buffer.
 232   auto testSwap = [](bool resetAFirst) {
 233     auto bufA = IOBuf::createCombined(256);
 234     EXPECT_FALSE(bufA->isShared());
 235     auto bufB = bufA->clone();
 236     EXPECT_TRUE(bufA->isShared());
 237     EXPECT_TRUE(bufB->isShared());
 238     bufA->unshare();
 239     EXPECT_FALSE(bufA->isShared());
 240     EXPECT_FALSE(bufB->isShared());
 241
 242     if (resetAFirst) {
 243       bufA.reset();
 244       bufB.reset();
 245     } else {
 246       bufB.reset();
 247       bufA.reset();
 248     }
 249   };
 250   testSwap(true);
 251   testSwap(false);
 252 }
 253
 254 void fillBuf(uint8_t* buf, uint32_t length, boost::mt19937& gen) {
 255   for (uint32_t n = 0; n < length; ++n) {
 256     buf[n] = static_cast<uint8_t>(gen() & 0xff);
 257   }
 258 }
 259
 260 void fillBuf(IOBuf* buf, boost::mt19937& gen) {
 261   buf->unshare();
 262   fillBuf(buf->writableData(), buf->length(), gen);
 263 }
 264
 265 void checkBuf(const uint8_t* buf, uint32_t length, boost::mt19937& gen) {
 266   // Rather than using EXPECT_EQ() to check each character,
 267   // count the number of differences and the first character that differs.
 268   // This way on error we'll report just that information, rather than tons of
 269   // failed checks for each byte in the buffer.
 270   uint32_t numDifferences = 0;
 271   uint32_t firstDiffIndex = 0;
 272   uint8_t firstDiffExpected = 0;
 273   for (uint32_t n = 0; n < length; ++n) {
 274     uint8_t expected = static_cast<uint8_t>(gen() & 0xff);
 275     if (buf[n] == expected) {
 276       continue;
 277     }
 278
 279     if (numDifferences == 0) {
 280       firstDiffIndex = n;
 281       firstDiffExpected = expected;
 282     }
 283     ++numDifferences;
 284   }
 285
 286   EXPECT_EQ(0, numDifferences);
 287   if (numDifferences > 0) {
 288     // Cast to int so it will be printed numerically
 289     // rather than as a char if the check fails
 290     EXPECT_EQ(static_cast<int>(buf[firstDiffIndex]),
 291               static_cast<int>(firstDiffExpected));
 292   }
 293 }
 294
 295 void checkBuf(IOBuf* buf, boost::mt19937& gen) {
 296   checkBuf(buf->data(), buf->length(), gen);
 297 }
 298
 299 void checkBuf(ByteRange buf, boost::mt19937& gen) {
 300   checkBuf(buf.data(), buf.size(), gen);
 301 }
 302
 303 void checkChain(IOBuf* buf, boost::mt19937& gen) {
 304   IOBuf *current = buf;
 305   do {
 306     checkBuf(current->data(), current->length(), gen);
 307     current = current->next();
 308   } while (current != buf);
 309 }
 310
 311 TEST(IOBuf, Chaining) {
 312   uint32_t fillSeed = 0x12345678;
 313   boost::mt19937 gen(fillSeed);
 314
 315   // An IOBuf with external storage
 316   uint32_t headroom = 123;
 317   unique_ptr<IOBuf> iob1(IOBuf::create(2048));
 318   iob1->advance(headroom);
 319   iob1->append(1500);
 320   fillBuf(iob1.get(), gen);
 321
 322   // An IOBuf with internal storage
 323   unique_ptr<IOBuf> iob2(IOBuf::create(20));
 324   iob2->append(20);
 325   fillBuf(iob2.get(), gen);
 326
 327   // An IOBuf around a buffer it doesn't own
 328   uint8_t localbuf[1234];
 329   fillBuf(localbuf, 1234, gen);
 330   unique_ptr<IOBuf> iob3(IOBuf::wrapBuffer(localbuf, sizeof(localbuf)));
 331
 332   // An IOBuf taking ownership of a user-supplied buffer
 333   uint32_t heapBufSize = 900;
 334   uint8_t* heapBuf = static_cast<uint8_t*>(malloc(heapBufSize));
 335   fillBuf(heapBuf, heapBufSize, gen);
 336   unique_ptr<IOBuf> iob4(IOBuf::takeOwnership(heapBuf, heapBufSize));
 337
 338   // An IOBuf taking ownership of a user-supplied buffer with
 339   // a custom free function
 340   uint32_t arrayBufSize = 321;
 341   uint8_t* arrayBuf = new uint8_t[arrayBufSize];
 342   fillBuf(arrayBuf, arrayBufSize, gen);
 343   uint32_t arrayBufFreeCount = 0;
 344   unique_ptr<IOBuf> iob5(IOBuf::takeOwnership(arrayBuf, arrayBufSize,
 345                                               deleteArrayBuffer,
 346                                               &arrayBufFreeCount));
 347
 348   EXPECT_FALSE(iob1->isChained());
 349   EXPECT_FALSE(iob2->isChained());
 350   EXPECT_FALSE(iob3->isChained());
 351   EXPECT_FALSE(iob4->isChained());
 352   EXPECT_FALSE(iob5->isChained());
 353
 354   EXPECT_FALSE(iob1->isSharedOne());
 355   EXPECT_FALSE(iob2->isSharedOne());
 356   EXPECT_TRUE(iob3->isSharedOne()); // since we own the buffer
 357   EXPECT_FALSE(iob4->isSharedOne());
 358   EXPECT_FALSE(iob5->isSharedOne());
 359
 360   // Chain the buffers all together
 361   // Since we are going to relinquish ownership of iob2-5 to the chain,
 362   // store raw pointers to them so we can reference them later.
 363   IOBuf* iob2ptr = iob2.get();
 364   IOBuf* iob3ptr = iob3.get();
 365   IOBuf* iob4ptr = iob4.get();
 366   IOBuf* iob5ptr = iob5.get();
 367
 368   iob1->prependChain(std::move(iob2));
 369   iob1->prependChain(std::move(iob4));
 370   iob2ptr->appendChain(std::move(iob3));
 371   iob1->prependChain(std::move(iob5));
 372
 373   EXPECT_EQ(iob2ptr, iob1->next());
 374   EXPECT_EQ(iob3ptr, iob2ptr->next());
 375   EXPECT_EQ(iob4ptr, iob3ptr->next());
 376   EXPECT_EQ(iob5ptr, iob4ptr->next());
 377   EXPECT_EQ(iob1.get(), iob5ptr->next());
 378
 379   EXPECT_EQ(iob5ptr, iob1->prev());
 380   EXPECT_EQ(iob1.get(), iob2ptr->prev());
 381   EXPECT_EQ(iob2ptr, iob3ptr->prev());
 382   EXPECT_EQ(iob3ptr, iob4ptr->prev());
 383   EXPECT_EQ(iob4ptr, iob5ptr->prev());
 384
 385   EXPECT_TRUE(iob1->isChained());
 386   EXPECT_TRUE(iob2ptr->isChained());
 387   EXPECT_TRUE(iob3ptr->isChained());
 388   EXPECT_TRUE(iob4ptr->isChained());
 389   EXPECT_TRUE(iob5ptr->isChained());
 390
 391   uint64_t fullLength = (iob1->length() + iob2ptr->length() +
 392                          iob3ptr->length() + iob4ptr->length() +
 393                         iob5ptr->length());
 394   EXPECT_EQ(5, iob1->countChainElements());
 395   EXPECT_EQ(fullLength, iob1->computeChainDataLength());
 396
 397   // Since iob3 is shared, the entire buffer should report itself as shared
 398   EXPECT_TRUE(iob1->isShared());
 399   // Unshare just iob3
 400   iob3ptr->unshareOne();
 401   EXPECT_FALSE(iob3ptr->isSharedOne());
 402   // Now everything in the chain should be unshared.
 403   // Check on all members of the chain just for good measure
 404   EXPECT_FALSE(iob1->isShared());
 405   EXPECT_FALSE(iob2ptr->isShared());
 406   EXPECT_FALSE(iob3ptr->isShared());
 407   EXPECT_FALSE(iob4ptr->isShared());
 408   EXPECT_FALSE(iob5ptr->isShared());
 409
 410   // Check iteration
 411   gen.seed(fillSeed);
 412   size_t count = 0;
 413   for (auto buf : *iob1) {
 414     checkBuf(buf, gen);
 415     ++count;
 416   }
 417   EXPECT_EQ(5, count);
 418
 419   // Clone one of the IOBufs in the chain
 420   unique_ptr<IOBuf> iob4clone = iob4ptr->cloneOne();
 421   gen.seed(fillSeed);
 422   checkBuf(iob1.get(), gen);
 423   checkBuf(iob2ptr, gen);
 424   checkBuf(iob3ptr, gen);
 425   checkBuf(iob4clone.get(), gen);
 426   checkBuf(iob5ptr, gen);
 427
 428   EXPECT_TRUE(iob1->isShared());
 429   EXPECT_TRUE(iob2ptr->isShared());
 430   EXPECT_TRUE(iob3ptr->isShared());
 431   EXPECT_TRUE(iob4ptr->isShared());
 432   EXPECT_TRUE(iob5ptr->isShared());
 433
 434   EXPECT_FALSE(iob1->isSharedOne());
 435   EXPECT_FALSE(iob2ptr->isSharedOne());
 436   EXPECT_FALSE(iob3ptr->isSharedOne());
 437   EXPECT_TRUE(iob4ptr->isSharedOne());
 438   EXPECT_FALSE(iob5ptr->isSharedOne());
 439
 440   // Unshare that clone
 441   EXPECT_TRUE(iob4clone->isSharedOne());
 442   iob4clone->unshare();
 443   EXPECT_FALSE(iob4clone->isSharedOne());
 444   EXPECT_FALSE(iob4ptr->isSharedOne());
 445   EXPECT_FALSE(iob1->isShared());
 446   iob4clone.reset();
 447
 448
 449   // Create a clone of a different IOBuf
 450   EXPECT_FALSE(iob1->isShared());
 451   EXPECT_FALSE(iob3ptr->isSharedOne());
 452
 453   unique_ptr<IOBuf> iob3clone = iob3ptr->cloneOne();
 454   gen.seed(fillSeed);
 455   checkBuf(iob1.get(), gen);
 456   checkBuf(iob2ptr, gen);
 457   checkBuf(iob3clone.get(), gen);
 458   checkBuf(iob4ptr, gen);
 459   checkBuf(iob5ptr, gen);
 460
 461   EXPECT_TRUE(iob1->isShared());
 462   EXPECT_TRUE(iob3ptr->isSharedOne());
 463   EXPECT_FALSE(iob1->isSharedOne());
 464
 465   // Delete the clone and make sure the original is unshared
 466   iob3clone.reset();
 467   EXPECT_FALSE(iob1->isShared());
 468   EXPECT_FALSE(iob3ptr->isSharedOne());
 469
 470
 471   // Clone the entire chain
 472   unique_ptr<IOBuf> chainClone = iob1->clone();
 473   // Verify that the data is correct.
 474   EXPECT_EQ(fullLength, chainClone->computeChainDataLength());
 475   gen.seed(fillSeed);
 476   checkChain(chainClone.get(), gen);
 477
 478   // Check that the buffers report sharing correctly
 479   EXPECT_TRUE(chainClone->isShared());
 480   EXPECT_TRUE(iob1->isShared());
 481
 482   EXPECT_TRUE(iob1->isSharedOne());
 483   EXPECT_TRUE(iob2ptr->isSharedOne());
 484   EXPECT_TRUE(iob3ptr->isSharedOne());
 485   EXPECT_TRUE(iob4ptr->isSharedOne());
 486   EXPECT_TRUE(iob5ptr->isSharedOne());
 487
 488   // Unshare the cloned chain
 489   chainClone->unshare();
 490   EXPECT_FALSE(chainClone->isShared());
 491   EXPECT_FALSE(iob1->isShared());
 492
 493   // Make sure the unshared result still has the same data
 494   EXPECT_EQ(fullLength, chainClone->computeChainDataLength());
 495   gen.seed(fillSeed);
 496   checkChain(chainClone.get(), gen);
 497
 498   // Destroy this chain
 499   chainClone.reset();
 500
 501
 502   // Clone a new chain
 503   EXPECT_FALSE(iob1->isShared());
 504   chainClone = iob1->clone();
 505   EXPECT_TRUE(iob1->isShared());
 506   EXPECT_TRUE(chainClone->isShared());
 507
 508   // Delete the original chain
 509   iob1.reset();
 510   EXPECT_FALSE(chainClone->isShared());
 511
 512   // Coalesce the chain
 513   //
 514   // Coalescing this chain will create a new buffer and release the last
 515   // refcount on the original buffers we created.  Also make sure
 516   // that arrayBufFreeCount increases to one to indicate that arrayBuf was
 517   // freed.
 518   EXPECT_EQ(5, chainClone->countChainElements());
 519   EXPECT_EQ(0, arrayBufFreeCount);
 520
 521   // Buffer lengths: 1500 20 1234 900 321
 522   // Attempting to gather more data than available should fail
 523   EXPECT_THROW(chainClone->gather(4000), std::overflow_error);
 524   // Coalesce the first 3 buffers
 525   chainClone->gather(1521);
 526   EXPECT_EQ(3, chainClone->countChainElements());
 527   EXPECT_EQ(0, arrayBufFreeCount);
 528
 529   // Make sure the data is still the same after coalescing
 530   EXPECT_EQ(fullLength, chainClone->computeChainDataLength());
 531   gen.seed(fillSeed);
 532   checkChain(chainClone.get(), gen);
 533
 534   // cloneCoalesced
 535   {
 536     auto chainCloneCoalesced = chainClone->cloneCoalesced();
 537     EXPECT_EQ(1, chainCloneCoalesced->countChainElements());
 538     EXPECT_EQ(fullLength, chainCloneCoalesced->computeChainDataLength());
 539     gen.seed(fillSeed);
 540     checkChain(chainCloneCoalesced.get(), gen);
 541   }
 542
 543   // Coalesce the entire chain
 544   chainClone->coalesce();
 545   EXPECT_EQ(1, chainClone->countChainElements());
 546   EXPECT_EQ(1, arrayBufFreeCount);
 547
 548   // Make sure the data is still the same after coalescing
 549   EXPECT_EQ(fullLength, chainClone->computeChainDataLength());
 550   gen.seed(fillSeed);
 551   checkChain(chainClone.get(), gen);
 552
 553   // Make a new chain to test the unlink and pop operations
 554   iob1 = IOBuf::create(1);
 555   iob1->append(1);
 556   IOBuf *iob1ptr = iob1.get();
 557   iob2 = IOBuf::create(3);
 558   iob2->append(3);
 559   iob2ptr = iob2.get();
 560   iob3 = IOBuf::create(5);
 561   iob3->append(5);
 562   iob3ptr = iob3.get();
 563   iob4 = IOBuf::create(7);
 564   iob4->append(7);
 565   iob4ptr = iob4.get();
 566   iob1->appendChain(std::move(iob2));
 567   iob1->prev()->appendChain(std::move(iob3));
 568   iob1->prev()->appendChain(std::move(iob4));
 569   EXPECT_EQ(4, iob1->countChainElements());
 570   EXPECT_EQ(16, iob1->computeChainDataLength());
 571
 572   // Unlink from the middle of the chain
 573   iob3 = iob3ptr->unlink();
 574   EXPECT_TRUE(iob3.get() == iob3ptr);
 575   EXPECT_EQ(3, iob1->countChainElements());
 576   EXPECT_EQ(11, iob1->computeChainDataLength());
 577
 578   // Unlink from the end of the chain
 579   iob4 = iob1->prev()->unlink();
 580   EXPECT_TRUE(iob4.get() == iob4ptr);
 581   EXPECT_EQ(2, iob1->countChainElements());
 582   EXPECT_TRUE(iob1->next() == iob2ptr);
 583   EXPECT_EQ(4, iob1->computeChainDataLength());
 584
 585   // Pop from the front of the chain
 586   iob2 = iob1->pop();
 587   EXPECT_TRUE(iob1.get() == iob1ptr);
 588   EXPECT_EQ(1, iob1->countChainElements());
 589   EXPECT_EQ(1, iob1->computeChainDataLength());
 590   EXPECT_TRUE(iob2.get() == iob2ptr);
 591   EXPECT_EQ(1, iob2->countChainElements());
 592   EXPECT_EQ(3, iob2->computeChainDataLength());
 593 }
 594
 595 void testFreeFn(void* buffer, void* ptr) {
 596   uint32_t* freeCount = static_cast<uint32_t*>(ptr);;
 597   delete[] static_cast<uint8_t*>(buffer);
 598   if (freeCount) {
 599     ++(*freeCount);
 600   }
 601 };
 602
 603 TEST(IOBuf, Reserve) {
 604   uint32_t fillSeed = 0x23456789;
 605   boost::mt19937 gen(fillSeed);
 606
 607   // Reserve does nothing if empty and doesn't have to grow the buffer
 608   {
 609     gen.seed(fillSeed);
 610     unique_ptr<IOBuf> iob(IOBuf::create(2000));
 611     EXPECT_EQ(0, iob->headroom());
 612     const void* p1 = iob->buffer();
 613     iob->reserve(5, 15);
 614     EXPECT_LE(5, iob->headroom());
 615     EXPECT_EQ(p1, iob->buffer());
 616   }
 617
 618   // Reserve doesn't reallocate if we have enough total room
 619   {
 620     gen.seed(fillSeed);
 621     unique_ptr<IOBuf> iob(IOBuf::create(2000));
 622     iob->append(100);
 623     fillBuf(iob.get(), gen);
 624     EXPECT_EQ(0, iob->headroom());
 625     EXPECT_EQ(100, iob->length());
 626     const void* p1 = iob->buffer();
 627     const uint8_t* d1 = iob->data();
 628     iob->reserve(100, 1800);
 629     EXPECT_LE(100, iob->headroom());
 630     EXPECT_EQ(p1, iob->buffer());
 631     EXPECT_EQ(d1 + 100, iob->data());
 632     gen.seed(fillSeed);
 633     checkBuf(iob.get(), gen);
 634   }
 635
 636   // Reserve reallocates if we don't have enough total room.
 637   // NOTE that, with jemalloc, we know that this won't reallocate in place
 638   // as the size is less than jemallocMinInPlaceExpanadable
 639   {
 640     gen.seed(fillSeed);
 641     unique_ptr<IOBuf> iob(IOBuf::create(2000));
 642     iob->append(100);
 643     fillBuf(iob.get(), gen);
 644     EXPECT_EQ(0, iob->headroom());
 645     EXPECT_EQ(100, iob->length());
 646     const void* p1 = iob->buffer();
 647     iob->reserve(100, 2512);  // allocation sizes are multiples of 256
 648     EXPECT_LE(100, iob->headroom());
 649     if (folly::usingJEMalloc()) {
 650       EXPECT_NE(p1, iob->buffer());
 651     }
 652     gen.seed(fillSeed);
 653     checkBuf(iob.get(), gen);
 654   }
 655
 656   // Test reserve from internal buffer, this used to segfault
 657   {
 658     unique_ptr<IOBuf> iob(IOBuf::create(0));
 659     iob->reserve(0, 2000);
 660     EXPECT_EQ(0, iob->headroom());
 661     EXPECT_LE(2000, iob->tailroom());
 662   }
 663
 664   // Test reserving from a user-allocated buffer.
 665   {
 666     uint8_t* buf = static_cast<uint8_t*>(malloc(100));
 667     auto iob = IOBuf::takeOwnership(buf, 100);
 668     iob->reserve(0, 2000);
 669     EXPECT_EQ(0, iob->headroom());
 670     EXPECT_LE(2000, iob->tailroom());
 671   }
 672
 673   // Test reserving from a user-allocated with a custom free function.
 674   {
 675     uint32_t freeCount{0};
 676     uint8_t* buf = new uint8_t[100];
 677     auto iob = IOBuf::takeOwnership(buf, 100, testFreeFn, &freeCount);
 678     iob->reserve(0, 2000);
 679     EXPECT_EQ(0, iob->headroom());
 680     EXPECT_LE(2000, iob->tailroom());
 681     EXPECT_EQ(1, freeCount);
 682   }
 683 }
 684
 685 TEST(IOBuf, copyBuffer) {
 686   std::string s("hello");
 687   auto buf = IOBuf::copyBuffer(s.data(), s.size(), 1, 2);
 688   EXPECT_EQ(1, buf->headroom());
 689   EXPECT_EQ(s, std::string(reinterpret_cast<const char*>(buf->data()),
 690                            buf->length()));
 691   EXPECT_LE(2, buf->tailroom());
 692
 693   buf = IOBuf::copyBuffer(s, 5, 7);
 694   EXPECT_EQ(5, buf->headroom());
 695   EXPECT_EQ(s, std::string(reinterpret_cast<const char*>(buf->data()),
 696                            buf->length()));
 697   EXPECT_LE(7, buf->tailroom());
 698
 699   std::string empty;
 700   buf = IOBuf::copyBuffer(empty, 3, 6);
 701   EXPECT_EQ(3, buf->headroom());
 702   EXPECT_EQ(0, buf->length());
 703   EXPECT_LE(6, buf->tailroom());
 704
 705   // A stack-allocated version
 706   IOBuf stackBuf(IOBuf::COPY_BUFFER, s, 1, 2);
 707   EXPECT_EQ(1, stackBuf.headroom());
 708   EXPECT_EQ(s, std::string(reinterpret_cast<const char*>(stackBuf.data()),
 709                            stackBuf.length()));
 710   EXPECT_LE(2, stackBuf.tailroom());
 711 }
 712
 713 TEST(IOBuf, maybeCopyBuffer) {
 714   std::string s("this is a test");
 715   auto buf = IOBuf::maybeCopyBuffer(s, 1, 2);
 716   EXPECT_EQ(1, buf->headroom());
 717   EXPECT_EQ(s, std::string(reinterpret_cast<const char*>(buf->data()),
 718                            buf->length()));
 719   EXPECT_LE(2, buf->tailroom());
 720
 721   std::string empty;
 722   buf = IOBuf::maybeCopyBuffer("", 5, 7);
 723   EXPECT_EQ(nullptr, buf.get());
 724
 725   buf = IOBuf::maybeCopyBuffer("");
 726   EXPECT_EQ(nullptr, buf.get());
 727 }
 728
 729 TEST(IOBuf, copyEmptyBuffer) {
 730   auto buf = IOBuf::copyBuffer(nullptr, 0);
 731   EXPECT_EQ(buf->length(), 0);
 732 }
 733
 734 namespace {
 735
 736 int customDeleterCount = 0;
 737 int destructorCount = 0;
 738 struct OwnershipTestClass {
 739   explicit OwnershipTestClass(int v = 0) : val(v) { }
 740   ~OwnershipTestClass() {
 741     ++destructorCount;
 742   }
 743   int val;
 744 };
 745
 746 typedef std::function<void(OwnershipTestClass*)> CustomDeleter;
 747
 748 void customDelete(OwnershipTestClass* p) {
 749   ++customDeleterCount;
 750   delete p;
 751 }
 752
 753 void customDeleteArray(OwnershipTestClass* p) {
 754   ++customDeleterCount;
 755   delete[] p;
 756 }
 757
 758 } // namespace
 759
 760 TEST(IOBuf, takeOwnershipUniquePtr) {
 761   destructorCount = 0;
 762   {
 763     std::unique_ptr<OwnershipTestClass> p(new OwnershipTestClass());
 764   }
 765   EXPECT_EQ(1, destructorCount);
 766
 767   destructorCount = 0;
 768   {
 769     std::unique_ptr<OwnershipTestClass[]> p(new OwnershipTestClass[2]);
 770   }
 771   EXPECT_EQ(2, destructorCount);
 772
 773   destructorCount = 0;
 774   {
 775     std::unique_ptr<OwnershipTestClass> p(new OwnershipTestClass());
 776     std::unique_ptr<IOBuf> buf(IOBuf::takeOwnership(std::move(p)));
 777     EXPECT_EQ(sizeof(OwnershipTestClass), buf->length());
 778     EXPECT_EQ(0, destructorCount);
 779   }
 780   EXPECT_EQ(1, destructorCount);
 781
 782   destructorCount = 0;
 783   {
 784     std::unique_ptr<OwnershipTestClass[]> p(new OwnershipTestClass[2]);
 785     std::unique_ptr<IOBuf> buf(IOBuf::takeOwnership(std::move(p), 2));
 786     EXPECT_EQ(2 * sizeof(OwnershipTestClass), buf->length());
 787     EXPECT_EQ(0, destructorCount);
 788   }
 789   EXPECT_EQ(2, destructorCount);
 790
 791   customDeleterCount = 0;
 792   destructorCount = 0;
 793   {
 794     std::unique_ptr<OwnershipTestClass, CustomDeleter>
 795       p(new OwnershipTestClass(), customDelete);
 796     std::unique_ptr<IOBuf> buf(IOBuf::takeOwnership(std::move(p)));
 797     EXPECT_EQ(sizeof(OwnershipTestClass), buf->length());
 798     EXPECT_EQ(0, destructorCount);
 799   }
 800   EXPECT_EQ(1, destructorCount);
 801   EXPECT_EQ(1, customDeleterCount);
 802
 803   customDeleterCount = 0;
 804   destructorCount = 0;
 805   {
 806     std::unique_ptr<OwnershipTestClass[], CustomDeleter>
 807       p(new OwnershipTestClass[2], CustomDeleter(customDeleteArray));
 808     std::unique_ptr<IOBuf> buf(IOBuf::takeOwnership(std::move(p), 2));
 809     EXPECT_EQ(2 * sizeof(OwnershipTestClass), buf->length());
 810     EXPECT_EQ(0, destructorCount);
 811   }
 812   EXPECT_EQ(2, destructorCount);
 813   EXPECT_EQ(1, customDeleterCount);
 814 }
 815
 816 TEST(IOBuf, Alignment) {
 817   size_t alignment = alignof(std::max_align_t);
 818
 819   std::vector<size_t> sizes {0, 1, 64, 256, 1024, 1 << 10};
 820   for (size_t size : sizes) {
 821     auto buf = IOBuf::create(size);
 822     uintptr_t p = reinterpret_cast<uintptr_t>(buf->data());
 823     EXPECT_EQ(0, p & (alignment - 1)) << "size=" << size;
 824   }
 825 }
 826
 827 TEST(TypedIOBuf, Simple) {
 828   auto buf = IOBuf::create(0);
 829   TypedIOBuf<uint64_t> typed(buf.get());
 830   const uint64_t n = 10000;
 831   typed.reserve(0, n);
 832   EXPECT_LE(n, typed.capacity());
 833   for (uint64_t i = 0; i < n; i++) {
 834     *typed.writableTail() = i;
 835     typed.append(1);
 836   }
 837   EXPECT_EQ(n, typed.length());
 838   for (uint64_t i = 0; i < n; i++) {
 839     EXPECT_EQ(i, typed.data()[i]);
 840   }
 841 }
 842 enum BufType {
 843   CREATE,
 844   TAKE_OWNERSHIP_MALLOC,
 845   TAKE_OWNERSHIP_CUSTOM,
 846   USER_OWNED,
 847 };
 848
 849 // chain element size, number of elements in chain, shared
 850 class MoveToFbStringTest
 851   : public ::testing::TestWithParam<std::tr1::tuple<int, int, bool, BufType>> {
 852  protected:
 853   void SetUp() override {
 854     elementSize_ = std::tr1::get<0>(GetParam());
 855     elementCount_ = std::tr1::get<1>(GetParam());
 856     shared_ = std::tr1::get<2>(GetParam());
 857     type_ = std::tr1::get<3>(GetParam());
 858
 859     buf_ = makeBuf();
 860     for (int i = 0; i < elementCount_ - 1; ++i) {
 861       buf_->prependChain(makeBuf());
 862     }
 863     EXPECT_EQ(elementCount_, buf_->countChainElements());
 864     EXPECT_EQ(elementCount_ * elementSize_, buf_->computeChainDataLength());
 865     if (shared_) {
 866       buf2_ = buf_->clone();
 867       EXPECT_EQ(elementCount_, buf2_->countChainElements());
 868       EXPECT_EQ(elementCount_ * elementSize_, buf2_->computeChainDataLength());
 869     }
 870   }
 871
 872   std::unique_ptr<IOBuf> makeBuf() {
 873     unique_ptr<IOBuf> buf;
 874     switch (type_) {
 875       case CREATE:
 876         buf = IOBuf::create(elementSize_);
 877         buf->append(elementSize_);
 878         break;
 879       case TAKE_OWNERSHIP_MALLOC: {
 880         void* data = malloc(elementSize_);
 881         if (!data) {
 882           throw std::bad_alloc();
 883         }
 884         buf = IOBuf::takeOwnership(data, elementSize_);
 885         break;
 886       }
 887       case TAKE_OWNERSHIP_CUSTOM: {
 888         uint8_t* data = new uint8_t[elementSize_];
 889         buf = IOBuf::takeOwnership(data, elementSize_, testFreeFn);
 890         break;
 891       }
 892       case USER_OWNED: {
 893         unique_ptr<uint8_t[]> data(new uint8_t[elementSize_]);
 894         buf = IOBuf::wrapBuffer(data.get(), elementSize_);
 895         ownedBuffers_.emplace_back(std::move(data));
 896         break;
 897       }
 898       default:
 899         throw std::invalid_argument("unexpected buffer type parameter");
 900     }
 901     memset(buf->writableData(), 'x', elementSize_);
 902     return buf;
 903   }
 904
 905   void check(std::unique_ptr<IOBuf>& buf) {
 906     fbstring str = buf->moveToFbString();
 907     EXPECT_EQ(elementCount_ * elementSize_, str.size());
 908     EXPECT_EQ(elementCount_ * elementSize_, strspn(str.c_str(), "x"));
 909     EXPECT_EQ(0, buf->length());
 910     EXPECT_EQ(1, buf->countChainElements());
 911     EXPECT_EQ(0, buf->computeChainDataLength());
 912     EXPECT_FALSE(buf->isChained());
 913   }
 914
 915   int elementSize_;
 916   int elementCount_;
 917   bool shared_;
 918   BufType type_;
 919   std::unique_ptr<IOBuf> buf_;
 920   std::unique_ptr<IOBuf> buf2_;
 921   std::vector<std::unique_ptr<uint8_t[]>> ownedBuffers_;
 922 };
 923
 924 TEST_P(MoveToFbStringTest, Simple) {
 925   check(buf_);
 926   if (shared_) {
 927     check(buf2_);
 928   }
 929 }
 930
 931 INSTANTIATE_TEST_CASE_P(
 932     MoveToFbString,
 933     MoveToFbStringTest,
 934     ::testing::Combine(
 935         ::testing::Values(0, 1, 24, 256, 1 << 10, 1 << 20),  // element size
 936         ::testing::Values(1, 2, 10),                         // element count
 937         ::testing::Bool(),                                   // shared
 938         ::testing::Values(CREATE, TAKE_OWNERSHIP_MALLOC,
 939                           TAKE_OWNERSHIP_CUSTOM, USER_OWNED)));
 940
 941 TEST(IOBuf, getIov) {
 942   uint32_t fillSeed = 0xdeadbeef;
 943   boost::mt19937 gen(fillSeed);
 944
 945   size_t len = 4096;
 946   size_t count = 32;
 947   auto buf = IOBuf::create(len + 1);
 948   buf->append(rand() % len + 1);
 949   fillBuf(buf.get(), gen);
 950
 951   for (size_t i = 0; i < count - 1; i++) {
 952     auto buf2 = IOBuf::create(len + 1);
 953     buf2->append(rand() % len + 1);
 954     fillBuf(buf2.get(), gen);
 955     buf->prependChain(std::move(buf2));
 956   }
 957   EXPECT_EQ(count, buf->countChainElements());
 958
 959   auto iov = buf->getIov();
 960   EXPECT_EQ(count, iov.size());
 961
 962   IOBuf const* p = buf.get();
 963   for (size_t i = 0; i < count; i++, p = p->next()) {
 964     EXPECT_EQ(p->data(), iov[i].iov_base);
 965     EXPECT_EQ(p->length(), iov[i].iov_len);
 966   }
 967
 968   // an empty buf should be skipped in the iov.
 969   buf->next()->clear();
 970   iov = buf->getIov();
 971   EXPECT_EQ(count - 1, iov.size());
 972   EXPECT_EQ(buf->next()->next()->data(), iov[1].iov_base);
 973
 974   // same for the first one being empty
 975   buf->clear();
 976   iov = buf->getIov();
 977   EXPECT_EQ(count - 2, iov.size());
 978   EXPECT_EQ(buf->next()->next()->data(), iov[0].iov_base);
 979
 980   // and the last one
 981   buf->prev()->clear();
 982   iov = buf->getIov();
 983   EXPECT_EQ(count - 3, iov.size());
 984
 985   // test appending to an existing iovec array
 986   iov.clear();
 987   const char localBuf[] = "hello";
 988   iov.push_back({(void*)localBuf, sizeof(localBuf)});
 989   iov.push_back({(void*)localBuf, sizeof(localBuf)});
 990   buf->appendToIov(&iov);
 991   EXPECT_EQ(count - 1, iov.size());
 992   EXPECT_EQ(localBuf, iov[0].iov_base);
 993   EXPECT_EQ(localBuf, iov[1].iov_base);
 994   // The first two IOBufs were cleared, so the next iov entry
 995   // should be the third IOBuf in the chain.
 996   EXPECT_EQ(buf->next()->next()->data(), iov[2].iov_base);
 997 }
 998
 999 TEST(IOBuf, move) {
1000   // Default allocate an IOBuf on the stack
1001   IOBuf outerBuf;
1002   char data[] = "foobar";
1003   uint32_t length = sizeof(data);
1004   uint32_t actualCapacity{0};
1005   const void* ptr{nullptr};
1006
1007   {
1008     // Create a small IOBuf on the stack.
1009     // Note that IOBufs created on the stack always use an external buffer.
1010     IOBuf b1(IOBuf::CREATE, 10);
1011     actualCapacity = b1.capacity();
1012     EXPECT_GE(actualCapacity, 10);
1013     EXPECT_EQ(0, b1.length());
1014     EXPECT_FALSE(b1.isShared());
1015     ptr = b1.data();
1016     ASSERT_TRUE(ptr != nullptr);
1017     memcpy(b1.writableTail(), data, length);
1018     b1.append(length);
1019     EXPECT_EQ(length, b1.length());
1020
1021     // Use the move constructor
1022     IOBuf b2(std::move(b1));
1023     EXPECT_EQ(ptr, b2.data());
1024     EXPECT_EQ(length, b2.length());
1025     EXPECT_EQ(actualCapacity, b2.capacity());
1026     EXPECT_FALSE(b2.isShared());
1027
1028     // Use the move assignment operator
1029     outerBuf = std::move(b2);
1030     // Close scope, destroying b1 and b2
1031     // (which are both be invalid now anyway after moving out of them)
1032   }
1033
1034   EXPECT_EQ(ptr, outerBuf.data());
1035   EXPECT_EQ(length, outerBuf.length());
1036   EXPECT_EQ(actualCapacity, outerBuf.capacity());
1037   EXPECT_FALSE(outerBuf.isShared());
1038 }
1039
1040 namespace {
1041 std::unique_ptr<IOBuf> fromStr(StringPiece sp) {
1042   return IOBuf::copyBuffer(ByteRange(sp));
1043 }
1044 } // namespace
1045
1046 TEST(IOBuf, HashAndEqual) {
1047   folly::IOBufEqual eq;
1048   folly::IOBufHash hash;
1049
1050   EXPECT_TRUE(eq(nullptr, nullptr));
1051   EXPECT_EQ(0, hash(nullptr));
1052
1053   auto empty = IOBuf::create(0);
1054
1055   EXPECT_TRUE(eq(*empty, *empty));
1056   EXPECT_TRUE(eq(empty, empty));
1057
1058   EXPECT_FALSE(eq(nullptr, empty));
1059   EXPECT_FALSE(eq(empty, nullptr));
1060
1061   EXPECT_EQ(hash(*empty), hash(empty));
1062   EXPECT_NE(0, hash(empty));
1063
1064   auto a = fromStr("hello");
1065
1066   EXPECT_TRUE(eq(*a, *a));
1067   EXPECT_TRUE(eq(a, a));
1068
1069   EXPECT_FALSE(eq(nullptr, a));
1070   EXPECT_FALSE(eq(a, nullptr));
1071
1072   EXPECT_EQ(hash(*a), hash(a));
1073   EXPECT_NE(0, hash(a));
1074
1075   auto b = fromStr("hello");
1076
1077   EXPECT_TRUE(eq(*a, *b));
1078   EXPECT_TRUE(eq(a, b));
1079
1080   EXPECT_EQ(hash(a), hash(b));
1081
1082   auto c = fromStr("hellow");
1083
1084   EXPECT_FALSE(eq(a, c));
1085   EXPECT_NE(hash(a), hash(c));
1086
1087   auto d = fromStr("world");
1088
1089   EXPECT_FALSE(eq(a, d));
1090   EXPECT_NE(hash(a), hash(d));
1091
1092   auto e = fromStr("helloworld");
1093   auto f = fromStr("hello");
1094   f->prependChain(fromStr("wo"));
1095   f->prependChain(fromStr("rld"));
1096
1097   EXPECT_TRUE(eq(e, f));
1098   EXPECT_EQ(hash(e), hash(f));
1099 }
1100
1101 // reserveSlow() had a bug when reallocating the buffer in place. It would
1102 // preserve old headroom if it's not too much (heuristically) but wouldn't
1103 // adjust the requested amount of memory to account for that; the end result
1104 // would be that reserve() would return with less tailroom than requested.
1105 TEST(IOBuf, ReserveWithHeadroom) {
1106   // This is assuming jemalloc, where we know that 4096 and 8192 bytes are
1107   // valid (and consecutive) allocation sizes. We're hoping that our
1108   // 4096-byte buffer can be expanded in place to 8192 (in practice, this
1109   // usually happens).
1110   const char data[] = "Lorem ipsum dolor sit amet, consectetur adipiscing elit";
1111   constexpr size_t reservedSize = 24;  // sizeof(SharedInfo)
1112   // chosen carefully so that the buffer is exactly 4096 bytes
1113   IOBuf buf(IOBuf::CREATE, 4096 - reservedSize);
1114   buf.advance(10);
1115   memcpy(buf.writableData(), data, sizeof(data));
1116   buf.append(sizeof(data));
1117   EXPECT_EQ(sizeof(data), buf.length());
1118
1119   // Grow the buffer (hopefully in place); this would incorrectly reserve
1120   // the 10 bytes of headroom, giving us 10 bytes less than requested.
1121   size_t tailroom = 8192 - reservedSize - sizeof(data);
1122   buf.reserve(0, tailroom);
1123   EXPECT_LE(tailroom, buf.tailroom());
1124   EXPECT_EQ(sizeof(data), buf.length());
1125   EXPECT_EQ(0, memcmp(data, buf.data(), sizeof(data)));
1126 }
1127
1128 TEST(IOBuf, CopyConstructorAndAssignmentOperator) {
1129   auto buf = IOBuf::create(4096);
1130   append(buf, "hello world");
1131   auto buf2 = IOBuf::create(4096);
1132   append(buf2, " goodbye");
1133   buf->prependChain(std::move(buf2));
1134   EXPECT_FALSE(buf->isShared());
1135
1136   {
1137     auto copy = *buf;
1138     EXPECT_TRUE(buf->isShared());
1139     EXPECT_TRUE(copy.isShared());
1140     EXPECT_EQ((void*)buf->data(), (void*)copy.data());
1141     EXPECT_NE(buf->next(), copy.next());  // actually different buffers
1142
1143     auto copy2 = *buf;
1144     copy2.coalesce();
1145     EXPECT_TRUE(buf->isShared());
1146     EXPECT_TRUE(copy.isShared());
1147     EXPECT_FALSE(copy2.isShared());
1148
1149     auto p = reinterpret_cast<const char*>(copy2.data());
1150     EXPECT_EQ("hello world goodbye", std::string(p, copy2.length()));
1151   }
1152
1153   EXPECT_FALSE(buf->isShared());
1154
1155   {
1156     folly::IOBuf newBuf(folly::IOBuf::CREATE, 4096);
1157     EXPECT_FALSE(newBuf.isShared());
1158
1159     auto newBufCopy = newBuf;
1160     EXPECT_TRUE(newBuf.isShared());
1161     EXPECT_TRUE(newBufCopy.isShared());
1162
1163     newBufCopy = *buf;
1164     EXPECT_TRUE(buf->isShared());
1165     EXPECT_FALSE(newBuf.isShared());
1166     EXPECT_TRUE(newBufCopy.isShared());
1167   }
1168
1169   EXPECT_FALSE(buf->isShared());
1170 }
1171
1172 TEST(IOBuf, CloneAsValue) {
1173   auto buf = IOBuf::create(4096);
1174   append(buf, "hello world");
1175   {
1176     auto buf2 = IOBuf::create(4096);
1177     append(buf2, " goodbye");
1178     buf->prependChain(std::move(buf2));
1179     EXPECT_FALSE(buf->isShared());
1180   }
1181
1182   {
1183     auto copy = buf->cloneOneAsValue();
1184     EXPECT_TRUE(buf->isShared());
1185     EXPECT_TRUE(copy.isShared());
1186     EXPECT_EQ((void*)buf->data(), (void*)copy.data());
1187     EXPECT_TRUE(buf->isChained());
1188     EXPECT_FALSE(copy.isChained());
1189
1190     auto copy2 = buf->cloneAsValue();
1191     EXPECT_TRUE(buf->isShared());
1192     EXPECT_TRUE(copy.isShared());
1193     EXPECT_TRUE(copy2.isShared());
1194     EXPECT_TRUE(buf->isChained());
1195     EXPECT_TRUE(copy2.isChained());
1196
1197     copy.unshareOne();
1198     EXPECT_TRUE(buf->isShared());
1199     EXPECT_FALSE(copy.isShared());
1200     EXPECT_NE((void*)buf->data(), (void*)copy.data());
1201     EXPECT_TRUE(copy2.isShared());
1202
1203     auto p = reinterpret_cast<const char*>(copy.data());
1204     EXPECT_EQ("hello world", std::string(p, copy.length()));
1205
1206     copy2.coalesce();
1207     EXPECT_FALSE(buf->isShared());
1208     EXPECT_FALSE(copy.isShared());
1209     EXPECT_FALSE(copy2.isShared());
1210     EXPECT_FALSE(copy2.isChained());
1211
1212     auto p2 = reinterpret_cast<const char*>(copy2.data());
1213     EXPECT_EQ("hello world goodbye", std::string(p2, copy2.length()));
1214   }
1215
1216   EXPECT_FALSE(buf->isShared());
1217 }
1218
1219 namespace {
1220 // Use with string literals only
1221 std::unique_ptr<IOBuf> wrap(const char* str) {
1222   return IOBuf::wrapBuffer(str, strlen(str));
1223 }
1224
1225 std::unique_ptr<IOBuf> copy(const char* str) {
1226   // At least 1KiB of tailroom, to ensure an external buffer
1227   return IOBuf::copyBuffer(str, strlen(str), 0, 1024);
1228 }
1229
1230 std::string toString(const folly::IOBuf& buf) {
1231   std::string result;
1232   result.reserve(buf.computeChainDataLength());
1233   for (auto& b : buf) {
1234     result.append(reinterpret_cast<const char*>(b.data()), b.size());
1235   }
1236   return result;
1237 }
1238
1239 char* writableStr(folly::IOBuf& buf) {
1240   return reinterpret_cast<char*>(buf.writableData());
1241 }
1242
1243 } // namespace
1244
1245 TEST(IOBuf, ExternallyShared) {
1246   struct Item {
1247     Item(const char* src, size_t len) : size(len) {
1248       CHECK_LE(len, sizeof(buffer));
1249       memcpy(buffer, src, len);
1250     }
1251     uint32_t refcount{0};
1252     uint8_t size;
1253     char buffer[256];
1254   };
1255
1256   auto hello = "hello";
1257   struct Item it(hello, strlen(hello));
1258
1259   {
1260     auto freeFn = [](void* /* unused */, void* userData) {
1261       auto it = static_cast<struct Item*>(userData);
1262       it->refcount--;
1263     };
1264     it.refcount++;
1265     auto buf1 = IOBuf::takeOwnership(it.buffer, it.size, freeFn, &it);
1266     EXPECT_TRUE(buf1->isManagedOne());
1267     EXPECT_FALSE(buf1->isSharedOne());
1268
1269     buf1->markExternallyShared();
1270     EXPECT_TRUE(buf1->isSharedOne());
1271
1272     {
1273       auto buf2 = buf1->clone();
1274       EXPECT_TRUE(buf2->isManagedOne());
1275       EXPECT_TRUE(buf2->isSharedOne());
1276       EXPECT_EQ(buf1->data(), buf2->data());
1277       EXPECT_EQ(it.refcount, 1);
1278     }
1279     EXPECT_EQ(it.refcount, 1);
1280   }
1281   EXPECT_EQ(it.refcount, 0);
1282 }
1283
1284 TEST(IOBuf, Managed) {
1285   auto hello = "hello";
1286   auto buf1UP = wrap(hello);
1287   auto buf1 = buf1UP.get();
1288   EXPECT_FALSE(buf1->isManagedOne());
1289   auto buf2UP = copy("world");
1290   auto buf2 = buf2UP.get();
1291   EXPECT_TRUE(buf2->isManagedOne());
1292   auto buf3UP = wrap(hello);
1293   auto buf3 = buf3UP.get();
1294   auto buf4UP = buf2->clone();
1295   auto buf4 = buf4UP.get();
1296
1297   // buf1 and buf3 share the same memory (but are unmanaged)
1298   EXPECT_FALSE(buf1->isManagedOne());
1299   EXPECT_FALSE(buf3->isManagedOne());
1300   EXPECT_TRUE(buf1->isSharedOne());
1301   EXPECT_TRUE(buf3->isSharedOne());
1302   EXPECT_EQ(buf1->data(), buf3->data());
1303
1304   // buf2 and buf4 share the same memory (but are managed)
1305   EXPECT_TRUE(buf2->isManagedOne());
1306   EXPECT_TRUE(buf4->isManagedOne());
1307   EXPECT_TRUE(buf2->isSharedOne());
1308   EXPECT_TRUE(buf4->isSharedOne());
1309   EXPECT_EQ(buf2->data(), buf4->data());
1310
1311   buf1->prependChain(std::move(buf2UP));
1312   buf1->prependChain(std::move(buf3UP));
1313   buf1->prependChain(std::move(buf4UP));
1314
1315   EXPECT_EQ("helloworldhelloworld", toString(*buf1));
1316   EXPECT_FALSE(buf1->isManaged());
1317
1318   buf1->makeManaged();
1319   EXPECT_TRUE(buf1->isManaged());
1320
1321   // buf1 and buf3 are now unshared (because they were unmanaged)
1322   EXPECT_TRUE(buf1->isManagedOne());
1323   EXPECT_TRUE(buf3->isManagedOne());
1324   EXPECT_FALSE(buf1->isSharedOne());
1325   EXPECT_FALSE(buf3->isSharedOne());
1326   EXPECT_NE(buf1->data(), buf3->data());
1327
1328   // buf2 and buf4 are still shared
1329   EXPECT_TRUE(buf2->isManagedOne());
1330   EXPECT_TRUE(buf4->isManagedOne());
1331   EXPECT_TRUE(buf2->isSharedOne());
1332   EXPECT_TRUE(buf4->isSharedOne());
1333   EXPECT_EQ(buf2->data(), buf4->data());
1334
1335   // And verify that the truth is what we expect: modify a byte in buf1 and
1336   // buf2, see that the change from buf1 is *not* reflected in buf3, but the
1337   // change from buf2 is reflected in buf4.
1338   writableStr(*buf1)[0] = 'j';
1339   writableStr(*buf2)[0] = 'x';
1340   EXPECT_EQ("jelloxorldhelloxorld", toString(*buf1));
1341 }
1342
1343 TEST(IOBuf, CoalesceEmptyBuffers) {
1344   auto b1 = IOBuf::takeOwnership(nullptr, 0);
1345   auto b2 = fromStr("hello");
1346   auto b3 = IOBuf::takeOwnership(nullptr, 0);
1347
1348   b2->appendChain(std::move(b3));
1349   b1->appendChain(std::move(b2));
1350
1351   auto br = b1->coalesce();
1352
1353   EXPECT_TRUE(ByteRange(StringPiece("hello")) == br);
1354 }
1355
1356 TEST(IOBuf, CloneCoalescedChain) {
1357   auto b = IOBuf::createChain(1000, 100);
1358   b->advance(10);
1359   const uint32_t fillSeed = 0x12345678;
1360   boost::mt19937 gen(fillSeed);
1361   {
1362     auto c = b.get();
1363     uint64_t length = c->tailroom();
1364     do {
1365       length = std::min(length, c->tailroom());
1366       c->append(length--);
1367       fillBuf(c, gen);
1368       c = c->next();
1369     } while (c != b.get());
1370   }
1371   auto c = b->cloneCoalescedAsValue();
1372   EXPECT_FALSE(c.isChained()); // Not chained
1373   EXPECT_FALSE(c.isSharedOne()); // Not shared
1374   EXPECT_EQ(b->headroom(), c.headroom()); // Preserves headroom
1375   EXPECT_LE(b->prev()->tailroom(), c.tailroom()); // Preserves minimum tailroom
1376   EXPECT_EQ(b->computeChainDataLength(), c.length()); // Same length
1377   gen.seed(fillSeed);
1378   checkBuf(&c, gen); // Same contents
1379 }
1380
1381 TEST(IOBuf, CloneCoalescedSingle) {
1382   auto b = IOBuf::create(1000);
1383   b->advance(10);
1384   b->append(900);
1385   const uint32_t fillSeed = 0x12345678;
1386   boost::mt19937 gen(fillSeed);
1387   fillBuf(b.get(), gen);
1388
1389   auto c = b->cloneCoalesced();
1390   EXPECT_FALSE(c->isChained()); // Not chained
1391   EXPECT_TRUE(c->isSharedOne()); // Shared
1392   EXPECT_EQ(b->buffer(), c->buffer());
1393   EXPECT_EQ(b->capacity(), c->capacity());
1394   EXPECT_EQ(b->data(), c->data());
1395   EXPECT_EQ(b->length(), c->length());
1396 }