fix flaky ConnectTFOTimeout and ConnectTFOFallbackTimeout tests

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

diff --git a/folly/io/async/AsyncUDPSocket.h b/folly/io/async/AsyncUDPSocket.h
index 1c5d3f0af2e891ce785b2708d78b2fc61bd7bedd..6e310feab1e501ad7a24179fff8a49f24b6574a1 100644 (file)
--- a/folly/io/async/AsyncUDPSocket.h
+++ b/folly/io/async/AsyncUDPSocket.h
@@ -1,5 +1,5 @@
 /*
- * Copyright 2014 Facebook, Inc.
+ * Copyright 2016 Facebook, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -19,6 +19,7 @@
 #include <folly/io/IOBuf.h>
 #include <folly/ScopeGuard.h>
 #include <folly/io/async/AsyncSocketException.h>
+#include <folly/io/async/AsyncSocketBase.h>
 #include <folly/io/async/EventHandler.h>
 #include <folly/io/async/EventBase.h>
 #include <folly/SocketAddress.h>
@@ -73,7 +74,7 @@ class AsyncUDPSocket : public EventHandler {
      */
     virtual void onReadClosed() noexcept = 0;
 
-    virtual ~ReadCallback() {}
+    virtual ~ReadCallback() = default;
   };
 
   /**
@@ -86,7 +87,7 @@ class AsyncUDPSocket : public EventHandler {
   /**
    * Returns the address server is listening on
    */
-  const folly::SocketAddress& address() const {
+  virtual const folly::SocketAddress& address() const {
     CHECK_NE(-1, fd_) << "Server not yet bound to an address";
     return localAddress_;
   }
@@ -97,7 +98,7 @@ class AsyncUDPSocket : public EventHandler {
    * use `address()` method above to get it after this method successfully
    * returns.
    */
-  void bind(const folly::SocketAddress& address);
+  virtual void bind(const folly::SocketAddress& address);
 
   /**
    * Use an already bound file descriptor. You can either transfer ownership
@@ -105,37 +106,58 @@ class AsyncUDPSocket : public EventHandler {
    * FDOwnership::SHARED. In case FD is shared, it will not be `close`d in
    * destructor.
    */
-  void setFD(int fd, FDOwnership ownership);
+  virtual void setFD(int fd, FDOwnership ownership);
 
   /**
    * Send the data in buffer to destination. Returns the return code from
-   * ::sendto.
+   * ::sendmsg.
    */
-  ssize_t write(const folly::SocketAddress& address,
-                const std::unique_ptr<folly::IOBuf>& buf);
+  virtual ssize_t write(const folly::SocketAddress& address,
+                        const std::unique_ptr<folly::IOBuf>& buf);
+
+  /**
+   * Send data in iovec to destination. Returns the return code from sendmsg.
+   */
+  virtual ssize_t writev(const folly::SocketAddress& address,
+                         const struct iovec* vec, size_t veclen);
 
   /**
    * Start reading datagrams
    */
-  void resumeRead(ReadCallback* cob);
+  virtual void resumeRead(ReadCallback* cob);
 
   /**
    * Pause reading datagrams
    */
-  void pauseRead();
+  virtual void pauseRead();
 
   /**
    * Stop listening on the socket.
    */
-  void close();
+  virtual void close();
 
   /**
    * Get internal FD used by this socket
    */
-  int getFD() const {
+  virtual int getFD() const {
     CHECK_NE(-1, fd_) << "Need to bind before getting FD out";
     return fd_;
   }
+
+  /**
+   * Set reuse port mode to call bind() on the same address multiple times
+   */
+  virtual void setReusePort(bool reusePort) {
+    reusePort_ = reusePort;
+  }
+
+  /**
+   * Set SO_REUSEADDR flag on the socket. Default is ON.
+   */
+  virtual void setReuseAddr(bool reuseAddr) {
+    reuseAddr_ = reuseAddr;
+  }
+
  private:
   AsyncUDPSocket(const AsyncUDPSocket&) = delete;
   AsyncUDPSocket& operator=(const AsyncUDPSocket&) = delete;
@@ -157,6 +179,9 @@ class AsyncUDPSocket : public EventHandler {
 
   // Non-null only when we are reading
   ReadCallback* readCallback_;
+
+  bool reuseAddr_{true};
+  bool reusePort_{false};
 };
 
 } // Namespace