IOBuf::getIov

Summary:
Generate an `fbvector` of `struct iovec` suitable for using with `writev` or
`sendmsg`.

This code is pretty straightforward, but Adam pointed out that something along
these lines has already been done in thrift, so I followed that code closely.
http://fburl.com/11586814

Test Plan:
fbmake runtests

I am using also this in a prototype and it's working there.

Reviewed By: agartrell@fb.com

FB internal diff: D744055

diff --git a/folly/io/IOBuf.cpp b/folly/io/IOBuf.cpp
index aac0be5154798b04bf44c36d8d43553bafe9a6f1..e9b9151b13dc93a29d1e31a715f4c4994bec7d48 100644 (file)
--- a/folly/io/IOBuf.cpp
+++ b/folly/io/IOBuf.cpp
@@ -643,4 +643,18 @@ IOBuf::Iterator IOBuf::cend() const {
   return Iterator(nullptr, nullptr);
 }
 
+folly::fbvector<struct iovec> IOBuf::getIov() const {
+  folly::fbvector<struct iovec> iov;
+  iov.reserve(countChainElements());
+  IOBuf const* p = this;
+  do {
+    // some code can get confused by empty iovs, so skip them
+    if (p->length() > 0) {
+      iov.push_back({(void*)p->data(), p->length()});
+    }
+    p = p->next();
+  } while (p != this);
+  return iov;
+}
+
 } // folly

diff --git a/folly/io/IOBuf.h b/folly/io/IOBuf.h
index 90dcb1e3cfc658de5f104e92bd297825077bc8f8..007c38414279c5a3e3e49bb16210f3a4a32a45ba 100644 (file)
--- a/folly/io/IOBuf.h
+++ b/folly/io/IOBuf.h
@@ -25,12 +25,14 @@
 #include <cstring>
 #include <memory>
 #include <limits>
+#include <sys/uio.h>
 #include <type_traits>
 
 #include <boost/iterator/iterator_facade.hpp>
 
 #include "folly/FBString.h"
 #include "folly/Range.h"
+#include "folly/FBVector.h"
 
 namespace folly {
 
@@ -922,6 +924,17 @@ class IOBuf {
    */
   std::unique_ptr<IOBuf> cloneOne() const;
 
+  /**
+   * Return an iovector suitable for e.g. writev()
+   *
+   *   auto iov = buf->getIov();
+   *   auto xfer = writev(fd, iov.data(), iov.size());
+   *
+   * Naturally, the returned iovector is invalid if you modify the buffer
+   * chain.
+   */
+  folly::fbvector<struct iovec> getIov() const;
+
   // Overridden operator new and delete.
   // These directly use malloc() and free() to allocate the space for IOBuf
   // objects.  This is needed since IOBuf::create() manually uses malloc when

diff --git a/folly/io/test/IOBufTest.cpp b/folly/io/test/IOBufTest.cpp
index 3bb86cf17927e17e10d84c95715486c01abffb5f..b17be70887072539d06b5fadbe6e8f258a719a2f 100644 (file)
--- a/folly/io/test/IOBufTest.cpp
+++ b/folly/io/test/IOBufTest.cpp
@@ -28,6 +28,7 @@
 #include "folly/Range.h"
 
 using folly::fbstring;
+using folly::fbvector;
 using folly::IOBuf;
 using folly::TypedIOBuf;
 using folly::StringPiece;
@@ -764,6 +765,51 @@ INSTANTIATE_TEST_CASE_P(
         ::testing::Values(1, 2, 10),                         // element count
         ::testing::Bool()));                                 // shared
 
+TEST(IOBuf, getIov) {
+  uint32_t fillSeed = 0xdeadbeef;
+  boost::mt19937 gen(fillSeed);
+
+  size_t len = 4096;
+  size_t count = 32;
+  auto buf = IOBuf::create(len + 1);
+  buf->append(rand() % len + 1);
+  fillBuf(buf.get(), gen);
+
+  for (size_t i = 0; i < count - 1; i++) {
+    auto buf2 = IOBuf::create(len + 1);
+    buf2->append(rand() % len + 1);
+    fillBuf(buf2.get(), gen);
+    buf->prependChain(std::move(buf2));
+  }
+  EXPECT_EQ(count, buf->countChainElements());
+
+  auto iov = buf->getIov();
+  EXPECT_EQ(count, iov.size());
+
+  IOBuf const* p = buf.get();
+  for (size_t i = 0; i < count; i++, p = p->next()) {
+    EXPECT_EQ(p->data(), iov[i].iov_base);
+    EXPECT_EQ(p->length(), iov[i].iov_len);
+  }
+
+  // an empty buf should be skipped in the iov.
+  buf->next()->clear();
+  iov = buf->getIov();
+  EXPECT_EQ(count - 1, iov.size());
+  EXPECT_EQ(buf->next()->next()->data(), iov[1].iov_base);
+
+  // same for the first one being empty
+  buf->clear();
+  iov = buf->getIov();
+  EXPECT_EQ(count - 2, iov.size());
+  EXPECT_EQ(buf->next()->next()->data(), iov[0].iov_base);
+
+  // and the last one
+  buf->prev()->clear();
+  iov = buf->getIov();
+  EXPECT_EQ(count - 3, iov.size());
+}
+
 int main(int argc, char** argv) {
   testing::InitGoogleTest(&argc, argv);
   google::ParseCommandLineFlags(&argc, &argv, true);