diff --git a/scripts/base/frameworks/sumstats/plugins/topk.bro b/scripts/base/frameworks/sumstats/plugins/topk.bro index 6107a252ae..a830b1c5ec 100644 --- a/scripts/base/frameworks/sumstats/plugins/topk.bro +++ b/scripts/base/frameworks/sumstats/plugins/topk.bro @@ -3,16 +3,27 @@ module SumStats; export { + redef record Reducer += { + ## number of elements to keep in the top-k list + topk_size: count &default=500; + }; + redef enum Calculation += { TOPK }; redef record ResultVal += { - topk: opaque of topk &default=topk_init(500); + topk: opaque of topk &optional; }; } +hook init_resultval_hook(r: Reducer, rv: ResultVal) + { + if ( TOPK in r$apply && ! rv?$topk ) + rv$topk = topk_init(r$topk_size); + } + hook observe_hook(r: Reducer, val: double, obs: Observation, rv: ResultVal) { if ( TOPK in r$apply ) @@ -24,7 +35,7 @@ hook observe_hook(r: Reducer, val: double, obs: Observation, rv: ResultVal) hook compose_resultvals_hook(result: ResultVal, rv1: ResultVal, rv2: ResultVal) { - result$topk = topk_init(500); + result$topk = topk_init(topk_size(rv1$topk)); topk_merge(result$topk, rv1$topk); topk_merge(result$topk, rv2$topk); diff --git a/src/Topk.h b/src/Topk.h index 30e87f7a99..51a2d75251 100644 --- a/src/Topk.h +++ b/src/Topk.h @@ -40,6 +40,7 @@ public: VectorVal* getTopK(int k) const; // returns vector uint64_t getCount(Val* value) const; uint64_t getEpsilon(Val* value) const; + uint64_t getSize() const { return size; } void Merge(const TopkVal* value); protected: diff --git a/src/bro.bif b/src/bro.bif index b6f101c025..4d0db54c8c 100644 --- a/src/bro.bif +++ b/src/bro.bif @@ -5684,6 +5684,13 @@ function topk_epsilon%(handle: opaque of topk, value: any%): count return new Val(h->getEpsilon(value), TYPE_COUNT); %} +function topk_size%(handle: opaque of topk%): count + %{ + assert(handle); + Topk::TopkVal* h = (Topk::TopkVal*) handle; + return new Val(h->getSize(), TYPE_COUNT); + %} + function topk_merge%(handle1: opaque of topk, handle2: opaque of topk%): any %{ assert(handle1);