2 * Copyright 2015 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 <netinet/in.h>
23 #include <netinet/tcp.h>
24 #include <sys/socket.h>
26 #include <glog/logging.h>
27 #include <gtest/gtest.h>
29 using folly::ShutdownSocketSet;
31 namespace folly { namespace test {
33 ShutdownSocketSet shutdownSocketSet;
39 void stop(bool abortive);
41 int port() const { return port_; }
42 int closeClients(bool abortive);
52 std::atomic<StopMode> stop_;
53 std::thread serverThread_;
54 std::vector<int> fds_;
61 acceptSocket_ = socket(PF_INET, SOCK_STREAM, 0);
62 CHECK_ERR(acceptSocket_);
63 shutdownSocketSet.add(acceptSocket_);
66 addr.sin_family = AF_INET;
68 addr.sin_addr.s_addr = INADDR_ANY;
69 CHECK_ERR(bind(acceptSocket_,
70 reinterpret_cast<const sockaddr*>(&addr),
73 CHECK_ERR(listen(acceptSocket_, 10));
75 socklen_t addrLen = sizeof(addr);
76 CHECK_ERR(getsockname(acceptSocket_,
77 reinterpret_cast<sockaddr*>(&addr),
80 port_ = ntohs(addr.sin_port);
82 serverThread_ = std::thread([this] {
83 while (stop_ == NO_STOP) {
85 socklen_t peerLen = sizeof(peer);
86 int fd = accept(acceptSocket_,
87 reinterpret_cast<sockaddr*>(&peer),
93 if (errno == EINVAL || errno == ENOTSOCK) { // socket broken
98 shutdownSocketSet.add(fd);
102 if (stop_ != NO_STOP) {
103 closeClients(stop_ == ABORTIVE);
106 shutdownSocketSet.close(acceptSocket_);
112 int Server::closeClients(bool abortive) {
113 for (int fd : fds_) {
115 struct linger l = {1, 0};
116 CHECK_ERR(setsockopt(fd, SOL_SOCKET, SO_LINGER, &l, sizeof(l)));
118 shutdownSocketSet.close(fd);
125 void Server::stop(bool abortive) {
126 stop_ = abortive ? ABORTIVE : ORDERLY;
127 shutdown(acceptSocket_, SHUT_RDWR);
130 void Server::join() {
131 serverThread_.join();
134 int createConnectedSocket(int port) {
135 int sock = socket(PF_INET, SOCK_STREAM, 0);
138 addr.sin_family = AF_INET;
139 addr.sin_port = htons(port);
140 addr.sin_addr.s_addr = htonl((127 << 24) | 1); // XXX
141 CHECK_ERR(connect(sock,
142 reinterpret_cast<const sockaddr*>(&addr),
147 void runCloseTest(bool abortive) {
150 int sock = createConnectedSocket(server.port());
152 std::thread stopper([&server, abortive] {
153 std::this_thread::sleep_for(std::chrono::milliseconds(200));
154 server.stop(abortive);
159 int r = read(sock, &c, 1);
163 EXPECT_EQ(ECONNRESET, e);
172 EXPECT_EQ(0, server.closeClients(false)); // closed by server when it exited
175 TEST(ShutdownSocketSetTest, OrderlyClose) {
179 TEST(ShutdownSocketSetTest, AbortiveClose) {
183 void runKillTest(bool abortive) {
186 int sock = createConnectedSocket(server.port());
188 std::thread killer([&server, abortive] {
189 std::this_thread::sleep_for(std::chrono::milliseconds(200));
190 shutdownSocketSet.shutdownAll(abortive);
195 int r = read(sock, &c, 1);
197 // "abortive" is just a hint for ShutdownSocketSet, so accept both
201 EXPECT_EQ(ECONNRESET, errno);
213 // NOT closed by server when it exited
214 EXPECT_EQ(1, server.closeClients(false));
217 TEST(ShutdownSocketSetTest, OrderlyKill) {
221 TEST(ShutdownSocketSetTest, AbortiveKill) {
227 int main(int argc, char *argv[]) {
228 testing::InitGoogleTest(&argc, argv);
229 gflags::ParseCommandLineFlags(&argc, &argv, true);
230 return RUN_ALL_TESTS();