Move DecoratedAsyncTransportWrapper and WriteChainAsyncTransportWrapper to folly.

Reviewed By: siyengar

Differential Revision: D3550430

fbshipit-source-id: 1489fe502f41e65ce4ce45f26de59db30c9874b8

folly/Makefile.am

@@ -197,6 +197,7 @@ nobase_follyinclude_HEADERS = \
 	io/async/AsyncSocketBase.h \
 	io/async/AsyncSSLSocket.h \
 	io/async/AsyncSocketException.h \
+	io/async/DecoratedAsyncTransportWrapper.h \
 	io/async/DelayedDestructionBase.h \
 	io/async/DelayedDestruction.h \
 	io/async/EventBase.h \
@@ -215,6 +216,7 @@ nobase_follyinclude_HEADERS = \
 	io/async/SSLContext.h \
 	io/async/ScopedEventBaseThread.h \
 	io/async/TimeoutManager.h \
+	io/async/WriteChainAsyncTransportWrapper.h \
 	io/async/test/AsyncSSLSocketTest.h \
 	io/async/test/BlockingSocket.h \
 	io/async/test/MockAsyncSocket.h \

folly/io/async/DecoratedAsyncTransportWrapper.h

+/*
+ * Copyright 2016 Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <folly/io/async/AsyncTransport.h>
+
+namespace folly {
+
+/**
+ * Convenience class so that AsyncTransportWrapper can be decorated without
+ * having to redefine every single method.
+ */
+template<class T>
+class DecoratedAsyncTransportWrapper : public folly::AsyncTransportWrapper {
+ public:
+  explicit DecoratedAsyncTransportWrapper(typename T::UniquePtr transport):
+    transport_(std::move(transport)) {}
+
+  const AsyncTransportWrapper* getWrappedTransport() const override {
+    return transport_.get();
+  }
+
+  // folly::AsyncTransportWrapper
+  virtual ReadCallback* getReadCallback() const override {
+    return transport_->getReadCallback();
+  }
+
+  virtual void setReadCB(
+      folly::AsyncTransportWrapper::ReadCallback* callback) override {
+    transport_->setReadCB(callback);
+  }
+
+  virtual void write(
+      folly::AsyncTransportWrapper::WriteCallback* callback,
+      const void* buf,
+      size_t bytes,
+      folly::WriteFlags flags = folly::WriteFlags::NONE) override {
+    transport_->write(callback, buf, bytes, flags);
+  }
+
+  virtual void writeChain(
+      folly::AsyncTransportWrapper::WriteCallback* callback,
+      std::unique_ptr<folly::IOBuf>&& buf,
+      folly::WriteFlags flags = folly::WriteFlags::NONE) override {
+    transport_->writeChain(callback, std::move(buf), flags);
+  }
+
+  virtual void writev(
+      folly::AsyncTransportWrapper::WriteCallback* callback,
+      const iovec* vec,
+      size_t bytes,
+      folly::WriteFlags flags = folly::WriteFlags::NONE) override {
+    transport_->writev(callback, vec, bytes, flags);
+  }
+
+  // folly::AsyncSocketBase
+  virtual folly::EventBase* getEventBase() const override {
+    return transport_->getEventBase();
+  }
+
+  // folly::AsyncTransport
+  virtual void attachEventBase(folly::EventBase* eventBase) override {
+    transport_->attachEventBase(eventBase);
+  }
+
+  virtual void close() override {
+    transport_->close();
+  }
+
+  virtual void closeNow() override {
+    transport_->closeNow();
+  }
+
+  virtual void closeWithReset() override {
+    transport_->closeWithReset();
+
+    // This will likely result in 2 closeNow() calls on the decorated transport,
+    // but otherwise it is very easy to miss the derived class's closeNow().
+    closeNow();
+  }
+
+  virtual bool connecting() const override {
+    return transport_->connecting();
+  }
+
+  virtual void detachEventBase() override {
+    transport_->detachEventBase();
+  }
+
+  virtual bool error() const override {
+    return transport_->error();
+  }
+
+  virtual size_t getAppBytesReceived() const override {
+    return transport_->getAppBytesReceived();
+  }
+
+  virtual size_t getAppBytesWritten() const override {
+    return transport_->getAppBytesWritten();
+  }
+
+  virtual void getLocalAddress(folly::SocketAddress* address) const override {
+    return transport_->getLocalAddress(address);
+  }
+
+  virtual void getPeerAddress(folly::SocketAddress* address) const override {
+    return transport_->getPeerAddress(address);
+  }
+
+  virtual folly::ssl::X509UniquePtr getPeerCert() const override {
+    return transport_->getPeerCert();
+  }
+
+  virtual size_t getRawBytesReceived() const override {
+    return transport_->getRawBytesReceived();
+  }
+
+  virtual size_t getRawBytesWritten() const override {
+    return transport_->getRawBytesWritten();
+  }
+
+  virtual uint32_t getSendTimeout() const override {
+    return transport_->getSendTimeout();
+  }
+
+  virtual bool good() const override {
+    return transport_->good();
+  }
+
+  virtual bool isDetachable() const override {
+    return transport_->isDetachable();
+  }
+
+  virtual bool isEorTrackingEnabled() const override {
+    return transport_->isEorTrackingEnabled();
+  }
+
+  virtual bool readable() const override {
+    return transport_->readable();
+  }
+
+  virtual void setEorTracking(bool track) override {
+    return transport_->setEorTracking(track);
+  }
+
+  virtual void setSendTimeout(uint32_t timeoutInMs) override {
+    transport_->setSendTimeout(timeoutInMs);
+  }
+
+  virtual void shutdownWrite() override {
+    transport_->shutdownWrite();
+  }
+
+  virtual void shutdownWriteNow() override {
+    transport_->shutdownWriteNow();
+  }
+
+  virtual std::string getApplicationProtocol() noexcept override {
+    return transport_->getApplicationProtocol();
+  }
+
+  virtual std::string getSecurityProtocol() const override {
+    return transport_->getSecurityProtocol();
+  }
+
+  virtual bool isReplaySafe() const override {
+    return transport_->isReplaySafe();
+  }
+
+  virtual void setReplaySafetyCallback(
+      folly::AsyncTransport::ReplaySafetyCallback* callback) override {
+    transport_->setReplaySafetyCallback(callback);
+  }
+
+ protected:
+  virtual ~DecoratedAsyncTransportWrapper() {}
+
+  typename T::UniquePtr transport_;
+};
+
+}

folly/io/async/WriteChainAsyncTransportWrapper.h

+/*
+ * Copyright 2016 Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <folly/io/IOBuf.h>
+#include <folly/io/async/AsyncTransport.h>
+#include <folly/io/async/DecoratedAsyncTransportWrapper.h>
+
+namespace folly {
+
+/**
+ * Helper class that redirects write() and writev() calls to writeChain().
+ */
+template <class T>
+class WriteChainAsyncTransportWrapper :
+  public DecoratedAsyncTransportWrapper<T> {
+ public:
+  using DecoratedAsyncTransportWrapper<T>::DecoratedAsyncTransportWrapper;
+
+  virtual void write(
+      folly::AsyncTransportWrapper::WriteCallback* callback,
+      const void* buf,
+      size_t bytes,
+      folly::WriteFlags flags = folly::WriteFlags::NONE) override {
+    auto ioBuf = folly::IOBuf::wrapBuffer(buf, bytes);
+    writeChain(callback, std::move(ioBuf), flags);
+  }
+
+  virtual void writev(
+      folly::AsyncTransportWrapper::WriteCallback* callback,
+      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));
+      }
+    }
+    if (writeBuffer) {
+      writeChain(callback, std::move(writeBuffer), flags);
+    }
+  }
+
+  /**
+   * It only makes sense to use this class if you override writeChain, so force
+   * derived classes to do that.
+   */
+  virtual void writeChain(
+      folly::AsyncTransportWrapper::WriteCallback* callback,
+      std::unique_ptr<folly::IOBuf>&& buf,
+      folly::WriteFlags flags = folly::WriteFlags::NONE) override = 0;
+};
+
+}

folly/io/async/test/WriteChainAsyncTransportWrapperTest.cpp

+/*
+ * Copyright 2016 Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <gtest/gtest.h>
+#include <gmock/gmock.h>
+
+#include <folly/io/async/AsyncTransport.h>
+#include <folly/io/async/WriteChainAsyncTransportWrapper.h>
+
+using namespace testing;
+using testing::_;
+
+namespace folly {
+namespace test {
+
+class TestWriteChainAsyncTransportWrapper :
+  public WriteChainAsyncTransportWrapper<folly::AsyncTransportWrapper> {
+ public:
+  TestWriteChainAsyncTransportWrapper() :
+    WriteChainAsyncTransportWrapper<folly::AsyncTransportWrapper>(nullptr) {}
+
+  MOCK_METHOD3(writeChain, void(
+        folly::AsyncTransportWrapper::WriteCallback*,
+        std::shared_ptr<folly::IOBuf>,
+        folly::WriteFlags));
+
+  // gmock doesn't work with the IOBuf&& so we have to wrap this.
+  void writeChain(WriteCallback* callback,
+                  std::unique_ptr<folly::IOBuf>&& iob,
+                  folly::WriteFlags flags = folly::WriteFlags::NONE) override {
+    writeChain(callback, std::shared_ptr<folly::IOBuf>(iob.release()), flags);
+  }
+
+  // Allow this to be constructed on the stack for easier testing.
+  virtual ~TestWriteChainAsyncTransportWrapper() {
+  }
+};
+
+MATCHER_P(BufMatches, expected, "") {
+  folly::IOBufEqual eq;
+  return eq(*arg, *expected);
+}
+
+TEST(WriteChainAsyncTransportWrapperTest, TestSimpleIov) {
+  TestWriteChainAsyncTransportWrapper transport;
+  auto buf = folly::IOBuf::copyBuffer("foo");
+
+  EXPECT_CALL(transport, writeChain(_, BufMatches(buf.get()), _));
+
+  auto iov = buf->getIov();
+  transport.writev(nullptr, iov.data(), iov.size());
+}
+
+TEST(WriteChainAsyncTransportWrapperTest, TestChainedIov) {
+  TestWriteChainAsyncTransportWrapper transport;
+  auto buf = folly::IOBuf::copyBuffer("hello");
+  buf->prependChain(folly::IOBuf::copyBuffer("world"));
+
+  EXPECT_CALL(transport, writeChain(_, BufMatches(buf.get()), _));
+
+  auto iov = buf->getIov();
+  transport.writev(nullptr, iov.data(), iov.size());
+}
+
+TEST(WriteChainAsyncTransportWrapperTest, TestSimpleBuf) {
+  TestWriteChainAsyncTransportWrapper transport;
+  auto buf = folly::IOBuf::copyBuffer("foobar");
+
+  EXPECT_CALL(transport, writeChain(_, BufMatches(buf.get()), _));
+
+  transport.write(nullptr, buf->data(), buf->length());
+}
+
+}}