Fix TSAN-reported race in ShutdownSocketSet test

Summary: [Folly] Fix TSAN-reported race in `ShutdownSocketSet` test where the shutdown does not wait for the server.

Reviewed By: yfeldblum

Differential Revision: D23057822

fbshipit-source-id: 4b682e6e80cff503e85cb2b663f8399b69265a77

folly/io/test/ShutdownSocketSetTest.cpp

@@ -25,12 +25,11 @@
 #include <folly/net/NetOps.h>
 #include <folly/net/NetworkSocket.h>
 #include <folly/portability/GTest.h>
+#include <folly/synchronization/Baton.h>
 
 namespace folly {
 namespace test {
 
-ShutdownSocketSet shutdownSocketSet;
-
 class Server {
  public:
   Server();
@@ -42,6 +41,8 @@ class Server {
   }
   int closeClients(bool abortive);
 
+  void shutdownAll(bool abortive);
+
  private:
   NetworkSocket acceptSocket_;
   int port_;
@@ -49,12 +50,14 @@ class Server {
   std::atomic<StopMode> stop_;
   std::thread serverThread_;
   std::vector<NetworkSocket> fds_;
+  folly::ShutdownSocketSet shutdownSocketSet_;
+  folly::Baton<> baton_;
 };
 
 Server::Server() : acceptSocket_(), port_(0), stop_(NO_STOP) {
   acceptSocket_ = netops::socket(PF_INET, SOCK_STREAM, 0);
   CHECK_NE(acceptSocket_, NetworkSocket());
-  shutdownSocketSet.add(acceptSocket_);
+  shutdownSocketSet_.add(acceptSocket_);
 
   sockaddr_in addr;
   addr.sin_family = AF_INET;
@@ -72,6 +75,7 @@ Server::Server() : acceptSocket_(), port_(0), stop_(NO_STOP) {
   port_ = ntohs(addr.sin_port);
 
   serverThread_ = std::thread([this] {
+    bool first = true;
     while (stop_ == NO_STOP) {
       sockaddr_in peer;
       socklen_t peerLen = sizeof(peer);
@@ -86,15 +90,18 @@ Server::Server() : acceptSocket_(), port_(0), stop_(NO_STOP) {
         }
       }
       CHECK_NE(fd, NetworkSocket());
-      shutdownSocketSet.add(fd);
+      shutdownSocketSet_.add(fd);
       fds_.push_back(fd);
+      CHECK(first);
+      first = false;
+      baton_.post();
     }
 
     if (stop_ != NO_STOP) {
       closeClients(stop_ == ABORTIVE);
     }
 
-    shutdownSocketSet.close(acceptSocket_);
+    shutdownSocketSet_.close(acceptSocket_);
     acceptSocket_ = NetworkSocket();
     port_ = 0;
   });
@@ -106,13 +113,18 @@ int Server::closeClients(bool abortive) {
       struct linger l = {1, 0};
       CHECK_ERR(netops::setsockopt(fd, SOL_SOCKET, SO_LINGER, &l, sizeof(l)));
     }
-    shutdownSocketSet.close(fd);
+    shutdownSocketSet_.close(fd);
   }
   int n = fds_.size();
   fds_.clear();
   return n;
 }
 
+void Server::shutdownAll(bool abortive) {
+  baton_.wait();
+  shutdownSocketSet_.shutdownAll(abortive);
+}
+
 void Server::stop(bool abortive) {
   stop_ = abortive ? ABORTIVE : ORDERLY;
   netops::shutdown(acceptSocket_, SHUT_RDWR);
@@ -176,8 +188,7 @@ void runKillTest(bool abortive) {
   auto sock = createConnectedSocket(server.port());
 
   std::thread killer([&server, abortive] {
-    std::this_thread::sleep_for(std::chrono::milliseconds(200));
-    shutdownSocketSet.shutdownAll(abortive);
+    server.shutdownAll(abortive);
     server.join();
   });