AsyncServerSocket.h

Go to the documentation of this file.

/*
 * Copyright 2014-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/SocketAddress.h>
#include <folly/String.h>
#include <folly/io/ShutdownSocketSet.h>
#include <folly/io/async/AsyncSocketBase.h>
#include <folly/io/async/AsyncTimeout.h>
#include <folly/io/async/DelayedDestruction.h>
#include <folly/io/async/EventBase.h>
#include <folly/io/async/EventHandler.h>
#include <folly/io/async/NotificationQueue.h>
#include <folly/portability/Sockets.h>

#include <limits.h>
#include <stddef.h>
#include <exception>
#include <memory>
#include <vector>

// Due to the way kernel headers are included, this may or may not be defined.
// Number pulled from 3.10 kernel headers.
#ifndef SO_REUSEPORT
#define SO_REUSEPORT 15
#endif

#if defined __linux__ && !defined SO_NO_TRANSPARENT_TLS
#define SO_NO_TRANSPARENT_TLS 200
#endif

namespace folly {

class AsyncServerSocket : public DelayedDestruction, public AsyncSocketBase {
 public:
  typedef std::unique_ptr<AsyncServerSocket, Destructor> UniquePtr;
  // Disallow copy, move, and default construction.
  AsyncServerSocket(AsyncServerSocket&&) = delete;

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

    virtual void onConnectionAccepted(
        const int socket,
        const SocketAddress& addr) noexcept = 0;

    virtual void onConnectionAcceptError(const int err) noexcept = 0;

    virtual void onConnectionDropped(
        const int socket,
        const SocketAddress& addr) noexcept = 0;

    virtual void onConnectionEnqueuedForAcceptorCallback(
        const int socket,
        const SocketAddress& addr) noexcept = 0;

    virtual void onConnectionDequeuedByAcceptorCallback(
        const int socket,
        const SocketAddress& addr) noexcept = 0;

    virtual void onBackoffStarted() noexcept = 0;

    virtual void onBackoffEnded() noexcept = 0;

    virtual void onBackoffError() noexcept = 0;
  };

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

    virtual void connectionAccepted(
        int fd,
        const SocketAddress& clientAddr) noexcept = 0;

    virtual void acceptError(const std::exception& ex) noexcept = 0;

    virtual void acceptStarted() noexcept {}

    virtual void acceptStopped() noexcept {}
  };

  static const uint32_t kDefaultMaxAcceptAtOnce = 30;
  static const uint32_t kDefaultCallbackAcceptAtOnce = 5;
  static const uint32_t kDefaultMaxMessagesInQueue = 1024;
  explicit AsyncServerSocket(EventBase* eventBase = nullptr);

  static std::shared_ptr<AsyncServerSocket> newSocket(
      EventBase* evb = nullptr) {
    return std::shared_ptr<AsyncServerSocket>(
        new AsyncServerSocket(evb), Destructor());
  }

  void setShutdownSocketSet(const std::weak_ptr<ShutdownSocketSet>& wNewSS);

  void destroy() override;

  void attachEventBase(EventBase* eventBase);

  void detachEventBase();

  EventBase* getEventBase() const override {
    return eventBase_;
  }

  void useExistingSocket(int fd);
  void useExistingSockets(const std::vector<int>& fds);

  std::vector<int> getSockets() const {
    std::vector<int> sockets;
    for (auto& handler : sockets_) {
      sockets.push_back(handler.socket_);
    }
    return sockets;
  }

  int getSocket() const {
    if (sockets_.size() > 1) {
      VLOG(2) << "Warning: getSocket can return multiple fds, "
              << "but getSockets was not called, so only returning the first";
    }
    if (sockets_.size() == 0) {
      return -1;
    } else {
      return sockets_[0].socket_;
    }
  }

  /* enable zerocopy support for the server sockets - the s = accept sockets
   * inherit it
   */
  bool setZeroCopy(bool enable);

  virtual void bind(const SocketAddress& address);

  virtual void bind(const std::vector<IPAddress>& ipAddresses, uint16_t port);

  virtual void bind(uint16_t port);

  void getAddress(SocketAddress* addressReturn) const override;

  SocketAddress getAddress() const {
    SocketAddress ret;
    getAddress(&ret);
    return ret;
  }

  std::vector<SocketAddress> getAddresses() const;

  virtual void listen(int backlog);

  virtual void addAcceptCallback(
      AcceptCallback* callback,
      EventBase* eventBase,
      uint32_t maxAtOnce = kDefaultCallbackAcceptAtOnce);

  void removeAcceptCallback(AcceptCallback* callback, EventBase* eventBase);

  virtual void startAccepting();

  void pauseAccepting();

  int stopAccepting(int shutdownFlags = -1);

  uint32_t getMaxAcceptAtOnce() const {
    return maxAcceptAtOnce_;
  }

  void setMaxAcceptAtOnce(uint32_t numConns) {
    maxAcceptAtOnce_ = numConns;
  }

  uint32_t getMaxNumMessagesInQueue() const {
    return maxNumMsgsInQueue_;
  }

  void setMaxNumMessagesInQueue(uint32_t num) {
    maxNumMsgsInQueue_ = num;
  }

  double getAcceptRateAdjustSpeed() const {
    return acceptRateAdjustSpeed_;
  }

  void setAcceptRateAdjustSpeed(double speed) {
    acceptRateAdjustSpeed_ = speed;
  }

  void setTosReflect(bool enable);

  bool getTosReflect() {
    return tosReflect_;
  }

  std::size_t getNumDroppedConnections() const {
    return numDroppedConnections_;
  }

  int64_t getNumPendingMessagesInQueue() const {
    if (eventBase_) {
      eventBase_->dcheckIsInEventBaseThread();
    }
    int64_t numMsgs = 0;
    for (const auto& callback : callbacks_) {
      if (callback.consumer) {
        numMsgs += callback.consumer->getQueue()->size();
      }
    }
    return numMsgs;
  }

  void setKeepAliveEnabled(bool enabled) {
    keepAliveEnabled_ = enabled;

    for (auto& handler : sockets_) {
      if (handler.socket_ < 0) {
        continue;
      }

      int val = (enabled) ? 1 : 0;
      if (setsockopt(
              handler.socket_, SOL_SOCKET, SO_KEEPALIVE, &val, sizeof(val)) !=
          0) {
        LOG(ERROR) << "failed to set SO_KEEPALIVE on async server socket: %s"
                   << errnoStr(errno);
      }
    }
  }

  bool getKeepAliveEnabled() const {
    return keepAliveEnabled_;
  }

  void setReusePortEnabled(bool enabled) {
    reusePortEnabled_ = enabled;

    for (auto& handler : sockets_) {
      if (handler.socket_ < 0) {
        continue;
      }

      int val = (enabled) ? 1 : 0;
      if (setsockopt(
              handler.socket_, SOL_SOCKET, SO_REUSEPORT, &val, sizeof(val)) !=
          0) {
        LOG(ERROR) << "failed to set SO_REUSEPORT on async server socket "
                   << errno;
        folly::throwSystemError(errno, "failed to bind to async server socket");
      }
    }
  }

  bool getReusePortEnabled_() const {
    return reusePortEnabled_;
  }

  void setCloseOnExec(bool closeOnExec) {
    closeOnExec_ = closeOnExec;
  }

  bool getCloseOnExec() const {
    return closeOnExec_;
  }

  void setTFOEnabled(bool enabled, uint32_t maxTFOQueueSize) {
    tfo_ = enabled;
    tfoMaxQueueSize_ = maxTFOQueueSize;
  }

  void disableTransparentTls() {
    noTransparentTls_ = true;
  }

  bool getAccepting() const {
    return accepting_;
  }

  void setConnectionEventCallback(
      ConnectionEventCallback* const connectionEventCallback) {
    connectionEventCallback_ = connectionEventCallback;
  }

  ConnectionEventCallback* getConnectionEventCallback() const {
    return connectionEventCallback_;
  }

 protected:
  ~AsyncServerSocket() override;

 private:
  enum class MessageType { MSG_NEW_CONN = 0, MSG_ERROR = 1 };

  struct QueueMessage {
    MessageType type;
    int fd;
    int err;
    SocketAddress address;
    std::string msg;
  };

  class RemoteAcceptor : private NotificationQueue<QueueMessage>::Consumer {
   public:
    explicit RemoteAcceptor(
        AcceptCallback* callback,
        ConnectionEventCallback* connectionEventCallback)
        : callback_(callback),
          connectionEventCallback_(connectionEventCallback) {}

    ~RemoteAcceptor() override = default;

    void start(EventBase* eventBase, uint32_t maxAtOnce, uint32_t maxInQueue);
    void stop(EventBase* eventBase, AcceptCallback* callback);

    void messageAvailable(QueueMessage&& message) noexcept override;

    NotificationQueue<QueueMessage>* getQueue() {
      return &queue_;
    }

   private:
    AcceptCallback* callback_;
    ConnectionEventCallback* connectionEventCallback_;

    NotificationQueue<QueueMessage> queue_;
  };

  struct CallbackInfo {
    CallbackInfo(AcceptCallback* cb, EventBase* evb)
        : callback(cb), eventBase(evb), consumer(nullptr) {}

    AcceptCallback* callback;
    EventBase* eventBase;

    RemoteAcceptor* consumer;
  };

  class BackoffTimeout;

  virtual void
  handlerReady(uint16_t events, int socket, sa_family_t family) noexcept;

  int createSocket(int family);
  void setupSocket(int fd, int family);
  void bindSocket(int fd, const SocketAddress& address, bool isExistingSocket);
  void dispatchSocket(int socket, SocketAddress&& address);
  void dispatchError(const char* msg, int errnoValue);
  void enterBackoff();
  void backoffTimeoutExpired();

  CallbackInfo* nextCallback() {
    CallbackInfo* info = &callbacks_[callbackIndex_];

    ++callbackIndex_;
    if (callbackIndex_ >= callbacks_.size()) {
      callbackIndex_ = 0;
    }

    return info;
  }

  struct ServerEventHandler : public EventHandler {
    ServerEventHandler(
        EventBase* eventBase,
        int socket,
        AsyncServerSocket* parent,
        sa_family_t addressFamily)
        : EventHandler(eventBase, socket),
          eventBase_(eventBase),
          socket_(socket),
          parent_(parent),
          addressFamily_(addressFamily) {}

    ServerEventHandler(const ServerEventHandler& other)
        : EventHandler(other.eventBase_, other.socket_),
          eventBase_(other.eventBase_),
          socket_(other.socket_),
          parent_(other.parent_),
          addressFamily_(other.addressFamily_) {}

    ServerEventHandler& operator=(const ServerEventHandler& other) {
      if (this != &other) {
        eventBase_ = other.eventBase_;
        socket_ = other.socket_;
        parent_ = other.parent_;
        addressFamily_ = other.addressFamily_;

        detachEventBase();
        attachEventBase(other.eventBase_);
        changeHandlerFD(other.socket_);
      }
      return *this;
    }

    // Inherited from EventHandler
    void handlerReady(uint16_t events) noexcept override {
      parent_->handlerReady(events, socket_, addressFamily_);
    }

    EventBase* eventBase_;
    int socket_;
    AsyncServerSocket* parent_;
    sa_family_t addressFamily_;
  };

  EventBase* eventBase_;
  std::vector<ServerEventHandler> sockets_;
  std::vector<int> pendingCloseSockets_;
  bool accepting_;
  uint32_t maxAcceptAtOnce_;
  uint32_t maxNumMsgsInQueue_;
  double acceptRateAdjustSpeed_; // 0 to disable auto adjust
  double acceptRate_;
  std::chrono::time_point<std::chrono::steady_clock> lastAccepTimestamp_;
  std::size_t numDroppedConnections_;
  uint32_t callbackIndex_;
  BackoffTimeout* backoffTimeout_;
  std::vector<CallbackInfo> callbacks_;
  bool keepAliveEnabled_;
  bool reusePortEnabled_{false};
  bool closeOnExec_;
  bool tfo_{false};
  bool noTransparentTls_{false};
  uint32_t tfoMaxQueueSize_{0};
  std::weak_ptr<ShutdownSocketSet> wShutdownSocketSet_;
  ConnectionEventCallback* connectionEventCallback_{nullptr};
  bool tosReflect_{false};
};

} // namespace folly