From 7bb7b1f8f2128061195212bd49f450676789f937 Mon Sep 17 00:00:00 2001
From: Tino Reichardt <milky-zfs@mcmilk.de>
Date: Mon, 13 Mar 2023 23:30:09 +0100
Subject: [PATCH] Colorize zpool iostat output

Use a bold header and colorize the space suffixes in iostat
by order of magnitude like this:
- K is green
- M is yellow
- G is red
- T is lightblue
- P is magenta
- E is cyan
- 0 space is colored gray

Signed-off-by: Ethan Coe-Renner <coerenner1@llnl.gov>
Signed-off-by: Tino Reichardt <milky-zfs@mcmilk.de>
Closes: #14459
---
 cmd/zpool/zpool_main.c | 35 ++++++++++++++++++++++++++++++++++-
 man/man8/zpool.8       |  2 ++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/cmd/zpool/zpool_main.c b/cmd/zpool/zpool_main.c
index d2614411df..ad37722829 100644
--- a/cmd/zpool/zpool_main.c
+++ b/cmd/zpool/zpool_main.c
@@ -4220,6 +4220,8 @@ print_iostat_header_impl(iostat_cbdata_t *cb, unsigned int force_column_width,
 	unsigned int namewidth;
 	const char *title;
 
+	color_start(ANSI_BOLD);
+
 	if (cb->cb_flags & IOS_ANYHISTO_M) {
 		title = histo_to_title[IOS_HISTO_IDX(cb->cb_flags)];
 	} else if (cb->cb_vdevs.cb_names_count) {
@@ -4253,6 +4255,8 @@ print_iostat_header_impl(iostat_cbdata_t *cb, unsigned int force_column_width,
 	if (cb->vcdl != NULL)
 		print_cmd_columns(cb->vcdl, 1);
 
+	color_end();
+
 	printf("\n");
 }
 
@@ -4262,6 +4266,35 @@ print_iostat_header(iostat_cbdata_t *cb)
 	print_iostat_header_impl(cb, 0, NULL);
 }
 
+/*
+ * Prints a size string (i.e. 120M) with the suffix ("M") colored
+ * by order of magnitude. Uses column_size to add padding.
+ */
+static void
+print_stat_color(char *statbuf, unsigned int column_size)
+{
+	fputs("  ", stdout);
+	if (*statbuf == '0') {
+		color_start(ANSI_GRAY);
+		fputc('0', stdout);
+		column_size--;
+	} else {
+		for (; *statbuf; statbuf++) {
+			if (*statbuf == 'K') color_start(ANSI_GREEN);
+			else if (*statbuf == 'M') color_start(ANSI_YELLOW);
+			else if (*statbuf == 'G') color_start(ANSI_RED);
+			else if (*statbuf == 'T') color_start(ANSI_BOLD_BLUE);
+			else if (*statbuf == 'P') color_start(ANSI_MAGENTA);
+			else if (*statbuf == 'E') color_start(ANSI_CYAN);
+			fputc(*statbuf, stdout);
+			if (--column_size <= 0)
+				break;
+		}
+	}
+	color_end();
+	for (; column_size > 0; column_size--)
+		fputc(' ', stdout);
+}
 
 /*
  * Display a single statistic.
@@ -4277,7 +4310,7 @@ print_one_stat(uint64_t value, enum zfs_nicenum_format format,
 	if (scripted)
 		printf("\t%s", buf);
 	else
-		printf("  %*s", column_size, buf);
+		print_stat_color(buf, column_size);
 }
 
 /*
diff --git a/man/man8/zpool.8 b/man/man8/zpool.8
index 48e4eb3c55..e8eadffa6f 100644
--- a/man/man8/zpool.8
+++ b/man/man8/zpool.8
@@ -445,6 +445,8 @@ to dump core on exit for the purposes of running
 .It Sy ZFS_COLOR
 Use ANSI color in
 .Nm zpool Cm status
+and
+.Nm zpool Cm iostat
 output.
 .It Sy ZPOOL_IMPORT_PATH
 The search path for devices or files to use with the pool.