/* getdelays.c
 *
 * Utility to get per-pid and per-tgid delay accounting statistics
 * Also illustrates usage of the taskstats interface
 *
 * Copyright (C) Shailabh Nagar, IBM Corp. 2005
 * Copyright (C) Balbir Singh, IBM Corp. 2006
 * Copyright (c) Jay Lan, SGI. 2006
 *
 */

#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <unistd.h>
#include <poll.h>
#include <string.h>
#include <fcntl.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <signal.h>

#include <linux/genetlink.h>
#include <taskstats.h>

/*
 * 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 *)(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)

#define err(code, fmt, arg...) do { printf(fmt, ##arg); exit(code); } while (0)
int done = 0;

int dbg=0, Delayacct=0;
__u64 stime, utime;
#define PRINTF(fmt, arg...) {			\
	    if (dbg) {				\
		printf(fmt, ##arg);		\
	    }					\
	}
/*
 * Create a raw netlink socket and bind
 */
static int create_nl_socket(int protocol, int groups)
{
    socklen_t addr_len;
    int fd;
    struct sockaddr_nl local;

    fd = socket(AF_NETLINK, SOCK_RAW, protocol);
    if (fd < 0)
	return -1;

    memset(&local, 0, sizeof(local));
    local.nl_family = AF_NETLINK;
    local.nl_groups = groups;

    if (bind(fd, (struct sockaddr *) &local, sizeof(local)) < 0)
	goto error;

    return fd;
  error:
    close(fd);
    return -1;
}

int sendto_fd(int s, const char *buf, int bufLen)
{
    struct sockaddr_nl nladdr;
    int r;

    memset(&nladdr, 0, sizeof(nladdr));
    nladdr.nl_family = AF_NETLINK;

    while ((r = sendto(s, buf, bufLen, 0, (struct sockaddr *) &nladdr,
		       sizeof(nladdr))) < bufLen) {
	if (r > 0) {
	    buf += r;
	    bufLen -= r;
	} else if (errno != EAGAIN)
	    return -1;
    }
    return 0;
}

/*
 * Probe the controller in genetlink to find the family id
 * for the TASKSTATS family
 */
int get_family_id(int sd)
{
    struct {
	struct nlmsghdr n;
	struct genlmsghdr g;
	char buf[256];
    } family_req;
    struct {
	struct nlmsghdr n;
	struct genlmsghdr g;
	char buf[256];
    } ans;

    int id;
    struct nlattr *na;
    int rep_len;

    /* Get family name */
    family_req.n.nlmsg_type = GENL_ID_CTRL;
    family_req.n.nlmsg_flags = NLM_F_REQUEST;
    family_req.n.nlmsg_seq = 0;
    family_req.n.nlmsg_pid = getpid();
    family_req.n.nlmsg_len = NLMSG_LENGTH(GENL_HDRLEN);
    family_req.g.cmd = CTRL_CMD_GETFAMILY;
    family_req.g.version = 0x1;
    na = (struct nlattr *) GENLMSG_DATA(&family_req);
    na->nla_type = CTRL_ATTR_FAMILY_NAME;
    na->nla_len = strlen(TASKSTATS_GENL_NAME) + 1 + NLA_HDRLEN;
    strcpy(NLA_DATA(na), TASKSTATS_GENL_NAME);
    family_req.n.nlmsg_len += NLMSG_ALIGN(na->nla_len);

    if (sendto_fd(sd, (char *) &family_req, family_req.n.nlmsg_len) < 0)
	err(1, "error sending message via Netlink\n");

    rep_len = recv(sd, &ans, sizeof(ans), 0);

    if (rep_len < 0)
	err(1, "error receiving reply message via Netlink\n");


    /* Validate response message */
    if (!NLMSG_OK((&ans.n), rep_len))
	err(1, "invalid reply message received via Netlink\n");

    if (ans.n.nlmsg_type == NLMSG_ERROR) {	/* error */
	printf("error received NACK - leaving\n");
	exit(1);
    }


    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;
}

void print_taskstats(struct taskstats *t)
{
    printf("\n\nCPU   %15s%15s%15s%15s\n"
	   "      %15llu%15llu%15llu%15llu\n"
	   "IO    %15s%15s\n"
	   "      %15llu%15llu\n"
	   "MEM   %15s%15s\n"
	   "      %15llu%15llu\n\n",
	   "count", "real total", "virtual total", "delay total",
	   t->cpu_count, t->cpu_run_real_total, t->cpu_run_virtual_total,
	   t->cpu_delay_total,
	   "count", "delay total",
	   t->blkio_count, t->blkio_delay_total,
	   "count", "delay total", t->swapin_count, t->swapin_delay_total);
}

void print_csa(struct taskstats *t)
{
    int sec, nsec;
    sec = t->exit_time.tv_sec - t->start_time.tv_sec;
    nsec = t->exit_time.tv_nsec - t->start_time.tv_nsec;
    printf("Command='%s'\n stime=%15llu, utime=%15llu, elapsed=%15llu msec\n",
    	t->ac_comm, t->ac_stime, t->ac_utime, sec*1000 + nsec/1000000);
    stime += t->ac_stime;
    utime += t->ac_utime;
}

void sigchld(int sig)
{
    done = 1;
}

int main(int argc, char *argv[])
{
    int rc;
    int sk_nl;
    struct nlmsghdr *nlh;
    struct genlmsghdr *genlhdr;
    __u16 id;
    struct nlattr *na;
    struct {
        struct nlmsghdr n;
	struct genlmsghdr g;
	char buf[800];
    } exitmsg;
    int	fd;

    /* For receiving */
    struct sockaddr_nl kern_nla, from_nla;
    socklen_t from_nla_len;
    int recv_len;

    int nl_sd = -1;
    int rep_len;
    int len = 0;
    int aggr_len, len2;
    struct sockaddr_nl nladdr;
    pid_t tid = 0;
    pid_t rtid = 0;
    int c;
    int count = 0, csa_summary=0;
    int write_file = 1;
    struct sigaction act = {
	.sa_handler = SIG_IGN,
    };
    struct sigaction tact = {
	.sa_handler = sigchld,
    };

    if (sigaction(SIGCHLD, &tact, NULL) < 0)
	err(1, "sigaction failed for SIGCHLD\n");

    while (1) {
    	c = getopt(argc, argv, "cdDw");
	if (c < 0)
	    break;

	switch (c) {
	    case 'c':
	    	printf("display csa\n");
		csa_summary = 1;
		break;

	    case 'd':
	    	printf("exit_recv: debug on\n");
	    	dbg = 1;
		break;

	    case 'D':
	    	printf("Delayacct summary ON\n");
		Delayacct = 1;
		break;

	    case 'w':		/* DON'T write data to a file */
	    	printf("exit_recv: write_file OFF\n");
	    	write_file = 0;
	    	break;

	    default: {
	    	printf("Unknown option %d\n", c);
		exit(-1);
	    }
	}
    }

    if (write_file)
	if ((fd = open("/var/csa/acct", O_WRONLY | O_CREAT | O_TRUNC, 
		S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH)) == -1) {
	    perror("Cannot open output file\n"); exit(1);
	}

    /* Open a NETLINK_GENERIC socket with TASKSTATS_LISTEN_GROUP */

    if ((nl_sd =
	 create_nl_socket(NETLINK_GENERIC, TASKSTATS_LISTEN_GROUP)) < 0)
	err(1, "error creating Netlink socket\n");

    if (sigaction(SIGINT, &act, NULL) < 0)
        err(1, "sigaction failed for SIGINT\n");

    do {
	int i;

	rep_len = recv(nl_sd, &exitmsg, sizeof(exitmsg), 0);
	PRINTF("\n\treceived %d bytes\n", rep_len);
	nladdr.nl_family = AF_NETLINK;
	nladdr.nl_groups = TASKSTATS_LISTEN_GROUP;

	if (exitmsg.n.nlmsg_type == NLMSG_ERROR) {	/* error */
	    printf("error received NACK - leaving\n");
	    exit(1);
	}

	if (rep_len < 0) {
	    printf("error receiving reply message via Netlink, rep_len=%d, errno=%d\n",
	    	rep_len, errno);
	    continue;
	}

	PRINTF("nlmsghdr size=%d, nlmsg_len=%d, rep_len=%d\n",
		sizeof(struct nlmsghdr), exitmsg.n.nlmsg_len, rep_len);
	/* Validate response message */
	if (!NLMSG_OK((&exitmsg.n), rep_len))
	    err(1, "invalid reply message received via Netlink\n");
/* #define NLMSG_OK(nlh,len) ((len) >= (int)sizeof(struct nlmsghdr) && \
                           (nlh)->nlmsg_len >= sizeof(struct nlmsghdr) && \
			                              (nlh)->nlmsg_len <= (len))
 */

	rep_len = GENLMSG_PAYLOAD(&exitmsg.n);

	na = (struct nlattr *) GENLMSG_DATA(&exitmsg);
	len = 0;
	i = 0;
	while (len < rep_len) {
	    len += NLA_ALIGN(na->nla_len);
	    switch (na->nla_type) {
	    case TASKSTATS_TYPE_AGGR_PID:
		/* Fall through */
	    case TASKSTATS_TYPE_AGGR_TGID:
		aggr_len = NLA_PAYLOAD(na->nla_len);
		len2 = 0;
		/* For nested attributes, na follows */
		na = (struct nlattr *) NLA_DATA(na);
		done = 0;
		while (len2 < aggr_len) {
		    switch (na->nla_type) {
		    case TASKSTATS_TYPE_PID:
			rtid = *(int *) NLA_DATA(na);
			PRINTF("PID\t%d\n", rtid);
			break;
		    case TASKSTATS_TYPE_TGID:
			rtid = *(int *) NLA_DATA(na);
			PRINTF("TGID\t%d\n", rtid);
			break;
		    case TASKSTATS_TYPE_STATS:
		    	count++;
		    	if (Delayacct)
			    print_taskstats((struct taskstats *) NLA_DATA(na));
			if (fd > 0) {
			    if (write(fd, NLA_DATA(na), na->nla_len) < 0) {
			    	err(1,"write error\n");
			    }
			}
			if (csa_summary)
			    print_csa((struct taskstats *)NLA_DATA(na));
			break;
		    default:
		    	printf("Unknown nested nla_type %d\n", na->nla_type);
			break;
		    }
		    len2 += NLA_ALIGN(na->nla_len);
		    na = (struct nlattr *) ((char *) na + len2);
		    if (done)
			break;
		}
		break;

	    default:
	    	printf("Unknown nla_type %d\n", na->nla_type);
		break;
	    }
	    na = (struct nlattr *) (GENLMSG_DATA(&exitmsg) + len);
	    if (done)
		break;
	}
	if (done)
	    break;
    }
    while (1);

    printf("Total taskstats STATS read %d, stime=%llu, utime=%llu\n",
    	count, stime, utime);
    close(nl_sd);
    if (fd > 0)
    	close(fd);
    return 0;
}