fix flaky ConnectTFOTimeout and ConnectTFOFallbackTimeout tests

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

diff --git a/folly/io/async/AsyncServerSocket.h b/folly/io/async/AsyncServerSocket.h
index 9b82ebd5597a0c903a627a8b4047581ee69d8c87..d2f5002b45a50fd0d8fa0cea61c3ef4fe3d27656 100644 (file)
--- a/folly/io/async/AsyncServerSocket.h
+++ b/folly/io/async/AsyncServerSocket.h
@@ -1,5 +1,5 @@
 /*
- * Copyright 2015 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.
@@ -16,21 +16,21 @@
 
 #pragma once
 
+#include <folly/SocketAddress.h>
+#include <folly/io/ShutdownSocketSet.h>
+#include <folly/io/async/AsyncSocketBase.h>
+#include <folly/io/async/AsyncTimeout.h>
 #include <folly/io/async/DelayedDestruction.h>
-#include <folly/io/async/EventHandler.h>
 #include <folly/io/async/EventBase.h>
+#include <folly/io/async/EventHandler.h>
 #include <folly/io/async/NotificationQueue.h>
-#include <folly/io/async/AsyncTimeout.h>
-#include <folly/io/async/AsyncSocketBase.h>
-#include <folly/io/ShutdownSocketSet.h>
-#include <folly/SocketAddress.h>
-#include <memory>
-#include <exception>
-#include <vector>
+#include <folly/portability/Sockets.h>
+
 #include <limits.h>
 #include <stddef.h>
-#include <sys/socket.h>
-
+#include <exception>
+#include <memory>
+#include <vector>
 
 // Due to the way kernel headers are included, this may or may not be defined.
 // Number pulled from 3.10 kernel headers.
@@ -200,7 +200,7 @@ class AsyncServerSocket : public DelayedDestruction
 
   static const uint32_t kDefaultMaxAcceptAtOnce = 30;
   static const uint32_t kDefaultCallbackAcceptAtOnce = 5;
-  static const uint32_t kDefaultMaxMessagesInQueue = 0;
+  static const uint32_t kDefaultMaxMessagesInQueue = 1024;
   /**
    * Create a new AsyncServerSocket with the specified EventBase.
    *
@@ -564,6 +564,23 @@ class AsyncServerSocket : public DelayedDestruction
     return numDroppedConnections_;
   }
 
+  /**
+   * Get the current number of unprocessed messages in NotificationQueue.
+   *
+   * This method must be invoked from the AsyncServerSocket's primary
+   * EventBase thread.  Use EventBase::runInEventBaseThread() to schedule the
+   * operation in the correct EventBase if your code is not in the server
+   * socket's primary EventBase.
+   */
+  int64_t getNumPendingMessagesInQueue() const {
+    assert(eventBase_ == nullptr || eventBase_->isInEventBaseThread());
+    int64_t numMsgs = 0;
+    for (const auto& callback : callbacks_) {
+      numMsgs += callback.consumer->getQueue()->size();
+    }
+    return numMsgs;
+  }
+
   /**
    * Set whether or not SO_KEEPALIVE should be enabled on the server socket
    * (and thus on all subsequently-accepted connections). By default, keepalive
@@ -642,6 +659,14 @@ class AsyncServerSocket : public DelayedDestruction
     return closeOnExec_;
   }
 
+  /**
+   * Tries to enable TFO if the machine supports it.
+   */
+  void setTFOEnabled(bool enabled, uint32_t maxTFOQueueSize) {
+    tfo_ = enabled;
+    tfoMaxQueueSize_ = maxTFOQueueSize;
+  }
+
   /**
    * Get whether or not the socket is accepting new connections
    */
@@ -743,7 +768,7 @@ class AsyncServerSocket : public DelayedDestruction
     uint16_t events, int socket, sa_family_t family) noexcept;
 
   int createSocket(int family);
-  void setupSocket(int fd);
+  void setupSocket(int fd, int family);
   void bindSocket(int fd, const SocketAddress& address, bool isExistingSocket);
   void dispatchSocket(int socket, SocketAddress&& address);
   void dispatchError(const char *msg, int errnoValue);
@@ -820,6 +845,8 @@ class AsyncServerSocket : public DelayedDestruction
   bool keepAliveEnabled_;
   bool reusePortEnabled_{false};
   bool closeOnExec_;
+  bool tfo_{false};
+  uint32_t tfoMaxQueueSize_{0};
   ShutdownSocketSet* shutdownSocketSet_;
   ConnectionEventCallback* connectionEventCallback_{nullptr};
 };