From 62adf177f1f8c36c85593b0a9330a099c797a5cf Mon Sep 17 00:00:00 2001
From: Christiano Haesbaert <haesbaert@elastic.co>
Date: Wed, 8 May 2024 15:40:12 +0200
Subject: [PATCH] Add ebpf lost event counter. Issue #38

Pretty straighforward, contrary to kprobes which we get the counter on the data
path, with ebpf we have to actually read it, so add a new ops for updating the
counter, we should caution users to not hammer the reading, as it's real
syscall.

Tested by hacking quark-mon away.
---
 bpf_queue.c    | 31 +++++++++++++++++++++++++------
 kprobe_queue.c | 15 ++++++++++++---
 quark.c        |  1 +
 quark.h        |  1 +
 4 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/bpf_queue.c b/bpf_queue.c
index 02dfc4a..a62e596 100644
--- a/bpf_queue.c
+++ b/bpf_queue.c
@@ -11,17 +11,19 @@
 #include "elastic-ebpf/GPL/Events/EbpfEventProto.h"
 
 struct bpf_queue {
-	struct bpf_prog			*prog;
-	struct ring_buffer		*ringbuf;
+	struct bpf_prog		*prog;
+	struct ring_buffer	*ringbuf;
 };
 
 static int	bpf_queue_populate(struct quark_queue *);
+static int	bpf_queue_update_stats(struct quark_queue *);
 static void	bpf_queue_close(struct quark_queue *);
 
 struct quark_queue_ops queue_ops_bpf = {
-	.open	  = bpf_queue_open,
-	.populate = bpf_queue_populate,
-	.close	  = bpf_queue_close,
+	.open	      = bpf_queue_open,
+	.populate     = bpf_queue_populate,
+	.update_stats = bpf_queue_update_stats,
+	.close	      = bpf_queue_close,
 };
 
 static int
@@ -262,7 +264,6 @@ bpf_queue_populate(struct quark_queue *qq)
 	struct bpf_queue	*bqq = qq->queue_be;
 	int			 npop, space_left;
 
-	npop = 0;
 	space_left = qq->length >= qq->max_length ?
 	    0 : qq->max_length - qq->length;
 	if (space_left == 0)
@@ -273,6 +274,24 @@ bpf_queue_populate(struct quark_queue *qq)
 	return (npop < 0 ? -1 : npop);
 }
 
+static int
+bpf_queue_update_stats(struct quark_queue *qq)
+{
+	struct bpf_queue	*bqq  = qq->queue_be;
+	struct ebpf_event_stats	 pcpu_ees[libbpf_num_possible_cpus()];
+	u32			 zero = 0;
+	int			 i;
+
+	if (bpf_map__lookup_elem(bqq->prog->maps.ringbuf_stats, &zero,
+	    sizeof(zero), pcpu_ees, sizeof(pcpu_ees), 0) != 0)
+		return (-1);
+
+	for (i = 0; i < libbpf_num_possible_cpus(); i++)
+		qq->stats.lost = pcpu_ees[i].lost;
+
+	return (0);
+}
+
 static void
 bpf_queue_close(struct quark_queue *qq)
 {
diff --git a/kprobe_queue.c b/kprobe_queue.c
index a111f94..0c3cd8f 100644
--- a/kprobe_queue.c
+++ b/kprobe_queue.c
@@ -266,12 +266,14 @@ struct kprobe_queue {
 };
 
 static int	kprobe_queue_populate(struct quark_queue *);
+static int	kprobe_queue_update_stats(struct quark_queue *);
 static void	kprobe_queue_close(struct quark_queue *);
 
 struct quark_queue_ops queue_ops_kprobe = {
-	.open	  = kprobe_queue_open,
-	.populate = kprobe_queue_populate,
-	.close	  = kprobe_queue_close,
+	.open	      = kprobe_queue_open,
+	.populate     = kprobe_queue_populate,
+	.update_stats = kprobe_queue_update_stats,
+	.close	      = kprobe_queue_close,
 };
 
 static char *
@@ -1298,6 +1300,13 @@ kprobe_queue_populate(struct quark_queue *qq)
 	return (npop);
 }
 
+static int
+kprobe_queue_update_stats(struct quark_queue *qq)
+{
+	/* NADA */
+	return (0);
+}
+
 static void
 kprobe_queue_close(struct quark_queue *qq)
 {
diff --git a/quark.c b/quark.c
index 9fdac60..24b6b09 100644
--- a/quark.c
+++ b/quark.c
@@ -1597,6 +1597,7 @@ quark_queue_get_epollfd(struct quark_queue *qq)
 void
 quark_queue_get_stats(struct quark_queue *qq, struct quark_queue_stats *qs)
 {
+	qq->queue_ops->update_stats(qq);
 	*qs = qq->stats;
 }
 
diff --git a/quark.h b/quark.h
index 03d3a09..6d56c87 100644
--- a/quark.h
+++ b/quark.h
@@ -320,6 +320,7 @@ struct quark_queue_stats {
 struct quark_queue_ops {
 	int	(*open)(struct quark_queue *);
 	int	(*populate)(struct quark_queue *);
+	int	(*update_stats)(struct quark_queue *);
 	void	(*close)(struct quark_queue *);
 };