folly/io/test/IOBufTest.cpp

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