[PATCH 12/20] libdtrace: index module address ranges across the handle

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

Build a sorted address-range index lazily and use prefix maxima to bound
searches across overlapping ranges. Retain module-list order when more
than one module covers an address. Fall back to the existing search if
index allocation fails.

Invalidate the index when symbols are added, a module is removed, or the
module list is reordered. Updates that discover nothing retain the index.

Expand discovery coverage for indexed text/data lookup, range boundaries,
gaps, overlaps, additive discovery, geometric range growth, and malformed
symbol input.
---
 libdtrace/dt_impl.h              |  10 +++
 libdtrace/dt_module.c            | 143 +++++++++++++++++++++++++++----
 libdtrace/dt_open.c              |   1 +
 test/triggers/module-discovery.c |  78 +++++++++++++++++
 4 files changed, 216 insertions(+), 16 deletions(-)

diff --git a/libdtrace/dt_impl.h b/libdtrace/dt_impl.h
index 5487bfcf..c83fedb2 100644
--- a/libdtrace/dt_impl.h
+++ b/libdtrace/dt_impl.h
@@ -188,6 +188,13 @@ typedef struct dt_module {
 	uint_t dm_aslen;	/* number of entries in dm_asmap */
 } dt_module_t;
 
+/* Handle-wide address index.  order preserves module-list precedence. */
+typedef struct dt_module_range {
+	GElf_Addr start, last, max_last;
+	dt_module_t *module;
+	size_t order;
+} dt_module_range_t;
+
 /*
  * The path to an actual kernel module residing on the disk.  Used only for
  * initialization of the corresponding dt_module, since modprobe -l is
@@ -308,6 +315,9 @@ struct dtrace_hdl {
 	uint64_t dt_maxframes;	/* system limit on stack traces */
 	dt_tstring_t *dt_tstrings; /* temporary string slots */
 	dt_list_t dt_modlist;	/* linked list of dt_module_t's */
+	dt_module_range_t *dt_modranges; /* sorted module address ranges */
+	size_t dt_modranges_count;
+	int dt_modranges_valid;
 	dt_htab_t *dt_mods;	/* hash table of dt_module_t's */
 	dt_htab_t *dt_kernsyms; /* htab of kernel symbol names */
 	int dt_ctfa_missing;	/* cached ENOENT/ENOTDIR for the archive */
diff --git a/libdtrace/dt_module.c b/libdtrace/dt_module.c
index f381bdcd..9d93b9d8 100644
--- a/libdtrace/dt_module.c
+++ b/libdtrace/dt_module.c
@@ -239,9 +239,19 @@ dt_module_cmp(const dt_module_t *p, const dt_module_t *q)
 
 DEFINE_HE_STD_LINK_FUNCS(dt_module, dt_module_t, dm_he)
 
+static void
+dt_module_ranges_invalidate(dtrace_hdl_t *dtp)
+{
+	free(dtp->dt_modranges);
+	dtp->dt_modranges = NULL;
+	dtp->dt_modranges_count = 0;
+	dtp->dt_modranges_valid = 0;
+}
+
 static void *
 dt_module_del_mod(dt_module_t *head, dt_module_t *dmp)
 {
+	dt_module_ranges_invalidate(dmp->dm_dtp);
 	head = dt_module_del(head, dmp);
 	dt_list_delete(&dmp->dm_dtp->dt_modlist, dmp);
 
@@ -1724,6 +1734,9 @@ dt_modsym_addsym(dtrace_hdl_t *dtp, dt_module_t *dmp, dt_kallsym_t *sym,
 	dtrace_addr_range_t	*range = NULL;
 	int			skip = 0;
 
+	if (dtp->dt_modranges_valid)
+		dt_module_ranges_invalidate(dtp);
+
 	/*
 	 * Add the symbol to the module's kernel symbol table.
 	 *
@@ -2171,6 +2184,7 @@ dt_module_shuffle_to_start(dtrace_hdl_t *dtp, const char *name)
 
 	dt_list_delete(&dtp->dt_modlist, dmp);
 	dt_list_prepend(&dtp->dt_modlist, dmp);
+	dt_module_ranges_invalidate(dtp);
 }
 
 static dt_module_t *
@@ -2338,6 +2352,118 @@ dtrace_lookup_by_name(dtrace_hdl_t *dtp, const char *object, const char *name,
 	return dt_set_errno(dtp, EDT_NOSYM);
 }
 
+static int
+dt_module_range_cmp(const void *a, const void *b)
+{
+	const dt_module_range_t *p = a, *q = b;
+
+	return p->start < q->start ? -1 : p->start > q->start;
+}
+
+/* Build lazily: handles that never symbolize addresses need no index. */
+static int
+dt_module_ranges_build(dtrace_hdl_t *dtp)
+{
+	dt_module_t *dmp;
+	dt_module_range_t *index;
+	size_t total = 0, n = 0, order = 0, i;
+	GElf_Addr max_last = 0;
+
+	for (dmp = dt_list_next(&dtp->dt_modlist); dmp;
+	     dmp = dt_list_next(dmp)) {
+		if (dmp->dm_text_addrs_size > SIZE_MAX - total)
+			return -1;
+		total += dmp->dm_text_addrs_size;
+		if (dmp->dm_data_addrs_size > SIZE_MAX - total)
+			return -1;
+		total += dmp->dm_data_addrs_size;
+	}
+	if (total > SIZE_MAX / sizeof(*index))
+		return -1;
+	index = total ? malloc(total * sizeof(*index)) : NULL;
+	if (total && index == NULL)
+		return -1;
+
+	for (dmp = dt_list_next(&dtp->dt_modlist); dmp;
+	     dmp = dt_list_next(dmp), order++) {
+		int text;
+
+		for (text = 0; text < 2; text++) {
+			dtrace_addr_range_t *ranges = text ? dmp->dm_text_addrs :
+							   dmp->dm_data_addrs;
+			size_t size = text ? dmp->dm_text_addrs_size :
+					    dmp->dm_data_addrs_size;
+
+			for (i = 0; i < size; i++) {
+				GElf_Addr start = ranges[i].dar_va;
+				GElf_Addr len = ranges[i].dar_size;
+
+				if (len == 0)
+					continue;
+				index[n].start = start;
+				index[n].last = len - 1 > UINT64_MAX - start ?
+						UINT64_MAX : start + len - 1;
+				index[n].module = dmp;
+				index[n++].order = order;
+			}
+		}
+	}
+	if (n)
+		qsort(index, n, sizeof(*index), dt_module_range_cmp);
+	for (i = 0; i < n; i++) {
+		if (index[i].last > max_last)
+			max_last = index[i].last;
+		index[i].max_last = max_last;
+	}
+	dtp->dt_modranges = index;
+	dtp->dt_modranges_count = n;
+	dtp->dt_modranges_valid = 1;
+	return 0;
+}
+
+static dt_module_t *
+dt_module_by_addr(dtrace_hdl_t *dtp, GElf_Addr addr)
+{
+	dt_module_range_t *index;
+	dt_module_t *dmp = NULL;
+	size_t lo = 0, hi, best = SIZE_MAX;
+
+	if (!dtp->dt_modranges_valid && dt_module_ranges_build(dtp) != 0) {
+		/* An optional index must not make a lookup fail under memory pressure. */
+		for (dmp = dt_list_next(&dtp->dt_modlist); dmp;
+		     dmp = dt_list_next(dmp)) {
+			if (bsearch(&addr, dmp->dm_text_addrs, dmp->dm_text_addrs_size,
+				    sizeof(dtrace_addr_range_t), dtrace_addr_range_cmp) ||
+			    bsearch(&addr, dmp->dm_data_addrs, dmp->dm_data_addrs_size,
+				    sizeof(dtrace_addr_range_t), dtrace_addr_range_cmp))
+				return dmp;
+		}
+		return NULL;
+	}
+	index = dtp->dt_modranges;
+	hi = dtp->dt_modranges_count;
+	while (lo < hi) {
+		size_t mid = lo + (hi - lo) / 2;
+
+		if (index[mid].start <= addr)
+			lo = mid + 1;
+		else
+			hi = mid;
+	}
+	/* Prefix maxima bound the backwards search, including overlapping ranges. */
+	while (lo && index[lo - 1].max_last >= addr) {
+		dt_module_range_t *r = &index[--lo];
+
+		if (r->last >= addr && r->order < best) {
+			dmp = r->module;
+			best = r->order;
+			if (best == 0)
+				break;
+		}
+	}
+	return dmp;
+}
+
 /*
  * Exported interface to look up a symbol by address.  We return the (possibly
  * partial) GElf_Sym and complete symbol information for the matching symbol.
@@ -2358,22 +2484,7 @@ dtrace_lookup_by_addr(dtrace_hdl_t *dtp, GElf_Addr addr,
 	if (v != NULL)
 		return v->dtv_lookup_by_addr(dtp->dt_varg, addr, symp, sip);
 
-	for (dmp = dt_list_next(&dtp->dt_modlist); dmp != NULL;
-	    dmp = dt_list_next(dmp)) {
-		void *i;
-
-		i = bsearch(&addr, dmp->dm_text_addrs, dmp->dm_text_addrs_size,
-		    sizeof(struct dtrace_addr_range), dtrace_addr_range_cmp);
-
-		if (i)
-			break;
-
-		i = bsearch(&addr, dmp->dm_data_addrs, dmp->dm_data_addrs_size,
-		    sizeof(struct dtrace_addr_range), dtrace_addr_range_cmp);
-
-		if (i)
-			break;
-	}
+	dmp = dt_module_by_addr(dtp, addr);
 
 	if (dmp == NULL) {
 		dt_dprintf("No module corresponds to %lx\n", addr);
diff --git a/libdtrace/dt_open.c b/libdtrace/dt_open.c
index 8291b87d..fbc97ccc 100644
--- a/libdtrace/dt_open.c
+++ b/libdtrace/dt_open.c
@@ -1312,6 +1312,7 @@ dtrace_close(dtrace_hdl_t *dtp)
 	dt_htab_destroy(dtp->dt_kernsyms);
 	dtp->dt_kernsyms = NULL;
 	dt_htab_destroy(dtp->dt_mods);
+	free(dtp->dt_modranges);
 	dt_htab_destroy(dtp->dt_kernpaths);
 	dt_htab_destroy(dtp->dt_btf_ids);
 
diff --git a/test/triggers/module-discovery.c b/test/triggers/module-discovery.c
index d05d2bab..ee9604b7 100644
--- a/test/triggers/module-discovery.c
+++ b/test/triggers/module-discovery.c
@@ -97,6 +97,7 @@ check_discovery(int fallback)
 	dtrace_addr_range_t *ranges, *kernel_ranges;
 	ctf_file_t *types;
 	GElf_Sym elf;
+	dtrace_syminfo_t si;
 	int err, i;
 
 	CHECK((dtp = dtrace_open(DTRACE_VERSION, DTRACE_O_LP64, &err)) != NULL);
@@ -144,6 +145,12 @@ check_discovery(int fallback)
 	CHECK(ranges->dar_va == 0x1000);
 	CHECK(ranges->dar_size == (fallback ? 0x11 : 0x20));
 	CHECK(dt_symbol_by_addr(first_syms, 0x1000) == sym);
+	first->dm_flags |= DT_DM_LOADED;
+	CHECK(dtrace_lookup_by_addr(dtp, 0x1000, &elf, &si) == 0);
+	CHECK(strcmp(si.object, first->dm_name) == 0);
+	CHECK(dtrace_lookup_by_addr(dtp, 0x1010, &elf, &si) == 0);
+	CHECK(dtrace_lookup_by_addr(dtp, fallback ? 0x1011 : 0x1020,
+				    &elf, &si) == -1);
 
 	/* Keep the old module while discovering another, including a duplicate name. */
 	btf_generation = 200;
@@ -152,15 +159,36 @@ check_discovery(int fallback)
 	    "1000 T discovery_symbol [discovery_first]\n"
 	    "1010 T discovery_tail [discovery_first]\n"
 	    "2000 T discovery_symbol [discovery_second]\n"
+	    "1008 T overlapping_symbol [discovery_overlap]\n"
 	    "3000 T discovery_live [discovery_offline]\n" :
 	    "100 1 T _end\n"
 	    "1000 10 T discovery_symbol [discovery_first]\n"
 	    "1010 10 T discovery_tail [discovery_first]\n"
 	    "2000 10 T discovery_symbol [discovery_second]\n"
+	    "1008 1 T overlapping_symbol [discovery_overlap]\n"
 	    "3000 10 T discovery_live [discovery_offline]\n";
 	CHECK(dtrace_update(dtp) == 0);
 	CHECK((second = dt_module_lookup_by_name(dtp, "discovery_second")) != NULL);
 	CHECK(first->dm_btf_id == 101 && second->dm_btf_id == 202);
+	second->dm_flags |= DT_DM_LOADED;
+	CHECK(dtrace_lookup_by_addr(dtp, 0x2000, &elf, &si) == 0);
+	CHECK(strcmp(si.object, second->dm_name) == 0);
+	{
+		dt_module_t *a = dt_module_lookup_by_name(dtp, "discovery_multi_a");
+		dt_module_t *b = dt_module_lookup_by_name(dtp, "discovery_multi_b");
+
+		a->dm_flags |= DT_DM_LOADED;
+		b->dm_flags |= DT_DM_LOADED;
+		dt_module_t *overlap = dt_module_lookup_by_name(dtp, "discovery_overlap");
+
+		CHECK(overlap != NULL);
+		overlap->dm_flags |= DT_DM_LOADED;
+		CHECK(dtrace_lookup_by_addr(dtp, 0x1008, &elf, &si) == 0);
+		CHECK(strcmp(si.object, first->dm_name) == 0);
+		/* Equal ranges retain the first module's precedence. */
+		CHECK(dtrace_lookup_by_addr(dtp, 0x4000, &elf, &si) == 0);
+		CHECK(strcmp(si.object, a->dm_name) == 0);
+	}
 	CHECK(dt_module_symbol_by_name(dtp, second, "discovery_symbol") != NULL);
 	CHECK(dt_module_symbol_by_name(dtp, offline, "discovery_live") != NULL);
 	CHECK(!(offline->dm_flags & DT_DM_KERN_UNLOADED));
@@ -192,6 +220,8 @@ check_discovery(int fallback)
 		CHECK(kernel->dm_kernsyms == kernel_syms);
 		CHECK(kernel->dm_text_addrs == kernel_ranges);
 		CHECK(first->dm_kernsyms == first_syms);
+		CHECK(dtrace_lookup_by_addr(dtp, 0x1000, &elf, &si) == 0);
+		CHECK(strcmp(si.object, first->dm_name) == 0);
 		CHECK(first->dm_btf_id == 101 && second->dm_btf_id == 202);
 		CHECK(first->dm_text_addrs == ranges);
 		CHECK(first->dm_text_addrs_size == 1);
@@ -205,6 +235,54 @@ check_discovery(int fallback)
 		CHECK(offline->dm_ctfp == types);
 	}
 
+	/* Alternating built-in text/data fragments cross both initial capacities. */
+	{
+		char snapshot[2048];
+		size_t used = 0;
+		dt_module_t *fragmented;
+
+		for (i = 0; i < 12; i++)
+			used += snprintf(snapshot + used, sizeof(snapshot) - used,
+				fallback ? "%x %c fragment_%d [discovery_fragmented]\n" :
+					   "%x 1 %c fragment_%d [discovery_fragmented]\n",
+				0x8000 + i * 0x10, i % 2 ? 'D' : 'T', i);
+		symbols = snapshot;
+		CHECK(dtrace_update(dtp) == 0);
+		fragmented = dt_module_lookup_by_name(dtp, "discovery_fragmented");
+		CHECK(fragmented != NULL);
+		CHECK(fragmented->dm_text_addrs_size == 6);
+		CHECK(fragmented->dm_data_addrs_size == 6);
+		fragmented->dm_flags |= DT_DM_LOADED;
+		for (i = 0; i < 12; i++) {
+			CHECK(dtrace_lookup_by_addr(dtp, 0x8000 + i * 0x10, &elf, &si) == 0);
+			CHECK(strcmp(si.object, fragmented->dm_name) == 0);
+			CHECK(dtrace_lookup_by_addr(dtp, 0x8001 + i * 0x10, &elf, &si) == -1);
+		}
+	}
+	{
+		char malformed[8192], longname[PATH_MAX + 1];
+		const char *rows[] = {
+			"10000000000000000 T bad [discovery_rejected]\n",
+			"1000 T bad [discovery_rejected\n",
+			"1000 ",
+			"1000 T ",
+			NULL
+		};
+
+		for (i = 0; rows[i]; i++) {
+			symbols = rows[i];
+			CHECK(dtrace_update(dtp) == 0);
+			CHECK(dt_module_lookup_by_name(dtp, "discovery_rejected") == NULL);
+		}
+		memset(longname, 'x', sizeof(longname) - 1);
+		longname[sizeof(longname) - 1] = '\0';
+		snprintf(malformed, sizeof(malformed), fallback ?
+			 "1000 T %s [discovery_rejected]\n" :
+			 "1000 1 T %s [discovery_rejected]\n", longname);
+		symbols = malformed;
+		CHECK(dtrace_update(dtp) == 0);
+		CHECK(dt_module_lookup_by_name(dtp, "discovery_rejected") == NULL);
+	}
 	symbols = NULL;
 	dtrace_close(dtp);
 	return 0;
-- 
2.52.0