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
28namespace folly {
29
30/**
31 * Set of sockets that allows immediate, take-no-prisoners abort.
32 */
33class 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