mirror of
https://github.com/ClickHouse/ClickHouse.git
synced 2024-10-09 18:10:48 +00:00
244 lines
6.1 KiB
C++
244 lines
6.1 KiB
C++
#include "TaskStatsInfoGetter.h"
|
|
#include <Common/Exception.h>
|
|
|
|
#include <asm/types.h>
|
|
#include <errno.h>
|
|
#include <linux/genetlink.h>
|
|
#include <linux/netlink.h>
|
|
#include <linux/taskstats.h>
|
|
#include <linux/unistd.h>
|
|
#include <pwd.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <sys/socket.h>
|
|
#include <sys/types.h>
|
|
#include <unistd.h>
|
|
#include <syscall.h>
|
|
|
|
/// Based on: https://github.com/Tomas-M/iotop/tree/master/src
|
|
|
|
/*
|
|
* Generic macros for dealing with netlink sockets. Might be duplicated
|
|
* elsewhere. It is recommended that commercial grade applications use
|
|
* libnl or libnetlink and use the interfaces provided by the library
|
|
*/
|
|
#define GENLMSG_DATA(glh) ((void *)((char*)NLMSG_DATA(glh) + GENL_HDRLEN))
|
|
#define GENLMSG_PAYLOAD(glh) (NLMSG_PAYLOAD(glh, 0) - GENL_HDRLEN)
|
|
#define NLA_DATA(na) ((void *)((char*)(na) + NLA_HDRLEN))
|
|
#define NLA_PAYLOAD(len) (len - NLA_HDRLEN)
|
|
|
|
|
|
namespace DB
|
|
{
|
|
|
|
namespace ErrorCodes
|
|
{
|
|
extern const int NETLINK_ERROR;
|
|
}
|
|
|
|
|
|
namespace
|
|
{
|
|
|
|
static size_t constexpr MAX_MSG_SIZE = 1024;
|
|
|
|
|
|
struct msgtemplate
|
|
{
|
|
struct nlmsghdr n;
|
|
struct genlmsghdr g;
|
|
char buf[MAX_MSG_SIZE];
|
|
};
|
|
|
|
|
|
#pragma GCC diagnostic push
|
|
#pragma GCC diagnostic ignored "-Wsign-compare"
|
|
|
|
int send_cmd(int sock_fd, __u16 nlmsg_type, __u32 nlmsg_pid,
|
|
__u8 genl_cmd, __u16 nla_type,
|
|
void *nla_data, int nla_len)
|
|
{
|
|
struct nlattr *na;
|
|
struct sockaddr_nl nladdr;
|
|
int r, buflen;
|
|
char *buf;
|
|
|
|
msgtemplate msg;
|
|
memset(&msg, 0, sizeof(msg));
|
|
|
|
msg.n.nlmsg_len = NLMSG_LENGTH(GENL_HDRLEN);
|
|
msg.n.nlmsg_type = nlmsg_type;
|
|
msg.n.nlmsg_flags = NLM_F_REQUEST;
|
|
msg.n.nlmsg_seq = 0;
|
|
msg.n.nlmsg_pid = nlmsg_pid;
|
|
msg.g.cmd = genl_cmd;
|
|
msg.g.version = 0x1;
|
|
|
|
na = (struct nlattr *) GENLMSG_DATA(&msg);
|
|
na->nla_type = nla_type;
|
|
na->nla_len = nla_len + 1 + NLA_HDRLEN;
|
|
|
|
memcpy(NLA_DATA(na), nla_data, nla_len);
|
|
msg.n.nlmsg_len += NLMSG_ALIGN(na->nla_len);
|
|
|
|
buf = (char *) &msg;
|
|
buflen = msg.n.nlmsg_len ;
|
|
memset(&nladdr, 0, sizeof(nladdr));
|
|
nladdr.nl_family = AF_NETLINK;
|
|
while ((r = sendto(sock_fd, buf, buflen, 0, (struct sockaddr *) &nladdr,
|
|
sizeof(nladdr))) < buflen)
|
|
{
|
|
if (r > 0)
|
|
{
|
|
buf += r;
|
|
buflen -= r;
|
|
}
|
|
else if (errno != EAGAIN)
|
|
return -1;
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
|
|
int get_family_id(int nl_sock_fd)
|
|
{
|
|
static char name[256];
|
|
|
|
struct
|
|
{
|
|
struct nlmsghdr n;
|
|
struct genlmsghdr g;
|
|
char buf[256];
|
|
} ans;
|
|
|
|
int id = 0;
|
|
struct nlattr *na;
|
|
int rep_len;
|
|
|
|
strcpy(name, TASKSTATS_GENL_NAME);
|
|
if (send_cmd(nl_sock_fd, GENL_ID_CTRL, getpid(), CTRL_CMD_GETFAMILY,
|
|
CTRL_ATTR_FAMILY_NAME, (void *) name,
|
|
strlen(TASKSTATS_GENL_NAME) + 1))
|
|
return 0;
|
|
|
|
rep_len = recv(nl_sock_fd, &ans, sizeof(ans), 0);
|
|
if (ans.n.nlmsg_type == NLMSG_ERROR
|
|
|| (rep_len < 0) || !NLMSG_OK((&ans.n), rep_len))
|
|
return 0;
|
|
|
|
na = (struct nlattr *) GENLMSG_DATA(&ans);
|
|
na = (struct nlattr *) ((char *) na + NLA_ALIGN(na->nla_len));
|
|
if (na->nla_type == CTRL_ATTR_FAMILY_ID)
|
|
id = *(__u16 *) NLA_DATA(na);
|
|
|
|
return id;
|
|
}
|
|
|
|
bool get_taskstats(int nl_sock_fd, int nl_family_id, pid_t xxxid, struct taskstats & out_stats, Exception * out_exception = nullptr)
|
|
{
|
|
if (send_cmd(nl_sock_fd, nl_family_id, xxxid, TASKSTATS_CMD_GET, TASKSTATS_CMD_ATTR_PID, &xxxid, sizeof(pid_t)))
|
|
throwFromErrno("Can't send a Netlink command");
|
|
|
|
msgtemplate msg;
|
|
int rv = recv(nl_sock_fd, &msg, sizeof(msg), 0);
|
|
|
|
if (msg.n.nlmsg_type == NLMSG_ERROR || !NLMSG_OK((&msg.n), rv))
|
|
{
|
|
struct nlmsgerr *err = static_cast<struct nlmsgerr *>(NLMSG_DATA(&msg));
|
|
Exception e("Can't get Netlink response, error=" + std::to_string(err->error), ErrorCodes::NETLINK_ERROR);
|
|
|
|
if (out_exception)
|
|
{
|
|
*out_exception = std::move(e);
|
|
return false;
|
|
}
|
|
|
|
throw Exception(std::move(e));
|
|
}
|
|
|
|
rv = GENLMSG_PAYLOAD(&msg.n);
|
|
|
|
struct nlattr *na = (struct nlattr *) GENLMSG_DATA(&msg);
|
|
int len = 0;
|
|
|
|
while (len < rv)
|
|
{
|
|
len += NLA_ALIGN(na->nla_len);
|
|
|
|
if (na->nla_type == TASKSTATS_TYPE_AGGR_TGID
|
|
|| na->nla_type == TASKSTATS_TYPE_AGGR_PID)
|
|
{
|
|
int aggr_len = NLA_PAYLOAD(na->nla_len);
|
|
int len2 = 0;
|
|
|
|
na = (struct nlattr *) NLA_DATA(na);
|
|
while (len2 < aggr_len)
|
|
{
|
|
if (na->nla_type == TASKSTATS_TYPE_STATS)
|
|
{
|
|
struct taskstats *ts = static_cast<struct taskstats *>(NLA_DATA(na));
|
|
out_stats = *ts;
|
|
}
|
|
len2 += NLA_ALIGN(na->nla_len);
|
|
na = (struct nlattr *) ((char *) na + len2);
|
|
}
|
|
}
|
|
na = (struct nlattr *) ((char *) GENLMSG_DATA(&msg) + len);
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
#pragma GCC diagnostic pop
|
|
|
|
}
|
|
|
|
|
|
TaskStatsInfoGetter::TaskStatsInfoGetter()
|
|
{
|
|
netlink_socket_fd = socket(PF_NETLINK, SOCK_RAW, NETLINK_GENERIC);
|
|
if (netlink_socket_fd < 0)
|
|
throwFromErrno("Can't create PF_NETLINK socket");
|
|
|
|
struct sockaddr_nl addr;
|
|
memset(&addr, 0, sizeof(addr));
|
|
addr.nl_family = AF_NETLINK;
|
|
|
|
if (bind(netlink_socket_fd, (struct sockaddr *) &addr, sizeof(addr)) < 0)
|
|
throwFromErrno("Can't bind PF_NETLINK socket");
|
|
|
|
netlink_family_id = get_family_id(netlink_socket_fd);
|
|
}
|
|
|
|
void TaskStatsInfoGetter::getStat(::taskstats & stat, int tid) const
|
|
{
|
|
if (tid < 0)
|
|
tid = getCurrentTID();
|
|
|
|
get_taskstats(netlink_socket_fd, netlink_family_id, tid, stat);
|
|
}
|
|
|
|
bool TaskStatsInfoGetter::tryGetStat(::taskstats & stat, int tid) const
|
|
{
|
|
if (tid < 0)
|
|
tid = getCurrentTID();
|
|
|
|
Exception e;
|
|
return get_taskstats(netlink_socket_fd, netlink_family_id, tid, stat, &e);
|
|
}
|
|
|
|
TaskStatsInfoGetter::~TaskStatsInfoGetter()
|
|
{
|
|
if (netlink_socket_fd > -1)
|
|
close(netlink_socket_fd);
|
|
}
|
|
|
|
int TaskStatsInfoGetter::getCurrentTID()
|
|
{
|
|
return static_cast<int>(syscall(SYS_gettid));
|
|
}
|
|
|
|
|
|
}
|