kernel-shark: Add KS_DEFINE_PLUGIN_CONTEXT macro

When we implement a KernelShark plugins, we often need a way to define
a structure that can hold context data that is visible only inside the
plugin and that has specific values for each data stream. The tricky
part here is that the total number of data streams and the IDs of the
active streams are dynamic quantities that can change as the user adds
or removes data streams. The macro defines an interface of functions
that will be useful for the plugin developer, helping to directly use
context objects, without caring for the complications due to the dynamic
configuration of active data streams.

Link: https://lore.kernel.org/linux-trace-devel/20210108143140.285037-4-y.karadz@gmail.com

Signed-off-by: Yordan Karadzhov (VMware) <y.karadz@gmail.com>
Signed-off-by: Steven Rostedt (VMware) <rostedt@goodmis.org>

2 files changed, 73 insertions, 0 deletions

diff --git a/src/libkshark-plugin.h b/src/libkshark-plugin.h
index 41043577..c62086f4 100644
--- a/src/libkshark-plugin.h
+++ b/src/libkshark-plugin.h
@@ -363,6 +363,47 @@ int kshark_handle_all_dpis(struct kshark_data_stream *stream,
 	__ok;								\
 })									\
 
+/** General purpose macro defining methods for adding plugin context. */
+#define KS_DEFINE_PLUGIN_CONTEXT(type)					\
+static type **__context_handler;					\
+static ssize_t __n_streams = -1;					\
+static inline type *__init(int sd)					\
+{									\
+	type *obj;							\
+	if (__n_streams < 0 && sd < KS_DEFAULT_NUM_STREAMS) {		\
+		__context_handler =					\
+			(type **) calloc(KS_DEFAULT_NUM_STREAMS,	\
+					 sizeof(*__context_handler));	\
+		if (!__context_handler)					\
+			return NULL;					\
+		__n_streams = KS_DEFAULT_NUM_STREAMS;			\
+	} else if (sd >= __n_streams) {					\
+		if (!KS_DOUBLE_SIZE(__context_handler,			\
+				    __n_streams))			\
+			return NULL;					\
+	}								\
+	assert(__context_handler[sd] == NULL);				\
+	obj = (type *) calloc(1, sizeof(*obj));				\
+	__context_handler[sd] = obj;					\
+	return obj;							\
+}									\
+static inline void __close(int sd)					\
+{									\
+	if (sd < 0) {							\
+		free(__context_handler);				\
+		__n_streams = -1;					\
+		return;							\
+	}								\
+	free(__context_handler[sd]);					\
+	__context_handler[sd] = NULL;					\
+}									\
+static inline type *__get_context(int sd)				\
+{									\
+	if (sd < 0 || sd >= __n_streams)				\
+		return NULL;						\
+	return __context_handler[sd];					\
+}									\
+
 #ifdef __cplusplus
 }
 #endif // __cplusplus
diff --git a/tests/libkshark-tests.cpp b/tests/libkshark-tests.cpp
index 8a5dcf10..780a3fae 100644
--- a/tests/libkshark-tests.cpp
+++ b/tests/libkshark-tests.cpp
@@ -10,6 +10,7 @@
 
 // KernelShark
 #include "libkshark.h"
+#include "libkshark-plugin.h"
 
 #define N_TEST_STREAMS	1000
 
@@ -103,3 +104,34 @@ BOOST_AUTO_TEST_CASE(fill_data_container)
 
 	kshark_free_data_container(data);
 }
+
+struct test_context {
+	int a;
+	char b;
+};
+
+KS_DEFINE_PLUGIN_CONTEXT(struct test_context);
+
+BOOST_AUTO_TEST_CASE(init_close_plugin)
+{
+	struct test_context *ctx;
+	int i;
+
+	for (i = 0; i < N_TEST_STREAMS; ++i) {
+		ctx = __init(i);
+		ctx->a = i * 10;
+		ctx->b = 'z';
+	}
+
+	for (i = 0; i < N_TEST_STREAMS; ++i) {
+		ctx = __get_context(i);
+		BOOST_REQUIRE(ctx != NULL);
+		BOOST_CHECK_EQUAL(ctx->a, i * 10);
+		BOOST_CHECK_EQUAL(ctx->b, 'z');
+
+		__close(i);
+		BOOST_REQUIRE(__get_context(i) == NULL);
+	}
+
+	__close(-1);
+}