poll_utils.h

前往此文件文档。

 
#pragma once
#include <xgboost/collective/result.h>
#include <xgboost/collective/socket.h>
 
#if defined(_WIN32)
#include <xgboost/windefs.h>
// 套接字 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)
/*
 * 在较新的 mingw 版本上，poll 应该受支持（有 bug）。请参阅：
 * https://stackoverflow.com/a/60623080
 *
 * 但目前 R 3.6 分发的 mingw 不支持它。
 * 所以我们只给出一个警告并提供虚拟实现以通过编译。
 * 否则我们将不得不为 RABIT 提供一个存根。
 * 即使在定义了这些结构和标志的 mingw 版本上，
 *
 * 诸如 `send` 和 `listen` 等函数也可能存在未解析的链接到
 * 它们的实现。因此，在撰写本文时支持 mingw 相当困难。
 * 
 * 
 */
#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) {
 // 对于 Windows 和 Linux，负超时表示无限超时。对于 freebsd，
 // 应该使用 INFTIM(-1)。
#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)
#undef POLLIN
#undef POLLPRI
#undef POLLOUT
#endif // IS_MINGW()