4a5366ddbe4cb90151c96201aa15becadfd7979c
[folly.git] / folly / io / test / IOBufCursorTest.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
19 #include <gflags/gflags.h>
20 #include <boost/random.hpp>
21 #include <gtest/gtest.h>
22 #include "folly/Benchmark.h"
23 #include "folly/Range.h"
24 #include "folly/io/Cursor.h"
25
26 DECLARE_bool(benchmark);
27
28 using folly::IOBuf;
29 using std::unique_ptr;
30 using namespace folly::io;
31
32 TEST(IOBuf, RWCursor) {
33   unique_ptr<IOBuf> iobuf1(IOBuf::create(20));
34   iobuf1->append(20);
35   unique_ptr<IOBuf> iobuf2(IOBuf::create(20));
36   iobuf2->append(20);
37
38   IOBuf* iob2ptr = iobuf2.get();
39   iobuf1->prependChain(std::move(iobuf2));
40
41   EXPECT_TRUE(iobuf1->isChained());
42
43   RWPrivateCursor wcursor(iobuf1.get());
44   Cursor rcursor(iobuf1.get());
45   wcursor.writeLE((uint64_t)1);
46   wcursor.writeLE((uint64_t)1);
47   wcursor.writeLE((uint64_t)1);
48   wcursor.write((uint8_t)1);
49
50   EXPECT_EQ(1, rcursor.readLE<uint64_t>());
51   rcursor.skip(8);
52   EXPECT_EQ(1, rcursor.readLE<uint32_t>());
53   rcursor.skip(0);
54   EXPECT_EQ(0, rcursor.read<uint8_t>());
55   EXPECT_EQ(0, rcursor.read<uint8_t>());
56   EXPECT_EQ(0, rcursor.read<uint8_t>());
57   EXPECT_EQ(0, rcursor.read<uint8_t>());
58   EXPECT_EQ(1, rcursor.read<uint8_t>());
59 }
60
61 TEST(IOBuf, skip) {
62   unique_ptr<IOBuf> iobuf1(IOBuf::create(20));
63   iobuf1->append(20);
64   RWPrivateCursor wcursor(iobuf1.get());
65   wcursor.write((uint8_t)1);
66   wcursor.write((uint8_t)2);
67   Cursor cursor(iobuf1.get());
68   cursor.skip(1);
69   EXPECT_EQ(2, cursor.read<uint8_t>());
70 }
71
72 TEST(IOBuf, reset) {
73   unique_ptr<IOBuf> iobuf1(IOBuf::create(20));
74   iobuf1->append(20);
75   RWPrivateCursor wcursor(iobuf1.get());
76   wcursor.write((uint8_t)1);
77   wcursor.write((uint8_t)2);
78   wcursor.reset(iobuf1.get());
79   EXPECT_EQ(1, wcursor.read<uint8_t>());
80 }
81
82 TEST(IOBuf, copy_assign_convert) {
83   unique_ptr<IOBuf> iobuf1(IOBuf::create(20));
84   iobuf1->append(20);
85   RWPrivateCursor wcursor(iobuf1.get());
86   RWPrivateCursor cursor2(wcursor);
87   RWPrivateCursor cursor3(iobuf1.get());
88
89   wcursor.write((uint8_t)1);
90   cursor3 = wcursor;
91   wcursor.write((uint8_t)2);
92   Cursor cursor4(wcursor);
93   RWPrivateCursor cursor5(wcursor);
94   wcursor.write((uint8_t)3);
95
96   EXPECT_EQ(1, cursor2.read<uint8_t>());
97   EXPECT_EQ(2, cursor3.read<uint8_t>());
98   EXPECT_EQ(3, cursor4.read<uint8_t>());
99 }
100
101 TEST(IOBuf, overloading) {
102   unique_ptr<IOBuf> iobuf1(IOBuf::create(20));
103   iobuf1->append(20);
104   RWPrivateCursor wcursor(iobuf1.get());
105   wcursor += 1;
106   wcursor.write((uint8_t)1);
107   Cursor cursor(iobuf1.get());
108   cursor += 1;
109   EXPECT_EQ(1, cursor.read<uint8_t>());
110 }
111
112 TEST(IOBuf, endian) {
113   unique_ptr<IOBuf> iobuf1(IOBuf::create(20));
114   iobuf1->append(20);
115   RWPrivateCursor wcursor(iobuf1.get());
116   Cursor rcursor(iobuf1.get());
117   uint16_t v = 1;
118   int16_t vu = -1;
119   wcursor.writeBE(v);
120   wcursor.writeBE(vu);
121   // Try a couple combinations to ensure they were generated correctly
122   wcursor.writeBE(vu);
123   wcursor.writeLE(vu);
124   wcursor.writeLE(vu);
125   wcursor.writeLE(v);
126   EXPECT_EQ(v, rcursor.readBE<uint16_t>());
127 }
128
129 TEST(IOBuf, Cursor) {
130   unique_ptr<IOBuf> iobuf1(IOBuf::create(1));
131   iobuf1->append(1);
132   RWPrivateCursor c(iobuf1.get());
133   c.write((uint8_t)40); // OK
134   try {
135     c.write((uint8_t)10); // Bad write, checked should except.
136     EXPECT_EQ(true, false);
137   } catch (...) {
138   }
139 }
140
141 TEST(IOBuf, UnshareCursor) {
142   uint8_t buf = 0;
143   unique_ptr<IOBuf> iobuf1(IOBuf::wrapBuffer(&buf, 1));
144   unique_ptr<IOBuf> iobuf2(IOBuf::wrapBuffer(&buf, 1));
145   RWUnshareCursor c1(iobuf1.get());
146   RWUnshareCursor c2(iobuf2.get());
147
148   c1.write((uint8_t)10); // This should duplicate the two buffers.
149   uint8_t t = c2.read<uint8_t>();
150   EXPECT_EQ(0, t);
151
152   iobuf1 = IOBuf::wrapBuffer(&buf, 1);
153   iobuf2 = IOBuf::wrapBuffer(&buf, 1);
154   RWPrivateCursor c3(iobuf1.get());
155   RWPrivateCursor c4(iobuf2.get());
156
157   c3.write((uint8_t)10); // This should _not_ duplicate the two buffers.
158   t = c4.read<uint8_t>();
159   EXPECT_EQ(10, t);
160 }
161
162 namespace {
163 void append(std::unique_ptr<IOBuf>& buf, folly::StringPiece data) {
164   EXPECT_LE(data.size(), buf->tailroom());
165   memcpy(buf->writableData(), data.data(), data.size());
166   buf->append(data.size());
167 }
168
169 void append(Appender& appender, folly::StringPiece data) {
170   appender.push(reinterpret_cast<const uint8_t*>(data.data()), data.size());
171 }
172
173 std::string toString(const IOBuf& buf) {
174   std::string str;
175   Cursor cursor(&buf);
176   std::pair<const uint8_t*, size_t> p;
177   while ((p = cursor.peek()).second) {
178     str.append(reinterpret_cast<const char*>(p.first), p.second);
179     cursor.skip(p.second);
180   }
181   return str;
182 }
183
184 }  // namespace
185
186 TEST(IOBuf, PullAndPeek) {
187   std::unique_ptr<IOBuf> iobuf1(IOBuf::create(10));
188   append(iobuf1, "he");
189   std::unique_ptr<IOBuf> iobuf2(IOBuf::create(10));
190   append(iobuf2, "llo ");
191   std::unique_ptr<IOBuf> iobuf3(IOBuf::create(10));
192   append(iobuf3, "world");
193   iobuf1->prependChain(std::move(iobuf2));
194   iobuf1->prependChain(std::move(iobuf3));
195   EXPECT_EQ(3, iobuf1->countChainElements());
196   EXPECT_EQ(11, iobuf1->computeChainDataLength());
197
198   char buf[12];
199   memset(buf, 0, sizeof(buf));
200   Cursor(iobuf1.get()).pull(buf, 11);
201   EXPECT_EQ("hello world", std::string(buf));
202
203   memset(buf, 0, sizeof(buf));
204   EXPECT_EQ(11, Cursor(iobuf1.get()).pullAtMost(buf, 20));
205   EXPECT_EQ("hello world", std::string(buf));
206
207   EXPECT_THROW({Cursor(iobuf1.get()).pull(buf, 20);},
208                std::out_of_range);
209
210   {
211     RWPrivateCursor cursor(iobuf1.get());
212     auto p = cursor.peek();
213     EXPECT_EQ("he", std::string(reinterpret_cast<const char*>(p.first),
214                                 p.second));
215     cursor.skip(p.second);
216     p = cursor.peek();
217     EXPECT_EQ("llo ", std::string(reinterpret_cast<const char*>(p.first),
218                                   p.second));
219     cursor.skip(p.second);
220     p = cursor.peek();
221     EXPECT_EQ("world", std::string(reinterpret_cast<const char*>(p.first),
222                                    p.second));
223     cursor.skip(p.second);
224     EXPECT_EQ(3, iobuf1->countChainElements());
225     EXPECT_EQ(11, iobuf1->computeChainDataLength());
226   }
227
228   {
229     RWPrivateCursor cursor(iobuf1.get());
230     cursor.gather(11);
231     auto p = cursor.peek();
232     EXPECT_EQ("hello world", std::string(reinterpret_cast<const
233                                          char*>(p.first), p.second));
234     EXPECT_EQ(1, iobuf1->countChainElements());
235     EXPECT_EQ(11, iobuf1->computeChainDataLength());
236   }
237 }
238
239 TEST(IOBuf, cloneAndInsert) {
240   std::unique_ptr<IOBuf> iobuf1(IOBuf::create(10));
241   append(iobuf1, "he");
242   std::unique_ptr<IOBuf> iobuf2(IOBuf::create(10));
243   append(iobuf2, "llo ");
244   std::unique_ptr<IOBuf> iobuf3(IOBuf::create(10));
245   append(iobuf3, "world");
246   iobuf1->prependChain(std::move(iobuf2));
247   iobuf1->prependChain(std::move(iobuf3));
248   EXPECT_EQ(3, iobuf1->countChainElements());
249   EXPECT_EQ(11, iobuf1->computeChainDataLength());
250
251   std::unique_ptr<IOBuf> cloned;
252
253   Cursor(iobuf1.get()).clone(cloned, 3);
254   EXPECT_EQ(2, cloned->countChainElements());
255   EXPECT_EQ(3, cloned->computeChainDataLength());
256
257
258   EXPECT_EQ(11, Cursor(iobuf1.get()).cloneAtMost(cloned, 20));
259   EXPECT_EQ(3, cloned->countChainElements());
260   EXPECT_EQ(11, cloned->computeChainDataLength());
261
262
263   EXPECT_THROW({Cursor(iobuf1.get()).clone(cloned, 20);},
264                std::out_of_range);
265
266   {
267     // Check that inserting in the middle of an iobuf splits
268     RWPrivateCursor cursor(iobuf1.get());
269     Cursor(iobuf1.get()).clone(cloned, 3);
270     EXPECT_EQ(2, cloned->countChainElements());
271     EXPECT_EQ(3, cloned->computeChainDataLength());
272
273     cursor.skip(1);
274
275     cursor.insert(std::move(cloned));
276     EXPECT_EQ(6, iobuf1->countChainElements());
277     EXPECT_EQ(14, iobuf1->computeChainDataLength());
278     // Check that nextBuf got set correctly
279     cursor.read<uint8_t>();
280   }
281
282   {
283     // Check that inserting at the end doesn't create empty buf
284     RWPrivateCursor cursor(iobuf1.get());
285     Cursor(iobuf1.get()).clone(cloned, 1);
286     EXPECT_EQ(1, cloned->countChainElements());
287     EXPECT_EQ(1, cloned->computeChainDataLength());
288
289     cursor.skip(1);
290
291     cursor.insert(std::move(cloned));
292     EXPECT_EQ(7, iobuf1->countChainElements());
293     EXPECT_EQ(15, iobuf1->computeChainDataLength());
294     // Check that nextBuf got set correctly
295     cursor.read<uint8_t>();
296   }
297   {
298     // Check that inserting at the beginning doesn't create empty buf
299     RWPrivateCursor cursor(iobuf1.get());
300     Cursor(iobuf1.get()).clone(cloned, 1);
301     EXPECT_EQ(1, cloned->countChainElements());
302     EXPECT_EQ(1, cloned->computeChainDataLength());
303
304     cursor.insert(std::move(cloned));
305     EXPECT_EQ(8, iobuf1->countChainElements());
306     EXPECT_EQ(16, iobuf1->computeChainDataLength());
307     // Check that nextBuf got set correctly
308     cursor.read<uint8_t>();
309   }
310 }
311
312 TEST(IOBuf, Appender) {
313   std::unique_ptr<IOBuf> head(IOBuf::create(10));
314   append(head, "hello");
315
316   Appender app(head.get(), 10);
317   uint32_t cap = head->capacity();
318   uint32_t len1 = app.length();
319   EXPECT_EQ(cap - 5, len1);
320   app.ensure(len1);  // won't grow
321   EXPECT_EQ(len1, app.length());
322   app.ensure(len1 + 1);  // will grow
323   EXPECT_LE(len1 + 1, app.length());
324
325   append(app, " world");
326   EXPECT_EQ("hello world", toString(*head));
327 }
328
329 TEST(IOBuf, QueueAppender) {
330   folly::IOBufQueue queue;
331
332   // Allocate 100 bytes at once, but don't grow past 1024
333   QueueAppender app(&queue, 100, 1024);
334   size_t n = 1024 / sizeof(uint32_t);
335   for (uint32_t i = 0; i < n; ++i) {
336     app.writeBE(i);
337   }
338
339   EXPECT_THROW({app.writeBE(0);}, std::out_of_range);
340
341   // There must be a goodMallocSize between 100 and 1024...
342   EXPECT_LT(1, queue.front()->countChainElements());
343   const IOBuf* buf = queue.front();
344   do {
345     EXPECT_LE(100, buf->capacity());
346     buf = buf->next();
347   } while (buf != queue.front());
348
349   Cursor cursor(queue.front());
350   for (uint32_t i = 0; i < n; ++i) {
351     EXPECT_EQ(i, cursor.readBE<uint32_t>());
352   }
353
354   EXPECT_THROW({cursor.readBE<uint32_t>();}, std::out_of_range);
355 }
356
357 TEST(IOBuf, CursorOperators) {
358   // Test operators on a single-item chain
359   {
360     std::unique_ptr<IOBuf> chain1(IOBuf::create(20));
361     chain1->append(10);
362
363     Cursor curs1(chain1.get());
364     EXPECT_EQ(0, curs1 - chain1.get());
365     curs1.skip(3);
366     EXPECT_EQ(3, curs1 - chain1.get());
367     curs1.skip(7);
368     EXPECT_EQ(10, curs1 - chain1.get());
369
370     Cursor curs2(chain1.get());
371     EXPECT_EQ(0, curs2 - chain1.get());
372     EXPECT_EQ(10, curs1 - curs2);
373     EXPECT_THROW(curs2 - curs1, std::out_of_range);
374   }
375
376   // Test cross-chain operations
377   {
378     std::unique_ptr<IOBuf> chain1(IOBuf::create(20));
379     chain1->append(10);
380     std::unique_ptr<IOBuf> chain2 = chain1->clone();
381
382     Cursor curs1(chain1.get());
383     Cursor curs2(chain2.get());
384     EXPECT_THROW(curs1 - curs2, std::out_of_range);
385     EXPECT_THROW(curs1 - chain2.get(), std::out_of_range);
386   }
387
388   // Test operations on multi-item chains
389   {
390     std::unique_ptr<IOBuf> chain(IOBuf::create(20));
391     chain->append(10);
392     chain->appendChain(chain->clone());
393     EXPECT_EQ(20, chain->computeChainDataLength());
394
395     Cursor curs1(chain.get());
396     curs1.skip(5);
397     Cursor curs2(chain.get());
398     curs2.skip(3);
399     EXPECT_EQ(2, curs1 - curs2);
400     EXPECT_EQ(5, curs1 - chain.get());
401     EXPECT_THROW(curs2 - curs1, std::out_of_range);
402
403     curs1.skip(7);
404     EXPECT_EQ(9, curs1 - curs2);
405     EXPECT_EQ(12, curs1 - chain.get());
406     EXPECT_THROW(curs2 - curs1, std::out_of_range);
407
408     curs2.skip(7);
409     EXPECT_EQ(2, curs1 - curs2);
410     EXPECT_THROW(curs2 - curs1, std::out_of_range);
411   }
412 }
413
414 TEST(IOBuf, StringOperations) {
415   // Test a single buffer with two null-terminated strings and an extra uint8_t
416   // at the end
417   {
418     std::unique_ptr<IOBuf> chain(IOBuf::create(16));
419     Appender app(chain.get(), 0);
420     app.pushAtMost(reinterpret_cast<const uint8_t*>("hello\0world\0\x01"), 13);
421
422     Cursor curs(chain.get());
423     EXPECT_STREQ("hello", curs.readTerminatedString().c_str());
424     EXPECT_STREQ("world", curs.readTerminatedString().c_str());
425     EXPECT_EQ(1, curs.read<uint8_t>());
426   }
427
428   // Test multiple buffers with a single null-terminated string spanning them
429   {
430     std::unique_ptr<IOBuf> chain(IOBuf::create(8));
431     chain->prependChain(IOBuf::create(8));
432     Appender app(chain.get(), 0);
433     app.pushAtMost(reinterpret_cast<const uint8_t*>("hello world\0"), 12);
434
435     Cursor curs(chain.get());
436     EXPECT_STREQ("hello world", curs.readTerminatedString().c_str());
437   }
438
439   // Test a reading a null-terminated string that's longer than the maximum
440   // allowable length
441   {
442     std::unique_ptr<IOBuf> chain(IOBuf::create(16));
443     Appender app(chain.get(), 0);
444     app.pushAtMost(reinterpret_cast<const uint8_t*>("hello world\0"), 12);
445
446     Cursor curs(chain.get());
447     EXPECT_THROW(curs.readTerminatedString('\0', 5), std::length_error);
448   }
449
450   // Test reading a null-termianted string from a chain with an empty buffer at
451   // the front
452   {
453     std::unique_ptr<IOBuf> buf(IOBuf::create(8));
454     Appender app(buf.get(), 0);
455     app.pushAtMost(reinterpret_cast<const uint8_t*>("hello\0"), 6);
456     std::unique_ptr<IOBuf> chain(IOBuf::create(8));
457     chain->prependChain(std::move(buf));
458
459     Cursor curs(chain.get());
460     EXPECT_STREQ("hello", curs.readTerminatedString().c_str());
461   }
462
463   // Test reading a two fixed-length strings from a single buffer with an extra
464   // uint8_t at the end
465   {
466     std::unique_ptr<IOBuf> chain(IOBuf::create(16));
467     Appender app(chain.get(), 0);
468     app.pushAtMost(reinterpret_cast<const uint8_t*>("helloworld\x01"), 11);
469
470     Cursor curs(chain.get());
471     EXPECT_STREQ("hello", curs.readFixedString(5).c_str());
472     EXPECT_STREQ("world", curs.readFixedString(5).c_str());
473     EXPECT_EQ(1, curs.read<uint8_t>());
474   }
475
476   // Test multiple buffers with a single fixed-length string spanning them
477   {
478     std::unique_ptr<IOBuf> chain(IOBuf::create(8));
479     chain->prependChain(IOBuf::create(8));
480     Appender app(chain.get(), 0);
481     app.pushAtMost(reinterpret_cast<const uint8_t*>("hello world"), 11);
482
483     Cursor curs(chain.get());
484     EXPECT_STREQ("hello world", curs.readFixedString(11).c_str());
485   }
486
487   // Test reading a fixed-length string from a chain with an empty buffer at
488   // the front
489   {
490     std::unique_ptr<IOBuf> buf(IOBuf::create(8));
491     Appender app(buf.get(), 0);
492     app.pushAtMost(reinterpret_cast<const uint8_t*>("hello"), 5);
493     std::unique_ptr<IOBuf> chain(IOBuf::create(8));
494     chain->prependChain(std::move(buf));
495
496     Cursor curs(chain.get());
497     EXPECT_STREQ("hello", curs.readFixedString(5).c_str());
498   }
499 }
500
501 int benchmark_size = 1000;
502 unique_ptr<IOBuf> iobuf_benchmark;
503
504 unique_ptr<IOBuf> iobuf_read_benchmark;
505
506 template <class CursClass>
507 void runBenchmark() {
508   CursClass c(iobuf_benchmark.get());
509
510   for(int i = 0; i < benchmark_size; i++) {
511     c.write((uint8_t)0);
512   }
513 }
514
515 BENCHMARK(rwPrivateCursorBenchmark, iters) {
516   while (--iters) {
517     runBenchmark<RWPrivateCursor>();
518   }
519 }
520
521 BENCHMARK(rwUnshareCursorBenchmark, iters) {
522   while (--iters) {
523     runBenchmark<RWUnshareCursor>();
524   }
525 }
526
527
528 BENCHMARK(cursorBenchmark, iters) {
529   while (--iters) {
530     Cursor c(iobuf_read_benchmark.get());
531     for(int i = 0; i < benchmark_size ; i++) {
532       c.read<uint8_t>();
533     }
534   }
535 }
536
537 BENCHMARK(skipBenchmark, iters) {
538   uint8_t buf;
539   while (--iters) {
540     Cursor c(iobuf_read_benchmark.get());
541     for(int i = 0; i < benchmark_size ; i++) {
542       c.peek();
543       c.skip(1);
544     }
545   }
546 }
547
548 // fbmake opt
549 // _bin/folly/experimental/io/test/iobuf_cursor_test -benchmark
550 //
551 // Benchmark                               Iters   Total t    t/iter iter/sec
552 // ---------------------------------------------------------------------------
553 // rwPrivateCursorBenchmark               100000  142.9 ms  1.429 us  683.5 k
554 // rwUnshareCursorBenchmark               100000  309.3 ms  3.093 us  315.7 k
555 // cursorBenchmark                        100000  741.4 ms  7.414 us  131.7 k
556 // skipBenchmark                          100000  738.9 ms  7.389 us  132.2 k
557 //
558 // uname -a:
559 //
560 // Linux dev2159.snc6.facebook.com 2.6.33-7_fbk15_104e4d0 #1 SMP
561 // Tue Oct 19 22:40:30 PDT 2010 x86_64 x86_64 x86_64 GNU/Linux
562 //
563 // 72GB RAM, 2 CPUs (Intel(R) Xeon(R) CPU L5630  @ 2.13GHz)
564 // hyperthreading disabled
565
566 int main(int argc, char** argv) {
567   testing::InitGoogleTest(&argc, argv);
568   google::ParseCommandLineFlags(&argc, &argv, true);
569
570   auto ret = RUN_ALL_TESTS();
571
572   if (ret == 0 && FLAGS_benchmark) {
573     iobuf_benchmark = IOBuf::create(benchmark_size);
574     iobuf_benchmark->append(benchmark_size);
575
576     iobuf_read_benchmark = IOBuf::create(1);
577     for (int i = 0; i < benchmark_size; i++) {
578       unique_ptr<IOBuf> iobuf2(IOBuf::create(1));
579       iobuf2->append(1);
580       iobuf_read_benchmark->prependChain(std::move(iobuf2));
581     }
582
583     folly::runBenchmarks();
584   }
585
586   return ret;
587 }