From 15fd86fa3c78f4c0060b221c5206b77eae331110 Mon Sep 17 00:00:00 2001
From: Dan Melnic <dmm@fb.com>
Date: Fri, 3 Nov 2017 11:16:10 -0700
Subject: [PATCH] Expose the zerocopy buf ID, change the AsyncSocket fd
 constructor to accept the id, Buff->Buf, make the cmsghdr methods private

Summary: Expose the zerocopy buf ID, change the AsyncScokey fd constructor to accept the id, Buff->Buf, make the cmsghdr methods private

Reviewed By: djwatson

Differential Revision: D6221324

fbshipit-source-id: d0fc4937adf6cf5790d11e406ffd3ec64c558b9c
---
 folly/io/async/AsyncSocket.cpp | 18 +++++++++++-------
 folly/io/async/AsyncSocket.h   | 18 ++++++++++++------
 2 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/folly/io/async/AsyncSocket.cpp b/folly/io/async/AsyncSocket.cpp
index a4c41969..0a78e404 100644
--- a/folly/io/async/AsyncSocket.cpp
+++ b/folly/io/async/AsyncSocket.cpp
@@ -272,13 +272,14 @@ AsyncSocket::AsyncSocket(EventBase* evb,
   connect(nullptr, ip, port, connectTimeout);
 }
 
-AsyncSocket::AsyncSocket(EventBase* evb, int fd)
-    : eventBase_(evb),
+AsyncSocket::AsyncSocket(EventBase* evb, int fd, uint32_t zeroCopyBufId)
+    : zeroCopyBufId_(zeroCopyBufId),
+      eventBase_(evb),
       writeTimeout_(this, evb),
       ioHandler_(this, evb, fd),
       immediateReadHandler_(this) {
-  VLOG(5) << "new AsyncSocket(" << this << ", evb=" << evb << ", fd="
-          << fd << ")";
+  VLOG(5) << "new AsyncSocket(" << this << ", evb=" << evb << ", fd=" << fd
+          << ", zeroCopyBufId=" << zeroCopyBufId << ")";
   init();
   fd_ = fd;
   setCloseOnExec();
@@ -286,7 +287,10 @@ AsyncSocket::AsyncSocket(EventBase* evb, int fd)
 }
 
 AsyncSocket::AsyncSocket(AsyncSocket::UniquePtr oldAsyncSocket)
-    : AsyncSocket(oldAsyncSocket->getEventBase(), oldAsyncSocket->detachFd()) {
+    : AsyncSocket(
+          oldAsyncSocket->getEventBase(),
+          oldAsyncSocket->detachFd(),
+          oldAsyncSocket->getZeroCopyBufId()) {
   preReceivedData_ = std::move(oldAsyncSocket->preReceivedData_);
 }
 
@@ -892,7 +896,7 @@ void AsyncSocket::adjustZeroCopyFlags(
 }
 
 void AsyncSocket::addZeroCopyBuf(std::unique_ptr<folly::IOBuf>&& buf) {
-  uint32_t id = getNextZeroCopyBuffId();
+  uint32_t id = getNextZeroCopyBufId();
   folly::IOBuf* ptr = buf.get();
 
   idZeroCopyBufPtrMap_[id] = ptr;
@@ -903,7 +907,7 @@ void AsyncSocket::addZeroCopyBuf(std::unique_ptr<folly::IOBuf>&& buf) {
 }
 
 void AsyncSocket::addZeroCopyBuf(folly::IOBuf* ptr) {
-  uint32_t id = getNextZeroCopyBuffId();
+  uint32_t id = getNextZeroCopyBufId();
   idZeroCopyBufPtrMap_[id] = ptr;
 
   idZeroCopyBufInfoMap_[ptr].count_++;
diff --git a/folly/io/async/AsyncSocket.h b/folly/io/async/AsyncSocket.h
index fb9d2c7f..22dbc394 100644
--- a/folly/io/async/AsyncSocket.h
+++ b/folly/io/async/AsyncSocket.h
@@ -261,8 +261,9 @@ class AsyncSocket : virtual public AsyncTransportWrapper {
    *
    * @param evb EventBase that will manage this socket.
    * @param fd  File descriptor to take over (should be a connected socket).
+   * @param zeroCopyBufId Zerocopy buf id to start with.
    */
-  AsyncSocket(EventBase* evb, int fd);
+  AsyncSocket(EventBase* evb, int fd, uint32_t zeroCopyBufId = 0);
 
   /**
    * Create an AsyncSocket from a different, already connected AsyncSocket.
@@ -516,8 +517,9 @@ class AsyncSocket : virtual public AsyncTransportWrapper {
     return zeroCopyWriteChainThreshold_;
   }
 
-  bool isZeroCopyMsg(const cmsghdr& cmsg) const;
-  void processZeroCopyMsg(const cmsghdr& cmsg);
+  uint32_t getZeroCopyBufId() const {
+    return zeroCopyBufId_;
+  }
 
   void write(WriteCallback* callback, const void* buf, size_t bytes,
              WriteFlags flags = WriteFlags::NONE) override;
@@ -1155,8 +1157,12 @@ class AsyncSocket : virtual public AsyncTransportWrapper {
   void cachePeerAddress() const;
 
   bool isZeroCopyRequest(WriteFlags flags);
-  uint32_t getNextZeroCopyBuffId() {
-    return zeroCopyBuffId_++;
+
+  bool isZeroCopyMsg(const cmsghdr& cmsg) const;
+  void processZeroCopyMsg(const cmsghdr& cmsg);
+
+  uint32_t getNextZeroCopyBufId() {
+    return zeroCopyBufId_++;
   }
   void adjustZeroCopyFlags(folly::IOBuf* buf, folly::WriteFlags& flags);
   void adjustZeroCopyFlags(
@@ -1173,7 +1179,7 @@ class AsyncSocket : virtual public AsyncTransportWrapper {
   // there is a that maps a buffer id to a raw folly::IOBuf ptr
   // and another one that adds a ref count for a folly::IOBuf that is either
   // the original ptr or nullptr
-  uint32_t zeroCopyBuffId_{0};
+  uint32_t zeroCopyBufId_{0};
 
   struct IOBufInfo {
     uint32_t count_{0};
-- 
2.34.1