ConnectionManager.h

Go to the documentation of this file.

/*
 * Copyright 2017-present Facebook, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once

#include <wangle/acceptor/ManagedConnection.h>

#include <chrono>
#include <iterator>
#include <utility>
#include <folly/Memory.h>
#include <folly/io/async/AsyncTimeout.h>
#include <folly/io/async/HHWheelTimer.h>
#include <folly/io/async/DelayedDestruction.h>
#include <folly/io/async/EventBase.h>

namespace wangle {

class ConnectionManager: public folly::DelayedDestruction,
                         private ManagedConnection::Callback {
 public:

  class Callback {
  public:
    virtual ~Callback() = default;

    virtual void onEmpty(const ConnectionManager& cm) = 0;

    virtual void onConnectionAdded(const ManagedConnection* conn) = 0;

    virtual void onConnectionRemoved(const ManagedConnection* conn) = 0;
  };

  typedef std::unique_ptr<ConnectionManager, Destructor> UniquePtr;

  typedef folly::CountedIntrusiveList<
    ManagedConnection, &ManagedConnection::listHook_>::iterator ConnectionIterator;
  template<typename... Args>
  static UniquePtr makeUnique(Args&&... args) {
    return UniquePtr(new ConnectionManager(std::forward<Args>(args)...));
  }

  ConnectionManager(folly::EventBase* eventBase,
                    std::chrono::milliseconds timeout,
                    Callback* callback = nullptr);

  void addConnection(ManagedConnection* connection,
      bool timeout = false);

  void scheduleTimeout(ManagedConnection* const connection,
                       std::chrono::milliseconds timeout);

  /*
   * Schedule a callback on the wheel timer
   */
  void scheduleTimeout(folly::HHWheelTimer::Callback* callback,
                       std::chrono::milliseconds timeout);

  void removeConnection(ManagedConnection* connection);

  /* Begin gracefully shutting down connections in this ConnectionManager.
   * Notify all connections of pending shutdown, and after idleGrace,
   * begin closing idle connections.
   */
  void initiateGracefulShutdown(std::chrono::milliseconds idleGrace);

  void drainConnections(double pct, std::chrono::milliseconds idleGrace);

  void dropAllConnections();

  void dropConnections(double pct);

  size_t getNumConnections() const { return conns_.size(); }

  template <typename F>
  void iterateConns(F func) {
    auto it = conns_.begin();
    while ( it != conns_.end()) {
      func(&(*it));
      it++;
    }
  }

  std::chrono::milliseconds getDefaultTimeout() const {
    return timeout_;
  }

  void setLoweredIdleTimeout(std::chrono::milliseconds timeout) {
    CHECK(timeout >= std::chrono::milliseconds(0));
    CHECK(timeout <= timeout_);
    idleConnEarlyDropThreshold_ = timeout;
  }

  size_t dropIdleConnections(size_t num);

  void onActivated(ManagedConnection& conn) override;

  void onDeactivated(ManagedConnection& conn) override;

 private:

   enum class ShutdownState : uint8_t {
     NONE = 0,
     // All ManagedConnections receive notifyPendingShutdown
     NOTIFY_PENDING_SHUTDOWN = 1,
     // All ManagedConnections have received notifyPendingShutdown
     NOTIFY_PENDING_SHUTDOWN_COMPLETE = 2,
     // All ManagedConnections receive closeWhenIdle
     CLOSE_WHEN_IDLE = 3,
     // All ManagedConnections have received closeWhenIdle
     CLOSE_WHEN_IDLE_COMPLETE = 4,
   };

  class DrainHelper :
      public folly::EventBase::LoopCallback,
      public folly::AsyncTimeout {
   public:
    explicit DrainHelper(ConnectionManager& manager)
        : folly::AsyncTimeout(manager.eventBase_),
          manager_(manager) {}

    ShutdownState getShutdownState() {
      // only cares about full shutdown state
      if (!all_) {
        return ShutdownState::NONE;
      }
      return shutdownState_;
    }

    void setShutdownState(ShutdownState state) {
      shutdownState_ = state;
    }

    void startDrainPartial(double pct, std::chrono::milliseconds idleGrace);
    void startDrainAll(std::chrono::milliseconds idleGrace);

    void runLoopCallback() noexcept override {
      VLOG(3) << "Draining more conns from loop callback";
      drainConnections();
    }

    void timeoutExpired() noexcept override {
      VLOG(3) << "Idle grace expired";
      idleGracefulTimeoutExpired();
    }

    void drainConnections();

    void idleGracefulTimeoutExpired();

    void startDrain(std::chrono::milliseconds idleGrace);

    ConnectionIterator drainStartIterator() const {
      if (all_) {
        return manager_.conns_.begin();
      }
      auto it = manager_.conns_.begin();
      const auto conns_size = manager_.conns_.size();
      const auto numToDrain =
        std::max<size_t>(0, std::min<size_t>(conns_size, conns_size * pct_));
      std::advance(it, conns_size - numToDrain);
      return it;
    }

   private:
    bool all_{true};
    double pct_{1.0};
    ConnectionManager& manager_;
    ShutdownState shutdownState_{ShutdownState::NONE};
  };

  ~ConnectionManager() override = default;

  ConnectionManager(const ConnectionManager&) = delete;
  ConnectionManager& operator=(ConnectionManager&) = delete;

  void drainAllConnections();

  void stopDrainingForShutdown();

  void idleGracefulTimeoutExpired();

  folly::CountedIntrusiveList<
    ManagedConnection,&ManagedConnection::listHook_> conns_;

  Callback* callback_;

  folly::EventBase* eventBase_;

  ConnectionIterator drainIterator_;
  ConnectionIterator idleIterator_;
  DrainHelper drainHelper_;
  bool notifyPendingShutdown_{true};

  std::chrono::milliseconds timeout_;

  std::chrono::milliseconds idleConnEarlyDropThreshold_;
};

} // wangle