Add LOCK_SHAREDMUTEX SSLLockType.

[folly.git] / folly / io / async / AsyncUDPServerSocket.h

diff --git a/folly/io/async/AsyncUDPServerSocket.h b/folly/io/async/AsyncUDPServerSocket.h
index f29353348344fb1bab118ca2fb3b4f832ccc4de9..651eb1cacd446fe6ded1e93c41d47cac839760c5 100644 (file)
--- a/folly/io/async/AsyncUDPServerSocket.h
+++ b/folly/io/async/AsyncUDPServerSocket.h
@@ -1,5 +1,5 @@
 /*
- * Copyright 2014 Facebook, Inc.
+ * Copyright 2017 Facebook, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,7 +16,6 @@
 
 #pragma once
 
-#include <folly/MoveWrapper.h>
 #include <folly/io/IOBufQueue.h>
 #include <folly/Memory.h>
 #include <folly/io/async/AsyncUDPSocket.h>
@@ -36,7 +35,8 @@ namespace folly {
  *       more than 1 packet will not work because they will end up with
  *       different event base to process.
  */
-class AsyncUDPServerSocket : private AsyncUDPSocket::ReadCallback {
+class AsyncUDPServerSocket : private AsyncUDPSocket::ReadCallback
+                           , public AsyncSocketBase {
  public:
   class Callback {
    public:
@@ -55,11 +55,13 @@ 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() {}
+    virtual ~Callback() = default;
   };
 
   /**
@@ -81,11 +83,16 @@ class AsyncUDPServerSocket : private AsyncUDPSocket::ReadCallback {
     }
   }
 
-  void bind(const folly::SocketAddress& address) {
+  void bind(const folly::SocketAddress& addy) {
     CHECK(!socket_);
 
-    socket_ = folly::make_unique<AsyncUDPSocket>(evb_);
-    socket_->bind(address);
+    socket_ = std::make_shared<AsyncUDPSocket>(evb_);
+    socket_->setReusePort(reusePort_);
+    socket_->bind(addy);
+  }
+
+  void setReusePort(bool reusePort) {
+    reusePort_ = reusePort;
   }
 
   folly::SocketAddress address() const {
@@ -93,6 +100,10 @@ class AsyncUDPServerSocket : private AsyncUDPSocket::ReadCallback {
     return socket_->address();
   }
 
+  void getAddress(SocketAddress* a) const {
+    *a = address();
+  }
+
   /**
    * Add a listener to the round robin list
    */
@@ -114,16 +125,21 @@ class AsyncUDPServerSocket : private AsyncUDPSocket::ReadCallback {
     socket_->resumeRead(this);
   }
 
-  int getFD() {
+  int getFD() const {
     CHECK(socket_) << "Need to bind before getting FD";
     return socket_->getFD();
   }
 
   void close() {
     CHECK(socket_) << "Need to bind before closing";
+    socket_->close();
     socket_.reset();
   }
 
+  EventBase* getEventBase() const {
+    return evb_;
+  }
+
  private:
   // AsyncUDPSocket::ReadCallback
   void getReadBuffer(void** buf, size_t* len) noexcept {
@@ -148,17 +164,21 @@ class AsyncUDPServerSocket : private AsyncUDPSocket::ReadCallback {
 
     auto client = clientAddress;
     auto callback = listeners_[nextListener_].second;
-    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);
+    auto f = [
+      socket,
+      client,
+      callback,
+      data = std::move(data),
+      truncated
+    ]() mutable {
+      callback->onDataAvailable(socket, client, std::move(data), truncated);
     };
 
-    listeners_[nextListener_].first->runInEventBaseThread(f);
+    listeners_[nextListener_].first->runInEventBaseThread(std::move(f));
     ++nextListener_;
   }
 
@@ -182,7 +202,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;
@@ -193,6 +213,8 @@ class AsyncUDPServerSocket : private AsyncUDPSocket::ReadCallback {
 
   // Temporary buffer for data
   folly::IOBufQueue buf_;
+
+  bool reusePort_{false};
 };
 
 } // Namespace