folly/io/test/IOBufTest.cpp

   1 /*
   2  * Copyright 2013 Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #include "folly/io/IOBuf.h"
  18 #include "folly/io/TypedIOBuf.h"
  19
  20 // googletest requires std::tr1::tuple, not std::tuple
  21 #include <tr1/tuple>
  22
  23 #include <gflags/gflags.h>
  24 #include <boost/random.hpp>
  25 #include <gtest/gtest.h>
  26
  27 #include "folly/Malloc.h"
  28 #include "folly/Range.h"
  29
  30 using folly::fbstring;
  31 using folly::fbvector;
  32 using folly::IOBuf;
  33 using folly::TypedIOBuf;
  34 using folly::StringPiece;
  35 using folly::ByteRange;
  36 using std::unique_ptr;
  37
  38 void append(std::unique_ptr<IOBuf>& buf, StringPiece str) {
  39   EXPECT_LE(str.size(), buf->tailroom());
  40   memcpy(buf->writableData(), str.data(), str.size());
  41   buf->append(str.size());
  42 }
  43
  44 void prepend(std::unique_ptr<IOBuf>& buf, StringPiece str) {
  45   EXPECT_LE(str.size(), buf->headroom());
  46   memcpy(buf->writableData() - str.size(), str.data(), str.size());
  47   buf->prepend(str.size());
  48 }
  49
  50 TEST(IOBuf, Simple) {
  51   unique_ptr<IOBuf> buf(IOBuf::create(100));
  52   uint32_t cap = buf->capacity();
  53   EXPECT_LE(100, cap);
  54   EXPECT_EQ(0, buf->headroom());
  55   EXPECT_EQ(0, buf->length());
  56   EXPECT_EQ(cap, buf->tailroom());
  57
  58   append(buf, "world");
  59   buf->advance(10);
  60   EXPECT_EQ(10, buf->headroom());
  61   EXPECT_EQ(5, buf->length());
  62   EXPECT_EQ(cap - 15, buf->tailroom());
  63
  64   prepend(buf, "hello ");
  65   EXPECT_EQ(4, buf->headroom());
  66   EXPECT_EQ(11, buf->length());
  67   EXPECT_EQ(cap - 15, buf->tailroom());
  68
  69   const char* p = reinterpret_cast<const char*>(buf->data());
  70   EXPECT_EQ("hello world", std::string(p, buf->length()));
  71
  72   buf->clear();
  73   EXPECT_EQ(0, buf->headroom());
  74   EXPECT_EQ(0, buf->length());
  75   EXPECT_EQ(cap, buf->tailroom());
  76 }
  77
  78
  79 void testAllocSize(uint32_t requestedCapacity) {
  80   unique_ptr<IOBuf> iobuf(IOBuf::create(requestedCapacity));
  81   EXPECT_GE(iobuf->capacity(), requestedCapacity);
  82 }
  83
  84 TEST(IOBuf, AllocSizes) {
  85   // Try with a small allocation size that should fit in the internal buffer
  86   testAllocSize(28);
  87
  88   // Try with a large allocation size that will require an external buffer.
  89   testAllocSize(9000);
  90
  91   // 220 bytes is currently the cutoff
  92   // (It would be nice to use the IOBuf::kMaxInternalDataSize constant,
  93   // but it's private and it doesn't seem worth making it public just for this
  94   // test code.)
  95   testAllocSize(220);
  96   testAllocSize(219);
  97   testAllocSize(221);
  98 }
  99
 100 void deleteArrayBuffer(void *buf, void* arg) {
 101   uint32_t* deleteCount = static_cast<uint32_t*>(arg);
 102   ++(*deleteCount);
 103   uint8_t* bufPtr = static_cast<uint8_t*>(buf);
 104   delete[] bufPtr;
 105 }
 106
 107 TEST(IOBuf, TakeOwnership) {
 108   uint32_t size1 = 99;
 109   uint8_t *buf1 = static_cast<uint8_t*>(malloc(size1));
 110   unique_ptr<IOBuf> iobuf1(IOBuf::takeOwnership(buf1, size1));
 111   EXPECT_EQ(buf1, iobuf1->data());
 112   EXPECT_EQ(size1, iobuf1->length());
 113   EXPECT_EQ(buf1, iobuf1->buffer());
 114   EXPECT_EQ(size1, iobuf1->capacity());
 115
 116   uint32_t deleteCount = 0;
 117   uint32_t size2 = 4321;
 118   uint8_t *buf2 = new uint8_t[size2];
 119   unique_ptr<IOBuf> iobuf2(IOBuf::takeOwnership(buf2, size2,
 120                                                 deleteArrayBuffer,
 121                                                 &deleteCount));
 122   EXPECT_EQ(buf2, iobuf2->data());
 123   EXPECT_EQ(size2, iobuf2->length());
 124   EXPECT_EQ(buf2, iobuf2->buffer());
 125   EXPECT_EQ(size2, iobuf2->capacity());
 126   EXPECT_EQ(0, deleteCount);
 127   iobuf2.reset();
 128   EXPECT_EQ(1, deleteCount);
 129
 130   deleteCount = 0;
 131   uint32_t size3 = 3456;
 132   uint8_t *buf3 = new uint8_t[size3];
 133   uint32_t length3 = 48;
 134   unique_ptr<IOBuf> iobuf3(IOBuf::takeOwnership(buf3, size3, length3,
 135                                                 deleteArrayBuffer,
 136                                                 &deleteCount));
 137   EXPECT_EQ(buf3, iobuf3->data());
 138   EXPECT_EQ(length3, iobuf3->length());
 139   EXPECT_EQ(buf3, iobuf3->buffer());
 140   EXPECT_EQ(size3, iobuf3->capacity());
 141   EXPECT_EQ(0, deleteCount);
 142   iobuf3.reset();
 143   EXPECT_EQ(1, deleteCount);
 144
 145   deleteCount = 0;
 146   {
 147     uint32_t size4 = 1234;
 148     uint8_t *buf4 = new uint8_t[size4];
 149     uint32_t length4 = 48;
 150     IOBuf iobuf4(IOBuf::TAKE_OWNERSHIP, buf4, size4, length4,
 151                  deleteArrayBuffer, &deleteCount);
 152     EXPECT_EQ(buf4, iobuf4.data());
 153     EXPECT_EQ(length4, iobuf4.length());
 154     EXPECT_EQ(buf4, iobuf4.buffer());
 155     EXPECT_EQ(size4, iobuf4.capacity());
 156
 157     IOBuf iobuf5 = std::move(iobuf4);
 158     EXPECT_EQ(buf4, iobuf5.data());
 159     EXPECT_EQ(length4, iobuf5.length());
 160     EXPECT_EQ(buf4, iobuf5.buffer());
 161     EXPECT_EQ(size4, iobuf5.capacity());
 162     EXPECT_EQ(0, deleteCount);
 163   }
 164   EXPECT_EQ(1, deleteCount);
 165 }
 166
 167 TEST(IOBuf, WrapBuffer) {
 168   const uint32_t size1 = 1234;
 169   uint8_t buf1[size1];
 170   unique_ptr<IOBuf> iobuf1(IOBuf::wrapBuffer(buf1, size1));
 171   EXPECT_EQ(buf1, iobuf1->data());
 172   EXPECT_EQ(size1, iobuf1->length());
 173   EXPECT_EQ(buf1, iobuf1->buffer());
 174   EXPECT_EQ(size1, iobuf1->capacity());
 175
 176   uint32_t size2 = 0x1234;
 177   unique_ptr<uint8_t[]> buf2(new uint8_t[size2]);
 178   unique_ptr<IOBuf> iobuf2(IOBuf::wrapBuffer(buf2.get(), size2));
 179   EXPECT_EQ(buf2.get(), iobuf2->data());
 180   EXPECT_EQ(size2, iobuf2->length());
 181   EXPECT_EQ(buf2.get(), iobuf2->buffer());
 182   EXPECT_EQ(size2, iobuf2->capacity());
 183
 184   uint32_t size3 = 4321;
 185   unique_ptr<uint8_t[]> buf3(new uint8_t[size3]);
 186   IOBuf iobuf3(IOBuf::WRAP_BUFFER, buf3.get(), size3);
 187   EXPECT_EQ(buf3.get(), iobuf3.data());
 188   EXPECT_EQ(size3, iobuf3.length());
 189   EXPECT_EQ(buf3.get(), iobuf3.buffer());
 190   EXPECT_EQ(size3, iobuf3.capacity());
 191 }
 192
 193 TEST(IOBuf, CreateCombined) {
 194   // Create a combined IOBuf, then destroy it.
 195   // The data buffer and IOBuf both become unused as part of the destruction
 196   {
 197     auto buf = IOBuf::createCombined(256);
 198     EXPECT_FALSE(buf->isShared());
 199   }
 200
 201   // Create a combined IOBuf, clone from it, and then destroy the original
 202   // IOBuf.  The data buffer cannot be deleted until the clone is also
 203   // destroyed.
 204   {
 205     auto bufA = IOBuf::createCombined(256);
 206     EXPECT_FALSE(bufA->isShared());
 207     auto bufB = bufA->clone();
 208     EXPECT_TRUE(bufA->isShared());
 209     EXPECT_TRUE(bufB->isShared());
 210     bufA.reset();
 211     EXPECT_FALSE(bufB->isShared());
 212   }
 213
 214   // Create a combined IOBuf, then call reserve() to get a larger buffer.
 215   // The IOBuf no longer points to the combined data buffer, but the
 216   // overall memory segment cannot be deleted until the IOBuf is also
 217   // destroyed.
 218   {
 219     auto buf = IOBuf::createCombined(256);
 220     buf->reserve(0, buf->capacity() + 100);
 221   }
 222
 223   // Create a combined IOBuf, clone from it, then call unshare() on the original
 224   // buffer.  This creates a situation where bufB is pointing at the combined
 225   // buffer associated with bufA, but bufA is now using a different buffer.
 226   auto testSwap = [](bool resetAFirst) {
 227     auto bufA = IOBuf::createCombined(256);
 228     EXPECT_FALSE(bufA->isShared());
 229     auto bufB = bufA->clone();
 230     EXPECT_TRUE(bufA->isShared());
 231     EXPECT_TRUE(bufB->isShared());
 232     bufA->unshare();
 233     EXPECT_FALSE(bufA->isShared());
 234     EXPECT_FALSE(bufB->isShared());
 235
 236     if (resetAFirst) {
 237       bufA.reset();
 238       bufB.reset();
 239     } else {
 240       bufB.reset();
 241       bufA.reset();
 242     }
 243   };
 244   testSwap(true);
 245   testSwap(false);
 246 }
 247
 248 void fillBuf(uint8_t* buf, uint32_t length, boost::mt19937& gen) {
 249   for (uint32_t n = 0; n < length; ++n) {
 250     buf[n] = static_cast<uint8_t>(gen() & 0xff);
 251   }
 252 }
 253
 254 void fillBuf(IOBuf* buf, boost::mt19937& gen) {
 255   buf->unshare();
 256   fillBuf(buf->writableData(), buf->length(), gen);
 257 }
 258
 259 void checkBuf(const uint8_t* buf, uint32_t length, boost::mt19937& gen) {
 260   // Rather than using EXPECT_EQ() to check each character,
 261   // count the number of differences and the first character that differs.
 262   // This way on error we'll report just that information, rather than tons of
 263   // failed checks for each byte in the buffer.
 264   uint32_t numDifferences = 0;
 265   uint32_t firstDiffIndex = 0;
 266   uint8_t firstDiffExpected = 0;
 267   for (uint32_t n = 0; n < length; ++n) {
 268     uint8_t expected = static_cast<uint8_t>(gen() & 0xff);
 269     if (buf[n] == expected) {
 270       continue;
 271     }
 272
 273     if (numDifferences == 0) {
 274       firstDiffIndex = n;
 275       firstDiffExpected = expected;
 276     }
 277     ++numDifferences;
 278   }
 279
 280   EXPECT_EQ(0, numDifferences);
 281   if (numDifferences > 0) {
 282     // Cast to int so it will be printed numerically
 283     // rather than as a char if the check fails
 284     EXPECT_EQ(static_cast<int>(buf[firstDiffIndex]),
 285               static_cast<int>(firstDiffExpected));
 286   }
 287 }
 288
 289 void checkBuf(IOBuf* buf, boost::mt19937& gen) {
 290   checkBuf(buf->data(), buf->length(), gen);
 291 }
 292
 293 void checkBuf(ByteRange buf, boost::mt19937& gen) {
 294   checkBuf(buf.data(), buf.size(), gen);
 295 }
 296
 297 void checkChain(IOBuf* buf, boost::mt19937& gen) {
 298   IOBuf *current = buf;
 299   do {
 300     checkBuf(current->data(), current->length(), gen);
 301     current = current->next();
 302   } while (current != buf);
 303 }
 304
 305 TEST(IOBuf, Chaining) {
 306   uint32_t fillSeed = 0x12345678;
 307   boost::mt19937 gen(fillSeed);
 308
 309   // An IOBuf with external storage
 310   uint32_t headroom = 123;
 311   unique_ptr<IOBuf> iob1(IOBuf::create(2048));
 312   iob1->advance(headroom);
 313   iob1->append(1500);
 314   fillBuf(iob1.get(), gen);
 315
 316   // An IOBuf with internal storage
 317   unique_ptr<IOBuf> iob2(IOBuf::create(20));
 318   iob2->append(20);
 319   fillBuf(iob2.get(), gen);
 320
 321   // An IOBuf around a buffer it doesn't own
 322   uint8_t localbuf[1234];
 323   fillBuf(localbuf, 1234, gen);
 324   unique_ptr<IOBuf> iob3(IOBuf::wrapBuffer(localbuf, sizeof(localbuf)));
 325
 326   // An IOBuf taking ownership of a user-supplied buffer
 327   uint32_t heapBufSize = 900;
 328   uint8_t* heapBuf = static_cast<uint8_t*>(malloc(heapBufSize));
 329   fillBuf(heapBuf, heapBufSize, gen);
 330   unique_ptr<IOBuf> iob4(IOBuf::takeOwnership(heapBuf, heapBufSize));
 331
 332   // An IOBuf taking ownership of a user-supplied buffer with
 333   // a custom free function
 334   uint32_t arrayBufSize = 321;
 335   uint8_t* arrayBuf = new uint8_t[arrayBufSize];
 336   fillBuf(arrayBuf, arrayBufSize, gen);
 337   uint32_t arrayBufFreeCount = 0;
 338   unique_ptr<IOBuf> iob5(IOBuf::takeOwnership(arrayBuf, arrayBufSize,
 339                                               deleteArrayBuffer,
 340                                               &arrayBufFreeCount));
 341
 342   EXPECT_FALSE(iob1->isChained());
 343   EXPECT_FALSE(iob2->isChained());
 344   EXPECT_FALSE(iob3->isChained());
 345   EXPECT_FALSE(iob4->isChained());
 346   EXPECT_FALSE(iob5->isChained());
 347
 348   EXPECT_FALSE(iob1->isSharedOne());
 349   EXPECT_FALSE(iob2->isSharedOne());
 350   EXPECT_TRUE(iob3->isSharedOne()); // since we own the buffer
 351   EXPECT_FALSE(iob4->isSharedOne());
 352   EXPECT_FALSE(iob5->isSharedOne());
 353
 354   // Chain the buffers all together
 355   // Since we are going to relinquish ownership of iob2-5 to the chain,
 356   // store raw pointers to them so we can reference them later.
 357   IOBuf* iob2ptr = iob2.get();
 358   IOBuf* iob3ptr = iob3.get();
 359   IOBuf* iob4ptr = iob4.get();
 360   IOBuf* iob5ptr = iob5.get();
 361
 362   iob1->prependChain(std::move(iob2));
 363   iob1->prependChain(std::move(iob4));
 364   iob2ptr->appendChain(std::move(iob3));
 365   iob1->prependChain(std::move(iob5));
 366
 367   EXPECT_EQ(iob2ptr, iob1->next());
 368   EXPECT_EQ(iob3ptr, iob2ptr->next());
 369   EXPECT_EQ(iob4ptr, iob3ptr->next());
 370   EXPECT_EQ(iob5ptr, iob4ptr->next());
 371   EXPECT_EQ(iob1.get(), iob5ptr->next());
 372
 373   EXPECT_EQ(iob5ptr, iob1->prev());
 374   EXPECT_EQ(iob1.get(), iob2ptr->prev());
 375   EXPECT_EQ(iob2ptr, iob3ptr->prev());
 376   EXPECT_EQ(iob3ptr, iob4ptr->prev());
 377   EXPECT_EQ(iob4ptr, iob5ptr->prev());
 378
 379   EXPECT_TRUE(iob1->isChained());
 380   EXPECT_TRUE(iob2ptr->isChained());
 381   EXPECT_TRUE(iob3ptr->isChained());
 382   EXPECT_TRUE(iob4ptr->isChained());
 383   EXPECT_TRUE(iob5ptr->isChained());
 384
 385   uint64_t fullLength = (iob1->length() + iob2ptr->length() +
 386                          iob3ptr->length() + iob4ptr->length() +
 387                         iob5ptr->length());
 388   EXPECT_EQ(5, iob1->countChainElements());
 389   EXPECT_EQ(fullLength, iob1->computeChainDataLength());
 390
 391   // Since iob3 is shared, the entire buffer should report itself as shared
 392   EXPECT_TRUE(iob1->isShared());
 393   // Unshare just iob3
 394   iob3ptr->unshareOne();
 395   EXPECT_FALSE(iob3ptr->isSharedOne());
 396   // Now everything in the chain should be unshared.
 397   // Check on all members of the chain just for good measure
 398   EXPECT_FALSE(iob1->isShared());
 399   EXPECT_FALSE(iob2ptr->isShared());
 400   EXPECT_FALSE(iob3ptr->isShared());
 401   EXPECT_FALSE(iob4ptr->isShared());
 402   EXPECT_FALSE(iob5ptr->isShared());
 403
 404   // Check iteration
 405   gen.seed(fillSeed);
 406   size_t count = 0;
 407   for (auto buf : *iob1) {
 408     checkBuf(buf, gen);
 409     ++count;
 410   }
 411   EXPECT_EQ(5, count);
 412
 413   // Clone one of the IOBufs in the chain
 414   unique_ptr<IOBuf> iob4clone = iob4ptr->cloneOne();
 415   gen.seed(fillSeed);
 416   checkBuf(iob1.get(), gen);
 417   checkBuf(iob2ptr, gen);
 418   checkBuf(iob3ptr, gen);
 419   checkBuf(iob4clone.get(), gen);
 420   checkBuf(iob5ptr, gen);
 421
 422   EXPECT_TRUE(iob1->isShared());
 423   EXPECT_TRUE(iob2ptr->isShared());
 424   EXPECT_TRUE(iob3ptr->isShared());
 425   EXPECT_TRUE(iob4ptr->isShared());
 426   EXPECT_TRUE(iob5ptr->isShared());
 427
 428   EXPECT_FALSE(iob1->isSharedOne());
 429   EXPECT_FALSE(iob2ptr->isSharedOne());
 430   EXPECT_FALSE(iob3ptr->isSharedOne());
 431   EXPECT_TRUE(iob4ptr->isSharedOne());
 432   EXPECT_FALSE(iob5ptr->isSharedOne());
 433
 434   // Unshare that clone
 435   EXPECT_TRUE(iob4clone->isSharedOne());
 436   iob4clone->unshare();
 437   EXPECT_FALSE(iob4clone->isSharedOne());
 438   EXPECT_FALSE(iob4ptr->isSharedOne());
 439   EXPECT_FALSE(iob1->isShared());
 440   iob4clone.reset();
 441
 442
 443   // Create a clone of a different IOBuf
 444   EXPECT_FALSE(iob1->isShared());
 445   EXPECT_FALSE(iob3ptr->isSharedOne());
 446
 447   unique_ptr<IOBuf> iob3clone = iob3ptr->cloneOne();
 448   gen.seed(fillSeed);
 449   checkBuf(iob1.get(), gen);
 450   checkBuf(iob2ptr, gen);
 451   checkBuf(iob3clone.get(), gen);
 452   checkBuf(iob4ptr, gen);
 453   checkBuf(iob5ptr, gen);
 454
 455   EXPECT_TRUE(iob1->isShared());
 456   EXPECT_TRUE(iob3ptr->isSharedOne());
 457   EXPECT_FALSE(iob1->isSharedOne());
 458
 459   // Delete the clone and make sure the original is unshared
 460   iob3clone.reset();
 461   EXPECT_FALSE(iob1->isShared());
 462   EXPECT_FALSE(iob3ptr->isSharedOne());
 463
 464
 465   // Clone the entire chain
 466   unique_ptr<IOBuf> chainClone = iob1->clone();
 467   // Verify that the data is correct.
 468   EXPECT_EQ(fullLength, chainClone->computeChainDataLength());
 469   gen.seed(fillSeed);
 470   checkChain(chainClone.get(), gen);
 471
 472   // Check that the buffers report sharing correctly
 473   EXPECT_TRUE(chainClone->isShared());
 474   EXPECT_TRUE(iob1->isShared());
 475
 476   EXPECT_TRUE(iob1->isSharedOne());
 477   EXPECT_TRUE(iob2ptr->isSharedOne());
 478   EXPECT_TRUE(iob3ptr->isSharedOne());
 479   EXPECT_TRUE(iob4ptr->isSharedOne());
 480   EXPECT_TRUE(iob5ptr->isSharedOne());
 481
 482   // Unshare the cloned chain
 483   chainClone->unshare();
 484   EXPECT_FALSE(chainClone->isShared());
 485   EXPECT_FALSE(iob1->isShared());
 486
 487   // Make sure the unshared result still has the same data
 488   EXPECT_EQ(fullLength, chainClone->computeChainDataLength());
 489   gen.seed(fillSeed);
 490   checkChain(chainClone.get(), gen);
 491
 492   // Destroy this chain
 493   chainClone.reset();
 494
 495
 496   // Clone a new chain
 497   EXPECT_FALSE(iob1->isShared());
 498   chainClone = iob1->clone();
 499   EXPECT_TRUE(iob1->isShared());
 500   EXPECT_TRUE(chainClone->isShared());
 501
 502   // Delete the original chain
 503   iob1.reset();
 504   EXPECT_FALSE(chainClone->isShared());
 505
 506   // Coalesce the chain
 507   //
 508   // Coalescing this chain will create a new buffer and release the last
 509   // refcount on the original buffers we created.  Also make sure
 510   // that arrayBufFreeCount increases to one to indicate that arrayBuf was
 511   // freed.
 512   EXPECT_EQ(5, chainClone->countChainElements());
 513   EXPECT_EQ(0, arrayBufFreeCount);
 514
 515   // Buffer lengths: 1500 20 1234 900 321
 516   // Attempting to gather more data than available should fail
 517   EXPECT_THROW(chainClone->gather(4000), std::overflow_error);
 518   // Coalesce the first 3 buffers
 519   chainClone->gather(1521);
 520   EXPECT_EQ(3, chainClone->countChainElements());
 521   EXPECT_EQ(0, arrayBufFreeCount);
 522
 523   // Make sure the data is still the same after coalescing
 524   EXPECT_EQ(fullLength, chainClone->computeChainDataLength());
 525   gen.seed(fillSeed);
 526   checkChain(chainClone.get(), gen);
 527
 528   // Coalesce the entire chain
 529   chainClone->coalesce();
 530   EXPECT_EQ(1, chainClone->countChainElements());
 531   EXPECT_EQ(1, arrayBufFreeCount);
 532
 533   // Make sure the data is still the same after coalescing
 534   EXPECT_EQ(fullLength, chainClone->computeChainDataLength());
 535   gen.seed(fillSeed);
 536   checkChain(chainClone.get(), gen);
 537
 538   // Make a new chain to test the unlink and pop operations
 539   iob1 = IOBuf::create(1);
 540   iob1->append(1);
 541   IOBuf *iob1ptr = iob1.get();
 542   iob2 = IOBuf::create(3);
 543   iob2->append(3);
 544   iob2ptr = iob2.get();
 545   iob3 = IOBuf::create(5);
 546   iob3->append(5);
 547   iob3ptr = iob3.get();
 548   iob4 = IOBuf::create(7);
 549   iob4->append(7);
 550   iob4ptr = iob4.get();
 551   iob1->appendChain(std::move(iob2));
 552   iob1->prev()->appendChain(std::move(iob3));
 553   iob1->prev()->appendChain(std::move(iob4));
 554   EXPECT_EQ(4, iob1->countChainElements());
 555   EXPECT_EQ(16, iob1->computeChainDataLength());
 556
 557   // Unlink from the middle of the chain
 558   iob3 = iob3ptr->unlink();
 559   EXPECT_TRUE(iob3.get() == iob3ptr);
 560   EXPECT_EQ(3, iob1->countChainElements());
 561   EXPECT_EQ(11, iob1->computeChainDataLength());
 562
 563   // Unlink from the end of the chain
 564   iob4 = iob1->prev()->unlink();
 565   EXPECT_TRUE(iob4.get() == iob4ptr);
 566   EXPECT_EQ(2, iob1->countChainElements());
 567   EXPECT_TRUE(iob1->next() == iob2ptr);
 568   EXPECT_EQ(4, iob1->computeChainDataLength());
 569
 570   // Pop from the front of the chain
 571   iob2 = iob1->pop();
 572   EXPECT_TRUE(iob1.get() == iob1ptr);
 573   EXPECT_EQ(1, iob1->countChainElements());
 574   EXPECT_EQ(1, iob1->computeChainDataLength());
 575   EXPECT_TRUE(iob2.get() == iob2ptr);
 576   EXPECT_EQ(1, iob2->countChainElements());
 577   EXPECT_EQ(3, iob2->computeChainDataLength());
 578 }
 579
 580 void testFreeFn(void* buffer, void* ptr) {
 581   uint32_t* freeCount = static_cast<uint32_t*>(ptr);;
 582   delete[] static_cast<uint8_t*>(buffer);
 583   if (freeCount) {
 584     ++(*freeCount);
 585   }
 586 };
 587
 588 TEST(IOBuf, Reserve) {
 589   uint32_t fillSeed = 0x23456789;
 590   boost::mt19937 gen(fillSeed);
 591
 592   // Reserve does nothing if empty and doesn't have to grow the buffer
 593   {
 594     gen.seed(fillSeed);
 595     unique_ptr<IOBuf> iob(IOBuf::create(2000));
 596     EXPECT_EQ(0, iob->headroom());
 597     const void* p1 = iob->buffer();
 598     iob->reserve(5, 15);
 599     EXPECT_LE(5, iob->headroom());
 600     EXPECT_EQ(p1, iob->buffer());
 601   }
 602
 603   // Reserve doesn't reallocate if we have enough total room
 604   {
 605     gen.seed(fillSeed);
 606     unique_ptr<IOBuf> iob(IOBuf::create(2000));
 607     iob->append(100);
 608     fillBuf(iob.get(), gen);
 609     EXPECT_EQ(0, iob->headroom());
 610     EXPECT_EQ(100, iob->length());
 611     const void* p1 = iob->buffer();
 612     const uint8_t* d1 = iob->data();
 613     iob->reserve(100, 1800);
 614     EXPECT_LE(100, iob->headroom());
 615     EXPECT_EQ(p1, iob->buffer());
 616     EXPECT_EQ(d1 + 100, iob->data());
 617     gen.seed(fillSeed);
 618     checkBuf(iob.get(), gen);
 619   }
 620
 621   // Reserve reallocates if we don't have enough total room.
 622   // NOTE that, with jemalloc, we know that this won't reallocate in place
 623   // as the size is less than jemallocMinInPlaceExpanadable
 624   {
 625     gen.seed(fillSeed);
 626     unique_ptr<IOBuf> iob(IOBuf::create(2000));
 627     iob->append(100);
 628     fillBuf(iob.get(), gen);
 629     EXPECT_EQ(0, iob->headroom());
 630     EXPECT_EQ(100, iob->length());
 631     const void* p1 = iob->buffer();
 632     const uint8_t* d1 = iob->data();
 633     iob->reserve(100, 2512);  // allocation sizes are multiples of 256
 634     EXPECT_LE(100, iob->headroom());
 635     if (folly::usingJEMalloc()) {
 636       EXPECT_NE(p1, iob->buffer());
 637     }
 638     gen.seed(fillSeed);
 639     checkBuf(iob.get(), gen);
 640   }
 641
 642   // Test reserve from internal buffer, this used to segfault
 643   {
 644     unique_ptr<IOBuf> iob(IOBuf::create(0));
 645     iob->reserve(0, 2000);
 646     EXPECT_EQ(0, iob->headroom());
 647     EXPECT_LE(2000, iob->tailroom());
 648   }
 649
 650   // Test reserving from a user-allocated buffer.
 651   {
 652     uint8_t* buf = static_cast<uint8_t*>(malloc(100));
 653     auto iob = IOBuf::takeOwnership(buf, 100);
 654     iob->reserve(0, 2000);
 655     EXPECT_EQ(0, iob->headroom());
 656     EXPECT_LE(2000, iob->tailroom());
 657   }
 658
 659   // Test reserving from a user-allocated with a custom free function.
 660   {
 661     uint32_t freeCount{0};
 662     uint8_t* buf = new uint8_t[100];
 663     auto iob = IOBuf::takeOwnership(buf, 100, testFreeFn, &freeCount);
 664     iob->reserve(0, 2000);
 665     EXPECT_EQ(0, iob->headroom());
 666     EXPECT_LE(2000, iob->tailroom());
 667     EXPECT_EQ(1, freeCount);
 668   }
 669 }
 670
 671 TEST(IOBuf, copyBuffer) {
 672   std::string s("hello");
 673   auto buf = IOBuf::copyBuffer(s.data(), s.size(), 1, 2);
 674   EXPECT_EQ(1, buf->headroom());
 675   EXPECT_EQ(s, std::string(reinterpret_cast<const char*>(buf->data()),
 676                            buf->length()));
 677   EXPECT_LE(2, buf->tailroom());
 678
 679   buf = IOBuf::copyBuffer(s, 5, 7);
 680   EXPECT_EQ(5, buf->headroom());
 681   EXPECT_EQ(s, std::string(reinterpret_cast<const char*>(buf->data()),
 682                            buf->length()));
 683   EXPECT_LE(7, buf->tailroom());
 684
 685   std::string empty;
 686   buf = IOBuf::copyBuffer(empty, 3, 6);
 687   EXPECT_EQ(3, buf->headroom());
 688   EXPECT_EQ(0, buf->length());
 689   EXPECT_LE(6, buf->tailroom());
 690
 691   // A stack-allocated version
 692   IOBuf stackBuf(IOBuf::COPY_BUFFER, s, 1, 2);
 693   EXPECT_EQ(1, stackBuf.headroom());
 694   EXPECT_EQ(s, std::string(reinterpret_cast<const char*>(stackBuf.data()),
 695                            stackBuf.length()));
 696   EXPECT_LE(2, stackBuf.tailroom());
 697 }
 698
 699 TEST(IOBuf, maybeCopyBuffer) {
 700   std::string s("this is a test");
 701   auto buf = IOBuf::maybeCopyBuffer(s, 1, 2);
 702   EXPECT_EQ(1, buf->headroom());
 703   EXPECT_EQ(s, std::string(reinterpret_cast<const char*>(buf->data()),
 704                            buf->length()));
 705   EXPECT_LE(2, buf->tailroom());
 706
 707   std::string empty;
 708   buf = IOBuf::maybeCopyBuffer("", 5, 7);
 709   EXPECT_EQ(nullptr, buf.get());
 710
 711   buf = IOBuf::maybeCopyBuffer("");
 712   EXPECT_EQ(nullptr, buf.get());
 713 }
 714
 715 namespace {
 716
 717 int customDeleterCount = 0;
 718 int destructorCount = 0;
 719 struct OwnershipTestClass {
 720   explicit OwnershipTestClass(int v = 0) : val(v) { }
 721   ~OwnershipTestClass() {
 722     ++destructorCount;
 723   }
 724   int val;
 725 };
 726
 727 typedef std::function<void(OwnershipTestClass*)> CustomDeleter;
 728
 729 void customDelete(OwnershipTestClass* p) {
 730   ++customDeleterCount;
 731   delete p;
 732 }
 733
 734 void customDeleteArray(OwnershipTestClass* p) {
 735   ++customDeleterCount;
 736   delete[] p;
 737 }
 738
 739 }  // namespace
 740
 741 TEST(IOBuf, takeOwnershipUniquePtr) {
 742   destructorCount = 0;
 743   {
 744     std::unique_ptr<OwnershipTestClass> p(new OwnershipTestClass());
 745   }
 746   EXPECT_EQ(1, destructorCount);
 747
 748   destructorCount = 0;
 749   {
 750     std::unique_ptr<OwnershipTestClass[]> p(new OwnershipTestClass[2]);
 751   }
 752   EXPECT_EQ(2, destructorCount);
 753
 754   destructorCount = 0;
 755   {
 756     std::unique_ptr<OwnershipTestClass> p(new OwnershipTestClass());
 757     std::unique_ptr<IOBuf> buf(IOBuf::takeOwnership(std::move(p)));
 758     EXPECT_EQ(sizeof(OwnershipTestClass), buf->length());
 759     EXPECT_EQ(0, destructorCount);
 760   }
 761   EXPECT_EQ(1, destructorCount);
 762
 763   destructorCount = 0;
 764   {
 765     std::unique_ptr<OwnershipTestClass[]> p(new OwnershipTestClass[2]);
 766     std::unique_ptr<IOBuf> buf(IOBuf::takeOwnership(std::move(p), 2));
 767     EXPECT_EQ(2 * sizeof(OwnershipTestClass), buf->length());
 768     EXPECT_EQ(0, destructorCount);
 769   }
 770   EXPECT_EQ(2, destructorCount);
 771
 772   customDeleterCount = 0;
 773   destructorCount = 0;
 774   {
 775     std::unique_ptr<OwnershipTestClass, CustomDeleter>
 776       p(new OwnershipTestClass(), customDelete);
 777     std::unique_ptr<IOBuf> buf(IOBuf::takeOwnership(std::move(p)));
 778     EXPECT_EQ(sizeof(OwnershipTestClass), buf->length());
 779     EXPECT_EQ(0, destructorCount);
 780   }
 781   EXPECT_EQ(1, destructorCount);
 782   EXPECT_EQ(1, customDeleterCount);
 783
 784   customDeleterCount = 0;
 785   destructorCount = 0;
 786   {
 787     std::unique_ptr<OwnershipTestClass[], CustomDeleter>
 788       p(new OwnershipTestClass[2], CustomDeleter(customDeleteArray));
 789     std::unique_ptr<IOBuf> buf(IOBuf::takeOwnership(std::move(p), 2));
 790     EXPECT_EQ(2 * sizeof(OwnershipTestClass), buf->length());
 791     EXPECT_EQ(0, destructorCount);
 792   }
 793   EXPECT_EQ(2, destructorCount);
 794   EXPECT_EQ(1, customDeleterCount);
 795 }
 796
 797 TEST(IOBuf, Alignment) {
 798   // max_align_t doesn't exist in gcc 4.6.2
 799   struct MaxAlign {
 800     char c;
 801   } __attribute__((aligned));
 802   size_t alignment = alignof(MaxAlign);
 803
 804   std::vector<size_t> sizes {0, 1, 64, 256, 1024, 1 << 10};
 805   for (size_t size : sizes) {
 806     auto buf = IOBuf::create(size);
 807     uintptr_t p = reinterpret_cast<uintptr_t>(buf->data());
 808     EXPECT_EQ(0, p & (alignment - 1)) << "size=" << size;
 809   }
 810 }
 811
 812 TEST(TypedIOBuf, Simple) {
 813   auto buf = IOBuf::create(0);
 814   TypedIOBuf<uint64_t> typed(buf.get());
 815   const uint64_t n = 10000;
 816   typed.reserve(0, n);
 817   EXPECT_LE(n, typed.capacity());
 818   for (uint64_t i = 0; i < n; i++) {
 819     *typed.writableTail() = i;
 820     typed.append(1);
 821   }
 822   EXPECT_EQ(n, typed.length());
 823   for (uint64_t i = 0; i < n; i++) {
 824     EXPECT_EQ(i, typed.data()[i]);
 825   }
 826 }
 827 enum BufType {
 828   CREATE,
 829   TAKE_OWNERSHIP_MALLOC,
 830   TAKE_OWNERSHIP_CUSTOM,
 831   USER_OWNED,
 832 };
 833
 834 // chain element size, number of elements in chain, shared
 835 class MoveToFbStringTest
 836   : public ::testing::TestWithParam<std::tr1::tuple<int, int, bool, BufType>> {
 837  protected:
 838   void SetUp() {
 839     std::tr1::tie(elementSize_, elementCount_, shared_, type_) = GetParam();
 840     buf_ = makeBuf();
 841     for (int i = 0; i < elementCount_ - 1; ++i) {
 842       buf_->prependChain(makeBuf());
 843     }
 844     EXPECT_EQ(elementCount_, buf_->countChainElements());
 845     EXPECT_EQ(elementCount_ * elementSize_, buf_->computeChainDataLength());
 846     if (shared_) {
 847       buf2_ = buf_->clone();
 848       EXPECT_EQ(elementCount_, buf2_->countChainElements());
 849       EXPECT_EQ(elementCount_ * elementSize_, buf2_->computeChainDataLength());
 850     }
 851   }
 852
 853   std::unique_ptr<IOBuf> makeBuf() {
 854     unique_ptr<IOBuf> buf;
 855     switch (type_) {
 856       case CREATE:
 857         buf = IOBuf::create(elementSize_);
 858         buf->append(elementSize_);
 859         break;
 860       case TAKE_OWNERSHIP_MALLOC: {
 861         void* data = malloc(elementSize_);
 862         if (!data) {
 863           throw std::bad_alloc();
 864         }
 865         buf = IOBuf::takeOwnership(data, elementSize_);
 866         break;
 867       }
 868       case TAKE_OWNERSHIP_CUSTOM: {
 869         uint8_t* data = new uint8_t[elementSize_];
 870         buf = IOBuf::takeOwnership(data, elementSize_, testFreeFn);
 871         break;
 872       }
 873       case USER_OWNED: {
 874         unique_ptr<uint8_t[]> data(new uint8_t[elementSize_]);
 875         buf = IOBuf::wrapBuffer(data.get(), elementSize_);
 876         ownedBuffers_.emplace_back(std::move(data));
 877         break;
 878       }
 879       default:
 880         throw std::invalid_argument("unexpected buffer type parameter");
 881         break;
 882     }
 883     memset(buf->writableData(), 'x', elementSize_);
 884     return buf;
 885   }
 886
 887   void check(std::unique_ptr<IOBuf>& buf) {
 888     fbstring str = buf->moveToFbString();
 889     EXPECT_EQ(elementCount_ * elementSize_, str.size());
 890     EXPECT_EQ(elementCount_ * elementSize_, strspn(str.c_str(), "x"));
 891     EXPECT_EQ(0, buf->length());
 892     EXPECT_EQ(1, buf->countChainElements());
 893     EXPECT_EQ(0, buf->computeChainDataLength());
 894     EXPECT_FALSE(buf->isChained());
 895   }
 896
 897   int elementSize_;
 898   int elementCount_;
 899   bool shared_;
 900   BufType type_;
 901   std::unique_ptr<IOBuf> buf_;
 902   std::unique_ptr<IOBuf> buf2_;
 903   std::vector<std::unique_ptr<uint8_t[]>> ownedBuffers_;
 904 };
 905
 906 TEST_P(MoveToFbStringTest, Simple) {
 907   check(buf_);
 908   if (shared_) {
 909     check(buf2_);
 910   }
 911 }
 912
 913 INSTANTIATE_TEST_CASE_P(
 914     MoveToFbString,
 915     MoveToFbStringTest,
 916     ::testing::Combine(
 917         ::testing::Values(0, 1, 24, 256, 1 << 10, 1 << 20),  // element size
 918         ::testing::Values(1, 2, 10),                         // element count
 919         ::testing::Bool(),                                   // shared
 920         ::testing::Values(CREATE, TAKE_OWNERSHIP_MALLOC,
 921                           TAKE_OWNERSHIP_CUSTOM, USER_OWNED)));
 922
 923 TEST(IOBuf, getIov) {
 924   uint32_t fillSeed = 0xdeadbeef;
 925   boost::mt19937 gen(fillSeed);
 926
 927   size_t len = 4096;
 928   size_t count = 32;
 929   auto buf = IOBuf::create(len + 1);
 930   buf->append(rand() % len + 1);
 931   fillBuf(buf.get(), gen);
 932
 933   for (size_t i = 0; i < count - 1; i++) {
 934     auto buf2 = IOBuf::create(len + 1);
 935     buf2->append(rand() % len + 1);
 936     fillBuf(buf2.get(), gen);
 937     buf->prependChain(std::move(buf2));
 938   }
 939   EXPECT_EQ(count, buf->countChainElements());
 940
 941   auto iov = buf->getIov();
 942   EXPECT_EQ(count, iov.size());
 943
 944   IOBuf const* p = buf.get();
 945   for (size_t i = 0; i < count; i++, p = p->next()) {
 946     EXPECT_EQ(p->data(), iov[i].iov_base);
 947     EXPECT_EQ(p->length(), iov[i].iov_len);
 948   }
 949
 950   // an empty buf should be skipped in the iov.
 951   buf->next()->clear();
 952   iov = buf->getIov();
 953   EXPECT_EQ(count - 1, iov.size());
 954   EXPECT_EQ(buf->next()->next()->data(), iov[1].iov_base);
 955
 956   // same for the first one being empty
 957   buf->clear();
 958   iov = buf->getIov();
 959   EXPECT_EQ(count - 2, iov.size());
 960   EXPECT_EQ(buf->next()->next()->data(), iov[0].iov_base);
 961
 962   // and the last one
 963   buf->prev()->clear();
 964   iov = buf->getIov();
 965   EXPECT_EQ(count - 3, iov.size());
 966 }
 967
 968 TEST(IOBuf, move) {
 969   // Default allocate an IOBuf on the stack
 970   IOBuf outerBuf;
 971   char data[] = "foobar";
 972   uint32_t length = sizeof(data);
 973   uint32_t actualCapacity{0};
 974   const void* ptr{nullptr};
 975
 976   {
 977     // Create a small IOBuf on the stack.
 978     // Note that IOBufs created on the stack always use an external buffer.
 979     IOBuf b1(IOBuf::CREATE, 10);
 980     actualCapacity = b1.capacity();
 981     EXPECT_GE(actualCapacity, 10);
 982     EXPECT_EQ(0, b1.length());
 983     EXPECT_FALSE(b1.isShared());
 984     ptr = b1.data();
 985     ASSERT_TRUE(ptr != nullptr);
 986     memcpy(b1.writableTail(), data, length);
 987     b1.append(length);
 988     EXPECT_EQ(length, b1.length());
 989
 990     // Use the move constructor
 991     IOBuf b2(std::move(b1));
 992     EXPECT_EQ(ptr, b2.data());
 993     EXPECT_EQ(length, b2.length());
 994     EXPECT_EQ(actualCapacity, b2.capacity());
 995     EXPECT_FALSE(b2.isShared());
 996
 997     // Use the move assignment operator
 998     outerBuf = std::move(b2);
 999     // Close scope, destroying b1 and b2
1000     // (which are both be invalid now anyway after moving out of them)
1001   }
1002
1003   EXPECT_EQ(ptr, outerBuf.data());
1004   EXPECT_EQ(length, outerBuf.length());
1005   EXPECT_EQ(actualCapacity, outerBuf.capacity());
1006   EXPECT_FALSE(outerBuf.isShared());
1007 }
1008
1009 int main(int argc, char** argv) {
1010   testing::InitGoogleTest(&argc, argv);
1011   google::ParseCommandLineFlags(&argc, &argv, true);
1012
1013   return RUN_ALL_TESTS();
1014 }