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
 146 }
 147
 148 TEST(IOBuf, WrapBuffer) {
 149   const uint32_t size1 = 1234;
 150   uint8_t buf1[size1];
 151   unique_ptr<IOBuf> iobuf1(IOBuf::wrapBuffer(buf1, size1));
 152   EXPECT_EQ(buf1, iobuf1->data());
 153   EXPECT_EQ(size1, iobuf1->length());
 154   EXPECT_EQ(buf1, iobuf1->buffer());
 155   EXPECT_EQ(size1, iobuf1->capacity());
 156
 157   uint32_t size2 = 0x1234;
 158   unique_ptr<uint8_t[]> buf2(new uint8_t[size2]);
 159   unique_ptr<IOBuf> iobuf2(IOBuf::wrapBuffer(buf2.get(), size2));
 160   EXPECT_EQ(buf2.get(), iobuf2->data());
 161   EXPECT_EQ(size2, iobuf2->length());
 162   EXPECT_EQ(buf2.get(), iobuf2->buffer());
 163   EXPECT_EQ(size2, iobuf2->capacity());
 164 }
 165
 166 TEST(IOBuf, CreateCombined) {
 167   // Create a combined IOBuf, then destroy it.
 168   // The data buffer and IOBuf both become unused as part of the destruction
 169   {
 170     auto buf = IOBuf::createCombined(256);
 171     EXPECT_FALSE(buf->isShared());
 172   }
 173
 174   // Create a combined IOBuf, clone from it, and then destroy the original
 175   // IOBuf.  The data buffer cannot be deleted until the clone is also
 176   // destroyed.
 177   {
 178     auto bufA = IOBuf::createCombined(256);
 179     EXPECT_FALSE(bufA->isShared());
 180     auto bufB = bufA->clone();
 181     EXPECT_TRUE(bufA->isShared());
 182     EXPECT_TRUE(bufB->isShared());
 183     bufA.reset();
 184     EXPECT_FALSE(bufB->isShared());
 185   }
 186
 187   // Create a combined IOBuf, then call reserve() to get a larger buffer.
 188   // The IOBuf no longer points to the combined data buffer, but the
 189   // overall memory segment cannot be deleted until the IOBuf is also
 190   // destroyed.
 191   {
 192     auto buf = IOBuf::createCombined(256);
 193     buf->reserve(0, buf->capacity() + 100);
 194   }
 195
 196   // Create a combined IOBuf, clone from it, then call unshare() on the original
 197   // buffer.  This creates a situation where bufB is pointing at the combined
 198   // buffer associated with bufA, but bufA is now using a different buffer.
 199   auto testSwap = [](bool resetAFirst) {
 200     auto bufA = IOBuf::createCombined(256);
 201     EXPECT_FALSE(bufA->isShared());
 202     auto bufB = bufA->clone();
 203     EXPECT_TRUE(bufA->isShared());
 204     EXPECT_TRUE(bufB->isShared());
 205     bufA->unshare();
 206     EXPECT_FALSE(bufA->isShared());
 207     EXPECT_FALSE(bufB->isShared());
 208
 209     if (resetAFirst) {
 210       bufA.reset();
 211       bufB.reset();
 212     } else {
 213       bufB.reset();
 214       bufA.reset();
 215     }
 216   };
 217   testSwap(true);
 218   testSwap(false);
 219 }
 220
 221 void fillBuf(uint8_t* buf, uint32_t length, boost::mt19937& gen) {
 222   for (uint32_t n = 0; n < length; ++n) {
 223     buf[n] = static_cast<uint8_t>(gen() & 0xff);
 224   }
 225 }
 226
 227 void fillBuf(IOBuf* buf, boost::mt19937& gen) {
 228   buf->unshare();
 229   fillBuf(buf->writableData(), buf->length(), gen);
 230 }
 231
 232 void checkBuf(const uint8_t* buf, uint32_t length, boost::mt19937& gen) {
 233   // Rather than using EXPECT_EQ() to check each character,
 234   // count the number of differences and the first character that differs.
 235   // This way on error we'll report just that information, rather than tons of
 236   // failed checks for each byte in the buffer.
 237   uint32_t numDifferences = 0;
 238   uint32_t firstDiffIndex = 0;
 239   uint8_t firstDiffExpected = 0;
 240   for (uint32_t n = 0; n < length; ++n) {
 241     uint8_t expected = static_cast<uint8_t>(gen() & 0xff);
 242     if (buf[n] == expected) {
 243       continue;
 244     }
 245
 246     if (numDifferences == 0) {
 247       firstDiffIndex = n;
 248       firstDiffExpected = expected;
 249     }
 250     ++numDifferences;
 251   }
 252
 253   EXPECT_EQ(0, numDifferences);
 254   if (numDifferences > 0) {
 255     // Cast to int so it will be printed numerically
 256     // rather than as a char if the check fails
 257     EXPECT_EQ(static_cast<int>(buf[firstDiffIndex]),
 258               static_cast<int>(firstDiffExpected));
 259   }
 260 }
 261
 262 void checkBuf(IOBuf* buf, boost::mt19937& gen) {
 263   checkBuf(buf->data(), buf->length(), gen);
 264 }
 265
 266 void checkBuf(ByteRange buf, boost::mt19937& gen) {
 267   checkBuf(buf.data(), buf.size(), gen);
 268 }
 269
 270 void checkChain(IOBuf* buf, boost::mt19937& gen) {
 271   IOBuf *current = buf;
 272   do {
 273     checkBuf(current->data(), current->length(), gen);
 274     current = current->next();
 275   } while (current != buf);
 276 }
 277
 278 TEST(IOBuf, Chaining) {
 279   uint32_t fillSeed = 0x12345678;
 280   boost::mt19937 gen(fillSeed);
 281
 282   // An IOBuf with external storage
 283   uint32_t headroom = 123;
 284   unique_ptr<IOBuf> iob1(IOBuf::create(2048));
 285   iob1->advance(headroom);
 286   iob1->append(1500);
 287   fillBuf(iob1.get(), gen);
 288
 289   // An IOBuf with internal storage
 290   unique_ptr<IOBuf> iob2(IOBuf::create(20));
 291   iob2->append(20);
 292   fillBuf(iob2.get(), gen);
 293
 294   // An IOBuf around a buffer it doesn't own
 295   uint8_t localbuf[1234];
 296   fillBuf(localbuf, 1234, gen);
 297   unique_ptr<IOBuf> iob3(IOBuf::wrapBuffer(localbuf, sizeof(localbuf)));
 298
 299   // An IOBuf taking ownership of a user-supplied buffer
 300   uint32_t heapBufSize = 900;
 301   uint8_t* heapBuf = static_cast<uint8_t*>(malloc(heapBufSize));
 302   fillBuf(heapBuf, heapBufSize, gen);
 303   unique_ptr<IOBuf> iob4(IOBuf::takeOwnership(heapBuf, heapBufSize));
 304
 305   // An IOBuf taking ownership of a user-supplied buffer with
 306   // a custom free function
 307   uint32_t arrayBufSize = 321;
 308   uint8_t* arrayBuf = new uint8_t[arrayBufSize];
 309   fillBuf(arrayBuf, arrayBufSize, gen);
 310   uint32_t arrayBufFreeCount = 0;
 311   unique_ptr<IOBuf> iob5(IOBuf::takeOwnership(arrayBuf, arrayBufSize,
 312                                               deleteArrayBuffer,
 313                                               &arrayBufFreeCount));
 314
 315   EXPECT_FALSE(iob1->isChained());
 316   EXPECT_FALSE(iob2->isChained());
 317   EXPECT_FALSE(iob3->isChained());
 318   EXPECT_FALSE(iob4->isChained());
 319   EXPECT_FALSE(iob5->isChained());
 320
 321   EXPECT_FALSE(iob1->isSharedOne());
 322   EXPECT_FALSE(iob2->isSharedOne());
 323   EXPECT_TRUE(iob3->isSharedOne()); // since we own the buffer
 324   EXPECT_FALSE(iob4->isSharedOne());
 325   EXPECT_FALSE(iob5->isSharedOne());
 326
 327   // Chain the buffers all together
 328   // Since we are going to relinquish ownership of iob2-5 to the chain,
 329   // store raw pointers to them so we can reference them later.
 330   IOBuf* iob2ptr = iob2.get();
 331   IOBuf* iob3ptr = iob3.get();
 332   IOBuf* iob4ptr = iob4.get();
 333   IOBuf* iob5ptr = iob5.get();
 334
 335   iob1->prependChain(std::move(iob2));
 336   iob1->prependChain(std::move(iob4));
 337   iob2ptr->appendChain(std::move(iob3));
 338   iob1->prependChain(std::move(iob5));
 339
 340   EXPECT_EQ(iob2ptr, iob1->next());
 341   EXPECT_EQ(iob3ptr, iob2ptr->next());
 342   EXPECT_EQ(iob4ptr, iob3ptr->next());
 343   EXPECT_EQ(iob5ptr, iob4ptr->next());
 344   EXPECT_EQ(iob1.get(), iob5ptr->next());
 345
 346   EXPECT_EQ(iob5ptr, iob1->prev());
 347   EXPECT_EQ(iob1.get(), iob2ptr->prev());
 348   EXPECT_EQ(iob2ptr, iob3ptr->prev());
 349   EXPECT_EQ(iob3ptr, iob4ptr->prev());
 350   EXPECT_EQ(iob4ptr, iob5ptr->prev());
 351
 352   EXPECT_TRUE(iob1->isChained());
 353   EXPECT_TRUE(iob2ptr->isChained());
 354   EXPECT_TRUE(iob3ptr->isChained());
 355   EXPECT_TRUE(iob4ptr->isChained());
 356   EXPECT_TRUE(iob5ptr->isChained());
 357
 358   uint64_t fullLength = (iob1->length() + iob2ptr->length() +
 359                          iob3ptr->length() + iob4ptr->length() +
 360                         iob5ptr->length());
 361   EXPECT_EQ(5, iob1->countChainElements());
 362   EXPECT_EQ(fullLength, iob1->computeChainDataLength());
 363
 364   // Since iob3 is shared, the entire buffer should report itself as shared
 365   EXPECT_TRUE(iob1->isShared());
 366   // Unshare just iob3
 367   iob3ptr->unshareOne();
 368   EXPECT_FALSE(iob3ptr->isSharedOne());
 369   // Now everything in the chain should be unshared.
 370   // Check on all members of the chain just for good measure
 371   EXPECT_FALSE(iob1->isShared());
 372   EXPECT_FALSE(iob2ptr->isShared());
 373   EXPECT_FALSE(iob3ptr->isShared());
 374   EXPECT_FALSE(iob4ptr->isShared());
 375   EXPECT_FALSE(iob5ptr->isShared());
 376
 377   // Check iteration
 378   gen.seed(fillSeed);
 379   size_t count = 0;
 380   for (auto buf : *iob1) {
 381     checkBuf(buf, gen);
 382     ++count;
 383   }
 384   EXPECT_EQ(5, count);
 385
 386   // Clone one of the IOBufs in the chain
 387   unique_ptr<IOBuf> iob4clone = iob4ptr->cloneOne();
 388   gen.seed(fillSeed);
 389   checkBuf(iob1.get(), gen);
 390   checkBuf(iob2ptr, gen);
 391   checkBuf(iob3ptr, gen);
 392   checkBuf(iob4clone.get(), gen);
 393   checkBuf(iob5ptr, gen);
 394
 395   EXPECT_TRUE(iob1->isShared());
 396   EXPECT_TRUE(iob2ptr->isShared());
 397   EXPECT_TRUE(iob3ptr->isShared());
 398   EXPECT_TRUE(iob4ptr->isShared());
 399   EXPECT_TRUE(iob5ptr->isShared());
 400
 401   EXPECT_FALSE(iob1->isSharedOne());
 402   EXPECT_FALSE(iob2ptr->isSharedOne());
 403   EXPECT_FALSE(iob3ptr->isSharedOne());
 404   EXPECT_TRUE(iob4ptr->isSharedOne());
 405   EXPECT_FALSE(iob5ptr->isSharedOne());
 406
 407   // Unshare that clone
 408   EXPECT_TRUE(iob4clone->isSharedOne());
 409   iob4clone->unshare();
 410   EXPECT_FALSE(iob4clone->isSharedOne());
 411   EXPECT_FALSE(iob4ptr->isSharedOne());
 412   EXPECT_FALSE(iob1->isShared());
 413   iob4clone.reset();
 414
 415
 416   // Create a clone of a different IOBuf
 417   EXPECT_FALSE(iob1->isShared());
 418   EXPECT_FALSE(iob3ptr->isSharedOne());
 419
 420   unique_ptr<IOBuf> iob3clone = iob3ptr->cloneOne();
 421   gen.seed(fillSeed);
 422   checkBuf(iob1.get(), gen);
 423   checkBuf(iob2ptr, gen);
 424   checkBuf(iob3clone.get(), gen);
 425   checkBuf(iob4ptr, gen);
 426   checkBuf(iob5ptr, gen);
 427
 428   EXPECT_TRUE(iob1->isShared());
 429   EXPECT_TRUE(iob3ptr->isSharedOne());
 430   EXPECT_FALSE(iob1->isSharedOne());
 431
 432   // Delete the clone and make sure the original is unshared
 433   iob3clone.reset();
 434   EXPECT_FALSE(iob1->isShared());
 435   EXPECT_FALSE(iob3ptr->isSharedOne());
 436
 437
 438   // Clone the entire chain
 439   unique_ptr<IOBuf> chainClone = iob1->clone();
 440   // Verify that the data is correct.
 441   EXPECT_EQ(fullLength, chainClone->computeChainDataLength());
 442   gen.seed(fillSeed);
 443   checkChain(chainClone.get(), gen);
 444
 445   // Check that the buffers report sharing correctly
 446   EXPECT_TRUE(chainClone->isShared());
 447   EXPECT_TRUE(iob1->isShared());
 448
 449   EXPECT_TRUE(iob1->isSharedOne());
 450   EXPECT_TRUE(iob2ptr->isSharedOne());
 451   EXPECT_TRUE(iob3ptr->isSharedOne());
 452   EXPECT_TRUE(iob4ptr->isSharedOne());
 453   EXPECT_TRUE(iob5ptr->isSharedOne());
 454
 455   // Unshare the cloned chain
 456   chainClone->unshare();
 457   EXPECT_FALSE(chainClone->isShared());
 458   EXPECT_FALSE(iob1->isShared());
 459
 460   // Make sure the unshared result still has the same data
 461   EXPECT_EQ(fullLength, chainClone->computeChainDataLength());
 462   gen.seed(fillSeed);
 463   checkChain(chainClone.get(), gen);
 464
 465   // Destroy this chain
 466   chainClone.reset();
 467
 468
 469   // Clone a new chain
 470   EXPECT_FALSE(iob1->isShared());
 471   chainClone = iob1->clone();
 472   EXPECT_TRUE(iob1->isShared());
 473   EXPECT_TRUE(chainClone->isShared());
 474
 475   // Delete the original chain
 476   iob1.reset();
 477   EXPECT_FALSE(chainClone->isShared());
 478
 479   // Coalesce the chain
 480   //
 481   // Coalescing this chain will create a new buffer and release the last
 482   // refcount on the original buffers we created.  Also make sure
 483   // that arrayBufFreeCount increases to one to indicate that arrayBuf was
 484   // freed.
 485   EXPECT_EQ(5, chainClone->countChainElements());
 486   EXPECT_EQ(0, arrayBufFreeCount);
 487
 488   // Buffer lengths: 1500 20 1234 900 321
 489   // Attempting to gather more data than available should fail
 490   EXPECT_THROW(chainClone->gather(4000), std::overflow_error);
 491   // Coalesce the first 3 buffers
 492   chainClone->gather(1521);
 493   EXPECT_EQ(3, chainClone->countChainElements());
 494   EXPECT_EQ(0, arrayBufFreeCount);
 495
 496   // Make sure the data is still the same after coalescing
 497   EXPECT_EQ(fullLength, chainClone->computeChainDataLength());
 498   gen.seed(fillSeed);
 499   checkChain(chainClone.get(), gen);
 500
 501   // Coalesce the entire chain
 502   chainClone->coalesce();
 503   EXPECT_EQ(1, chainClone->countChainElements());
 504   EXPECT_EQ(1, arrayBufFreeCount);
 505
 506   // Make sure the data is still the same after coalescing
 507   EXPECT_EQ(fullLength, chainClone->computeChainDataLength());
 508   gen.seed(fillSeed);
 509   checkChain(chainClone.get(), gen);
 510
 511   // Make a new chain to test the unlink and pop operations
 512   iob1 = IOBuf::create(1);
 513   iob1->append(1);
 514   IOBuf *iob1ptr = iob1.get();
 515   iob2 = IOBuf::create(3);
 516   iob2->append(3);
 517   iob2ptr = iob2.get();
 518   iob3 = IOBuf::create(5);
 519   iob3->append(5);
 520   iob3ptr = iob3.get();
 521   iob4 = IOBuf::create(7);
 522   iob4->append(7);
 523   iob4ptr = iob4.get();
 524   iob1->appendChain(std::move(iob2));
 525   iob1->prev()->appendChain(std::move(iob3));
 526   iob1->prev()->appendChain(std::move(iob4));
 527   EXPECT_EQ(4, iob1->countChainElements());
 528   EXPECT_EQ(16, iob1->computeChainDataLength());
 529
 530   // Unlink from the middle of the chain
 531   iob3 = iob3ptr->unlink();
 532   EXPECT_TRUE(iob3.get() == iob3ptr);
 533   EXPECT_EQ(3, iob1->countChainElements());
 534   EXPECT_EQ(11, iob1->computeChainDataLength());
 535
 536   // Unlink from the end of the chain
 537   iob4 = iob1->prev()->unlink();
 538   EXPECT_TRUE(iob4.get() == iob4ptr);
 539   EXPECT_EQ(2, iob1->countChainElements());
 540   EXPECT_TRUE(iob1->next() == iob2ptr);
 541   EXPECT_EQ(4, iob1->computeChainDataLength());
 542
 543   // Pop from the front of the chain
 544   iob2 = iob1->pop();
 545   EXPECT_TRUE(iob1.get() == iob1ptr);
 546   EXPECT_EQ(1, iob1->countChainElements());
 547   EXPECT_EQ(1, iob1->computeChainDataLength());
 548   EXPECT_TRUE(iob2.get() == iob2ptr);
 549   EXPECT_EQ(1, iob2->countChainElements());
 550   EXPECT_EQ(3, iob2->computeChainDataLength());
 551 }
 552
 553 void testFreeFn(void* buffer, void* ptr) {
 554   uint32_t* freeCount = static_cast<uint32_t*>(ptr);;
 555   delete[] static_cast<uint8_t*>(buffer);
 556   if (freeCount) {
 557     ++(*freeCount);
 558   }
 559 };
 560
 561 TEST(IOBuf, Reserve) {
 562   uint32_t fillSeed = 0x23456789;
 563   boost::mt19937 gen(fillSeed);
 564
 565   // Reserve does nothing if empty and doesn't have to grow the buffer
 566   {
 567     gen.seed(fillSeed);
 568     unique_ptr<IOBuf> iob(IOBuf::create(2000));
 569     EXPECT_EQ(0, iob->headroom());
 570     const void* p1 = iob->buffer();
 571     iob->reserve(5, 15);
 572     EXPECT_LE(5, iob->headroom());
 573     EXPECT_EQ(p1, iob->buffer());
 574   }
 575
 576   // Reserve doesn't reallocate if we have enough total room
 577   {
 578     gen.seed(fillSeed);
 579     unique_ptr<IOBuf> iob(IOBuf::create(2000));
 580     iob->append(100);
 581     fillBuf(iob.get(), gen);
 582     EXPECT_EQ(0, iob->headroom());
 583     EXPECT_EQ(100, iob->length());
 584     const void* p1 = iob->buffer();
 585     const uint8_t* d1 = iob->data();
 586     iob->reserve(100, 1800);
 587     EXPECT_LE(100, iob->headroom());
 588     EXPECT_EQ(p1, iob->buffer());
 589     EXPECT_EQ(d1 + 100, iob->data());
 590     gen.seed(fillSeed);
 591     checkBuf(iob.get(), gen);
 592   }
 593
 594   // Reserve reallocates if we don't have enough total room.
 595   // NOTE that, with jemalloc, we know that this won't reallocate in place
 596   // as the size is less than jemallocMinInPlaceExpanadable
 597   {
 598     gen.seed(fillSeed);
 599     unique_ptr<IOBuf> iob(IOBuf::create(2000));
 600     iob->append(100);
 601     fillBuf(iob.get(), gen);
 602     EXPECT_EQ(0, iob->headroom());
 603     EXPECT_EQ(100, iob->length());
 604     const void* p1 = iob->buffer();
 605     const uint8_t* d1 = iob->data();
 606     iob->reserve(100, 2512);  // allocation sizes are multiples of 256
 607     EXPECT_LE(100, iob->headroom());
 608     if (folly::usingJEMalloc()) {
 609       EXPECT_NE(p1, iob->buffer());
 610     }
 611     gen.seed(fillSeed);
 612     checkBuf(iob.get(), gen);
 613   }
 614
 615   // Test reserve from internal buffer, this used to segfault
 616   {
 617     unique_ptr<IOBuf> iob(IOBuf::create(0));
 618     iob->reserve(0, 2000);
 619     EXPECT_EQ(0, iob->headroom());
 620     EXPECT_LE(2000, iob->tailroom());
 621   }
 622
 623   // Test reserving from a user-allocated buffer.
 624   {
 625     uint8_t* buf = static_cast<uint8_t*>(malloc(100));
 626     auto iob = IOBuf::takeOwnership(buf, 100);
 627     iob->reserve(0, 2000);
 628     EXPECT_EQ(0, iob->headroom());
 629     EXPECT_LE(2000, iob->tailroom());
 630   }
 631
 632   // Test reserving from a user-allocated with a custom free function.
 633   {
 634     uint32_t freeCount{0};
 635     uint8_t* buf = new uint8_t[100];
 636     auto iob = IOBuf::takeOwnership(buf, 100, testFreeFn, &freeCount);
 637     iob->reserve(0, 2000);
 638     EXPECT_EQ(0, iob->headroom());
 639     EXPECT_LE(2000, iob->tailroom());
 640     EXPECT_EQ(1, freeCount);
 641   }
 642 }
 643
 644 TEST(IOBuf, copyBuffer) {
 645   std::string s("hello");
 646   auto buf = IOBuf::copyBuffer(s.data(), s.size(), 1, 2);
 647   EXPECT_EQ(1, buf->headroom());
 648   EXPECT_EQ(s, std::string(reinterpret_cast<const char*>(buf->data()),
 649                            buf->length()));
 650   EXPECT_LE(2, buf->tailroom());
 651
 652   buf = IOBuf::copyBuffer(s, 5, 7);
 653   EXPECT_EQ(5, buf->headroom());
 654   EXPECT_EQ(s, std::string(reinterpret_cast<const char*>(buf->data()),
 655                            buf->length()));
 656   EXPECT_LE(7, buf->tailroom());
 657
 658   std::string empty;
 659   buf = IOBuf::copyBuffer(empty, 3, 6);
 660   EXPECT_EQ(3, buf->headroom());
 661   EXPECT_EQ(0, buf->length());
 662   EXPECT_LE(6, buf->tailroom());
 663 }
 664
 665 TEST(IOBuf, maybeCopyBuffer) {
 666   std::string s("this is a test");
 667   auto buf = IOBuf::maybeCopyBuffer(s, 1, 2);
 668   EXPECT_EQ(1, buf->headroom());
 669   EXPECT_EQ(s, std::string(reinterpret_cast<const char*>(buf->data()),
 670                            buf->length()));
 671   EXPECT_LE(2, buf->tailroom());
 672
 673   std::string empty;
 674   buf = IOBuf::maybeCopyBuffer("", 5, 7);
 675   EXPECT_EQ(nullptr, buf.get());
 676
 677   buf = IOBuf::maybeCopyBuffer("");
 678   EXPECT_EQ(nullptr, buf.get());
 679 }
 680
 681 namespace {
 682
 683 int customDeleterCount = 0;
 684 int destructorCount = 0;
 685 struct OwnershipTestClass {
 686   explicit OwnershipTestClass(int v = 0) : val(v) { }
 687   ~OwnershipTestClass() {
 688     ++destructorCount;
 689   }
 690   int val;
 691 };
 692
 693 typedef std::function<void(OwnershipTestClass*)> CustomDeleter;
 694
 695 void customDelete(OwnershipTestClass* p) {
 696   ++customDeleterCount;
 697   delete p;
 698 }
 699
 700 void customDeleteArray(OwnershipTestClass* p) {
 701   ++customDeleterCount;
 702   delete[] p;
 703 }
 704
 705 }  // namespace
 706
 707 TEST(IOBuf, takeOwnershipUniquePtr) {
 708   destructorCount = 0;
 709   {
 710     std::unique_ptr<OwnershipTestClass> p(new OwnershipTestClass());
 711   }
 712   EXPECT_EQ(1, destructorCount);
 713
 714   destructorCount = 0;
 715   {
 716     std::unique_ptr<OwnershipTestClass[]> p(new OwnershipTestClass[2]);
 717   }
 718   EXPECT_EQ(2, destructorCount);
 719
 720   destructorCount = 0;
 721   {
 722     std::unique_ptr<OwnershipTestClass> p(new OwnershipTestClass());
 723     std::unique_ptr<IOBuf> buf(IOBuf::takeOwnership(std::move(p)));
 724     EXPECT_EQ(sizeof(OwnershipTestClass), buf->length());
 725     EXPECT_EQ(0, destructorCount);
 726   }
 727   EXPECT_EQ(1, destructorCount);
 728
 729   destructorCount = 0;
 730   {
 731     std::unique_ptr<OwnershipTestClass[]> p(new OwnershipTestClass[2]);
 732     std::unique_ptr<IOBuf> buf(IOBuf::takeOwnership(std::move(p), 2));
 733     EXPECT_EQ(2 * sizeof(OwnershipTestClass), buf->length());
 734     EXPECT_EQ(0, destructorCount);
 735   }
 736   EXPECT_EQ(2, destructorCount);
 737
 738   customDeleterCount = 0;
 739   destructorCount = 0;
 740   {
 741     std::unique_ptr<OwnershipTestClass, CustomDeleter>
 742       p(new OwnershipTestClass(), customDelete);
 743     std::unique_ptr<IOBuf> buf(IOBuf::takeOwnership(std::move(p)));
 744     EXPECT_EQ(sizeof(OwnershipTestClass), buf->length());
 745     EXPECT_EQ(0, destructorCount);
 746   }
 747   EXPECT_EQ(1, destructorCount);
 748   EXPECT_EQ(1, customDeleterCount);
 749
 750   customDeleterCount = 0;
 751   destructorCount = 0;
 752   {
 753     std::unique_ptr<OwnershipTestClass[], CustomDeleter>
 754       p(new OwnershipTestClass[2], CustomDeleter(customDeleteArray));
 755     std::unique_ptr<IOBuf> buf(IOBuf::takeOwnership(std::move(p), 2));
 756     EXPECT_EQ(2 * sizeof(OwnershipTestClass), buf->length());
 757     EXPECT_EQ(0, destructorCount);
 758   }
 759   EXPECT_EQ(2, destructorCount);
 760   EXPECT_EQ(1, customDeleterCount);
 761 }
 762
 763 TEST(IOBuf, Alignment) {
 764   // max_align_t doesn't exist in gcc 4.6.2
 765   struct MaxAlign {
 766     char c;
 767   } __attribute__((aligned));
 768   size_t alignment = alignof(MaxAlign);
 769
 770   std::vector<size_t> sizes {0, 1, 64, 256, 1024, 1 << 10};
 771   for (size_t size : sizes) {
 772     auto buf = IOBuf::create(size);
 773     uintptr_t p = reinterpret_cast<uintptr_t>(buf->data());
 774     EXPECT_EQ(0, p & (alignment - 1)) << "size=" << size;
 775   }
 776 }
 777
 778 TEST(TypedIOBuf, Simple) {
 779   auto buf = IOBuf::create(0);
 780   TypedIOBuf<uint64_t> typed(buf.get());
 781   const uint64_t n = 10000;
 782   typed.reserve(0, n);
 783   EXPECT_LE(n, typed.capacity());
 784   for (uint64_t i = 0; i < n; i++) {
 785     *typed.writableTail() = i;
 786     typed.append(1);
 787   }
 788   EXPECT_EQ(n, typed.length());
 789   for (uint64_t i = 0; i < n; i++) {
 790     EXPECT_EQ(i, typed.data()[i]);
 791   }
 792 }
 793 enum BufType {
 794   CREATE,
 795   TAKE_OWNERSHIP_MALLOC,
 796   TAKE_OWNERSHIP_CUSTOM,
 797   USER_OWNED,
 798 };
 799
 800 // chain element size, number of elements in chain, shared
 801 class MoveToFbStringTest
 802   : public ::testing::TestWithParam<std::tr1::tuple<int, int, bool, BufType>> {
 803  protected:
 804   void SetUp() {
 805     std::tr1::tie(elementSize_, elementCount_, shared_, type_) = GetParam();
 806     buf_ = makeBuf();
 807     for (int i = 0; i < elementCount_ - 1; ++i) {
 808       buf_->prependChain(makeBuf());
 809     }
 810     EXPECT_EQ(elementCount_, buf_->countChainElements());
 811     EXPECT_EQ(elementCount_ * elementSize_, buf_->computeChainDataLength());
 812     if (shared_) {
 813       buf2_ = buf_->clone();
 814       EXPECT_EQ(elementCount_, buf2_->countChainElements());
 815       EXPECT_EQ(elementCount_ * elementSize_, buf2_->computeChainDataLength());
 816     }
 817   }
 818
 819   std::unique_ptr<IOBuf> makeBuf() {
 820     unique_ptr<IOBuf> buf;
 821     switch (type_) {
 822       case CREATE:
 823         buf = IOBuf::create(elementSize_);
 824         buf->append(elementSize_);
 825         break;
 826       case TAKE_OWNERSHIP_MALLOC: {
 827         void* data = malloc(elementSize_);
 828         if (!data) {
 829           throw std::bad_alloc();
 830         }
 831         buf = IOBuf::takeOwnership(data, elementSize_);
 832         break;
 833       }
 834       case TAKE_OWNERSHIP_CUSTOM: {
 835         uint8_t* data = new uint8_t[elementSize_];
 836         buf = IOBuf::takeOwnership(data, elementSize_, testFreeFn);
 837         break;
 838       }
 839       case USER_OWNED: {
 840         unique_ptr<uint8_t[]> data(new uint8_t[elementSize_]);
 841         buf = IOBuf::wrapBuffer(data.get(), elementSize_);
 842         ownedBuffers_.emplace_back(std::move(data));
 843         break;
 844       }
 845       default:
 846         throw std::invalid_argument("unexpected buffer type parameter");
 847         break;
 848     }
 849     memset(buf->writableData(), 'x', elementSize_);
 850     return buf;
 851   }
 852
 853   void check(std::unique_ptr<IOBuf>& buf) {
 854     fbstring str = buf->moveToFbString();
 855     EXPECT_EQ(elementCount_ * elementSize_, str.size());
 856     EXPECT_EQ(elementCount_ * elementSize_, strspn(str.c_str(), "x"));
 857     EXPECT_EQ(0, buf->length());
 858     EXPECT_EQ(1, buf->countChainElements());
 859     EXPECT_EQ(0, buf->computeChainDataLength());
 860     EXPECT_FALSE(buf->isChained());
 861   }
 862
 863   int elementSize_;
 864   int elementCount_;
 865   bool shared_;
 866   BufType type_;
 867   std::unique_ptr<IOBuf> buf_;
 868   std::unique_ptr<IOBuf> buf2_;
 869   std::vector<std::unique_ptr<uint8_t[]>> ownedBuffers_;
 870 };
 871
 872 TEST_P(MoveToFbStringTest, Simple) {
 873   check(buf_);
 874   if (shared_) {
 875     check(buf2_);
 876   }
 877 }
 878
 879 INSTANTIATE_TEST_CASE_P(
 880     MoveToFbString,
 881     MoveToFbStringTest,
 882     ::testing::Combine(
 883         ::testing::Values(0, 1, 24, 256, 1 << 10, 1 << 20),  // element size
 884         ::testing::Values(1, 2, 10),                         // element count
 885         ::testing::Bool(),                                   // shared
 886         ::testing::Values(CREATE, TAKE_OWNERSHIP_MALLOC,
 887                           TAKE_OWNERSHIP_CUSTOM, USER_OWNED)));
 888
 889 TEST(IOBuf, getIov) {
 890   uint32_t fillSeed = 0xdeadbeef;
 891   boost::mt19937 gen(fillSeed);
 892
 893   size_t len = 4096;
 894   size_t count = 32;
 895   auto buf = IOBuf::create(len + 1);
 896   buf->append(rand() % len + 1);
 897   fillBuf(buf.get(), gen);
 898
 899   for (size_t i = 0; i < count - 1; i++) {
 900     auto buf2 = IOBuf::create(len + 1);
 901     buf2->append(rand() % len + 1);
 902     fillBuf(buf2.get(), gen);
 903     buf->prependChain(std::move(buf2));
 904   }
 905   EXPECT_EQ(count, buf->countChainElements());
 906
 907   auto iov = buf->getIov();
 908   EXPECT_EQ(count, iov.size());
 909
 910   IOBuf const* p = buf.get();
 911   for (size_t i = 0; i < count; i++, p = p->next()) {
 912     EXPECT_EQ(p->data(), iov[i].iov_base);
 913     EXPECT_EQ(p->length(), iov[i].iov_len);
 914   }
 915
 916   // an empty buf should be skipped in the iov.
 917   buf->next()->clear();
 918   iov = buf->getIov();
 919   EXPECT_EQ(count - 1, iov.size());
 920   EXPECT_EQ(buf->next()->next()->data(), iov[1].iov_base);
 921
 922   // same for the first one being empty
 923   buf->clear();
 924   iov = buf->getIov();
 925   EXPECT_EQ(count - 2, iov.size());
 926   EXPECT_EQ(buf->next()->next()->data(), iov[0].iov_base);
 927
 928   // and the last one
 929   buf->prev()->clear();
 930   iov = buf->getIov();
 931   EXPECT_EQ(count - 3, iov.size());
 932 }
 933
 934 int main(int argc, char** argv) {
 935   testing::InitGoogleTest(&argc, argv);
 936   google::ParseCommandLineFlags(&argc, &argv, true);
 937
 938   return RUN_ALL_TESTS();
 939 }