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   virtual void onConnectionAccepted(
  38       const int /* socket */,
  39       const SocketAddress& /* addr */) noexcept override {
  40     folly::RWSpinLock::WriteHolder holder(spinLock_);
  41     connectionAccepted_++;
  42   }
  43
  44   virtual void onConnectionAcceptError(const int /* err */) noexcept override {
  45     folly::RWSpinLock::WriteHolder holder(spinLock_);
  46     connectionAcceptedError_++;
  47   }
  48
  49   virtual void onConnectionDropped(
  50       const int /* socket */,
  51       const SocketAddress& /* addr */) noexcept override {
  52     folly::RWSpinLock::WriteHolder holder(spinLock_);
  53     connectionDropped_++;
  54   }
  55
  56   virtual void onConnectionEnqueuedForAcceptorCallback(
  57       const int /* socket */,
  58       const SocketAddress& /* addr */) noexcept override {
  59     folly::RWSpinLock::WriteHolder holder(spinLock_);
  60     connectionEnqueuedForAcceptCallback_++;
  61   }
  62
  63   virtual void onConnectionDequeuedByAcceptorCallback(
  64       const int /* socket */,
  65       const SocketAddress& /* addr */) noexcept override {
  66     folly::RWSpinLock::WriteHolder holder(spinLock_);
  67     connectionDequeuedByAcceptCallback_++;
  68   }
  69
  70   virtual void onBackoffStarted() noexcept override {
  71     folly::RWSpinLock::WriteHolder holder(spinLock_);
  72     backoffStarted_++;
  73   }
  74
  75   virtual void onBackoffEnded() noexcept override {
  76     folly::RWSpinLock::WriteHolder holder(spinLock_);
  77     backoffEnded_++;
  78   }
  79
  80   virtual 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 }
 224 }