[PATCH bpf-next 4/5] bpftool: support optional 'task main_thread_only' argument

@@ -17,15 +17,16 @@ SYNOPSIS
 ITER COMMANDS
 ===================
 
-|	**bpftool** **iter pin** *OBJ* *PATH* [**map** *MAP*]
+|	**bpftool** **iter pin** *OBJ* *PATH* [**map** *MAP* | **task** *TASK_OPT*]
 |	**bpftool** **iter help**
 |
 |	*OBJ* := /a/file/of/bpf_iter_target.o
 |	*MAP* := { **id** *MAP_ID* | **pinned** *FILE* }
+|	*TASK_OPT* := { **main_thread_only** }
 
 DESCRIPTION
 ===========
-	**bpftool iter pin** *OBJ* *PATH* [**map** *MAP*]
+	**bpftool iter pin** *OBJ* *PATH* [**map** *MAP* | **task** *TASK_OPT*]
 		  A bpf iterator combines a kernel iterating of
 		  particular kernel data (e.g., tasks, bpf_maps, etc.)
 		  and a bpf program called for each kernel data object

@@ -44,6 +45,11 @@ DESCRIPTION
 		  with each map element, do checking, filtering, aggregation,
 		  etc. without copying data to user space.
 
+		  The task or task_file bpf iterator can have an optional
+		  parameter *TASK_OPT*. The current supported value is
+		  **main_thread_only** which supports to iterate only main
+		  threads of each process.
+
 		  User can then *cat PATH* to see the bpf iterator output.
 
 	**bpftool iter help**

@@ -78,6 +84,13 @@ EXAMPLES
    Create a file-based bpf iterator from bpf_iter_hashmap.o and map with
    id 20, and pin it to /sys/fs/bpf/my_hashmap
 
+**# bpftool iter pin bpf_iter_task.o /sys/fs/bpf/my_task task main_thread_only**
+
+::
+
+   Create a file-based bpf iterator from bpf_iter_task.o which iterates main
+   threads of processes only, and pin it to /sys/fs/bpf/my_hashmap
+
 SEE ALSO
 ========
 	**bpf**\ (2),

@@ -613,6 +613,7 @@ _bpftool()
             esac
             ;;
         iter)
+            local TARGET_TYPE='map task'
             case $command in
                 pin)
                     case $prev in

@@ -628,9 +629,15 @@ _bpftool()
                         pinned)
                             _filedir
                             ;;
-                        *)
+                        task)
+                            _bpftool_one_of_list 'main_thread_only'
+                            ;;
+                        map)
                             _bpftool_one_of_list $MAP_TYPE
                             ;;
+                        *)
+                            _bpftool_one_of_list $TARGET_TYPE
+                            ;;
                     esac
                     return 0
                     ;;

@@ -26,6 +26,7 @@ static int do_pin(int argc, char **argv)
 
 	/* optional arguments */
 	if (argc) {
+		memset(&linfo, 0, sizeof(linfo));
 		if (is_prefix(*argv, "map")) {
 			NEXT_ARG();
 

@@ -38,11 +39,29 @@ static int do_pin(int argc, char **argv)
 			if (map_fd < 0)
 				return -1;
 
-			memset(&linfo, 0, sizeof(linfo));
 			linfo.map.map_fd = map_fd;
-			iter_opts.link_info = &linfo;
-			iter_opts.link_info_len = sizeof(linfo);
+		} else if (is_prefix(*argv, "task")) {
+			NEXT_ARG();
+
+			if (!REQ_ARGS(1)) {
+				p_err("incorrect task spec");
+				return -1;
+			}
+
+			if (strcmp(*argv, "main_thread_only") != 0) {
+				p_err("incorrect task spec");
+				return -1;
+			}
+
+			linfo.task.main_thread_only = true;
+		} else {
+			p_err("expected no more arguments, 'map' or 'task', got: '%s'?",
+			      *argv);
+			return -1;
 		}
+
+		iter_opts.link_info = &linfo;
+		iter_opts.link_info_len = sizeof(linfo);
 	}
 
 	obj = bpf_object__open(objfile);

@@ -95,9 +114,10 @@ static int do_pin(int argc, char **argv)
 static int do_help(int argc, char **argv)
 {
 	fprintf(stderr,
-		"Usage: %1$s %2$s pin OBJ PATH [map MAP]\n"
+		"Usage: %1$s %2$s pin OBJ PATH [map MAP | task TASK_OPT]\n"
 		"       %1$s %2$s help\n"
 		"       " HELP_SPEC_MAP "\n"
+		"       TASK_OPT := { main_thread_only }\n"
 		"",
 		bin_name, "iter");
 

@@ -83,6 +83,12 @@ static bool is_iter_map_target(const char *target_name)
 	       strcmp(target_name, "bpf_sk_storage_map") == 0;
 }
 
+static bool is_iter_task_target(const char *target_name)
+{
+	return strcmp(target_name, "task") == 0 ||
+	       strcmp(target_name, "task_file") == 0;
+}
+
 static void show_iter_json(struct bpf_link_info *info, json_writer_t *wtr)
 {
 	const char *target_name = u64_to_ptr(info->iter.target_name);

@@ -91,6 +97,9 @@ static void show_iter_json(struct bpf_link_info *info, json_writer_t *wtr)
 
 	if (is_iter_map_target(target_name))
 		jsonw_uint_field(wtr, "map_id", info->iter.map.map_id);
+	else if (is_iter_task_target(target_name))
+		jsonw_uint_field(wtr, "main_thread_only",
+				 info->iter.task.main_thread_only);
 }
 
 static int get_prog_info(int prog_id, struct bpf_prog_info *info)

@@ -202,6 +211,9 @@ static void show_iter_plain(struct bpf_link_info *info)
 
 	if (is_iter_map_target(target_name))
 		printf("map_id %u  ", info->iter.map.map_id);
+	else if (is_iter_task_target(target_name))
+		printf("main_thread_only %u  ",
+		       info->iter.task.main_thread_only);
 }
 
 static int show_link_close_plain(int fd, struct bpf_link_info *info)