2 * Copyright 2014-present Facebook, Inc.
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
8 * http://www.apache.org/licenses/LICENSE-2.0
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.
17 #include <folly/io/ShutdownSocketSet.h>
22 #include <glog/logging.h>
24 #include <folly/FileUtil.h>
25 #include <folly/Singleton.h>
26 #include <folly/portability/Sockets.h>
32 folly::Singleton<folly::ShutdownSocketSet, PrivateTag> singleton;
35 ShutdownSocketSet::ShutdownSocketSet(int 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) {}
41 std::shared_ptr<ShutdownSocketSet> ShutdownSocketSet::getInstance() {
42 return singleton.try_get();
45 void ShutdownSocketSet::add(int fd) {
46 // Silently ignore any fds >= maxFd_, very unlikely
52 auto& sref = data_[size_t(fd)];
53 uint8_t prevState = FREE;
54 CHECK(sref.compare_exchange_strong(
55 prevState, IN_USE, std::memory_order_relaxed))
56 << "Invalid prev state for fd " << fd << ": " << int(prevState);
59 void ShutdownSocketSet::remove(int fd) {
65 auto& sref = data_[size_t(fd)];
66 uint8_t prevState = 0;
68 prevState = sref.load(std::memory_order_relaxed);
72 std::this_thread::sleep_for(std::chrono::milliseconds(1));
73 prevState = sref.load(std::memory_order_relaxed);
76 LOG(FATAL) << "Invalid prev state for fd " << fd << ": "
80 !sref.compare_exchange_weak(prevState, FREE, std::memory_order_relaxed));
83 int ShutdownSocketSet::close(int fd) {
86 return folly::closeNoInt(fd);
89 auto& sref = data_[size_t(fd)];
90 uint8_t prevState = sref.load(std::memory_order_relaxed);
100 newState = MUST_CLOSE;
103 LOG(FATAL) << "Invalid prev state for fd " << fd << ": "
106 } while (!sref.compare_exchange_weak(
107 prevState, newState, std::memory_order_relaxed));
109 return newState == FREE ? folly::closeNoInt(fd) : 0;
112 void ShutdownSocketSet::shutdown(int fd, bool abortive) {
115 doShutdown(fd, abortive);
119 auto& sref = data_[size_t(fd)];
120 uint8_t prevState = IN_USE;
121 if (!sref.compare_exchange_strong(
122 prevState, IN_SHUTDOWN, std::memory_order_relaxed)) {
126 doShutdown(fd, abortive);
128 prevState = IN_SHUTDOWN;
129 if (sref.compare_exchange_strong(
130 prevState, SHUT_DOWN, std::memory_order_relaxed)) {
134 CHECK_EQ(prevState, MUST_CLOSE)
135 << "Invalid prev state for fd " << fd << ": " << int(prevState);
137 folly::closeNoInt(fd); // ignore errors, nothing to do
140 sref.compare_exchange_strong(prevState, FREE, std::memory_order_relaxed))
141 << "Invalid prev state for fd " << fd << ": " << int(prevState);
144 void ShutdownSocketSet::shutdownAll(bool abortive) {
145 for (int i = 0; i < maxFd_; ++i) {
146 auto& sref = data_[size_t(i)];
147 if (sref.load(std::memory_order_relaxed) == IN_USE) {
148 shutdown(i, abortive);
153 void ShutdownSocketSet::doShutdown(int fd, bool abortive) {
154 // shutdown() the socket first, to awaken any threads blocked on the fd
155 // (subsequent IO will fail because it's been shutdown); close()ing the
156 // socket does not wake up blockers, see
157 // http://stackoverflow.com/a/3624545/1736339
158 folly::shutdownNoInt(fd, SHUT_RDWR);
160 // If abortive shutdown is desired, we'll set the SO_LINGER option on
161 // the socket with a timeout of 0; this will cause RST to be sent on
164 struct linger l = {1, 0};
165 if (setsockopt(fd, SOL_SOCKET, SO_LINGER, &l, sizeof(l)) != 0) {
166 // Probably not a socket, ignore.
171 // We can't close() the socket, as that would be dangerous; a new file
172 // could be opened and get the same file descriptor, and then code assuming
173 // the old fd would do IO in the wrong place. We'll (atomically) dup2
174 // /dev/null onto the fd instead.
175 folly::dup2NoInt(nullFile_.fd(), fd);