ServerBootstrap-inl.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 <folly/ExceptionWrapper.h>
#include <folly/SharedMutex.h>
#include <folly/io/async/DelayedDestruction.h>
#include <folly/io/async/EventBaseManager.h>
#include <wangle/acceptor/Acceptor.h>
#include <wangle/acceptor/ManagedConnection.h>
#include <wangle/bootstrap/ServerSocketFactory.h>
#include <wangle/channel/Handler.h>
#include <wangle/channel/Pipeline.h>
#include <folly/executors/IOThreadPoolExecutor.h>
#include <wangle/ssl/SSLStats.h>

namespace wangle {

class AcceptorException : public std::runtime_error {
 public:
  enum class ExceptionType {
    UNKNOWN = 0,
    TIMED_OUT = 1,
    DROPPED = 2,
    ACCEPT_STOPPED = 3,
    DRAIN_CONN_PCT = 4,
    DROP_CONN_PCT = 5,
    FORCE_STOP = 6,
    INTERNAL_ERROR = 7,
  };

  explicit AcceptorException(ExceptionType type) :
      std::runtime_error(""), type_(type), pct_(0.0) {}

  explicit AcceptorException(ExceptionType type, const std::string& message) :
      std::runtime_error(message), type_(type), pct_(0.0) {}

  explicit AcceptorException(ExceptionType type, const std::string& message,
                             double pct) :
      std::runtime_error(message), type_(type), pct_(pct) {}

  ExceptionType getType() const noexcept { return type_; }
  double getPct() const noexcept { return pct_; }

 protected:
  const ExceptionType type_;
  // the percentage of connections to be drained or dropped during the shutdown
  const double pct_;
};

template <typename Pipeline>
class ServerAcceptor
    : public Acceptor
    , public wangle::InboundHandler<AcceptPipelineType> {
 public:
  class ServerConnection : public wangle::ManagedConnection,
                           public wangle::PipelineManager {
   public:
    explicit ServerConnection(typename Pipeline::Ptr pipeline)
        : pipeline_(std::move(pipeline)) {
      pipeline_->setPipelineManager(this);
    }

    void timeoutExpired() noexcept override {
      auto ew = folly::make_exception_wrapper<AcceptorException>(
          AcceptorException::ExceptionType::TIMED_OUT, "timeout");
      pipeline_->readException(ew);
    }

    void describe(std::ostream&) const override {}
    bool isBusy() const override {
      return true;
    }
    void notifyPendingShutdown() override {}
    void closeWhenIdle() override {}
    void dropConnection() override {
      auto ew = folly::make_exception_wrapper<AcceptorException>(
          AcceptorException::ExceptionType::DROPPED, "dropped");
      pipeline_->readException(ew);
    }
    void dumpConnectionState(uint8_t /* loglevel */) override {}

    void deletePipeline(wangle::PipelineBase* p) override {
      CHECK(p == pipeline_.get());
      destroy();
    }

    void init() {
      pipeline_->transportActive();
    }

    void refreshTimeout() override {
      resetTimeout();
    }

   private:
    ~ServerConnection() override {
      pipeline_->setPipelineManager(nullptr);
    }
    typename Pipeline::Ptr pipeline_;
  };

  explicit ServerAcceptor(
      std::shared_ptr<AcceptPipelineFactory> acceptPipelineFactory,
      std::shared_ptr<PipelineFactory<Pipeline>> childPipelineFactory,
      const ServerSocketConfig& accConfig)
      : Acceptor(accConfig),
        acceptPipelineFactory_(acceptPipelineFactory),
        childPipelineFactory_(childPipelineFactory) {
  }

  void init(folly::AsyncServerSocket* serverSocket,
            folly::EventBase* eventBase,
            SSLStats* stats = nullptr) override {
    Acceptor::init(serverSocket, eventBase, stats);

    acceptPipeline_ = acceptPipelineFactory_->newPipeline(this);

    if (childPipelineFactory_) {
      // This means a custom AcceptPipelineFactory was not passed in via
      // pipeline() and we're using the DefaultAcceptPipelineFactory.
      // Add the default inbound handler here.
      acceptPipeline_->addBack(this);
    }
    acceptPipeline_->finalize();
  }

  void read(Context*, AcceptPipelineType conn) override {
    if (conn.type() != typeid(ConnInfo&)) {
      return;
    }

    auto connInfo = boost::get<ConnInfo&>(conn);
    folly::AsyncTransportWrapper::UniquePtr transport(connInfo.sock);

    // Setup local and remote addresses
    auto tInfoPtr = std::make_shared<TransportInfo>(connInfo.tinfo);
    tInfoPtr->localAddr =
      std::make_shared<folly::SocketAddress>(accConfig_.bindAddress);
    transport->getLocalAddress(tInfoPtr->localAddr.get());
    tInfoPtr->remoteAddr =
      std::make_shared<folly::SocketAddress>(*connInfo.clientAddr);
    tInfoPtr->appProtocol =
      std::make_shared<std::string>(connInfo.nextProtoName);

    auto pipeline = childPipelineFactory_->newPipeline(
      std::shared_ptr<folly::AsyncTransportWrapper>(
        transport.release(), folly::DelayedDestruction::Destructor()));
    pipeline->setTransportInfo(tInfoPtr);
    auto connection = new ServerConnection(std::move(pipeline));
    Acceptor::addConnection(connection);
    connection->init();
  }

  // Null implementation to terminate the call in this handler
  // and suppress warnings
  void readEOF(Context*) override {}
  void readException(Context*, folly::exception_wrapper) override {}

  /* See Acceptor::onNewConnection for details */
  void onNewConnection(folly::AsyncTransportWrapper::UniquePtr transport,
                       const folly::SocketAddress* clientAddr,
                       const std::string& nextProtocolName,
                       SecureTransportType secureTransportType,
                       const TransportInfo& tinfo) override {
    ConnInfo connInfo = {transport.release(), clientAddr, nextProtocolName,
                         secureTransportType, tinfo};
    acceptPipeline_->read(connInfo);
  }

  // notify the acceptors in the acceptPipeline to drain & drop conns
  void acceptStopped() noexcept override {
    auto ew = folly::make_exception_wrapper<AcceptorException>(
      AcceptorException::ExceptionType::ACCEPT_STOPPED,
      "graceful shutdown timeout");

    acceptPipeline_->readException(ew);
    Acceptor::acceptStopped();
  }

  void drainConnections(double pct) noexcept override {
    auto ew = folly::make_exception_wrapper<AcceptorException>(
      AcceptorException::ExceptionType::DRAIN_CONN_PCT,
      "draining some connections", pct);

    acceptPipeline_->readException(ew);
    Acceptor::drainConnections(pct);
  }

  void dropConnections(double pct) noexcept override {
    auto ew = folly::make_exception_wrapper<AcceptorException>(
      AcceptorException::ExceptionType::DROP_CONN_PCT,
      "dropping some connections", pct);

    acceptPipeline_->readException(ew);
    Acceptor::dropConnections(pct);
  }

  void forceStop() noexcept override {
    auto ew = folly::make_exception_wrapper<AcceptorException>(
      AcceptorException::ExceptionType::FORCE_STOP,
      "hard shutdown timeout");

    acceptPipeline_->readException(ew);
    Acceptor::forceStop();
  }

  // UDP thunk
  void onDataAvailable(std::shared_ptr<folly::AsyncUDPSocket> socket,
                       const folly::SocketAddress& addr,
                       std::unique_ptr<folly::IOBuf> buf,
                       bool /* truncated */) noexcept override {
    acceptPipeline_->read(
        AcceptPipelineType(make_tuple(buf.release(), socket, addr)));
  }

  void onConnectionAdded(const ManagedConnection*) override {
    acceptPipeline_->read(ConnEvent::CONN_ADDED);
  }

  void onConnectionRemoved(const ManagedConnection*) override {
    acceptPipeline_->read(ConnEvent::CONN_REMOVED);
  }

  void sslConnectionError(const folly::exception_wrapper& ex) override {
    acceptPipeline_->readException(ex);
    Acceptor::sslConnectionError(ex);
  }

 private:
  std::shared_ptr<AcceptPipelineFactory> acceptPipelineFactory_;
  std::shared_ptr<AcceptPipeline> acceptPipeline_;
  std::shared_ptr<PipelineFactory<Pipeline>> childPipelineFactory_;
};

template <typename Pipeline>
class ServerAcceptorFactory : public AcceptorFactory {
 public:
  explicit ServerAcceptorFactory(
      std::shared_ptr<AcceptPipelineFactory> acceptPipelineFactory,
      std::shared_ptr<PipelineFactory<Pipeline>> childPipelineFactory,
      const ServerSocketConfig& accConfig)
      : acceptPipelineFactory_(acceptPipelineFactory),
        childPipelineFactory_(childPipelineFactory),
        accConfig_(accConfig) {}

  std::shared_ptr<Acceptor> newAcceptor(folly::EventBase* base) override {
    auto acceptor = std::make_shared<ServerAcceptor<Pipeline>>(
        acceptPipelineFactory_, childPipelineFactory_, accConfig_);
    acceptor->init(nullptr, base, nullptr);
    return acceptor;
  }

 private:
  std::shared_ptr<AcceptPipelineFactory> acceptPipelineFactory_;
  std::shared_ptr<PipelineFactory<Pipeline>> childPipelineFactory_;
  ServerSocketConfig accConfig_;
};

class ServerWorkerPool : public folly::ThreadPoolExecutor::Observer {
 public:
  explicit ServerWorkerPool(
    std::shared_ptr<AcceptorFactory> acceptorFactory,
    folly::IOThreadPoolExecutor* exec,
    std::shared_ptr<std::vector<std::shared_ptr<folly::AsyncSocketBase>>> sockets,
    std::shared_ptr<ServerSocketFactory> socketFactory)
      : workers_(std::make_shared<WorkerMap>())
      , workersMutex_(std::make_shared<Mutex>())
      , acceptorFactory_(acceptorFactory)
      , exec_(exec)
      , sockets_(sockets)
      , socketFactory_(socketFactory) {
    CHECK(exec);
  }

  template <typename F>
  void forEachWorker(F&& f) const;

  void threadStarted(folly::ThreadPoolExecutor::ThreadHandle*) override;
  void threadStopped(folly::ThreadPoolExecutor::ThreadHandle*) override;
  void threadPreviouslyStarted(
      folly::ThreadPoolExecutor::ThreadHandle* thread) override {
    threadStarted(thread);
  }
  void threadNotYetStopped(
      folly::ThreadPoolExecutor::ThreadHandle* thread) override {
    threadStopped(thread);
  }

 private:
  using WorkerMap = std::map<folly::ThreadPoolExecutor::ThreadHandle*,
        std::shared_ptr<Acceptor>>;
  using Mutex = folly::SharedMutexReadPriority;

  std::shared_ptr<WorkerMap> workers_;
  std::shared_ptr<Mutex> workersMutex_;
  std::shared_ptr<AcceptorFactory> acceptorFactory_;
  folly::IOThreadPoolExecutor* exec_{nullptr};
  std::shared_ptr<std::vector<std::shared_ptr<folly::AsyncSocketBase>>>
      sockets_;
  std::shared_ptr<ServerSocketFactory> socketFactory_;
};

template <typename F>
void ServerWorkerPool::forEachWorker(F&& f) const {
  Mutex::ReadHolder holder(workersMutex_.get());
  for (const auto& kv : *workers_) {
    f(kv.second.get());
  }
}

class DefaultAcceptPipelineFactory : public AcceptPipelineFactory {
 public:
  typename AcceptPipeline::Ptr newPipeline(Acceptor*) override {
    return AcceptPipeline::create();
  }
};

} // namespace wangle