folly/wangle/acceptor/ConnectionManager.h

   1 /*
   2  * Copyright 2014 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* connection);
  94
  95   /*
  96    * Schedule a callback on the wheel timer
  97    */
  98   void scheduleTimeout(folly::HHWheelTimer::Callback* callback,
  99                        std::chrono::milliseconds timeout);
 100
 101   /**
 102    * Remove a connection from this ConnectionManager and, if
 103    * applicable, cancel the pending timeout callback that the
 104    * ConnectionManager has scheduled for the connection.
 105    *
 106    * @note This method does NOT destroy the connection.
 107    */
 108   void removeConnection(ManagedConnection* connection);
 109
 110   /* Begin gracefully shutting down connections in this ConnectionManager.
 111    * Notify all connections of pending shutdown, and after idleGrace,
 112    * begin closing idle connections.
 113    */
 114   void initiateGracefulShutdown(std::chrono::milliseconds idleGrace);
 115
 116   /**
 117    * Destroy all connections Managed by this ConnectionManager, even
 118    * the ones that are busy.
 119    */
 120   void dropAllConnections();
 121
 122   size_t getNumConnections() const { return conns_.size(); }
 123
 124   template <typename F>
 125   void iterateConns(F func) {
 126     auto it = conns_.begin();
 127     while ( it != conns_.end()) {
 128       func(&(*it));
 129       it++;
 130     }
 131   }
 132
 133  private:
 134   class CloseIdleConnsCallback :
 135       public folly::EventBase::LoopCallback,
 136       public folly::AsyncTimeout {
 137    public:
 138     explicit CloseIdleConnsCallback(ConnectionManager* manager)
 139         : folly::AsyncTimeout(manager->eventBase_),
 140           manager_(manager) {}
 141
 142     void runLoopCallback() noexcept override {
 143       VLOG(3) << "Draining more conns from loop callback";
 144       manager_->drainAllConnections();
 145     }
 146
 147     void timeoutExpired() noexcept override {
 148       VLOG(3) << "Idle grace expired";
 149       manager_->drainAllConnections();
 150     }
 151
 152    private:
 153     ConnectionManager* manager_;
 154   };
 155
 156   enum class ShutdownAction : uint8_t {
 157     /**
 158      * Drain part 1: inform remote that you will soon reject new requests.
 159      */
 160     DRAIN1 = 0,
 161     /**
 162      * Drain part 2: start rejecting new requests.
 163      */
 164     DRAIN2 = 1,
 165   };
 166
 167   ~ConnectionManager() {}
 168
 169   ConnectionManager(const ConnectionManager&) = delete;
 170   ConnectionManager& operator=(ConnectionManager&) = delete;
 171
 172   /**
 173    * Destroy all connections managed by this ConnectionManager that
 174    * are currently idle, as determined by a call to each ManagedConnection's
 175    * isBusy() method.
 176    */
 177   void drainAllConnections();
 178
 179   /** All connections */
 180   folly::CountedIntrusiveList<
 181     ManagedConnection,&ManagedConnection::listHook_> conns_;
 182
 183   /** Connections that currently are registered for timeouts */
 184   folly::HHWheelTimer::UniquePtr connTimeouts_;
 185
 186   /** Optional callback to notify of state changes */
 187   Callback* callback_;
 188
 189   /** Event base in which we run */
 190   folly::EventBase* eventBase_;
 191
 192   /** Iterator to the next connection to shed; used by drainAllConnections() */
 193   folly::CountedIntrusiveList<
 194     ManagedConnection,&ManagedConnection::listHook_>::iterator idleIterator_;
 195   CloseIdleConnsCallback idleLoopCallback_;
 196   ShutdownAction action_{ShutdownAction::DRAIN1};
 197   std::chrono::milliseconds timeout_;
 198 };
 199
 200 }} // folly::wangle