#!/usr/sbin/dtrace -s

#pragma D option quiet
#pragma D option switchrate=10hz

dtrace:::BEGIN
{
    in = 0;
    out = 0;
}

fbt::udp_attach:entry
{
    self->so = args[0];
}

fbt::udp_attach:return
/args[1] == 0 && self->so != NULL/
{
    procs[(uintptr_t)self->so->so_pcb] = execname;
    self->so = NULL;
}

fbt::sosend_dgram:entry, fbt::soreceive:entry
/args[0]->so_proto->pr_protocol == IPPROTO_UDP/
{
    procs[(uintptr_t)args[0]->so_pcb] = execname;
}

fbt::in_pcbdetach:entry
{
    procs[(uintptr_t)args[0]] = 0;
}

udp:::send
/procs[args[1]->cs_cid] != ""/
{
    /* Subtract UDP header size. */
    this->bytes = args[4]->udp_length - 8;

    out += this->bytes;
    @bytes[procs[args[1]->cs_cid], args[2]->ip_saddr, args[4]->udp_sport,
        args[2]->ip_daddr, args[4]->udp_dport] = sum(this->bytes);
}

udp:::receive
/procs[args[1]->cs_cid] != ""/
{
    /* Subtract UDP header size. */
    this->bytes = args[4]->udp_length - 8;

    in += this->bytes;
    @bytes[procs[args[1]->cs_cid], args[2]->ip_daddr, args[4]->udp_dport,
        args[2]->ip_saddr, args[4]->udp_sport] = sum(this->bytes);
}

profile:::tick-2sec
{
    out /= 1024;
    in /= 1024;

    printf("%Y, UDP in: %6dKB, UDP out: %6dKB, UDP total: %6dKB\n",
            walltimestamp, in, out, in + out);

    printf("%-12s %-15s %5s %-15s %5s %9s\n",
        "PROC", "LADDR", "LPORT", "RADDR", "RPORT", "SIZE");

    printa("%-12s %-15s %5d %-15s %5d %@9d\n", @bytes);
    printf("\n");

    trunc(@bytes);
    in = 0;
    out = 0;
}