From f51158b61b5f4e6f75a3c721e90a4078bde1c29f Mon Sep 17 00:00:00 2001 From: Kyle Nekritz Date: Tue, 12 Jul 2016 18:04:37 -0700 Subject: [PATCH] Move DecoratedAsyncTransportWrapper and WriteChainAsyncTransportWrapper to folly. Reviewed By: siyengar Differential Revision: D3550430 fbshipit-source-id: 1489fe502f41e65ce4ce45f26de59db30c9874b8 --- folly/Makefile.am | 2 + .../io/async/DecoratedAsyncTransportWrapper.h | 194 ++++++++++++++++++ .../async/WriteChainAsyncTransportWrapper.h | 74 +++++++ .../WriteChainAsyncTransportWrapperTest.cpp | 86 ++++++++ 4 files changed, 356 insertions(+) create mode 100644 folly/io/async/DecoratedAsyncTransportWrapper.h create mode 100644 folly/io/async/WriteChainAsyncTransportWrapper.h create mode 100644 folly/io/async/test/WriteChainAsyncTransportWrapperTest.cpp diff --git a/folly/Makefile.am b/folly/Makefile.am index 69559225..a987eb74 100644 --- a/folly/Makefile.am +++ b/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 \ diff --git a/folly/io/async/DecoratedAsyncTransportWrapper.h b/folly/io/async/DecoratedAsyncTransportWrapper.h new file mode 100644 index 00000000..db96390d --- /dev/null +++ b/folly/io/async/DecoratedAsyncTransportWrapper.h @@ -0,0 +1,194 @@ +/* + * 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 + +namespace folly { + +/** + * Convenience class so that AsyncTransportWrapper can be decorated without + * having to redefine every single method. + */ +template +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&& 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_; +}; + +} diff --git a/folly/io/async/WriteChainAsyncTransportWrapper.h b/folly/io/async/WriteChainAsyncTransportWrapper.h new file mode 100644 index 00000000..797be780 --- /dev/null +++ b/folly/io/async/WriteChainAsyncTransportWrapper.h @@ -0,0 +1,74 @@ +/* + * 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 +#include +#include + +namespace folly { + +/** + * Helper class that redirects write() and writev() calls to writeChain(). + */ +template +class WriteChainAsyncTransportWrapper : + public DecoratedAsyncTransportWrapper { + public: + using DecoratedAsyncTransportWrapper::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 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&& buf, + folly::WriteFlags flags = folly::WriteFlags::NONE) override = 0; +}; + +} diff --git a/folly/io/async/test/WriteChainAsyncTransportWrapperTest.cpp b/folly/io/async/test/WriteChainAsyncTransportWrapperTest.cpp new file mode 100644 index 00000000..6f8bbc7b --- /dev/null +++ b/folly/io/async/test/WriteChainAsyncTransportWrapperTest.cpp @@ -0,0 +1,86 @@ +/* + * 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 +#include + +#include +#include + +using namespace testing; +using testing::_; + +namespace folly { +namespace test { + +class TestWriteChainAsyncTransportWrapper : + public WriteChainAsyncTransportWrapper { + public: + TestWriteChainAsyncTransportWrapper() : + WriteChainAsyncTransportWrapper(nullptr) {} + + MOCK_METHOD3(writeChain, void( + folly::AsyncTransportWrapper::WriteCallback*, + std::shared_ptr, + folly::WriteFlags)); + + // gmock doesn't work with the IOBuf&& so we have to wrap this. + void writeChain(WriteCallback* callback, + std::unique_ptr&& iob, + folly::WriteFlags flags = folly::WriteFlags::NONE) override { + writeChain(callback, std::shared_ptr(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()); +} + +}} -- 2.34.1