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