folly/wangle/acceptor/ConnectionManager.h

   1 /*
   2  * Copyright 2015 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
  17 #pragma once
  18
  19 #include <folly/wangle/acceptor/ManagedConnection.h>
  20
  21 #include <chrono>
  22 #include <folly/Memory.h>
  23 #include <folly/io/async/AsyncTimeout.h>
  24 #include <folly/io/async/HHWheelTimer.h>
  25 #include <folly/io/async/DelayedDestruction.h>
  26 #include <folly/io/async/EventBase.h>
  27
  28 namespace folly { namespace wangle {
  29
  30 /**
  31  * A ConnectionManager keeps track of ManagedConnections.
  32  */
  33 class ConnectionManager: public folly::DelayedDestruction {
  34  public:
  35
  36   /**
  37    * Interface for an optional observer that's notified about
  38    * various events in a ConnectionManager
  39    */
  40   class Callback {
  41   public:
  42     virtual ~Callback() {}
  43
  44     /**
  45      * Invoked when the number of connections managed by the
  46      * ConnectionManager changes from nonzero to zero.
  47      */
  48     virtual void onEmpty(const ConnectionManager& cm) = 0;
  49
  50     /**
  51      * Invoked when a connection is added to the ConnectionManager.
  52      */
  53     virtual void onConnectionAdded(const ConnectionManager& cm) = 0;
  54
  55     /**
  56      * Invoked when a connection is removed from the ConnectionManager.
  57      */
  58     virtual void onConnectionRemoved(const ConnectionManager& cm) = 0;
  59   };
  60
  61   typedef std::unique_ptr<ConnectionManager, Destructor> UniquePtr;
  62
  63   /**
  64    * Returns a new instance of ConnectionManager wrapped in a unique_ptr
  65    */
  66   template<typename... Args>
  67   static UniquePtr makeUnique(Args&&... args) {
  68     return folly::make_unique<ConnectionManager, Destructor>(
  69       std::forward<Args>(args)...);
  70   }
  71
  72   /**
  73    * Constructor not to be used by itself.
  74    */
  75   ConnectionManager(folly::EventBase* eventBase,
  76                     std::chrono::milliseconds timeout,
  77                     Callback* callback = nullptr);
  78
  79   /**
  80    * Add a connection to the set of connections managed by this
  81    * ConnectionManager.
  82    *
  83    * @param connection     The connection to add.
  84    * @param timeout        Whether to immediately register this connection
  85    *                         for an idle timeout callback.
  86    */
  87   void addConnection(ManagedConnection* connection,
  88       bool timeout = false);
  89
  90   /**
  91    * Schedule a timeout callback for a connection.
  92    */
  93   void scheduleTimeout(ManagedConnection* const connection,
  94                        std::chrono::milliseconds timeout);
  95
  96   /*
  97    * Schedule a callback on the wheel timer
  98    */
  99   void scheduleTimeout(folly::HHWheelTimer::Callback* callback,
 100                        std::chrono::milliseconds timeout);
 101
 102   /**
 103    * Remove a connection from this ConnectionManager and, if
 104    * applicable, cancel the pending timeout callback that the
 105    * ConnectionManager has scheduled for the connection.
 106    *
 107    * @note This method does NOT destroy the connection.
 108    */
 109   void removeConnection(ManagedConnection* connection);
 110
 111   /* Begin gracefully shutting down connections in this ConnectionManager.
 112    * Notify all connections of pending shutdown, and after idleGrace,
 113    * begin closing idle connections.
 114    */
 115   void initiateGracefulShutdown(std::chrono::milliseconds idleGrace);
 116
 117   /**
 118    * Destroy all connections Managed by this ConnectionManager, even
 119    * the ones that are busy.
 120    */
 121   void dropAllConnections();
 122
 123   size_t getNumConnections() const { return conns_.size(); }
 124
 125   template <typename F>
 126   void iterateConns(F func) {
 127     auto it = conns_.begin();
 128     while ( it != conns_.end()) {
 129       func(&(*it));
 130       it++;
 131     }
 132   }
 133
 134   std::chrono::milliseconds getDefaultTimeout() const {
 135     return timeout_;
 136   }
 137
 138  private:
 139   class CloseIdleConnsCallback :
 140       public folly::EventBase::LoopCallback,
 141       public folly::AsyncTimeout {
 142    public:
 143     explicit CloseIdleConnsCallback(ConnectionManager* manager)
 144         : folly::AsyncTimeout(manager->eventBase_),
 145           manager_(manager) {}
 146
 147     void runLoopCallback() noexcept override {
 148       VLOG(3) << "Draining more conns from loop callback";
 149       manager_->drainAllConnections();
 150     }
 151
 152     void timeoutExpired() noexcept override {
 153       VLOG(3) << "Idle grace expired";
 154       manager_->drainAllConnections();
 155     }
 156
 157    private:
 158     ConnectionManager* manager_;
 159   };
 160
 161   enum class ShutdownAction : uint8_t {
 162     /**
 163      * Drain part 1: inform remote that you will soon reject new requests.
 164      */
 165     DRAIN1 = 0,
 166     /**
 167      * Drain part 2: start rejecting new requests.
 168      */
 169     DRAIN2 = 1,
 170   };
 171
 172   ~ConnectionManager() {}
 173
 174   ConnectionManager(const ConnectionManager&) = delete;
 175   ConnectionManager& operator=(ConnectionManager&) = delete;
 176
 177   /**
 178    * Destroy all connections managed by this ConnectionManager that
 179    * are currently idle, as determined by a call to each ManagedConnection's
 180    * isBusy() method.
 181    */
 182   void drainAllConnections();
 183
 184   /** All connections */
 185   folly::CountedIntrusiveList<
 186     ManagedConnection,&ManagedConnection::listHook_> conns_;
 187
 188   /** Connections that currently are registered for timeouts */
 189   folly::HHWheelTimer::UniquePtr connTimeouts_;
 190
 191   /** Optional callback to notify of state changes */
 192   Callback* callback_;
 193
 194   /** Event base in which we run */
 195   folly::EventBase* eventBase_;
 196
 197   /** Iterator to the next connection to shed; used by drainAllConnections() */
 198   folly::CountedIntrusiveList<
 199     ManagedConnection,&ManagedConnection::listHook_>::iterator idleIterator_;
 200   CloseIdleConnsCallback idleLoopCallback_;
 201   ShutdownAction action_{ShutdownAction::DRAIN1};
 202   std::chrono::milliseconds timeout_;
 203 };
 204
 205 }} // folly::wangle