2018-05-17 16:01:41 +00:00
|
|
|
#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
|
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
|
|
|
|
|
|
|
/*
|
|
|
|
* 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;
|
|
|
|
|
|
|
|
|
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];
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
|
|
#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,
|
2018-06-04 14:16:27 +00:00
|
|
|
void * nla_data, int nla_len) noexcept
|
2018-05-17 16:01:41 +00:00
|
|
|
{
|
2018-06-04 14:16:27 +00:00
|
|
|
NetlinkMessage msg;
|
2018-05-17 16:01:41 +00:00
|
|
|
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;
|
|
|
|
|
2018-06-04 14:16:27 +00:00
|
|
|
::nlattr * na = static_cast<::nlattr *>(GENLMSG_DATA(&msg));
|
2018-05-17 16:01:41 +00:00
|
|
|
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);
|
|
|
|
|
2018-06-04 14:16:27 +00:00
|
|
|
char * buf = reinterpret_cast<char *>(&msg);
|
|
|
|
ssize_t buflen = msg.n.nlmsg_len;
|
|
|
|
|
|
|
|
::sockaddr_nl nladdr;
|
2018-05-17 16:01:41 +00:00
|
|
|
memset(&nladdr, 0, sizeof(nladdr));
|
|
|
|
nladdr.nl_family = AF_NETLINK;
|
2018-06-04 14:16:27 +00:00
|
|
|
|
|
|
|
while (true)
|
2018-05-17 16:01:41 +00:00
|
|
|
{
|
2018-06-04 14:16:27 +00:00
|
|
|
ssize_t r = ::sendto(sock_fd, buf, buflen, 0, reinterpret_cast<::sockaddr *>(&nladdr), sizeof(nladdr));
|
|
|
|
|
|
|
|
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-06-04 14:16:27 +00:00
|
|
|
int get_family_id(int nl_sock_fd) noexcept
|
2018-05-17 16:01:41 +00:00
|
|
|
{
|
|
|
|
struct
|
|
|
|
{
|
2018-06-04 14:16:27 +00:00
|
|
|
::nlmsghdr n;
|
|
|
|
::genlmsghdr g;
|
2018-05-17 16:01:41 +00:00
|
|
|
char buf[256];
|
|
|
|
} ans;
|
|
|
|
|
2018-06-04 14:16:27 +00:00
|
|
|
static char name[] = TASKSTATS_GENL_NAME;
|
2018-05-17 16:01:41 +00:00
|
|
|
|
|
|
|
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;
|
|
|
|
|
2018-06-04 14:16:27 +00:00
|
|
|
int id = 0;
|
|
|
|
ssize_t 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))
|
2018-05-17 16:01:41 +00:00
|
|
|
return 0;
|
|
|
|
|
2018-06-04 14:16:27 +00:00
|
|
|
::nlattr * na;
|
|
|
|
na = static_cast<::nlattr *>(GENLMSG_DATA(&ans));
|
|
|
|
na = reinterpret_cast<::nlattr *>((char *) na + NLA_ALIGN(na->nla_len));
|
2018-05-17 16:01:41 +00:00
|
|
|
if (na->nla_type == CTRL_ATTR_FAMILY_ID)
|
2018-06-04 14:16:27 +00:00
|
|
|
id = *static_cast<__u16 *>(NLA_DATA(na));
|
2018-05-17 16:01:41 +00:00
|
|
|
|
|
|
|
return id;
|
|
|
|
}
|
|
|
|
|
2018-06-04 14:16:27 +00:00
|
|
|
bool get_taskstats(int nl_sock_fd, int nl_family_id, pid_t xxxid, ::taskstats & out_stats, Exception * out_exception = nullptr)
|
2018-05-17 16:01:41 +00:00
|
|
|
{
|
|
|
|
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");
|
|
|
|
|
2018-06-04 14:16:27 +00:00
|
|
|
NetlinkMessage msg;
|
|
|
|
int rv = ::recv(nl_sock_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-06-04 14:16:27 +00:00
|
|
|
::nlmsgerr * err = static_cast<::nlmsgerr *>(NLMSG_DATA(&msg));
|
2018-05-17 16:01:41 +00:00
|
|
|
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);
|
|
|
|
|
2018-06-04 14:16:27 +00:00
|
|
|
::nlattr * na = static_cast<::nlattr *>(GENLMSG_DATA(&msg));
|
2018-05-17 16:01:41 +00:00
|
|
|
int len = 0;
|
|
|
|
|
|
|
|
while (len < rv)
|
|
|
|
{
|
|
|
|
len += NLA_ALIGN(na->nla_len);
|
|
|
|
|
2018-06-04 14:16:27 +00:00
|
|
|
if (na->nla_type == TASKSTATS_TYPE_AGGR_TGID || na->nla_type == TASKSTATS_TYPE_AGGR_PID)
|
2018-05-17 16:01:41 +00:00
|
|
|
{
|
|
|
|
int aggr_len = NLA_PAYLOAD(na->nla_len);
|
|
|
|
int len2 = 0;
|
|
|
|
|
2018-06-04 14:16:27 +00:00
|
|
|
na = static_cast<::nlattr *>(NLA_DATA(na));
|
2018-05-17 16:01:41 +00:00
|
|
|
while (len2 < aggr_len)
|
|
|
|
{
|
|
|
|
if (na->nla_type == TASKSTATS_TYPE_STATS)
|
|
|
|
{
|
2018-06-04 14:16:27 +00:00
|
|
|
::taskstats *ts = static_cast<::taskstats *>(NLA_DATA(na));
|
2018-05-17 16:01:41 +00:00
|
|
|
out_stats = *ts;
|
|
|
|
}
|
2018-06-04 14:16:27 +00:00
|
|
|
|
2018-05-17 16:01:41 +00:00
|
|
|
len2 += NLA_ALIGN(na->nla_len);
|
2018-06-04 14:16:27 +00:00
|
|
|
na = reinterpret_cast<::nlattr *>((char *) na + len2);
|
2018-05-17 16:01:41 +00:00
|
|
|
}
|
|
|
|
}
|
2018-06-04 14:16:27 +00:00
|
|
|
|
|
|
|
na = reinterpret_cast<::nlattr *>((char *) GENLMSG_DATA(&msg) + len);
|
2018-05-17 16:01:41 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
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");
|
|
|
|
|
2018-06-04 14:16:27 +00:00
|
|
|
::sockaddr_nl addr;
|
2018-05-17 16:01:41 +00:00
|
|
|
memset(&addr, 0, sizeof(addr));
|
|
|
|
addr.nl_family = AF_NETLINK;
|
|
|
|
|
2018-06-04 14:16:27 +00:00
|
|
|
if (::bind(netlink_socket_fd, reinterpret_cast<::sockaddr *>(&addr), sizeof(addr)) < 0)
|
2018-05-17 16:01:41 +00:00
|
|
|
throwFromErrno("Can't bind PF_NETLINK socket");
|
|
|
|
|
|
|
|
netlink_family_id = get_family_id(netlink_socket_fd);
|
2018-05-29 18:14:31 +00:00
|
|
|
|
|
|
|
initial_tid = getCurrentTID();
|
2018-05-17 16:01:41 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
void TaskStatsInfoGetter::getStat(::taskstats & stat, int tid) const
|
|
|
|
{
|
|
|
|
if (tid < 0)
|
2018-05-29 18:14:31 +00:00
|
|
|
tid = initial_tid;
|
2018-05-17 16:01:41 +00:00
|
|
|
|
|
|
|
get_taskstats(netlink_socket_fd, netlink_family_id, tid, stat);
|
|
|
|
}
|
|
|
|
|
|
|
|
bool TaskStatsInfoGetter::tryGetStat(::taskstats & stat, int tid) const
|
|
|
|
{
|
|
|
|
if (tid < 0)
|
2018-05-29 18:14:31 +00:00
|
|
|
tid = initial_tid;
|
2018-05-17 16:01:41 +00:00
|
|
|
|
|
|
|
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));
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
}
|