dtracy: add support for aggregations

8 files changed, 347 insertions, 8 deletions

diff --git a/sys/src/cmd/dtracy/act.c b/sys/src/cmd/dtracy/act.c
index 335d1bd23..063be6cb9 100644
--- a/sys/src/cmd/dtracy/act.c
+++ b/sys/src/cmd/dtracy/act.c
@@ -55,6 +55,27 @@ addprobe(char *s)
 	clause->probs[clause->nprob++] = strdup(s);
 }
 
+static char *aggtypes[] = {
+	[AGGCNT] "count",
+	[AGGMIN] "min",
+	[AGGMAX] "max",
+	[AGGSUM] "sum",
+	[AGGAVG] "avg",
+	[AGGSTD] "std",
+};
+
+int
+aggtype(Symbol *s)
+{
+	int i;
+
+	for(i = 0; i < nelem(aggtypes); i++)
+		if(strcmp(s->name, aggtypes[i]) == 0)
+			return i;
+	error("%s unknown aggregation type", s->name);
+	return 0;
+}
+
 void
 addstat(int type, ...)
 {
@@ -73,6 +94,19 @@ addstat(int type, ...)
 	case STATPRINT:
 	case STATPRINTF:
 		break;
+	case STATAGG:
+		s->agg.name = va_arg(va, Symbol *);
+		s->agg.key = va_arg(va, Node *);
+		s->agg.type = aggtype(va_arg(va, Symbol *));
+		s->agg.value = va_arg(va, Node *);
+		if(s->agg.type == AGGCNT){
+			if(s->agg.value != nil)
+				error("too many arguments for count()");
+		}else{
+			if(s->agg.value == nil)
+				error("need argument for %s()", aggtypes[s->agg.type]);
+		}
+		break;
 	default:
 		sysfatal("addstat: unknown type %d", type);
 	}
@@ -158,12 +192,26 @@ prepprintf(Node **arg, int narg, DTActGr *g, int *recoff)
 	(*arg)->str = fmtstrflush(&f);
 }
 
+int aggid;
+
+int
+allagg(Clause *c)
+{
+	Stat *s;
+
+	for(s = c->stats; s < c->stats + c->nstats; s++)
+		if(s->type != STATAGG)
+			return 0;
+	return 1;
+}
+
 DTClause *
 mkdtclause(Clause *c)
 {
 	DTClause *d;
 	Stat *s;
 	int recoff, i;
+	Node *n;
 	
 	d = emalloc(sizeof(DTClause));
 	d->nprob = c->nprob;
@@ -175,7 +223,7 @@ mkdtclause(Clause *c)
 	for(s = c->stats; s < c->stats + c->nstats; s++)
 		switch(s->type){
 		case STATEXPR:
-			actgradd(d->gr, (DTAct){ACTTRACE, codegen(s->n), 0});
+			actgradd(d->gr, (DTAct){ACTTRACE, codegen(s->n), 0, noagg});
 			break;
 		case STATPRINT:
 			for(i = 0; i < s->narg; i++)
@@ -184,7 +232,22 @@ mkdtclause(Clause *c)
 		case STATPRINTF:
 			prepprintf(s->arg, s->narg, d->gr, &recoff);
 			break;
+		case STATAGG: {
+			DTAgg agg = {.id = s->agg.type << 28 | 1 << 16 | aggid++};
+			assert(dtaunpackid(&agg) >= 0);
+			aggs = realloc(aggs, sizeof(Agg) * aggid);
+			memset(&aggs[aggid-1], 0, sizeof(Agg));
+			aggs[aggid-1].DTAgg = agg;
+			aggs[aggid-1].name = strdup(s->agg.name == nil ? "" : s->agg.name->name);
+			actgradd(d->gr, (DTAct){ACTAGGKEY, codegen(s->agg.key), 8, agg});
+			n = s->agg.value;
+			if(n == nil) n = node(ONUM, 0ULL);
+			actgradd(d->gr, (DTAct){ACTAGGVAL, codegen(n), 8, agg});
+			break;
 		}
+		}
+	if(allagg(c))
+		actgradd(d->gr, (DTAct){ACTCANCEL, codegen(node(ONUM, 0)), 0, noagg});
 	return d;
 }
 
@@ -392,6 +455,7 @@ parseclause(Clause *cl, uchar *p, uchar *e, Enab *en, Biobuf *bp)
 		case STATPRINTF:
 			execprintf(s->arg, s->narg, p, e, en);
 			break;
+		case STATAGG: break;
 		default:
 			sysfatal("parseclause: unknown type %d", s->type);
 		}
@@ -546,6 +610,17 @@ dump(void)
 				print("\t\ttrace string (%d bytes)\n", a->size);
 				dumpexpr(a->p, "\t\t\t");
 				break;
+			case ACTAGGKEY:
+				print("\t\taggregation key (%s,%d,%d)\n", a->agg.type >= nelem(aggtypes) ? "???" : aggtypes[a->agg.type], a->agg.keysize, (u16int)a->agg.id);
+				dumpexpr(a->p, "\t\t\t");
+				break;
+			case ACTAGGVAL:
+				print("\t\taggregation value (%s,%d,%d)\n", a->agg.type >= nelem(aggtypes) ? "???" : aggtypes[a->agg.type], a->agg.keysize, (u16int)a->agg.id);
+				dumpexpr(a->p, "\t\t\t");
+				break;
+			case ACTCANCEL:
+				print("\t\tcancel record\n");
+				break;
 			default:
 				print("\t\t??? %d\n", a->type);
 			}
@@ -564,6 +639,8 @@ dump(void)
 				for(j = 0; j < s->narg; j++)
 					print("\t\t\targ %ε\n", s->arg[j]);
 				break;
+			case STATAGG:
+				break;
 			default:
 				print("\t\t??? %d\n", s->type);
 			}
diff --git a/sys/src/cmd/dtracy/agg.c b/sys/src/cmd/dtracy/agg.c
new file mode 100644
index 000000000..c54636479
--- /dev/null
+++ b/sys/src/cmd/dtracy/agg.c
@@ -0,0 +1,199 @@
+#include <u.h>
+#include <libc.h>
+#include <dtracy.h>
+#include <bio.h>
+#include <avl.h>
+#include <mp.h>
+#include "dat.h"
+#include "fns.h"
+
+typedef struct ANode ANode;
+
+struct ANode {
+	Avl;
+	s64int val, cnt;
+	u64int sq[2];
+	int keysize;
+	uchar key[1];
+};
+
+Agg *aggs;
+static Avltree **trees;
+static ANode *key;
+int interrupted;
+
+static int
+aggcmp(Avl *ap, Avl *bp)
+{
+	ANode *a, *b;
+	
+	a = (ANode *) ap;
+	b = (ANode *) bp;
+	return memcmp(a->key, b->key, a->keysize);
+}
+
+static void
+createrecord(int type, ANode *n, s64int *q)
+{
+	switch(type){
+	case AGGCNT: n->cnt = q[0]; break;
+	case AGGSUM: case AGGMIN: case AGGMAX: n->val = q[0]; break;
+	case AGGAVG: n->cnt = q[1]; n->val = q[0]; break;
+	case AGGSTD: n->cnt = q[1]; n->val = q[0]; n->sq[0] = q[2]; n->sq[1] = q[3]; break;
+	default: abort();
+	}
+}
+
+static void
+updaterecord(int type, ANode *n, s64int *q)
+{
+	u64int r;
+
+	switch(type){
+	case AGGCNT: n->cnt += q[0]; break;
+	case AGGSUM: n->val += q[0]; break;
+	case AGGAVG: n->cnt += q[1]; n->val += q[0]; break;
+	case AGGSTD:
+		n->cnt += q[1];
+		n->val += q[0];
+		r = n->sq[0] + q[2];
+		if(r < q[2]) n->sq[1]++;
+		n->sq[0] = r;
+		n->sq[1] += q[3];
+		break;
+	default: abort();
+	}
+}
+
+
+int
+aggparsebuf(uchar *p, int n)
+{
+	uchar *e;
+	Agg *a;
+	u32int id;
+	Avltree *tp;
+	ANode *np;
+	
+	e = p + n;
+	for(; p + 8 < e; p += a->recsize){
+		id = *(u32int*)&p[4];
+		if((u16int)id >= aggid){
+		inval:
+			fprint(2, "invalid record in aggregation buffer\n");
+			return -1;
+		}
+		a = &aggs[(u16int)id];
+		if(a->type != id>>28) goto inval;
+		if(a->keysize != (id>>13&0x7ff8)) goto inval;
+		if(p + a->recsize > e) goto inval;
+		tp = trees[(u16int)id];
+		key->keysize = a->keysize;
+		memcpy(key->key, &p[8], a->keysize);
+		np = (ANode *) avllookup(tp, key, 0);
+		if(np == nil){
+			np = emalloc(sizeof(ANode) - 1 + a->keysize);
+			*np = *key;
+			createrecord(a->type, np, (s64int*)&p[8+a->keysize]);
+			avlinsert(tp, np);
+		}else
+			updaterecord(a->type, np, (s64int*)&p[8+a->keysize]);
+	}
+	return 0;
+}
+
+void
+agginit(void)
+{
+	int i, m;
+	
+	trees = emalloc(sizeof(Avltree *) * aggid);
+	m = 0;
+	for(i = 0; i < aggid; i++){
+		trees[i] = avlcreate(aggcmp);
+		if(aggs[i].keysize > m)
+			m = aggs[i].keysize;
+	}
+	key = emalloc(sizeof(ANode) - 1 + m);
+}
+
+int
+aggnote(void *, char *note)
+{
+	if(strcmp(note, "interrupt") != 0 || interrupted)
+		return 0;
+	interrupted = 1;
+	return 1;
+}
+
+void
+aggkeyprint(Fmt *f, Agg *, ANode *a)
+{
+	fmtprint(f, "%20lld ", *(u64int*)a->key);
+}
+
+static double
+variance(ANode *a)
+{
+	mpint *x, *y, *z;
+	double r;
+	
+	x = vtomp(a->val, nil);
+	y = uvtomp(a->sq[0], nil);
+	z = vtomp(a->sq[1], nil);
+	mpleft(z, 64, z);
+	mpadd(z, y, y);
+	vtomp(a->cnt, z);
+	mpmul(x, x, x);
+	mpmul(y, z, y);
+	mpsub(y, x, x);
+	r = mptod(x) / a->cnt;
+	mpfree(x);
+	mpfree(y);
+	mpfree(z);
+	return r;
+}
+
+void
+aggvalprint(Fmt *f, int type, ANode *a)
+{
+	double x, s;
+
+	switch(type){
+	case AGGCNT: fmtprint(f, "%20lld", a->cnt); break;
+	case AGGSUM: case AGGMIN: case AGGMAX: fmtprint(f, "%20lld", a->val); break;
+	case AGGAVG: fmtprint(f, "%20g", (double)a->val / a->cnt); break;
+	case AGGSTD:
+		x = (double)a->val / a->cnt;
+		s = variance(a);
+		if(s < 0)
+			fmtprint(f, "%20g %20s", x, "NaN");
+		else{
+			fmtprint(f, "%20g %20g", x, sqrt(s));
+		}
+		break;
+	default:
+		abort();
+	}
+}
+
+void
+aggdump(void)
+{
+	Fmt f;
+	char buf[8192];
+	int i;
+	ANode *a;
+	
+	fmtfdinit(&f, 1, buf, sizeof(buf));
+	for(i = 0; i < aggid; i++){
+		a = (ANode *) avlmin(trees[i]);
+		for(; a != nil; a = (ANode *) avlnext(a)){
+			fmtprint(&f, "%s\t", aggs[i].name);
+			aggkeyprint(&f, &aggs[i], a);
+			aggvalprint(&f, aggs[i].type, a);
+			fmtprint(&f, "\n");
+		}
+	}
+	fmtfdflush(&f);
+}
diff --git a/sys/src/cmd/dtracy/cgen.c b/sys/src/cmd/dtracy/cgen.c
index 6b4f77d19..ace7f63f2 100644
--- a/sys/src/cmd/dtracy/cgen.c
+++ b/sys/src/cmd/dtracy/cgen.c
@@ -296,10 +296,10 @@ tracegen(Node *n, DTActGr *g, int *recoff)
 	case ORECORD:
 		switch(n->typ->type){
 		case TYPINT:
-			actgradd(g, (DTAct){ACTTRACE, codegen(n->n1), n->typ->size});
+			actgradd(g, (DTAct){ACTTRACE, codegen(n->n1), n->typ->size, noagg});
 			break;
 		case TYPSTRING:
-			actgradd(g, (DTAct){ACTTRACESTR, codegen(n->n1), n->typ->size});
+			actgradd(g, (DTAct){ACTTRACESTR, codegen(n->n1), n->typ->size, noagg});
 			break;
 		default:
 			sysfatal("tracegen: don't know how to record %τ", n->typ);
diff --git a/sys/src/cmd/dtracy/dat.h b/sys/src/cmd/dtracy/dat.h
index fc6737b4c..5087a609e 100644
--- a/sys/src/cmd/dtracy/dat.h
+++ b/sys/src/cmd/dtracy/dat.h
@@ -5,6 +5,7 @@ typedef struct Clause Clause;
 typedef struct Enab Enab;
 typedef struct Stat Stat;
 typedef struct Type Type;
+typedef struct Agg Agg;
 
 enum {
 	SYMHASH = 256,
@@ -89,10 +90,19 @@ struct Stat {
 		STATEXPR,
 		STATPRINT,
 		STATPRINTF,
+		STATAGG,
 	} type;
+	/* STATEXPR */
 	Node *n;
+	/* STATPRINT, STATPRINTF */
 	int narg;
 	Node **arg;
+	/* STATAGG */
+	struct {
+		Symbol *name;
+		int type;
+		Node *key, *value;
+	} agg;
 };
 
 struct Clause {
@@ -112,6 +122,11 @@ struct Enab {
 	Enab *next;
 };
 
+struct Agg {
+	DTAgg;
+	char *name;
+};
+
 extern int errors;
 
 #pragma	varargck	type	"α"	int
@@ -121,3 +136,6 @@ extern int errors;
 #pragma varargck	argpos error 1
 
 extern int dflag;
+extern DTAgg noagg;
+extern int aggid;
+extern Agg *aggs;
diff --git a/sys/src/cmd/dtracy/dtracy.c b/sys/src/cmd/dtracy/dtracy.c
index 8585180f2..540ed79bd 100644
--- a/sys/src/cmd/dtracy/dtracy.c
+++ b/sys/src/cmd/dtracy/dtracy.c
@@ -5,6 +5,8 @@
 #include "dat.h"
 #include "fns.h"
 
+DTAgg noagg;
+
 char *dtracyroot = "#Δ";
 int dtracyno;
 int ctlfd, buffd;
@@ -160,23 +162,56 @@ err:
 	
 }
 
-void
+int
 bufread(Biobuf *bp)
 {
 	static uchar buf[65536];
 	int n;
 	
 	n = read(buffd, buf, sizeof(buf));
-	if(n < 0) sysfatal("bufread: %r");
+	if(n < 0)
+		sysfatal("bufread: %r");
 	if(parsebuf(buf, n, bp) < 0)
 		sysfatal("parsebuf: %r");
 	Bflush(bp);
+	return 0;
+}
+
+void
+aggproc(void)
+{
+	char buf[65536];
+	int buffd, n;
+	extern int interrupted;
+
+	switch(rfork(RFPROC|RFMEM)){
+	case -1: sysfatal("rfork: %r");
+	case 0: return;
+	default: break;
+	}
+	snprint(buf, sizeof(buf), "%s/%d/aggbuf", dtracyroot, dtracyno);
+	buffd = open(buf, OREAD);
+	if(buffd < 0) sysfatal("open: %r");
+	agginit();
+	atnotify(aggnote, 1);
+	while(!interrupted){
+		n = read(buffd, buf, sizeof(buf));
+		if(n < 0){
+			if(interrupted)
+				break;
+			sysfatal("aggbufread: %r");
+		}
+		if(aggparsebuf((uchar *) buf, n) < 0)
+			exits("error");
+	}
+	aggdump();
+	exits(nil);
 }
 
 static void
 usage(void)
 {
-	fprint(2, "usage: %s [ -cd ] script\n", argv0);
+	fprint(2, "usage: %s [ -d ] script\n", argv0);
 	exits("usage");
 }
 
@@ -217,6 +252,8 @@ main(int argc, char **argv)
 		fprint(ctlfd, "go");
 		out = Bfdopen(1, OWRITE);
 		if(out == nil) sysfatal("Bfdopen: %r");
+		if(aggid > 0)
+			aggproc();
 		for(;;)
 			bufread(out);
 	}
diff --git a/sys/src/cmd/dtracy/fns.h b/sys/src/cmd/dtracy/fns.h
index f418e4f34..83a677025 100644
--- a/sys/src/cmd/dtracy/fns.h
+++ b/sys/src/cmd/dtracy/fns.h
@@ -33,3 +33,7 @@ Type *type(int, ...);
 int min(int, int);
 int max(int, int);
 Node *addtype(Type *, Node *);
+int aggparsebuf(uchar *, int);
+int aggnote(void *, char *);
+void aggdump(void);
+void agginit(void);
diff --git a/sys/src/cmd/dtracy/mkfile b/sys/src/cmd/dtracy/mkfile
index 7d1b4b45c..0c0d36afa 100644
--- a/sys/src/cmd/dtracy/mkfile
+++ b/sys/src/cmd/dtracy/mkfile
@@ -9,6 +9,7 @@ OFILES=\
 	cgen.$O\
 	act.$O\
 	type.$O\
+	agg.$O\
 
 YFILES=parse.y
 
diff --git a/sys/src/cmd/dtracy/parse.y b/sys/src/cmd/dtracy/parse.y
index 9f231dbc9..81a57f79b 100644
--- a/sys/src/cmd/dtracy/parse.y
+++ b/sys/src/cmd/dtracy/parse.y
@@ -16,7 +16,8 @@
 	Type *t;
 }
 
-%type <n> expr
+%type <n> expr optexpr
+%type <sym> optsym
 %type <t> type
 
 %token <sym> TSYM
@@ -63,7 +64,9 @@ stats0: stat | stats0 ';' stat
 stat: expr { addstat(STATEXPR, exprcheck($1, 0)); }
 | TPRINT { addstat(STATPRINT); } pelist
 | TPRINTF { addstat(STATPRINTF); } pelist
-
+| '@' optsym '[' expr ']' '=' TSYM '(' optexpr ')' { addstat(STATAGG, $2, $4, $7, $9); }
+optsym: TSYM | { $$ = nil; }
+optexpr: expr | { $$ = nil; }
 
 pelist:
 	'(' ')'