folly/io/ShutdownSocketSet.cpp

   1 /*
   2  * Copyright 2017 Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #include <folly/io/ShutdownSocketSet.h>
  18
  19 #include <chrono>
  20 #include <thread>
  21
  22 #include <glog/logging.h>
  23
  24 #include <folly/FileUtil.h>
  25 #include <folly/Singleton.h>
  26 #include <folly/portability/Sockets.h>
  27
  28 namespace folly {
  29
  30 namespace {
  31 struct PrivateTag {};
  32 folly::Singleton<folly::ShutdownSocketSet, PrivateTag> singleton;
  33 } // namespace
  34
  35 ShutdownSocketSet::ShutdownSocketSet(int maxFd)
  36     : maxFd_(maxFd),
  37       data_(static_cast<std::atomic<uint8_t>*>(
  38           folly::checkedCalloc(size_t(maxFd), sizeof(std::atomic<uint8_t>)))),
  39       nullFile_("/dev/null", O_RDWR) {}
  40
  41 std::shared_ptr<ShutdownSocketSet> ShutdownSocketSet::getInstance() {
  42   return singleton.try_get();
  43 }
  44
  45 void ShutdownSocketSet::add(int fd) {
  46   // Silently ignore any fds >= maxFd_, very unlikely
  47   DCHECK_GE(fd, 0);
  48   if (fd >= maxFd_) {
  49     return;
  50   }
  51
  52   auto& sref = data_[size_t(fd)];
  53   uint8_t prevState = FREE;
  54   CHECK(sref.compare_exchange_strong(prevState,
  55                                      IN_USE,
  56                                      std::memory_order_acq_rel))
  57     << "Invalid prev state for fd " << fd << ": " << int(prevState);
  58 }
  59
  60 void ShutdownSocketSet::remove(int fd) {
  61   DCHECK_GE(fd, 0);
  62   if (fd >= maxFd_) {
  63     return;
  64   }
  65
  66   auto& sref = data_[size_t(fd)];
  67   uint8_t prevState = 0;
  68
  69 retry_load:
  70   prevState = sref.load(std::memory_order_relaxed);
  71
  72 retry:
  73   switch (prevState) {
  74   case IN_SHUTDOWN:
  75     std::this_thread::sleep_for(std::chrono::milliseconds(1));
  76     goto retry_load;
  77   case FREE:
  78     LOG(FATAL) << "Invalid prev state for fd " << fd << ": " << int(prevState);
  79   }
  80
  81   if (!sref.compare_exchange_weak(prevState,
  82                                   FREE,
  83                                   std::memory_order_acq_rel)) {
  84     goto retry;
  85   }
  86 }
  87
  88 int ShutdownSocketSet::close(int fd) {
  89   DCHECK_GE(fd, 0);
  90   if (fd >= maxFd_) {
  91     return folly::closeNoInt(fd);
  92   }
  93
  94   auto& sref = data_[size_t(fd)];
  95   uint8_t prevState = sref.load(std::memory_order_relaxed);
  96   uint8_t newState = 0;
  97
  98 retry:
  99   switch (prevState) {
 100   case IN_USE:
 101   case SHUT_DOWN:
 102     newState = FREE;
 103     break;
 104   case IN_SHUTDOWN:
 105     newState = MUST_CLOSE;
 106     break;
 107   default:
 108     LOG(FATAL) << "Invalid prev state for fd " << fd << ": " << int(prevState);
 109   }
 110
 111   if (!sref.compare_exchange_weak(prevState,
 112                                   newState,
 113                                   std::memory_order_acq_rel)) {
 114     goto retry;
 115   }
 116
 117   return newState == FREE ? folly::closeNoInt(fd) : 0;
 118 }
 119
 120 void ShutdownSocketSet::shutdown(int fd, bool abortive) {
 121   DCHECK_GE(fd, 0);
 122   if (fd >= maxFd_) {
 123     doShutdown(fd, abortive);
 124     return;
 125   }
 126
 127   auto& sref = data_[size_t(fd)];
 128   uint8_t prevState = IN_USE;
 129   if (!sref.compare_exchange_strong(prevState,
 130                                     IN_SHUTDOWN,
 131                                     std::memory_order_acq_rel)) {
 132     return;
 133   }
 134
 135   doShutdown(fd, abortive);
 136
 137   prevState = IN_SHUTDOWN;
 138   if (sref.compare_exchange_strong(prevState,
 139                                    SHUT_DOWN,
 140                                    std::memory_order_acq_rel)) {
 141     return;
 142   }
 143
 144   CHECK_EQ(prevState, MUST_CLOSE)
 145     << "Invalid prev state for fd " << fd << ": " << int(prevState);
 146
 147   folly::closeNoInt(fd);  // ignore errors, nothing to do
 148
 149   CHECK(sref.compare_exchange_strong(prevState,
 150                                      FREE,
 151                                      std::memory_order_acq_rel))
 152     << "Invalid prev state for fd " << fd << ": " << int(prevState);
 153 }
 154
 155 void ShutdownSocketSet::shutdownAll(bool abortive) {
 156   for (int i = 0; i < maxFd_; ++i) {
 157     auto& sref = data_[size_t(i)];
 158     if (sref.load(std::memory_order_acquire) == IN_USE) {
 159       shutdown(i, abortive);
 160     }
 161   }
 162 }
 163
 164 void ShutdownSocketSet::doShutdown(int fd, bool abortive) {
 165   // shutdown() the socket first, to awaken any threads blocked on the fd
 166   // (subsequent IO will fail because it's been shutdown); close()ing the
 167   // socket does not wake up blockers, see
 168   // http://stackoverflow.com/a/3624545/1736339
 169   folly::shutdownNoInt(fd, SHUT_RDWR);
 170
 171   // If abortive shutdown is desired, we'll set the SO_LINGER option on
 172   // the socket with a timeout of 0; this will cause RST to be sent on
 173   // close.
 174   if (abortive) {
 175     struct linger l = {1, 0};
 176     if (setsockopt(fd, SOL_SOCKET, SO_LINGER, &l, sizeof(l)) != 0) {
 177       // Probably not a socket, ignore.
 178       return;
 179     }
 180   }
 181
 182   // We can't close() the socket, as that would be dangerous; a new file
 183   // could be opened and get the same file descriptor, and then code assuming
 184   // the old fd would do IO in the wrong place. We'll (atomically) dup2
 185   // /dev/null onto the fd instead.
 186   folly::dup2NoInt(nullFile_.fd(), fd);
 187 }
 188
 189 } // namespace folly