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