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