torch/csrc/distributed/c10d/socket.h

// Copyright (c) Meta Platforms, Inc. and its affiliates.
// 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.

#pragma once

#include <chrono>
#include <cstdint>
#include <memory>
#include <string>

#include <c10/macros/Macros.h>
#include <torch/csrc/distributed/c10d/exception.h>

namespace c10d {
namespace detail {

class SocketOptions {
 public:
  SocketOptions& prefer_ipv6(bool value) noexcept {
    prefer_ipv6_ = value;

    return *this;
  }

  bool prefer_ipv6() const noexcept {
    return prefer_ipv6_;
  }

  SocketOptions& connect_timeout(std::chrono::seconds value) noexcept {
    connect_timeout_ = value;

    return *this;
  }

  std::chrono::seconds connect_timeout() const noexcept {
    return connect_timeout_;
  }

 private:
  bool prefer_ipv6_ = true;
  std::chrono::seconds connect_timeout_{30};
};

class SocketImpl;

class Socket {
 public:
  // This function initializes the underlying socket library and must be called
  // before any other socket function.
  static void initialize();

  static Socket listen(std::uint16_t port, const SocketOptions& opts = {});

  static Socket connect(
      const std::string& host,
      std::uint16_t port,
      const SocketOptions& opts = {});

  Socket() noexcept = default;

  Socket(const Socket& other) = delete;

  Socket& operator=(const Socket& other) = delete;

  Socket(Socket&& other) noexcept;

  Socket& operator=(Socket&& other) noexcept;

  ~Socket();

  Socket accept() const;

  int handle() const noexcept;

  std::uint16_t port() const;

 private:
  explicit Socket(std::unique_ptr<SocketImpl>&& impl) noexcept;

  std::unique_ptr<SocketImpl> impl_;
};

} // namespace detail

class TORCH_API SocketError : public C10dError {
 public:
  using C10dError::C10dError;

  SocketError(const SocketError&) = default;

  SocketError& operator=(const SocketError&) = default;

  SocketError(SocketError&&) = default;

  SocketError& operator=(SocketError&&) = default;

  ~SocketError() override;
};

} // namespace c10d