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 04682297fcccc4fcaceeb8a86c2cb0c46637b388..1a5201b9ac4103b78aa732dd90967d20ae5501a2 100644 (file)
--- 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 b19eeb3f42090a996f08cb7e876e5e6f08a1b5d5..b9f99af1f1aa70b05fc712f1d15325945e6aad57 100644 (file)
--- 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 9ad992fc03995909fb12ab024ff34fdccb13301c..11bb3992fa29479d315bb908ab1e166dd9a237db 100644 (file)
--- 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 ecc1372169d5b5a67dfc540d198337fc2f80bb12..e1909265eb6e99204d35d5ad06fbc1677e23e2c8 100644 (file)
--- 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());