From: Dave Watson <davejwatson@fb.com>
Date: Tue, 7 Apr 2015 16:53:44 +0000 (-0700)
Subject: AsyncUDPServerSocket passes socket in callback
X-Git-Tag: v0.35.0~9
X-Git-Url: http://plrg.eecs.uci.edu/git/?a=commitdiff_plain;h=f697ba4934d867fb1a59a80650b3a8082659e976;p=folly.git

AsyncUDPServerSocket passes socket in callback

Summary: AsyncUDPServerSocket doesn't make it easy to write to the same socket you read from.  Add the socket as a callback param, similar to AsyncServerSocket

Test Plan:
fbconfig -r folly; fbmake dbg

Will fixup any other spots contbuild finds

Reviewed By: hans@fb.com

Subscribers: bmatheny, doug, fugalh, folly-diffs@, jsedgwick, yfeldblum, chalfant

FB internal diff: D1948936

Signature: t1:1948936:1427841651:20d13d73c06d31c75056624f051a6fd35b9701fb
---

diff --git a/folly/io/async/AsyncUDPServerSocket.h b/folly/io/async/AsyncUDPServerSocket.h
index 04682297..1a5201b9 100644
--- a/folly/io/async/AsyncUDPServerSocket.h
+++ b/folly/io/async/AsyncUDPServerSocket.h
@@ -56,9 +56,11 @@ class AsyncUDPServerSocket : private AsyncUDPSocket::ReadCallback
     /**
      * Invoked when a new packet is received
      */
-    virtual void onDataAvailable(const folly::SocketAddress& addr,
-                                 std::unique_ptr<folly::IOBuf> buf,
-                                 bool truncated) noexcept = 0;
+    virtual void onDataAvailable(
+      std::shared_ptr<AsyncUDPSocket> socket,
+      const folly::SocketAddress& addr,
+      std::unique_ptr<folly::IOBuf> buf,
+      bool truncated) noexcept = 0;
 
     virtual ~Callback() {}
   };
@@ -85,7 +87,7 @@ class AsyncUDPServerSocket : private AsyncUDPSocket::ReadCallback
   void bind(const folly::SocketAddress& addy) {
     CHECK(!socket_);
 
-    socket_ = folly::make_unique<AsyncUDPSocket>(evb_);
+    socket_ = std::make_shared<AsyncUDPSocket>(evb_);
     socket_->setReusePort(reusePort_);
     socket_->bind(addy);
   }
@@ -131,6 +133,7 @@ class AsyncUDPServerSocket : private AsyncUDPSocket::ReadCallback
 
   void close() {
     CHECK(socket_) << "Need to bind before closing";
+    socket_->close();
     socket_.reset();
   }
 
@@ -165,11 +168,12 @@ class AsyncUDPServerSocket : private AsyncUDPSocket::ReadCallback
     auto mvp =
         folly::MoveWrapper<
             std::unique_ptr<folly::IOBuf>>(std::move(data));
+    auto socket = socket_;
 
     // Schedule it in the listener's eventbase
     // XXX: Speed this up
-    std::function<void()> f = [client, callback, mvp, truncated] () mutable {
-      callback->onDataAvailable(client, std::move(*mvp), truncated);
+    std::function<void()> f = [socket, client, callback, mvp, truncated] () mutable {
+      callback->onDataAvailable(socket, client, std::move(*mvp), truncated);
     };
 
     listeners_[nextListener_].first->runInEventBaseThread(f);
@@ -196,7 +200,7 @@ class AsyncUDPServerSocket : private AsyncUDPSocket::ReadCallback
   EventBase* const evb_;
   const size_t packetSize_;
 
-  std::unique_ptr<AsyncUDPSocket> socket_;
+  std::shared_ptr<AsyncUDPSocket> socket_;
 
   // List of listener to distribute packets among
   typedef std::pair<EventBase*, Callback*> Listener;
diff --git a/folly/io/async/test/AsyncUDPSocketTest.cpp b/folly/io/async/test/AsyncUDPSocketTest.cpp
index b19eeb3f..b9f99af1 100644
--- a/folly/io/async/test/AsyncUDPSocketTest.cpp
+++ b/folly/io/async/test/AsyncUDPSocketTest.cpp
@@ -47,7 +47,8 @@ class UDPAcceptor
   void onListenStopped() noexcept {
   }
 
-  void onDataAvailable(const folly::SocketAddress& client,
+  void onDataAvailable(std::shared_ptr<folly::AsyncUDPSocket> socket,
+                       const folly::SocketAddress& client,
                        std::unique_ptr<folly::IOBuf> data,
                        bool truncated) noexcept {
 
diff --git a/folly/wangle/acceptor/Acceptor.h b/folly/wangle/acceptor/Acceptor.h
index 9ad992fc..11bb3992 100644
--- a/folly/wangle/acceptor/Acceptor.h
+++ b/folly/wangle/acceptor/Acceptor.h
@@ -233,7 +233,10 @@ class Acceptor :
 
   void onListenStarted() noexcept {}
   void onListenStopped() noexcept {}
-  void onDataAvailable(const SocketAddress&, std::unique_ptr<IOBuf>, bool) noexcept {}
+  void onDataAvailable(
+    std::shared_ptr<AsyncUDPSocket> socket,
+    const SocketAddress&,
+    std::unique_ptr<IOBuf>, bool) noexcept {}
 
   virtual AsyncSocket::UniquePtr makeNewAsyncSocket(EventBase* base, int fd) {
     return AsyncSocket::UniquePtr(new AsyncSocket(base, fd));
diff --git a/folly/wangle/bootstrap/ServerBootstrap-inl.h b/folly/wangle/bootstrap/ServerBootstrap-inl.h
index ecc13721..e1909265 100644
--- a/folly/wangle/bootstrap/ServerBootstrap-inl.h
+++ b/folly/wangle/bootstrap/ServerBootstrap-inl.h
@@ -98,7 +98,8 @@ class ServerAcceptor
   }
 
   // UDP thunk
-  void onDataAvailable(const folly::SocketAddress& addr,
+  void onDataAvailable(std::shared_ptr<AsyncUDPSocket> socket,
+                       const folly::SocketAddress& addr,
                        std::unique_ptr<folly::IOBuf> buf,
                        bool truncated) noexcept {
     acceptorPipeline_->read(buf.release());