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