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 | |
28 | namespace folly { |
29 | |
30 | ShutdownSocketSet::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 | |
36 | void 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 | |
50 | void 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 | |
74 | int 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 | |
103 | void 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 | |
135 | void 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 | |
146 | void 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 | |