folly/wangle/acceptor/ConnectionManager.cpp

   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 #include <folly/wangle/acceptor/ConnectionManager.h>
  18
  19 #include <glog/logging.h>
  20 #include <folly/io/async/EventBase.h>
  21
  22 using folly::HHWheelTimer;
  23 using std::chrono::milliseconds;
  24
  25 namespace folly { namespace wangle {
  26
  27 ConnectionManager::ConnectionManager(EventBase* eventBase,
  28     milliseconds timeout, Callback* callback)
  29   : connTimeouts_(new HHWheelTimer(eventBase)),
  30     callback_(callback),
  31     eventBase_(eventBase),
  32     idleIterator_(conns_.end()),
  33     idleLoopCallback_(this),
  34     timeout_(timeout) {
  35
  36 }
  37
  38 void
  39 ConnectionManager::addConnection(ManagedConnection* connection,
  40     bool timeout) {
  41   CHECK_NOTNULL(connection);
  42   ConnectionManager* oldMgr = connection->getConnectionManager();
  43   if (oldMgr != this) {
  44     if (oldMgr) {
  45       // 'connection' was being previously managed in a different thread.
  46       // We must remove it from that manager before adding it to this one.
  47       oldMgr->removeConnection(connection);
  48     }
  49     conns_.push_back(*connection);
  50     connection->setConnectionManager(this);
  51     if (callback_) {
  52       callback_->onConnectionAdded(*this);
  53     }
  54   }
  55   if (timeout) {
  56     scheduleTimeout(connection, timeout_);
  57   }
  58 }
  59
  60 void
  61 ConnectionManager::scheduleTimeout(ManagedConnection* const connection,
  62     std::chrono::milliseconds timeout) {
  63   if (timeout > std::chrono::milliseconds(0)) {
  64     connTimeouts_->scheduleTimeout(connection, timeout);
  65   }
  66 }
  67
  68 void ConnectionManager::scheduleTimeout(
  69   folly::HHWheelTimer::Callback* callback,
  70   std::chrono::milliseconds timeout) {
  71   connTimeouts_->scheduleTimeout(callback, timeout);
  72 }
  73
  74 void
  75 ConnectionManager::removeConnection(ManagedConnection* connection) {
  76   if (connection->getConnectionManager() == this) {
  77     connection->cancelTimeout();
  78     connection->setConnectionManager(nullptr);
  79
  80     // Un-link the connection from our list, being careful to keep the iterator
  81     // that we're using for idle shedding valid
  82     auto it = conns_.iterator_to(*connection);
  83     if (it == idleIterator_) {
  84       ++idleIterator_;
  85     }
  86     conns_.erase(it);
  87
  88     if (callback_) {
  89       callback_->onConnectionRemoved(*this);
  90       if (getNumConnections() == 0) {
  91         callback_->onEmpty(*this);
  92       }
  93     }
  94   }
  95 }
  96
  97 void
  98 ConnectionManager::initiateGracefulShutdown(
  99   std::chrono::milliseconds idleGrace) {
 100   if (idleGrace.count() > 0) {
 101     idleLoopCallback_.scheduleTimeout(idleGrace);
 102     VLOG(3) << "Scheduling idle grace period of " << idleGrace.count() << "ms";
 103   } else {
 104     action_ = ShutdownAction::DRAIN2;
 105     VLOG(3) << "proceeding directly to closing idle connections";
 106   }
 107   drainAllConnections();
 108 }
 109
 110 void
 111 ConnectionManager::drainAllConnections() {
 112   DestructorGuard g(this);
 113   size_t numCleared = 0;
 114   size_t numKept = 0;
 115
 116   auto it = idleIterator_ == conns_.end() ?
 117     conns_.begin() : idleIterator_;
 118
 119   while (it != conns_.end() && (numKept + numCleared) < 64) {
 120     ManagedConnection& conn = *it++;
 121     if (action_ == ShutdownAction::DRAIN1) {
 122       conn.notifyPendingShutdown();
 123     } else {
 124       // Second time around: close idle sessions. If they aren't idle yet,
 125       // have them close when they are idle
 126       if (conn.isBusy()) {
 127         numKept++;
 128       } else {
 129         numCleared++;
 130       }
 131       conn.closeWhenIdle();
 132     }
 133   }
 134
 135   if (action_ == ShutdownAction::DRAIN2) {
 136     VLOG(2) << "Idle connections cleared: " << numCleared <<
 137       ", busy conns kept: " << numKept;
 138   }
 139   if (it != conns_.end()) {
 140     idleIterator_ = it;
 141     eventBase_->runInLoop(&idleLoopCallback_);
 142   } else {
 143     action_ = ShutdownAction::DRAIN2;
 144   }
 145 }
 146
 147 void
 148 ConnectionManager::dropAllConnections() {
 149   DestructorGuard g(this);
 150
 151   // Iterate through our connection list, and drop each connection.
 152   VLOG(3) << "connections to drop: " << conns_.size();
 153   idleLoopCallback_.cancelTimeout();
 154   unsigned i = 0;
 155   while (!conns_.empty()) {
 156     ManagedConnection& conn = conns_.front();
 157     conns_.pop_front();
 158     conn.cancelTimeout();
 159     conn.setConnectionManager(nullptr);
 160     // For debugging purposes, dump information about the first few
 161     // connections.
 162     static const unsigned MAX_CONNS_TO_DUMP = 2;
 163     if (++i <= MAX_CONNS_TO_DUMP) {
 164       conn.dumpConnectionState(3);
 165     }
 166     conn.dropConnection();
 167   }
 168   idleIterator_ = conns_.end();
 169   idleLoopCallback_.cancelLoopCallback();
 170
 171   if (callback_) {
 172     callback_->onEmpty(*this);
 173   }
 174 }
 175
 176 }} // folly::wangle