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