poll_utils.h

前往此文件的文档。

#pragma once
#include <xgboost/collective/result.h>
#include <xgboost/collective/socket.h>

#if defined(_WIN32)
#include <xgboost/windefs.h>
// Socket API
#include <winsock2.h>
#include <ws2tcpip.h>
#else

#include <arpa/inet.h>
#include <fcntl.h>
#include <netdb.h>
#include <netinet/in.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <unistd.h>

#include <cerrno>

#endif // defined(_WIN32)

#include <chrono>
#include <cstring>
#include <string>
#include <system_error> // make_error_code, errc
#include <unordered_map>
#include <vector>

#if !defined(_WIN32)

#include <poll.h>

using SOCKET = int;
using sock_size_t = size_t; // NOLINT
#endif // !defined(_WIN32)

#define IS_MINGW() defined(__MINGW32__)

#if IS_MINGW() && !defined(POLLRDNORM) && !defined(POLLRDBAND)
/*
 * On later mingw versions poll should be supported (with bugs). See:
 * https://stackoverflow.com/a/60623080
 *
 * But right now the mingw distributed with R 3.6 doesn't support it.
 * So we just give a warning and provide dummy implementation to get
 * compilation passed. Otherwise we will have to provide a stub for
 * RABIT.
 *
 * Even on mingw version that has these structures and flags defined,
 * functions like `send` and `listen` might have unresolved linkage to
 * their implementation. So supporting mingw is quite difficult at
 * the time of writing.
 */
#pragma message("mingw上不支持分布式训练。")
typedef struct pollfd {
 SOCKET fd;
 short events; // NOLINT
 short revents; // NOLINT
} WSAPOLLFD, *PWSAPOLLFD, *LPWSAPOLLFD;

// POLLRDNORM | POLLRDBAND
#define POLLIN (0x0100 | 0x0200)
#define POLLPRI 0x0400
// POLLWRNORM
#define POLLOUT 0x0010

#endif // IS_MINGW() && !defined(POLLRDNORM) && !defined(POLLRDBAND)

namespace rabit {
namespace utils {

template <typename PollFD>
int PollImpl(PollFD* pfd, int nfds, std::chrono::seconds timeout) noexcept(true) {
 // For Windows and Linux, negative timeout means infinite timeout. For freebsd,
 // INFTIM(-1) should be used instead.
#if defined(_WIN32)

#if IS_MINGW()
 xgboost::MingWError();
 return -1;
#else
 return WSAPoll(pfd, nfds, std::chrono::milliseconds(timeout).count());
#endif // IS_MINGW()

#else
 return poll(pfd, nfds, timeout.count() < 0 ? -1 : std::chrono::milliseconds(timeout).count());
#endif // IS_MINGW()
}

template <typename E>
std::enable_if_t<std::is_integral_v<E>, xgboost::collective::Result> PollError(E const& revents) {
 if ((revents & POLLERR) != 0) {
 auto err = errno;
 auto str = strerror(err);
 return xgboost::system::FailWithCode(std::string{"轮询错误条件:"} + // NOLINT
 std::string{str} + // NOLINT
 " 代码:" + std::to_string(err));
 }
 if ((revents & POLLNVAL) != 0) {
 return xgboost::system::FailWithCode("无效的轮询请求。");
 }
 if ((revents & POLLHUP) != 0) {
 // 摘自 Linux 手册:
 //
 // 请注意，当从通道（例如管道或流套接字）读取时，此事件仅表示对端关闭了其通道端点。
 // 只有在通道中所有未处理的数据都被消耗后，后续从通道的读取才会返回 0（文件末尾）。
 //
 // 我们通常没有工作节点退出的障碍，一端退出而另一端仍在读取数据是正常的。
 //
 return xgboost::collective::Success();
 }
#if defined(POLLRDHUP)
 // 仅限 Linux 的标志
 if ((revents & POLLRDHUP) != 0) {
 return xgboost::system::FailWithCode("轮询到对端挂断。");
 }
#endif // defined(POLLRDHUP)
 return xgboost::collective::Success();
}


struct PollHelper {
 public
 inline void WatchRead(SOCKET fd) {
 // 添加文件描述符以监视读取
 auto& pfd = fds[fd];
 pfd.fd = fd;
 pfd.events |= POLLIN;
 }
 void WatchRead(xgboost::collective::TCPSocket const &socket) { this->WatchRead(socket.Handle()); }

 inline void WatchWrite(SOCKET fd) {
 // 添加文件描述符以监视写入
 auto& pfd = fds[fd];
 pfd.fd = fd;
 pfd.events |= POLLOUT;
 }
 void WatchWrite(xgboost::collective::TCPSocket const &socket) {
 this->WatchWrite(socket.Handle());
 }
 inline void WatchException(SOCKET fd) {
 // 添加文件描述符以监视异常
 auto& pfd = fds[fd];
 pfd.fd = fd;
 pfd.events |= POLLPRI;
 }
 void WatchException(xgboost::collective::TCPSocket const &socket) {
 this->WatchException(socket.Handle());
 }
 [[nodiscard]] bool CheckRead(SOCKET fd) const {
 // 检查描述符是否已准备好读取。
 const auto& pfd = fds.find(fd);
 return pfd != fds.end() && ((pfd->second.events & POLLIN) != 0);
 }
 [[nodiscard]] bool CheckRead(xgboost::collective::TCPSocket const& socket) const {
 return this->CheckRead(socket.Handle());
 }

 [[nodiscard]] bool CheckWrite(SOCKET fd) const {
 // 检查描述符是否已准备好写入。
 const auto& pfd = fds.find(fd);
 return pfd != fds.end() && ((pfd->second.events & POLLOUT) != 0);
 }
 [[nodiscard]] bool CheckWrite(xgboost::collective::TCPSocket const& socket) const {
 return this->CheckWrite(socket.Handle());
 }
 [[nodiscard]] xgboost::collective::Result Poll(std::chrono::seconds timeout,
 bool check_error = true) {
 // 在定义的集合上执行轮询，检查读取、写入、异常
 std::vector<pollfd> fdset;
 fdset.reserve(fds.size());
 for (auto kv : fds) {
 fdset.push_back(kv.second);
 }
 std::int32_t ret = PollImpl(fdset.data(), fdset.size(), timeout);
 if (ret == 0) {
 return xgboost::collective::Fail(
 "轮询超时:" + std::to_string(timeout.count()) + " 秒。",
 std::make_error_code(std::errc::timed_out));
 } else if (ret < 0) {
 return xgboost::system::FailWithCode("轮询失败, nfds:" + std::to_string(fdset.size()));
 }

 for (auto& pfd : fdset) {
 auto result = PollError(pfd.revents);
 if (check_error && !result.OK()) {
 return result;
 }

 auto revents = pfd.revents & pfd.events;
 fds[pfd.fd].events = revents;
 }
 return xgboost::collective::Success();
 }

 std::unordered_map<SOCKET, pollfd> fds;
};
} // namespace utils
} // namespace rabit

#if IS_MINGW() && !defined(POLLRDNORM) && !defined(POLLRDBAND)