132875a10df5cea31e922b87130599bf0f01b04f
[folly.git] / folly / io / async / test / AsyncSocketTest2.h
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 #pragma once
17
18 #include <deque>
19 #include <exception>
20 #include <functional>
21 #include <string>
22
23 #include <folly/RWSpinLock.h>
24
25 #include <folly/io/async/AsyncServerSocket.h>
26
27 namespace folly {
28 namespace test {
29
30 /**
31  * Helper ConnectionEventCallback class for the test code.
32  * It maintains counters protected by a spin lock.
33  */
34 class TestConnectionEventCallback
35     : public AsyncServerSocket::ConnectionEventCallback {
36  public:
37   void onConnectionAccepted(
38       const int /* socket */,
39       const SocketAddress& /* addr */) noexcept override {
40     folly::RWSpinLock::WriteHolder holder(spinLock_);
41     connectionAccepted_++;
42   }
43
44   void onConnectionAcceptError(const int /* err */) noexcept override {
45     folly::RWSpinLock::WriteHolder holder(spinLock_);
46     connectionAcceptedError_++;
47   }
48
49   void onConnectionDropped(
50       const int /* socket */,
51       const SocketAddress& /* addr */) noexcept override {
52     folly::RWSpinLock::WriteHolder holder(spinLock_);
53     connectionDropped_++;
54   }
55
56   void onConnectionEnqueuedForAcceptorCallback(
57       const int /* socket */,
58       const SocketAddress& /* addr */) noexcept override {
59     folly::RWSpinLock::WriteHolder holder(spinLock_);
60     connectionEnqueuedForAcceptCallback_++;
61   }
62
63   void onConnectionDequeuedByAcceptorCallback(
64       const int /* socket */,
65       const SocketAddress& /* addr */) noexcept override {
66     folly::RWSpinLock::WriteHolder holder(spinLock_);
67     connectionDequeuedByAcceptCallback_++;
68   }
69
70   void onBackoffStarted() noexcept override {
71     folly::RWSpinLock::WriteHolder holder(spinLock_);
72     backoffStarted_++;
73   }
74
75   void onBackoffEnded() noexcept override {
76     folly::RWSpinLock::WriteHolder holder(spinLock_);
77     backoffEnded_++;
78   }
79
80   void onBackoffError() noexcept override {
81     folly::RWSpinLock::WriteHolder holder(spinLock_);
82     backoffError_++;
83   }
84
85   unsigned int getConnectionAccepted() const {
86     folly::RWSpinLock::ReadHolder holder(spinLock_);
87     return connectionAccepted_;
88   }
89
90   unsigned int getConnectionAcceptedError() const {
91     folly::RWSpinLock::ReadHolder holder(spinLock_);
92     return connectionAcceptedError_;
93   }
94
95   unsigned int getConnectionDropped() const {
96     folly::RWSpinLock::ReadHolder holder(spinLock_);
97     return connectionDropped_;
98   }
99
100   unsigned int getConnectionEnqueuedForAcceptCallback() const {
101     folly::RWSpinLock::ReadHolder holder(spinLock_);
102     return connectionEnqueuedForAcceptCallback_;
103   }
104
105   unsigned int getConnectionDequeuedByAcceptCallback() const {
106     folly::RWSpinLock::ReadHolder holder(spinLock_);
107     return connectionDequeuedByAcceptCallback_;
108   }
109
110   unsigned int getBackoffStarted() const {
111     folly::RWSpinLock::ReadHolder holder(spinLock_);
112     return backoffStarted_;
113   }
114
115   unsigned int getBackoffEnded() const {
116     folly::RWSpinLock::ReadHolder holder(spinLock_);
117     return backoffEnded_;
118   }
119
120   unsigned int getBackoffError() const {
121     folly::RWSpinLock::ReadHolder holder(spinLock_);
122     return backoffError_;
123   }
124
125  private:
126   mutable folly::RWSpinLock spinLock_;
127   unsigned int connectionAccepted_{0};
128   unsigned int connectionAcceptedError_{0};
129   unsigned int connectionDropped_{0};
130   unsigned int connectionEnqueuedForAcceptCallback_{0};
131   unsigned int connectionDequeuedByAcceptCallback_{0};
132   unsigned int backoffStarted_{0};
133   unsigned int backoffEnded_{0};
134   unsigned int backoffError_{0};
135 };
136
137 /**
138  * Helper AcceptCallback class for the test code
139  * It records the callbacks that were invoked, and also supports calling
140  * generic std::function objects in each callback.
141  */
142 class TestAcceptCallback : public AsyncServerSocket::AcceptCallback {
143  public:
144   enum EventType { TYPE_START, TYPE_ACCEPT, TYPE_ERROR, TYPE_STOP };
145   struct EventInfo {
146     EventInfo(int fd, const folly::SocketAddress& addr)
147         : type(TYPE_ACCEPT), fd(fd), address(addr), errorMsg() {}
148     explicit EventInfo(const std::string& msg)
149         : type(TYPE_ERROR), fd(-1), address(), errorMsg(msg) {}
150     explicit EventInfo(EventType et)
151         : type(et), fd(-1), address(), errorMsg() {}
152
153     EventType type;
154     int fd; // valid for TYPE_ACCEPT
155     folly::SocketAddress address; // valid for TYPE_ACCEPT
156     std::string errorMsg; // valid for TYPE_ERROR
157   };
158   typedef std::deque<EventInfo> EventList;
159
160   TestAcceptCallback()
161       : connectionAcceptedFn_(),
162         acceptErrorFn_(),
163         acceptStoppedFn_(),
164         events_() {}
165
166   std::deque<EventInfo>* getEvents() {
167     return &events_;
168   }
169
170   void setConnectionAcceptedFn(
171       const std::function<void(int, const folly::SocketAddress&)>& fn) {
172     connectionAcceptedFn_ = fn;
173   }
174   void setAcceptErrorFn(const std::function<void(const std::exception&)>& fn) {
175     acceptErrorFn_ = fn;
176   }
177   void setAcceptStartedFn(const std::function<void()>& fn) {
178     acceptStartedFn_ = fn;
179   }
180   void setAcceptStoppedFn(const std::function<void()>& fn) {
181     acceptStoppedFn_ = fn;
182   }
183
184   void connectionAccepted(
185       int fd,
186       const folly::SocketAddress& clientAddr) noexcept override {
187     events_.emplace_back(fd, clientAddr);
188
189     if (connectionAcceptedFn_) {
190       connectionAcceptedFn_(fd, clientAddr);
191     }
192   }
193   void acceptError(const std::exception& ex) noexcept override {
194     events_.emplace_back(ex.what());
195
196     if (acceptErrorFn_) {
197       acceptErrorFn_(ex);
198     }
199   }
200   void acceptStarted() noexcept override {
201     events_.emplace_back(TYPE_START);
202
203     if (acceptStartedFn_) {
204       acceptStartedFn_();
205     }
206   }
207   void acceptStopped() noexcept override {
208     events_.emplace_back(TYPE_STOP);
209
210     if (acceptStoppedFn_) {
211       acceptStoppedFn_();
212     }
213   }
214
215  private:
216   std::function<void(int, const folly::SocketAddress&)> connectionAcceptedFn_;
217   std::function<void(const std::exception&)> acceptErrorFn_;
218   std::function<void()> acceptStartedFn_;
219   std::function<void()> acceptStoppedFn_;
220
221   std::deque<EventInfo> events_;
222 };
223 } // namespace test
224 } // namespace folly