AsyncSocketTest2.h 6.4 KB

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