2 * Copyright 2016 Facebook, Inc.
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
8 * http://www.apache.org/licenses/LICENSE-2.0
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.
16 #include <folly/io/ShutdownSocketSet.h>
22 #include <glog/logging.h>
23 #include <gtest/gtest.h>
25 #include <folly/portability/Sockets.h>
27 using folly::ShutdownSocketSet;
29 namespace folly { namespace test {
31 ShutdownSocketSet shutdownSocketSet;
37 void stop(bool abortive);
39 int port() const { return port_; }
40 int closeClients(bool abortive);
50 std::atomic<StopMode> stop_;
51 std::thread serverThread_;
52 std::vector<int> fds_;
59 acceptSocket_ = socket(PF_INET, SOCK_STREAM, 0);
60 CHECK_ERR(acceptSocket_);
61 shutdownSocketSet.add(acceptSocket_);
64 addr.sin_family = AF_INET;
66 addr.sin_addr.s_addr = INADDR_ANY;
67 CHECK_ERR(bind(acceptSocket_,
68 reinterpret_cast<const sockaddr*>(&addr),
71 CHECK_ERR(listen(acceptSocket_, 10));
73 socklen_t addrLen = sizeof(addr);
74 CHECK_ERR(getsockname(acceptSocket_,
75 reinterpret_cast<sockaddr*>(&addr),
78 port_ = ntohs(addr.sin_port);
80 serverThread_ = std::thread([this] {
81 while (stop_ == NO_STOP) {
83 socklen_t peerLen = sizeof(peer);
84 int fd = accept(acceptSocket_,
85 reinterpret_cast<sockaddr*>(&peer),
91 if (errno == EINVAL || errno == ENOTSOCK) { // socket broken
96 shutdownSocketSet.add(fd);
100 if (stop_ != NO_STOP) {
101 closeClients(stop_ == ABORTIVE);
104 shutdownSocketSet.close(acceptSocket_);
110 int Server::closeClients(bool abortive) {
111 for (int fd : fds_) {
113 struct linger l = {1, 0};
114 CHECK_ERR(setsockopt(fd, SOL_SOCKET, SO_LINGER, &l, sizeof(l)));
116 shutdownSocketSet.close(fd);
123 void Server::stop(bool abortive) {
124 stop_ = abortive ? ABORTIVE : ORDERLY;
125 shutdown(acceptSocket_, SHUT_RDWR);
128 void Server::join() {
129 serverThread_.join();
132 int createConnectedSocket(int port) {
133 int sock = socket(PF_INET, SOCK_STREAM, 0);
136 addr.sin_family = AF_INET;
137 addr.sin_port = htons(port);
138 addr.sin_addr.s_addr = htonl((127 << 24) | 1); // XXX
139 CHECK_ERR(connect(sock,
140 reinterpret_cast<const sockaddr*>(&addr),
145 void runCloseTest(bool abortive) {
148 int sock = createConnectedSocket(server.port());
150 std::thread stopper([&server, abortive] {
151 std::this_thread::sleep_for(std::chrono::milliseconds(200));
152 server.stop(abortive);
157 int r = read(sock, &c, 1);
161 EXPECT_EQ(ECONNRESET, e);
170 EXPECT_EQ(0, server.closeClients(false)); // closed by server when it exited
173 TEST(ShutdownSocketSetTest, OrderlyClose) {
177 TEST(ShutdownSocketSetTest, AbortiveClose) {
181 void runKillTest(bool abortive) {
184 int sock = createConnectedSocket(server.port());
186 std::thread killer([&server, abortive] {
187 std::this_thread::sleep_for(std::chrono::milliseconds(200));
188 shutdownSocketSet.shutdownAll(abortive);
193 int r = read(sock, &c, 1);
195 // "abortive" is just a hint for ShutdownSocketSet, so accept both
199 EXPECT_EQ(ECONNRESET, errno);
211 // NOT closed by server when it exited
212 EXPECT_EQ(1, server.closeClients(false));
215 TEST(ShutdownSocketSetTest, OrderlyKill) {
219 TEST(ShutdownSocketSetTest, AbortiveKill) {