Add tryRead() and endian variants

Summary:
Add a `tryRead()`, and endian variants, which try to read into an arithmetic type, and if there isn't enough data they return false.
One use case is to quickly check if an IOBuf starts with a certain prefix, benchmarks show that using `tryReadLE()` is 6x faster than using `pullAtMost()` and `memcmp()`.

Reviewed By: yfeldblum

Differential Revision: D4767855

fbshipit-source-id: feb8c61092772933d4b8496b27d464559ff8b827

diff --git a/folly/io/Cursor.h b/folly/io/Cursor.h
index 32d65605262905be0b2be9f400ee15ce4b767683..516b10393be3dae1c050651676584706d5a36e57 100644 (file)
--- a/folly/io/Cursor.h
+++ b/folly/io/Cursor.h
@@ -200,14 +200,36 @@ class CursorBase {
   }
 
   template <class T>
-  typename std::enable_if<std::is_arithmetic<T>::value, T>::type read() {
-    T val;
+  typename std::enable_if<std::is_arithmetic<T>::value, bool>::type tryRead(
+      T& val) {
     if (LIKELY(length() >= sizeof(T))) {
       val = loadUnaligned<T>(data());
       offset_ += sizeof(T);
       advanceBufferIfEmpty();
-    } else {
-      pullSlow(&val, sizeof(T));
+      return true;
+    }
+    return pullAtMostSlow(&val, sizeof(T)) == sizeof(T);
+  }
+
+  template <class T>
+  bool tryReadBE(T& val) {
+    const bool result = tryRead(val);
+    val = Endian::big(val);
+    return result;
+  }
+
+  template <class T>
+  bool tryReadLE(T& val) {
+    const bool result = tryRead(val);
+    val = Endian::little(val);
+    return result;
+  }
+
+  template <class T>
+  T read() {
+    T val;
+    if (!tryRead(val)) {
+      std::__throw_out_of_range("underflow");
     }
     return val;
   }

diff --git a/folly/io/test/IOBufCursorBenchmark.cpp b/folly/io/test/IOBufCursorBenchmark.cpp
index c828ba64189056880a24942d59c006f7ab250a2c..3811a612783f4fe9acee55acd2fa322d433ae278 100644 (file)
--- a/folly/io/test/IOBufCursorBenchmark.cpp
+++ b/folly/io/test/IOBufCursorBenchmark.cpp
@@ -85,23 +85,83 @@ BENCHMARK(cloneBenchmark, iters) {
   }
 }
 
-// fbmake opt
-// _bin/folly/experimental/io/test/iobuf_cursor_test -benchmark
-//
-// Benchmark                               Iters   Total t    t/iter iter/sec
-// ---------------------------------------------------------------------------
-// rwPrivateCursorBenchmark               100000  142.9 ms  1.429 us  683.5 k
-// rwUnshareCursorBenchmark               100000  309.3 ms  3.093 us  315.7 k
-// cursorBenchmark                        100000  741.4 ms  7.414 us  131.7 k
-// skipBenchmark                          100000  738.9 ms  7.389 us  132.2 k
-//
-// uname -a:
-//
-// Linux dev2159.snc6.facebook.com 2.6.33-7_fbk15_104e4d0 #1 SMP
-// Tue Oct 19 22:40:30 PDT 2010 x86_64 x86_64 x86_64 GNU/Linux
-//
-// 72GB RAM, 2 CPUs (Intel(R) Xeon(R) CPU L5630  @ 2.13GHz)
-// hyperthreading disabled
+BENCHMARK(read, iters) {
+  while (iters--) {
+    Cursor c(iobuf_read_benchmark.get());
+    for (int i = 0; i < benchmark_size; ++i) {
+      const auto val = c.read<uint8_t>();
+      folly::doNotOptimizeAway(val);
+    }
+  }
+}
+
+BENCHMARK(readSlow, iters) {
+  while (iters--) {
+    Cursor c(iobuf_read_benchmark.get());
+    const int size = benchmark_size / 2;
+    for (int i = 0; i < size; ++i) {
+      const auto val = c.read<uint16_t>();
+      folly::doNotOptimizeAway(val);
+    }
+  }
+}
+
+bool prefixBaseline(Cursor& c, const std::array<uint8_t, 4>& expected) {
+  std::array<uint8_t, 4> actual;
+  if (c.pullAtMost(actual.data(), actual.size()) != actual.size()) {
+    return false;
+  }
+  return memcmp(actual.data(), expected.data(), actual.size()) == 0;
+}
+
+bool prefix(Cursor& c, uint32_t expected) {
+  uint32_t actual;
+  if (!c.tryReadLE(actual)) {
+    return false;
+  }
+  return actual == expected;
+}
+
+BENCHMARK(prefixBaseline, iters) {
+  IOBuf buf{IOBuf::CREATE, 10};
+  buf.append(10);
+  constexpr std::array<uint8_t, 4> prefix = {{0x01, 0x02, 0x03, 0x04}};
+  while (iters--) {
+    for (int i = 0; i < benchmark_size; ++i) {
+      Cursor c(&buf);
+      bool result = prefixBaseline(c, prefix);
+      folly::doNotOptimizeAway(result);
+    }
+  }
+}
+
+BENCHMARK_RELATIVE(prefix, iters) {
+  IOBuf buf{IOBuf::CREATE, 10};
+  buf.append(10);
+  while (iters--) {
+    for (int i = 0; i < benchmark_size; ++i) {
+      Cursor c(&buf);
+      bool result = prefix(c, 0x01020304);
+      folly::doNotOptimizeAway(result);
+    }
+  }
+}
+
+/**
+ * ============================================================================
+ * folly/io/test/IOBufCursorBenchmark.cpp          relative  time/iter  iters/s
+ * ============================================================================
+ * rwPrivateCursorBenchmark                                     1.01us  985.85K
+ * rwUnshareCursorBenchmark                                     1.01us  986.70K
+ * cursorBenchmark                                              4.77us  209.61K
+ * skipBenchmark                                                4.78us  209.42K
+ * cloneBenchmark                                              26.65us   37.52K
+ * read                                                         4.35us  230.07K
+ * readSlow                                                     5.45us  183.48K
+ * prefixBaseline                                               6.44us  155.24K
+ * prefix                                           589.31%     1.09us  914.87K
+ * ============================================================================
+ */
 
 int main(int argc, char** argv) {
   gflags::ParseCommandLineFlags(&argc, &argv, true);

diff --git a/folly/io/test/IOBufCursorTest.cpp b/folly/io/test/IOBufCursorTest.cpp
index 36e3ab554b1aee6025ddbbccbbd29a8fcaa305ac..36002b2a7105af2147f52da85571a7ce51742160 100644 (file)
--- a/folly/io/test/IOBufCursorTest.cpp
+++ b/folly/io/test/IOBufCursorTest.cpp
@@ -1032,3 +1032,86 @@ TEST(IOBuf, TestRetreatOperators) {
   EXPECT_EQ(retreated.totalLength(), 5);
   EXPECT_EQ(curs.totalLength(), 0);
 }
+
+TEST(IOBuf, tryRead) {
+  unique_ptr<IOBuf> iobuf1(IOBuf::create(6));
+  iobuf1->append(6);
+  unique_ptr<IOBuf> iobuf2(IOBuf::create(24));
+  iobuf2->append(24);
+
+  iobuf1->prependChain(std::move(iobuf2));
+
+  EXPECT_TRUE(iobuf1->isChained());
+
+  RWPrivateCursor wcursor(iobuf1.get());
+  Cursor rcursor(iobuf1.get());
+  wcursor.writeLE((uint32_t)1);
+  wcursor.writeLE((uint64_t)1);
+  wcursor.writeLE((uint64_t)1);
+  wcursor.writeLE((uint64_t)1);
+  wcursor.writeLE((uint16_t)1);
+  EXPECT_EQ(0, wcursor.totalLength());
+
+  EXPECT_EQ(1u, rcursor.readLE<uint32_t>());
+
+  EXPECT_EQ(1u, rcursor.readLE<uint32_t>());
+  EXPECT_EQ(0u, rcursor.readLE<uint32_t>());
+
+  EXPECT_EQ(1u, rcursor.readLE<uint32_t>());
+  rcursor.skip(4);
+
+  uint32_t val;
+  EXPECT_TRUE(rcursor.tryRead(val));
+  EXPECT_EQ(1, val);
+  EXPECT_TRUE(rcursor.tryRead(val));
+
+  EXPECT_EQ(0, val);
+  EXPECT_FALSE(rcursor.tryRead(val));
+}
+
+TEST(IOBuf, tryReadLE) {
+  IOBuf buf{IOBuf::CREATE, 4};
+  buf.append(4);
+
+  RWPrivateCursor wcursor(&buf);
+  Cursor rcursor(&buf);
+
+  const uint32_t expected = 0x01020304;
+  wcursor.writeLE(expected);
+  uint32_t actual;
+  EXPECT_TRUE(rcursor.tryReadLE(actual));
+  EXPECT_EQ(expected, actual);
+}
+
+TEST(IOBuf, tryReadBE) {
+  IOBuf buf{IOBuf::CREATE, 4};
+  buf.append(4);
+
+  RWPrivateCursor wcursor(&buf);
+  Cursor rcursor(&buf);
+
+  const uint32_t expected = 0x01020304;
+  wcursor.writeBE(expected);
+  uint32_t actual;
+  EXPECT_TRUE(rcursor.tryReadBE(actual));
+  EXPECT_EQ(expected, actual);
+}
+
+TEST(IOBuf, tryReadConsumesAllInputOnFailure) {
+  IOBuf buf{IOBuf::CREATE, 2};
+  buf.append(2);
+
+  Cursor rcursor(&buf);
+  uint32_t val;
+  EXPECT_FALSE(rcursor.tryRead(val));
+  EXPECT_EQ(0, rcursor.totalLength());
+}
+
+TEST(IOBuf, readConsumesAllInputOnFailure) {
+  IOBuf buf{IOBuf::CREATE, 2};
+  buf.append(2);
+
+  Cursor rcursor(&buf);
+  EXPECT_THROW(rcursor.read<uint32_t>(), std::out_of_range);
+  EXPECT_EQ(0, rcursor.totalLength());
+}