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