ServiceWorker.h

Go to the documentation of this file.

/*
 *  Copyright (c) 2015-present, Facebook, Inc.
 *  All rights reserved.
 *
 *  This source code is licensed under the BSD-style license found in the
 *  LICENSE file in the root directory of this source tree. An additional grant
 *  of patent rights can be found in the PATENTS file in the same directory.
 *
 */
#pragma once

#include <folly/io/async/AsyncServerSocket.h>
#include <list>
#include <memory>
#include <proxygen/lib/utils/AcceptorAddress.h>
#include <wangle/acceptor/Acceptor.h>
#include <wangle/acceptor/ConnectionCounter.h>

namespace proxygen {

class Service;
class RequestWorker;

class ServiceWorker {
 public:
  using AcceptorMap =
      std::map<AcceptorAddress, std::unique_ptr<wangle::Acceptor>>;

  using NamedAddressMap = std::map<std::string, AcceptorAddress>;

  ServiceWorker(Service* service, RequestWorker* worker)
      : service_(service), worker_(worker) {
  }

  virtual ~ServiceWorker() {
  }

  Service* getService() const {
    return service_;
  }

  void addServiceAcceptor(const folly::SocketAddress& address,
                          std::unique_ptr<wangle::Acceptor> acceptor) {
    addServiceAcceptor(
        AcceptorAddress(address, AcceptorAddress::AcceptorType::TCP),
        std::move(acceptor));
  }

  void addServiceAcceptor(const AcceptorAddress& accAddress,
                          std::unique_ptr<wangle::Acceptor> acceptor) {
    namedAddress_.emplace(acceptor->getName(), accAddress);
    addAcceptor(accAddress, std::move(acceptor), acceptors_);
  }

  void drainServiceAcceptor(const folly::SocketAddress& address) {
    drainServiceAcceptor(
        AcceptorAddress(address, AcceptorAddress::AcceptorType::TCP));
  }

  void drainServiceAcceptor(const AcceptorAddress& accAddress) {
    // Move the old acceptor to drainingAcceptors_ if present
    const auto& it = acceptors_.find(accAddress);
    if (it != acceptors_.end()) {
      auto name = it->second->getName();
      addAcceptor(accAddress, std::move(it->second), drainingAcceptors_);
      acceptors_.erase(it);
      namedAddress_.erase(name);
    }
  }

  RequestWorker* getRequestWorker() const {
    return worker_;
  }

  const AcceptorMap& getAcceptors() {
    return acceptors_;
  }

  wangle::Acceptor* getAcceptorByName(std::string name) const {
    auto it = namedAddress_.find(name);
    if (it != namedAddress_.end()) {
      auto it2 = acceptors_.find(it->second);
      if (it2 != acceptors_.end()) {
        return it2->second.get();
      }
    }
    return nullptr;
  }

  const AcceptorMap& getDrainingAcceptors() {
    return drainingAcceptors_;
  }

  // Flush any thread-local stats that the service is tracking
  virtual void flushStats() {
  }

  // Destruct all the acceptors
  virtual void clearAcceptors() {
    acceptors_.clear();
    namedAddress_.clear();
    drainingAcceptors_.clear();
  }

  virtual void clearDrainingAcceptors() {
    drainingAcceptors_.clear();
  }

  wangle::IConnectionCounter* getConnectionCounter() {
    return &connectionCounter_;
  }

  virtual void forceStop() {
  }

 protected:
  wangle::SimpleConnectionCounter connectionCounter_;

 private:
  // Forbidden copy constructor and assignment operator
  ServiceWorker(ServiceWorker const&) = delete;
  ServiceWorker& operator=(ServiceWorker const&) = delete;

  void addAcceptor(const AcceptorAddress& accAddress,
                   std::unique_ptr<wangle::Acceptor> acceptor,
                   AcceptorMap& acceptors) {
    CHECK(acceptors.find(accAddress) == acceptors.end());
    acceptors.insert(std::make_pair(accAddress, std::move(acceptor)));
  }

  Service* service_;

  RequestWorker* worker_;

  AcceptorMap acceptors_;

  NamedAddressMap namedAddress_;

  AcceptorMap drainingAcceptors_;
};

} // namespace proxygen