[PATCH 10/20] libdtrace: retain descriptors for cached module BTF IDs

[Kris Van Hees <kris.van.hees@oracle.com>]

---
 libdtrace/dt_btf.c               | 79 ++++++++++++++++++++++++++++++--
 libdtrace/dt_impl.h              |  1 +
 libdtrace/dt_open.c              |  1 +
 test/triggers/module-discovery.c | 11 +++++
 4 files changed, 88 insertions(+), 4 deletions(-)

diff --git a/libdtrace/dt_btf.c b/libdtrace/dt_btf.c
index 617b0128..727e5f54 100644
--- a/libdtrace/dt_btf.c
+++ b/libdtrace/dt_btf.c
@@ -1479,17 +1479,67 @@ dt_btf_lookup_name_kind(dtrace_hdl_t *dtp, dt_module_t *dmp, const char *name,
 	return id;
 }
 
+/*
+ * Retain a descriptor as well as the ID: otherwise an unloaded object's ID
+ * could be reused and a later enumeration could mistake it for a cached one.
+ * Only IDs assigned to modules are retained, bounding descriptor use by the
+ * number of modules discovered by this handle.
+ */
+typedef struct dt_btf_id {
+	dt_hentry_t he;
+	uint32_t id;
+	int fd;
+} dt_btf_id_t;
+
+static uint32_t
+dt_btf_id_hval(const dt_btf_id_t *entry)
+{
+	return entry->id;
+}
+
+static int
+dt_btf_id_cmp(const dt_btf_id_t *a, const dt_btf_id_t *b)
+{
+	return a->id != b->id;
+}
+
+DEFINE_HE_STD_LINK_FUNCS(dt_btf_id, dt_btf_id_t, he)
+
+static void *
+dt_btf_id_release(dt_btf_id_t *head, dt_btf_id_t *entry)
+{
+	head = dt_btf_id_del(head, entry);
+	close(entry->fd);
+	free(entry);
+	return head;
+}
+
+static dt_htab_ops_t dt_btf_id_ops = {
+	.hval = (htab_hval_fn)dt_btf_id_hval,
+	.cmp = (htab_cmp_fn)dt_btf_id_cmp,
+	.add = (htab_add_fn)dt_btf_id_add,
+	.del = (htab_del_fn)dt_btf_id_release,
+	.next = (htab_next_fn)dt_btf_id_next
+};
+
 int
 dt_btf_get_module_ids(dtrace_hdl_t *dtp)
 {
 	uint32_t	id = 0;
 
+	if (dtp->dt_btf_ids == NULL) {
+		dtp->dt_btf_ids = dt_htab_create(&dt_btf_id_ops);
+		if (dtp->dt_btf_ids == NULL)
+			return dt_set_errno(dtp, EDT_NOMEM);
+	}
+
 	for (;;) {
 		int		rc, fd;
 		btf_info_t	info;
 		char		name[PATH_MAX];
 		uint32_t	size = sizeof(info);
 		dt_module_t	*dmp;
+		dt_btf_id_t	tmpl, *entry;
 
 		rc = dt_bpf_btf_get_next_id(id, &id);
 		if (rc) {
@@ -1501,6 +1551,10 @@ dt_btf_get_module_ids(dtrace_hdl_t *dtp)
 			return dt_set_errno(dtp, errno);
 		}
 
+		tmpl.id = id;
+		if (dt_htab_lookup(dtp->dt_btf_ids, &tmpl) != NULL)
+			continue;
+
 		fd = dt_bpf_btf_get_fd_by_id(id);
 		if (fd < 0) {
 			if (errno == ENOENT)
@@ -1513,14 +1567,31 @@ dt_btf_get_module_ids(dtrace_hdl_t *dtp)
 		info.name = (uint64_t)name;
 		info.name_len = sizeof(name);
 		rc = dt_bpf_btf_get_info_by_fd(fd, &info, &size);
-		close(fd);
-		if (rc)
-			return dt_set_errno(dtp, errno);
+		if (rc) {
+			int err = errno;
+
+			close(fd);
+			return dt_set_errno(dtp, err);
+		}
 
 		dmp = dt_module_lookup_by_name(dtp, name);
-		if (dmp == NULL || dmp->dm_btf_id != 0)
+		if (!info.kernel_btf || dmp == NULL || dmp->dm_btf_id != 0) {
+			close(fd);
 			continue;
+		}
 
+		entry = calloc(1, sizeof(*entry));
+		if (entry == NULL) {
+			close(fd);
+			return dt_set_errno(dtp, EDT_NOMEM);
+		}
+		entry->id = id;
+		entry->fd = fd;
+		if (dt_htab_insert(dtp->dt_btf_ids, entry) != 0) {
+			close(fd);
+			free(entry);
+			return dt_set_errno(dtp, EDT_NOMEM);
+		}
 		dmp->dm_btf_id = id;
 	}
 }
diff --git a/libdtrace/dt_impl.h b/libdtrace/dt_impl.h
index f89ce2bc..972fceaf 100644
--- a/libdtrace/dt_impl.h
+++ b/libdtrace/dt_impl.h
@@ -314,6 +314,7 @@ struct dtrace_hdl {
 	ctf_archive_t *dt_ctfa; /* ctf archive for the entire kernel tree */
 	void *dt_ctfa_data;	/* backing data for dt_ctfa */
 	char *dt_btf_path;	/* path to vmlinux.btf */
+	dt_htab_t *dt_btf_ids; /* retained module BTF IDs and descriptors */
 	dt_htab_t *dt_kernpaths; /* hash table of dt_kern_path_t's */
 	dt_module_t *dt_exec;	/* pointer to executable module */
 	dt_module_t *dt_cdefs;	/* pointer to C dynamic type module */
diff --git a/libdtrace/dt_open.c b/libdtrace/dt_open.c
index 245ccb70..8291b87d 100644
--- a/libdtrace/dt_open.c
+++ b/libdtrace/dt_open.c
@@ -1313,6 +1313,7 @@ dtrace_close(dtrace_hdl_t *dtp)
 	dtp->dt_kernsyms = NULL;
 	dt_htab_destroy(dtp->dt_mods);
 	dt_htab_destroy(dtp->dt_kernpaths);
+	dt_htab_destroy(dtp->dt_btf_ids);
 
 	if (dtp->dt_ctfa != NULL)
 		ctf_arc_close(dtp->dt_ctfa);
diff --git a/test/triggers/module-discovery.c b/test/triggers/module-discovery.c
index 0b09ae2b..d05d2bab 100644
--- a/test/triggers/module-discovery.c
+++ b/test/triggers/module-discovery.c
@@ -22,6 +22,7 @@
 static const char *symbols;
 static int kallsyms;
 static uint32_t btf_generation, btf_id;
+static int btf_queries;
 
 extern FILE *__real_fopen64(const char *, const char *);
 
@@ -78,6 +79,8 @@ __wrap_dt_bpf_btf_get_info_by_fd(int fd, btf_info_t *info, uint32_t *size)
 
 	if (symbols == NULL)
 		return __real_dt_bpf_btf_get_info_by_fd(fd, info, size);
+	btf_queries++;
+	info->kernel_btf = 1;
 	snprintf((char *)(uintptr_t)info->name, info->name_len, "%s",
 		 names[btf_id - btf_generation - 1]);
 	return 0;
@@ -171,6 +174,14 @@ check_discovery(int fallback)
 	}
 	CHECK(count == 2 && found_first == 1 && found_second == 1);
 
+	/* Cached IDs are pinned; enumeration only queries unassigned reload IDs. */
+	btf_queries = 0;
+	CHECK(dtrace_update(dtp) == 0);
+	CHECK(btf_queries == 2);
+	btf_queries = 0;
+	CHECK(dtrace_update(dtp) == 0);
+	CHECK(btf_queries == 2);
+
 	/* Disappearance and reappearance at a new address do not replace metadata. */
 	for (i = 0; i < 4; i++) {
 		symbols = i % 2 == 0 ? "" : fallback ?
-- 
2.52.0