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/portability/Sockets.h>
  26
  27 namespace folly {
  28
  29 ShutdownSocketSet::ShutdownSocketSet(int maxFd)
  30     : maxFd_(maxFd),
  31       data_(static_cast<std::atomic<uint8_t>*>(
  32           folly::checkedCalloc(size_t(maxFd), sizeof(std::atomic<uint8_t>)))),
  33       nullFile_("/dev/null", O_RDWR) {}
  34
  35 void ShutdownSocketSet::add(int fd) {
  36   // Silently ignore any fds >= maxFd_, very unlikely
  37   DCHECK_GE(fd, 0);
  38   if (fd >= maxFd_) {
  39     return;
  40   }
  41
  42   auto& sref = data_[size_t(fd)];
  43   uint8_t prevState = FREE;
  44   CHECK(sref.compare_exchange_strong(prevState,
  45                                      IN_USE,
  46                                      std::memory_order_acq_rel))
  47     << "Invalid prev state for fd " << fd << ": " << int(prevState);
  48 }
  49
  50 void ShutdownSocketSet::remove(int fd) {
  51   DCHECK_GE(fd, 0);
  52   if (fd >= maxFd_) {
  53     return;
  54   }
  55
  56   auto& sref = data_[size_t(fd)];
  57   uint8_t prevState = 0;
  58
  59 retry_load:
  60   prevState = sref.load(std::memory_order_relaxed);
  61
  62 retry:
  63   switch (prevState) {
  64   case IN_SHUTDOWN:
  65     std::this_thread::sleep_for(std::chrono::milliseconds(1));
  66     goto retry_load;
  67   case FREE:
  68     LOG(FATAL) << "Invalid prev state for fd " << fd << ": " << int(prevState);
  69   }
  70
  71   if (!sref.compare_exchange_weak(prevState,
  72                                   FREE,
  73                                   std::memory_order_acq_rel)) {
  74     goto retry;
  75   }
  76 }
  77
  78 int ShutdownSocketSet::close(int fd) {
  79   DCHECK_GE(fd, 0);
  80   if (fd >= maxFd_) {
  81     return folly::closeNoInt(fd);
  82   }
  83
  84   auto& sref = data_[size_t(fd)];
  85   uint8_t prevState = sref.load(std::memory_order_relaxed);
  86   uint8_t newState = 0;
  87
  88 retry:
  89   switch (prevState) {
  90   case IN_USE:
  91   case SHUT_DOWN:
  92     newState = FREE;
  93     break;
  94   case IN_SHUTDOWN:
  95     newState = MUST_CLOSE;
  96     break;
  97   default:
  98     LOG(FATAL) << "Invalid prev state for fd " << fd << ": " << int(prevState);
  99   }
 100
 101   if (!sref.compare_exchange_weak(prevState,
 102                                   newState,
 103                                   std::memory_order_acq_rel)) {
 104     goto retry;
 105   }
 106
 107   return newState == FREE ? folly::closeNoInt(fd) : 0;
 108 }
 109
 110 void ShutdownSocketSet::shutdown(int fd, bool abortive) {
 111   DCHECK_GE(fd, 0);
 112   if (fd >= maxFd_) {
 113     doShutdown(fd, abortive);
 114     return;
 115   }
 116
 117   auto& sref = data_[size_t(fd)];
 118   uint8_t prevState = IN_USE;
 119   if (!sref.compare_exchange_strong(prevState,
 120                                     IN_SHUTDOWN,
 121                                     std::memory_order_acq_rel)) {
 122     return;
 123   }
 124
 125   doShutdown(fd, abortive);
 126
 127   prevState = IN_SHUTDOWN;
 128   if (sref.compare_exchange_strong(prevState,
 129                                    SHUT_DOWN,
 130                                    std::memory_order_acq_rel)) {
 131     return;
 132   }
 133
 134   CHECK_EQ(prevState, MUST_CLOSE)
 135     << "Invalid prev state for fd " << fd << ": " << int(prevState);
 136
 137   folly::closeNoInt(fd);  // ignore errors, nothing to do
 138
 139   CHECK(sref.compare_exchange_strong(prevState,
 140                                      FREE,
 141                                      std::memory_order_acq_rel))
 142     << "Invalid prev state for fd " << fd << ": " << int(prevState);
 143 }
 144
 145 void ShutdownSocketSet::shutdownAll(bool abortive) {
 146   for (int i = 0; i < maxFd_; ++i) {
 147     auto& sref = data_[size_t(i)];
 148     if (sref.load(std::memory_order_acquire) == IN_USE) {
 149       shutdown(i, abortive);
 150     }
 151   }
 152 }
 153
 154 void ShutdownSocketSet::doShutdown(int fd, bool abortive) {
 155   // shutdown() the socket first, to awaken any threads blocked on the fd
 156   // (subsequent IO will fail because it's been shutdown); close()ing the
 157   // socket does not wake up blockers, see
 158   // http://stackoverflow.com/a/3624545/1736339
 159   folly::shutdownNoInt(fd, SHUT_RDWR);
 160
 161   // If abortive shutdown is desired, we'll set the SO_LINGER option on
 162   // the socket with a timeout of 0; this will cause RST to be sent on
 163   // close.
 164   if (abortive) {
 165     struct linger l = {1, 0};
 166     if (setsockopt(fd, SOL_SOCKET, SO_LINGER, &l, sizeof(l)) != 0) {
 167       // Probably not a socket, ignore.
 168       return;
 169     }
 170   }
 171
 172   // We can't close() the socket, as that would be dangerous; a new file
 173   // could be opened and get the same file descriptor, and then code assuming
 174   // the old fd would do IO in the wrong place. We'll (atomically) dup2
 175   // /dev/null onto the fd instead.
 176   folly::dup2NoInt(nullFile_.fd(), fd);
 177 }
 178
 179 }  // namespaces