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 | #pragma once |
18 | |
19 | #include <atomic> |
20 | #include <cstdlib> |
21 | #include <memory> |
22 | |
23 | #include <boost/noncopyable.hpp> |
24 | |
25 | #include <folly/File.h> |
26 | #include <folly/net/NetworkSocket.h> |
27 | |
28 | namespace folly { |
29 | |
30 | /** |
31 | * Set of sockets that allows immediate, take-no-prisoners abort. |
32 | */ |
33 | class ShutdownSocketSet : private boost::noncopyable { |
34 | public: |
35 | /** |
36 | * Create a socket set that can handle file descriptors < maxFd. |
37 | * The default value (256Ki) is high enough for just about all |
38 | * applications, even if you increased the number of file descriptors |
39 | * on your system. |
40 | */ |
41 | explicit ShutdownSocketSet(int maxFd = 1 << 18); |
42 | |
43 | /** |
44 | * Add an already open socket to the list of sockets managed by |
45 | * ShutdownSocketSet. You MUST close the socket by calling |
46 | * ShutdownSocketSet::close (which will, as a side effect, also handle EINTR |
47 | * properly) and not by calling close() on the file descriptor. |
48 | */ |
49 | void add(NetworkSocket fd); |
50 | |
51 | /** |
52 | * Remove a socket from the list of sockets managed by ShutdownSocketSet. |
53 | * Note that remove() might block! (which we lamely implement by |
54 | * sleep()-ing) in the extremely rare case that the fd is currently |
55 | * being shutdown(). |
56 | */ |
57 | void remove(NetworkSocket fd); |
58 | |
59 | /** |
60 | * Close a socket managed by ShutdownSocketSet. Returns the same return code |
61 | * as ::close() (and sets errno accordingly). |
62 | */ |
63 | int close(NetworkSocket fd); |
64 | |
65 | /** |
66 | * Shut down a socket. If abortive is true, we perform an abortive |
67 | * shutdown (send RST rather than FIN). Note that we might still end up |
68 | * sending FIN due to the rather interesting implementation. |
69 | * |
70 | * This is async-signal-safe and ignores errors. Obviously, subsequent |
71 | * read() and write() operations to the socket will fail. During normal |
72 | * operation, just call ::shutdown() on the socket. |
73 | */ |
74 | void shutdown(NetworkSocket fd, bool abortive = false); |
75 | |
76 | /** |
77 | * Immediate shutdown of all connections. This is a hard-hitting hammer; |
78 | * all reads and writes will return errors and no new connections will |
79 | * be accepted. |
80 | * |
81 | * To be used only in dire situations. We're using it from the failure |
82 | * signal handler to close all connections quickly, even though the server |
83 | * might take multiple seconds to finish crashing. |
84 | * |
85 | * The optional bool parameter indicates whether to set the active |
86 | * connections in to not linger. The effect of that includes RST packets |
87 | * being immediately sent to clients which will result |
88 | * in errors (and not normal EOF) on the client side. This also causes |
89 | * the local (ip, tcp port number) tuple to be reusable immediately, instead |
90 | * of having to wait the standard amount of time. For full details see |
91 | * the `shutdown` method of `ShutdownSocketSet` (incl. notes about the |
92 | * `abortive` parameter). |
93 | * |
94 | * This is async-signal-safe and ignores errors. |
95 | */ |
96 | void shutdownAll(bool abortive = false); |
97 | |
98 | private: |
99 | void doShutdown(NetworkSocket fd, bool abortive); |
100 | |
101 | // State transitions: |
102 | // add(): |
103 | // FREE -> IN_USE |
104 | // |
105 | // close(): |
106 | // IN_USE -> (::close()) -> FREE |
107 | // SHUT_DOWN -> (::close()) -> FREE |
108 | // IN_SHUTDOWN -> MUST_CLOSE |
109 | // (If the socket is currently being shut down, we must defer the |
110 | // ::close() until the shutdown completes) |
111 | // |
112 | // shutdown(): |
113 | // IN_USE -> IN_SHUTDOWN |
114 | // (::shutdown()) |
115 | // IN_SHUTDOWN -> SHUT_DOWN |
116 | // MUST_CLOSE -> (::close()) -> FREE |
117 | // |
118 | // State atomic operation memory orders: |
119 | // All atomic operations on per-socket states use std::memory_order_relaxed |
120 | // because there is no associated per-socket data guarded by the state and |
121 | // the states for different sockets are unrelated. If there were associated |
122 | // per-socket data, acquire and release orders would be desired; and if the |
123 | // states for different sockets were related, it could be that sequential |
124 | // consistent orders would be desired. |
125 | enum State : uint8_t { |
126 | FREE = 0, |
127 | IN_USE, |
128 | IN_SHUTDOWN, |
129 | SHUT_DOWN, |
130 | MUST_CLOSE, |
131 | }; |
132 | |
133 | struct Free { |
134 | template <class T> |
135 | void operator()(T* ptr) const { |
136 | ::free(ptr); |
137 | } |
138 | }; |
139 | |
140 | const int maxFd_; |
141 | std::unique_ptr<std::atomic<uint8_t>[], Free> data_; |
142 | folly::File nullFile_; |
143 | }; |
144 | |
145 | } // namespace folly |
146 | |