socket.h 2.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100
  1. // Copyright (c) Meta Platforms, Inc. and its affiliates.
  2. // All rights reserved.
  3. //
  4. // This source code is licensed under the BSD-style license found in the
  5. // LICENSE file in the root directory of this source tree.
  6. #pragma once
  7. #include <chrono>
  8. #include <cstdint>
  9. #include <memory>
  10. #include <string>
  11. #include <c10/macros/Macros.h>
  12. #include <c10d/exception.h>
  13. namespace c10d {
  14. namespace detail {
  15. class SocketOptions {
  16. public:
  17. SocketOptions& prefer_ipv6(bool value) noexcept {
  18. prefer_ipv6_ = value;
  19. return *this;
  20. }
  21. bool prefer_ipv6() const noexcept {
  22. return prefer_ipv6_;
  23. }
  24. SocketOptions& connect_timeout(std::chrono::seconds value) noexcept {
  25. connect_timeout_ = value;
  26. return *this;
  27. }
  28. std::chrono::seconds connect_timeout() const noexcept {
  29. return connect_timeout_;
  30. }
  31. private:
  32. bool prefer_ipv6_ = true;
  33. std::chrono::seconds connect_timeout_{30};
  34. };
  35. class SocketImpl;
  36. class Socket {
  37. public:
  38. // This function initializes the underlying socket library and must be called
  39. // before any other socket function.
  40. static void initialize();
  41. static Socket listen(std::uint16_t port, const SocketOptions& opts = {});
  42. static Socket connect(const std::string& host, std::uint16_t port, const SocketOptions& opts = {});
  43. Socket() noexcept = default;
  44. Socket(const Socket& other) = delete;
  45. Socket& operator=(const Socket& other) = delete;
  46. Socket(Socket&& other) noexcept;
  47. Socket& operator=(Socket&& other) noexcept;
  48. ~Socket();
  49. Socket accept() const;
  50. int handle() const noexcept;
  51. std::uint16_t port() const;
  52. private:
  53. explicit Socket(std::unique_ptr<SocketImpl>&& impl) noexcept;
  54. std::unique_ptr<SocketImpl> impl_;
  55. };
  56. } // namespace detail
  57. class TORCH_API SocketError : public C10dError {
  58. public:
  59. using C10dError::C10dError;
  60. SocketError(const SocketError&) = default;
  61. SocketError& operator=(const SocketError&) = default;
  62. SocketError(SocketError&&) = default;
  63. SocketError& operator=(SocketError&&) = default;
  64. ~SocketError() override;
  65. };
  66. } // namespace c10d