folly/io/ShutdownSocketSet.cpp

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