2 * Copyright 2015 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>
19 #include <sys/socket.h>
20 #include <sys/types.h>
25 #include <glog/logging.h>
27 #include <folly/FileUtil.h>
28 #include <folly/Malloc.h>
32 ShutdownSocketSet::ShutdownSocketSet(size_t maxFd)
34 data_(static_cast<std::atomic<uint8_t>*>(
35 folly::checkedCalloc(maxFd, sizeof(std::atomic<uint8_t>)))),
36 nullFile_("/dev/null", O_RDWR) {
39 void ShutdownSocketSet::add(int fd) {
40 // Silently ignore any fds >= maxFd_, very unlikely
42 if (size_t(fd) >= maxFd_) {
46 auto& sref = data_[fd];
47 uint8_t prevState = FREE;
48 CHECK(sref.compare_exchange_strong(prevState,
50 std::memory_order_acq_rel))
51 << "Invalid prev state for fd " << fd << ": " << int(prevState);
54 void ShutdownSocketSet::remove(int fd) {
56 if (size_t(fd) >= maxFd_) {
60 auto& sref = data_[fd];
61 uint8_t prevState = 0;
64 prevState = sref.load(std::memory_order_relaxed);
69 std::this_thread::sleep_for(std::chrono::milliseconds(1));
72 LOG(FATAL) << "Invalid prev state for fd " << fd << ": " << int(prevState);
75 if (!sref.compare_exchange_weak(prevState,
77 std::memory_order_acq_rel)) {
82 int ShutdownSocketSet::close(int fd) {
84 if (size_t(fd) >= maxFd_) {
85 return folly::closeNoInt(fd);
88 auto& sref = data_[fd];
89 uint8_t prevState = sref.load(std::memory_order_relaxed);
99 newState = MUST_CLOSE;
102 LOG(FATAL) << "Invalid prev state for fd " << fd << ": " << int(prevState);
105 if (!sref.compare_exchange_weak(prevState,
107 std::memory_order_acq_rel)) {
111 return newState == FREE ? folly::closeNoInt(fd) : 0;
114 void ShutdownSocketSet::shutdown(int fd, bool abortive) {
116 if (fd >= 0 && size_t(fd) >= maxFd_) {
117 doShutdown(fd, abortive);
121 auto& sref = data_[fd];
122 uint8_t prevState = IN_USE;
123 if (!sref.compare_exchange_strong(prevState,
125 std::memory_order_acq_rel)) {
129 doShutdown(fd, abortive);
131 prevState = IN_SHUTDOWN;
132 if (sref.compare_exchange_strong(prevState,
134 std::memory_order_acq_rel)) {
138 CHECK_EQ(prevState, MUST_CLOSE)
139 << "Invalid prev state for fd " << fd << ": " << int(prevState);
141 folly::closeNoInt(fd); // ignore errors, nothing to do
143 CHECK(sref.compare_exchange_strong(prevState,
145 std::memory_order_acq_rel))
146 << "Invalid prev state for fd " << fd << ": " << int(prevState);
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);
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);
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
169 struct linger l = {1, 0};
170 if (setsockopt(fd, SOL_SOCKET, SO_LINGER, &l, sizeof(l)) != 0) {
171 // Probably not a socket, ignore.
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);