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

folly/io/async/AsyncUDPServerSocket.h

@@ -56,7 +56,9 @@ class AsyncUDPServerSocket : private AsyncUDPSocket::ReadCallback
     /**
      * Invoked when a new packet is received
      */
-    virtual void onDataAvailable(const folly::SocketAddress& addr,
+    virtual void onDataAvailable(
+      std::shared_ptr<AsyncUDPSocket> socket,
+      const folly::SocketAddress& addr,
       std::unique_ptr<folly::IOBuf> buf,
       bool truncated) noexcept = 0;
 
@@ -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;

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 {
 

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));

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());