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