Add helper to wrap iovecs

Summary:
Add utility to IOBuf to wrap iovecs and have WriteChainAsyncTransportWrapper use this.

In the future, custom transports could use this method specialized handling of writev without subclassing WriteChainAsyncTransportWrapper.

Reviewed By: yfeldblum

Differential Revision: D9297258

fbshipit-source-id: 9afc17f2ac2c46a49ce384e7c4d8c3a82afa1eec

folly/io/IOBuf.cpp

@@ -1015,6 +1015,23 @@ void IOBuf::appendToIov(folly::fbvector<struct iovec>* iov) const {
   } while (p != this);
 }
 
+unique_ptr<IOBuf> IOBuf::wrapIov(const iovec* vec, size_t count) {
+  unique_ptr<IOBuf> result = nullptr;
+  for (size_t i = 0; i < count; ++i) {
+    size_t len = vec[i].iov_len;
+    void* data = vec[i].iov_base;
+    if (len > 0) {
+      auto buf = folly::IOBuf::wrapBuffer(data, len);
+      if (!result) {
+        result = std::move(buf);
+      } else {
+        result->prependChain(std::move(buf));
+      }
+    }
+  }
+  return result;
+}
+
 size_t IOBuf::fillIov(struct iovec* iov, size_t len) const {
   IOBuf const* p = this;
   size_t i = 0;

folly/io/IOBuf.h

@@ -1223,6 +1223,12 @@ class IOBuf {
    */
   size_t fillIov(struct iovec* iov, size_t len) const;
 
+  /**
+   * A helper that wraps a number of iovecs into an IOBuf chain.  Empty iovecs
+   * are ignored.  If no data is wrapped, nullptr is returned.
+   */
+  static std::unique_ptr<IOBuf> wrapIov(const iovec* vec, size_t count);
+
   /*
    * Overridden operator new and delete.
    * These perform specialized memory management to help support

folly/io/async/WriteChainAsyncTransportWrapper.h

@@ -44,18 +44,7 @@ class WriteChainAsyncTransportWrapper
       const iovec* vec,
       size_t count,
       folly::WriteFlags flags = folly::WriteFlags::NONE) override {
-    std::unique_ptr<folly::IOBuf> writeBuffer;
-
-    for (size_t i = 0; i < count; ++i) {
-      size_t len = vec[i].iov_len;
-      void* data = vec[i].iov_base;
-      auto buf = folly::IOBuf::wrapBuffer(data, len);
-      if (i == 0) {
-        writeBuffer = std::move(buf);
-      } else {
-        writeBuffer->prependChain(std::move(buf));
-      }
-    }
+    auto writeBuffer = folly::IOBuf::wrapIov(vec, count);
     if (writeBuffer) {
       writeChain(callback, std::move(writeBuffer), flags);
     }

folly/io/test/IOBufTest.cpp

@@ -1001,6 +1001,32 @@ TEST(IOBuf, getIov) {
   EXPECT_EQ(buf->next()->next()->data(), iov[2].iov_base);
 }
 
+TEST(IOBuf, wrapIov) {
+  // Test wrapping IOVs
+  constexpr folly::StringPiece hello = "hello";
+  constexpr folly::StringPiece world = "world!";
+  folly::fbvector<struct iovec> iov;
+  iov.push_back({nullptr, 0});
+  iov.push_back({(void*)hello.data(), hello.size()});
+  iov.push_back({(void*)world.data(), world.size()});
+  auto wrapped = IOBuf::wrapIov(iov.data(), iov.size());
+  EXPECT_EQ(iov.size() - 1, wrapped->countChainElements());
+  IOBuf const* w = wrapped.get();
+  // skip the first iovec, which is empty/null, as it is ignored by IOBuf::wrapIov
+  for (size_t i = 0; i < wrapped->countChainElements(); ++i, w = w->next()) {
+    EXPECT_EQ(w->data(), iov[i + 1].iov_base);
+    EXPECT_EQ(w->length(), iov[i + 1].iov_len);
+  }
+}
+
+TEST(IOBuf, wrapZeroLenIov) {
+  folly::fbvector<struct iovec> iov;
+  iov.push_back({nullptr, 0});
+  iov.push_back({nullptr, 0});
+  auto wrapped = IOBuf::wrapIov(iov.data(), iov.size());
+  EXPECT_EQ(nullptr, wrapped);
+}
+
 TEST(IOBuf, move) {
   // Default allocate an IOBuf on the stack
   IOBuf outerBuf;