2018-08-19 04:43:58 +00:00
|
|
|
#include <Common/TaskStatsInfoGetter.h>
|
2018-05-17 16:01:41 +00:00
|
|
|
#include <Common/Exception.h>
|
2018-08-19 04:43:58 +00:00
|
|
|
#include <Core/Types.h>
|
2018-05-17 16:01:41 +00:00
|
|
|
|
|
|
|
#include <errno.h>
|
|
|
|
#include <linux/genetlink.h>
|
|
|
|
#include <linux/netlink.h>
|
|
|
|
#include <linux/taskstats.h>
|
|
|
|
#include <stdio.h>
|
|
|
|
#include <stdlib.h>
|
|
|
|
#include <string.h>
|
|
|
|
#include <sys/socket.h>
|
|
|
|
#include <sys/types.h>
|
|
|
|
#include <unistd.h>
|
|
|
|
#include <syscall.h>
|
|
|
|
|
2018-08-19 04:43:58 +00:00
|
|
|
/// Basic idea is motivated by "iotop" tool.
|
2018-06-04 14:16:27 +00:00
|
|
|
/// More info: https://www.kernel.org/doc/Documentation/accounting/taskstats.txt
|
2018-05-17 16:01:41 +00:00
|
|
|
|
|
|
|
#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;
|
|
|
|
|
|
|
|
|
2018-06-04 14:16:27 +00:00
|
|
|
struct NetlinkMessage
|
2018-05-17 16:01:41 +00:00
|
|
|
{
|
2018-06-04 14:16:27 +00:00
|
|
|
::nlmsghdr n;
|
|
|
|
::genlmsghdr g;
|
2018-05-17 16:01:41 +00:00
|
|
|
char buf[MAX_MSG_SIZE];
|
|
|
|
};
|
|
|
|
|
|
|
|
|
2018-08-19 04:43:58 +00:00
|
|
|
int sendCommand(
|
|
|
|
int sock_fd,
|
|
|
|
UInt16 nlmsg_type,
|
|
|
|
UInt32 nlmsg_pid,
|
|
|
|
UInt8 genl_cmd,
|
|
|
|
UInt16 nla_type,
|
|
|
|
void * nla_data,
|
|
|
|
int nla_len) noexcept
|
2018-05-17 16:01:41 +00:00
|
|
|
{
|
2018-08-19 04:43:58 +00:00
|
|
|
NetlinkMessage msg{};
|
2018-05-17 16:01:41 +00:00
|
|
|
|
|
|
|
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;
|
2018-08-19 04:43:58 +00:00
|
|
|
msg.g.version = 1;
|
2018-05-17 16:01:41 +00:00
|
|
|
|
2018-08-19 04:43:58 +00:00
|
|
|
::nlattr * attr = static_cast<::nlattr *>(GENLMSG_DATA(&msg));
|
|
|
|
attr->nla_type = nla_type;
|
|
|
|
attr->nla_len = nla_len + 1 + NLA_HDRLEN;
|
2018-05-17 16:01:41 +00:00
|
|
|
|
2018-08-19 04:43:58 +00:00
|
|
|
memcpy(NLA_DATA(attr), nla_data, nla_len);
|
|
|
|
msg.n.nlmsg_len += NLMSG_ALIGN(attr->nla_len);
|
2018-05-17 16:01:41 +00:00
|
|
|
|
2018-06-04 14:16:27 +00:00
|
|
|
char * buf = reinterpret_cast<char *>(&msg);
|
|
|
|
ssize_t buflen = msg.n.nlmsg_len;
|
|
|
|
|
2018-08-19 04:43:58 +00:00
|
|
|
::sockaddr_nl nladdr{};
|
2018-05-17 16:01:41 +00:00
|
|
|
nladdr.nl_family = AF_NETLINK;
|
2018-06-04 14:16:27 +00:00
|
|
|
|
|
|
|
while (true)
|
2018-05-17 16:01:41 +00:00
|
|
|
{
|
2018-08-19 04:43:58 +00:00
|
|
|
ssize_t r = ::sendto(sock_fd, buf, buflen, 0, reinterpret_cast<const ::sockaddr *>(&nladdr), sizeof(nladdr));
|
2018-06-04 14:16:27 +00:00
|
|
|
|
|
|
|
if (r >= buflen)
|
|
|
|
break;
|
|
|
|
|
2018-05-17 16:01:41 +00:00
|
|
|
if (r > 0)
|
|
|
|
{
|
|
|
|
buf += r;
|
|
|
|
buflen -= r;
|
|
|
|
}
|
|
|
|
else if (errno != EAGAIN)
|
|
|
|
return -1;
|
|
|
|
}
|
2018-06-04 14:16:27 +00:00
|
|
|
|
2018-05-17 16:01:41 +00:00
|
|
|
return 0;
|
|
|
|
}
|
|
|
|
|
|
|
|
|
2018-08-19 04:43:58 +00:00
|
|
|
UInt16 getFamilyId(int nl_sock_fd) noexcept
|
2018-05-17 16:01:41 +00:00
|
|
|
{
|
|
|
|
struct
|
|
|
|
{
|
2018-08-19 04:43:58 +00:00
|
|
|
::nlmsghdr header;
|
|
|
|
::genlmsghdr ge_header;
|
2018-05-17 16:01:41 +00:00
|
|
|
char buf[256];
|
2018-08-19 04:43:58 +00:00
|
|
|
} answer;
|
2018-05-17 16:01:41 +00:00
|
|
|
|
2018-06-04 14:16:27 +00:00
|
|
|
static char name[] = TASKSTATS_GENL_NAME;
|
2018-05-17 16:01:41 +00:00
|
|
|
|
2018-08-19 04:43:58 +00:00
|
|
|
if (sendCommand(
|
|
|
|
nl_sock_fd, GENL_ID_CTRL, getpid(), CTRL_CMD_GETFAMILY,
|
|
|
|
CTRL_ATTR_FAMILY_NAME, (void *) name,
|
|
|
|
strlen(TASKSTATS_GENL_NAME) + 1))
|
2018-05-17 16:01:41 +00:00
|
|
|
return 0;
|
|
|
|
|
2018-06-14 14:29:42 +00:00
|
|
|
UInt16 id = 0;
|
2018-08-19 04:43:58 +00:00
|
|
|
ssize_t rep_len = ::recv(nl_sock_fd, &answer, sizeof(answer), 0);
|
|
|
|
if (answer.header.nlmsg_type == NLMSG_ERROR || (rep_len < 0) || !NLMSG_OK((&answer.header), rep_len))
|
2018-05-17 16:01:41 +00:00
|
|
|
return 0;
|
|
|
|
|
2018-08-19 04:43:58 +00:00
|
|
|
const ::nlattr * attr;
|
|
|
|
attr = static_cast<const ::nlattr *>(GENLMSG_DATA(&answer));
|
|
|
|
attr = reinterpret_cast<const ::nlattr *>(reinterpret_cast<const char *>(attr) + NLA_ALIGN(attr->nla_len));
|
|
|
|
if (attr->nla_type == CTRL_ATTR_FAMILY_ID)
|
|
|
|
id = *static_cast<const UInt16 *>(NLA_DATA(attr));
|
2018-05-17 16:01:41 +00:00
|
|
|
|
|
|
|
return id;
|
|
|
|
}
|
|
|
|
|
2018-06-14 14:29:42 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
TaskStatsInfoGetter::TaskStatsInfoGetter() = default;
|
|
|
|
|
|
|
|
void TaskStatsInfoGetter::init()
|
|
|
|
{
|
|
|
|
if (netlink_socket_fd >= 0)
|
|
|
|
return;
|
|
|
|
|
|
|
|
netlink_socket_fd = ::socket(PF_NETLINK, SOCK_RAW, NETLINK_GENERIC);
|
|
|
|
if (netlink_socket_fd < 0)
|
|
|
|
throwFromErrno("Can't create PF_NETLINK socket");
|
|
|
|
|
2018-08-19 04:43:58 +00:00
|
|
|
::sockaddr_nl addr{};
|
2018-06-14 14:29:42 +00:00
|
|
|
addr.nl_family = AF_NETLINK;
|
|
|
|
|
2018-08-19 04:43:58 +00:00
|
|
|
if (::bind(netlink_socket_fd, reinterpret_cast<const ::sockaddr *>(&addr), sizeof(addr)) < 0)
|
2018-06-14 14:29:42 +00:00
|
|
|
throwFromErrno("Can't bind PF_NETLINK socket");
|
|
|
|
|
2018-08-19 04:43:58 +00:00
|
|
|
netlink_family_id = getFamilyId(netlink_socket_fd);
|
2018-06-14 14:29:42 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
bool TaskStatsInfoGetter::getStatImpl(int tid, ::taskstats & out_stats, bool throw_on_error)
|
2018-05-17 16:01:41 +00:00
|
|
|
{
|
2018-06-14 14:29:42 +00:00
|
|
|
init();
|
|
|
|
|
2018-08-19 04:43:58 +00:00
|
|
|
if (sendCommand(netlink_socket_fd, netlink_family_id, tid, TASKSTATS_CMD_GET, TASKSTATS_CMD_ATTR_PID, &tid, sizeof(pid_t)))
|
2018-05-17 16:01:41 +00:00
|
|
|
throwFromErrno("Can't send a Netlink command");
|
|
|
|
|
2018-06-04 14:16:27 +00:00
|
|
|
NetlinkMessage msg;
|
2018-06-14 14:29:42 +00:00
|
|
|
ssize_t rv = ::recv(netlink_socket_fd, &msg, sizeof(msg), 0);
|
2018-05-17 16:01:41 +00:00
|
|
|
|
|
|
|
if (msg.n.nlmsg_type == NLMSG_ERROR || !NLMSG_OK((&msg.n), rv))
|
|
|
|
{
|
2018-08-19 04:43:58 +00:00
|
|
|
const ::nlmsgerr * err = static_cast<const ::nlmsgerr *>(NLMSG_DATA(&msg));
|
2018-06-14 14:29:42 +00:00
|
|
|
if (throw_on_error)
|
2018-08-17 18:57:07 +00:00
|
|
|
throw Exception("Can't get Netlink response, error: " + std::to_string(err->error), ErrorCodes::NETLINK_ERROR);
|
2018-06-14 14:29:42 +00:00
|
|
|
else
|
2018-05-17 16:01:41 +00:00
|
|
|
return false;
|
|
|
|
}
|
|
|
|
|
|
|
|
rv = GENLMSG_PAYLOAD(&msg.n);
|
|
|
|
|
2018-08-19 04:43:58 +00:00
|
|
|
const ::nlattr * attr = static_cast<const ::nlattr *>(GENLMSG_DATA(&msg));
|
2018-06-14 14:29:42 +00:00
|
|
|
ssize_t len = 0;
|
2018-05-17 16:01:41 +00:00
|
|
|
|
|
|
|
while (len < rv)
|
|
|
|
{
|
2018-08-19 04:43:58 +00:00
|
|
|
len += NLA_ALIGN(attr->nla_len);
|
2018-05-17 16:01:41 +00:00
|
|
|
|
2018-08-19 04:43:58 +00:00
|
|
|
if (attr->nla_type == TASKSTATS_TYPE_AGGR_TGID || attr->nla_type == TASKSTATS_TYPE_AGGR_PID)
|
2018-05-17 16:01:41 +00:00
|
|
|
{
|
2018-08-19 04:43:58 +00:00
|
|
|
int aggr_len = NLA_PAYLOAD(attr->nla_len);
|
2018-05-17 16:01:41 +00:00
|
|
|
int len2 = 0;
|
|
|
|
|
2018-08-19 04:43:58 +00:00
|
|
|
attr = static_cast<const ::nlattr *>(NLA_DATA(attr));
|
2018-05-17 16:01:41 +00:00
|
|
|
while (len2 < aggr_len)
|
|
|
|
{
|
2018-08-19 04:43:58 +00:00
|
|
|
if (attr->nla_type == TASKSTATS_TYPE_STATS)
|
2018-05-17 16:01:41 +00:00
|
|
|
{
|
2018-08-19 04:43:58 +00:00
|
|
|
const ::taskstats * ts = static_cast<const ::taskstats *>(NLA_DATA(attr));
|
2018-05-17 16:01:41 +00:00
|
|
|
out_stats = *ts;
|
|
|
|
}
|
2018-06-04 14:16:27 +00:00
|
|
|
|
2018-08-19 04:43:58 +00:00
|
|
|
len2 += NLA_ALIGN(attr->nla_len);
|
|
|
|
attr = reinterpret_cast<const ::nlattr *>(reinterpret_cast<const char *>(attr) + len2);
|
2018-05-17 16:01:41 +00:00
|
|
|
}
|
|
|
|
}
|
2018-06-04 14:16:27 +00:00
|
|
|
|
2018-08-19 04:43:58 +00:00
|
|
|
attr = reinterpret_cast<const ::nlattr *>(reinterpret_cast<const char *>(GENLMSG_DATA(&msg)) + len);
|
2018-05-17 16:01:41 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
return true;
|
|
|
|
}
|
|
|
|
|
2018-06-14 14:29:42 +00:00
|
|
|
void TaskStatsInfoGetter::getStat(::taskstats & stat, int tid)
|
2018-05-17 16:01:41 +00:00
|
|
|
{
|
2018-08-17 18:57:07 +00:00
|
|
|
tid = tid < 0 ? getDefaultTID() : tid;
|
2018-06-14 14:29:42 +00:00
|
|
|
getStatImpl(tid, stat, true);
|
2018-05-17 16:01:41 +00:00
|
|
|
}
|
|
|
|
|
2018-06-14 14:29:42 +00:00
|
|
|
bool TaskStatsInfoGetter::tryGetStat(::taskstats & stat, int tid)
|
2018-05-17 16:01:41 +00:00
|
|
|
{
|
2018-08-17 18:57:07 +00:00
|
|
|
tid = tid < 0 ? getDefaultTID() : tid;
|
2018-06-14 14:29:42 +00:00
|
|
|
return getStatImpl(tid, stat, false);
|
2018-05-17 16:01:41 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
TaskStatsInfoGetter::~TaskStatsInfoGetter()
|
|
|
|
{
|
2018-06-14 14:29:42 +00:00
|
|
|
if (netlink_socket_fd >= 0)
|
2018-05-17 16:01:41 +00:00
|
|
|
close(netlink_socket_fd);
|
|
|
|
}
|
|
|
|
|
|
|
|
int TaskStatsInfoGetter::getCurrentTID()
|
|
|
|
{
|
2018-08-17 18:57:07 +00:00
|
|
|
/// This call is always successful. - man gettid
|
2018-05-17 16:01:41 +00:00
|
|
|
return static_cast<int>(syscall(SYS_gettid));
|
|
|
|
}
|
|
|
|
|
2018-08-17 18:57:07 +00:00
|
|
|
int TaskStatsInfoGetter::getDefaultTID()
|
2018-06-14 14:29:42 +00:00
|
|
|
{
|
|
|
|
if (default_tid < 0)
|
|
|
|
default_tid = getCurrentTID();
|
|
|
|
|
|
|
|
return default_tid;
|
|
|
|
}
|
|
|
|
|
2018-06-18 13:50:20 +00:00
|
|
|
static bool tryGetTaskStats()
|
2018-06-14 14:29:42 +00:00
|
|
|
{
|
2018-06-18 13:50:20 +00:00
|
|
|
TaskStatsInfoGetter getter;
|
|
|
|
::taskstats stat;
|
|
|
|
return getter.tryGetStat(stat);
|
|
|
|
}
|
2018-06-14 14:29:42 +00:00
|
|
|
|
2018-06-18 13:50:20 +00:00
|
|
|
bool TaskStatsInfoGetter::checkProcessHasRequiredPermissions()
|
|
|
|
{
|
|
|
|
/// It is thread- and exception- safe since C++11
|
|
|
|
static bool res = tryGetTaskStats();
|
|
|
|
return res;
|
2018-06-14 14:29:42 +00:00
|
|
|
}
|
2018-05-17 16:01:41 +00:00
|
|
|
|
|
|
|
}
|