libbpf-tools/profile: Use batch method if available

Use bpf_map_lookup_batch if available

Signed-off-by: Eunseon Lee <[email protected]>

libbpf-tools/profile.c

@@ -53,6 +53,8 @@ static struct env {
 	.cpu = -1,
 };
 
+#define warn(...) fprintf(stderr, __VA_ARGS__)
+
 const char *argp_program_version = "profile 0.1";
 const char *argp_program_bug_address =
 	"https://github.com/iovisor/bcc/tree/master/libbpf-tools";
@@ -250,13 +252,65 @@ static int cmp_counts(const void *dx, const void *dy)
 	return x > y ? -1 : !(x == y);
 }
 
+static bool batch_map_ops = true; /* hope for the best */
+
+static bool read_batch_counts_map(int fd, struct key_ext_t *items, __u32 *count)
+{
+	void *in = NULL, *out;
+	__u32 i, n, n_read = 0;
+	int err = 0;
+	__u32 vals[*count];
+	struct key_t keys[*count];
+
+	while (n_read < *count && !err) {
+		n = *count - n_read;
+		err = bpf_map_lookup_batch(fd, &in, &out, keys + n_read,
+					   vals + n_read, &n, NULL);
+		if (err && errno != ENOENT) {
+			/* we want to propagate EINVAL upper, so that
+			 * the batch_map_ops flag is set to false */
+			if (errno != EINVAL)
+				warn("bpf_map_lookup_batch: %s\n",
+				     strerror(-err));
+			return false;
+		}
+		n_read += n;
+		in = out;
+	}
+
+	for (i = 0; i < n_read; i++) {
+		items[i].k.pid = keys[i].pid;
+		items[i].k.kernel_ip = keys[i].kernel_ip;
+		items[i].k.user_stack_id = keys[i].user_stack_id;
+		items[i].k.kern_stack_id = keys[i].kern_stack_id;
+		strncpy(items[i].k.name, keys[i].name, TASK_COMM_LEN);
+		items[i].v = vals[i];
+	}
+
+	*count = n_read;
+	return true;
+}
+
 static bool read_counts_map(int fd, struct key_ext_t *items, __u32 *count)
 {
 	struct key_t empty = {};
 	struct key_t *lookup_key = &empty;
 	int i = 0;
 	int err;
 
+	if (batch_map_ops) {
+		bool ok = read_batch_counts_map(fd, items, count);
+		if (!ok && errno == EINVAL) {
+			/* fall back to a racy variant */
+			batch_map_ops = false;
+		} else {
+			return ok;
+		}
+	}
+
+	if (!items || !count || !*count)
+		return true;
+
 	while (!bpf_map_get_next_key(fd, lookup_key, &items[i].k)) {
 
 		err = bpf_map_lookup_elem(fd, &items[i].k, &items[i].v);