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

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;
   }

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);

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());
+}