Colorize zpool status output

include/libzutil.h

@@ -146,6 +146,18 @@ extern int zpool_history_unpack(char *, uint64_t, uint64_t *, nvlist_t ***,
 struct zfs_cmd;
 int zfs_ioctl_fd(int fd, unsigned long request, struct zfs_cmd *zc);
 
+/*
+ * List of colors to use
+ */
+#define	ANSI_RED	"\033[0;31m"
+#define	ANSI_YELLOW	"\033[0;33m"
+#define	ANSI_RESET	"\033[0m"
+#define	ANSI_BOLD	"\033[1m"
+
+void color_start(char *color);
+void color_end(void);
+int printf_color(char *color, char *format, ...);
+
 #ifdef	__cplusplus
 }
 #endif

lib/libzfs/libzfs_util.c

@@ -1886,3 +1886,95 @@ zfs_version_print(void)
 
 	return (0);
 }
+
+/*
+ * Return 1 if the user requested ANSI color output, and our terminal supports
+ * it.  Return 0 for no color.
+ */
+static int
+use_color(void)
+{
+	static int use_color = -1;
+	char *term;
+
+	/*
+	 * Optimization:
+	 *
+	 * For each zpool invocation, we do a single check to see if we should
+	 * be using color or not, and cache that value for the lifetime of the
+	 * the zpool command.  That makes it cheap to call use_color() when
+	 * we're printing with color.  We assume that the settings are not going
+	 * to change during the invocation of a zpool command (the user isn't
+	 * going to change the ZFS_COLOR value while zpool is running, for
+	 * example).
+	 */
+	if (use_color != -1) {
+		/*
+		 * We've already figured out if we should be using color or
+		 * not.  Return the cached value.
+		 */
+		return (use_color);
+	}
+
+	term = getenv("TERM");
+	/*
+	 * The user sets the ZFS_COLOR env var set to enable zpool ANSI color
+	 * output.  However if NO_COLOR is set (https://no-color.org/) then
+	 * don't use it.  Also, don't use color if terminal doesn't support
+	 * it.
+	 */
+	if (libzfs_envvar_is_set("ZFS_COLOR") &&
+	    !libzfs_envvar_is_set("NO_COLOR") &&
+	    isatty(STDOUT_FILENO) && term && strcmp("dumb", term) != 0 &&
+	    strcmp("unknown", term) != 0) {
+		/* Color supported */
+		use_color = 1;
+	} else {
+		use_color = 0;
+	}
+
+	return (use_color);
+}
+
+/*
+ * color_start() and color_end() are used for when you want to colorize a block
+ * of text.  For example:
+ *
+ * color_start(ANSI_RED_FG)
+ * printf("hello");
+ * printf("world");
+ * color_end();
+ */
+void
+color_start(char *color)
+{
+	if (use_color())
+		printf("%s", color);
+}
+
+void
+color_end(void)
+{
+	if (use_color())
+		printf(ANSI_RESET);
+}
+
+/* printf() with a color.  If color is NULL, then do a normal printf. */
+int
+printf_color(char *color, char *format, ...)
+{
+	va_list aptr;
+	int rc;
+
+	if (color)
+		color_start(color);
+
+	va_start(aptr, format);
+	rc = vprintf(format, aptr);
+	va_end(aptr);
+
+	if (color)
+		color_end();
+
+	return (rc);
+}

man/man8/zpool.8

@@ -453,6 +453,12 @@ Cause
 to dump core on exit for the purposes of running
 .Sy ::findleaks .
 .El
+.Bl -tag -width "ZFS_COLOR"
+.It Ev ZFS_COLOR
+Use ANSI color in
+.Nm zpool status
+output.
+.El
 .Bl -tag -width "ZPOOL_IMPORT_PATH"
 .It Ev ZPOOL_IMPORT_PATH
 The search path for devices or files to use with the pool. This is a colon-separated list of directories in which

tests/runfiles/common.run

@@ -276,7 +276,7 @@ tests = ['zfs_upgrade_001_pos', 'zfs_upgrade_002_pos', 'zfs_upgrade_003_pos',
 tags = ['functional', 'cli_root', 'zfs_upgrade']
 
 [tests/functional/cli_root/zpool]
-tests = ['zpool_001_neg', 'zpool_002_pos', 'zpool_003_pos']
+tests = ['zpool_001_neg', 'zpool_002_pos', 'zpool_003_pos', 'zpool_colors']
 tags = ['functional', 'cli_root', 'zpool']
 
 [tests/functional/cli_root/zpool_add]

tests/zfs-tests/include/commands.cfg

@@ -77,6 +77,7 @@ export SYSTEM_FILES_COMMON='arp
     rm
     rmdir
     scp
+    script
     sed
     seq
     setfacl

tests/zfs-tests/include/libtest.shlib

@@ -3916,3 +3916,18 @@ function stat_size #<path>
 		;;
 	esac
 }
+
+# Run a command as if it was being run in a TTY.
+#
+# Usage:
+#
+#    faketty command
+#
+function faketty
+{
+    if is_freebsd; then
+        script -q /dev/null "$@"
+    else
+        script --return --quiet -c "$*" /dev/null
+    fi
+}

tests/zfs-tests/tests/functional/cli_root/zpool/Makefile.am

@@ -4,4 +4,5 @@ dist_pkgdata_SCRIPTS = \
 	cleanup.ksh \
 	zpool_001_neg.ksh \
 	zpool_002_pos.ksh \
-	zpool_003_pos.ksh
+	zpool_003_pos.ksh \
+	zpool_colors.ksh

tests/zfs-tests/tests/functional/cli_root/zpool/setup.ksh

@@ -29,4 +29,4 @@
 
 DISK=${DISKS%% *}
 
-default_setup $DISK
+default_mirror_setup $DISKS

tests/zfs-tests/tests/functional/cli_root/zpool/zpool_colors.ksh

@@ -0,0 +1,91 @@
+#!/bin/ksh -p
+#
+# CDDL HEADER START
+#
+# This file and its contents are supplied under the terms of the
+# Common Development and Distribution License ("CDDL"), version 1.0.
+# You may only use this file in accordance with the terms of version
+# 1.0 of the CDDL.
+#
+# A full copy of the text of the CDDL should have accompanied this
+# source.  A copy of the CDDL is also available via the Internet at
+# http://www.illumos.org/license/CDDL.
+#
+# CDDL HEADER END
+#
+# Copyright (c) 2019 Lawrence Livermore National Security, LLC.
+
+. $STF_SUITE/include/libtest.shlib
+
+#
+# DESCRIPTION:
+#	Test that zpool status colored output works.
+#
+# STRATEGY:
+# 1. Create a pool with a bunch of errors and force fault one of the vdevs.
+# 2. Look for 'pool:' in bold.
+# 3. Look for 'DEGRADED' in yellow
+# 3. Look for 'FAULTED' in red
+#
+
+verify_runnable "both"
+
+function cleanup
+{
+	zinject -c all
+}
+
+log_onexit cleanup
+
+log_assert "Test colorized zpool status output"
+
+DISK2="$(echo $DISKS | cut -d' ' -f2)"
+DISK3="$(echo $DISKS | cut -d' ' -f3)"
+
+log_must dd if=/dev/urandom of=/$TESTDIR/testfile bs=10M count=1
+
+log_must zpool sync
+
+log_must zpool offline -f $TESTPOOL $DISK3
+log_must wait_for_degraded $TESTPOOL
+log_must zinject -d $DISK2 -e io -T read -f 20 $TESTPOOL
+log_must zinject -d $DISK2 -e io -T write -f 20 $TESTPOOL
+
+
+log_must zpool scrub -w $TESTPOOL
+log_must zinject -c all
+
+
+# Use 'script' to fake zpool status into thinking it's running in a tty.
+# Log the output here in case it's needed for postmortem.
+log_note "$(faketty TERM=xterm-256color ZFS_COLOR=1 zpool status)"
+
+# Replace the escape codes with "ESC" so they're easier to grep
+out="$(faketty TERM=xterm-256color ZFS_COLOR=1 zpool status | \
+    grep -E 'pool:|DEGRADED' | \
+    sed -r 's/\s+//g;'$(echo -e 's/\033/ESC/g'))"
+
+log_note "$(echo $out)"
+
+log_note "Look for 'pool:' in bold"
+log_must eval "echo \"$out\" | grep -q 'ESC\[1mpool:ESC\[0m' "
+
+log_note "Look for 'DEGRADED' in yellow"
+log_must eval "echo \"$out\" | grep -q 'ESC\[0;33mDEGRADEDESC\[0m'"
+
+#
+# The escape code for 'FAULTED' is a little more tricky.  The line starts like
+# this:
+#
+# <start red escape code> loop2  FAULTED <end escape code>
+#
+# Luckily, awk counts the start and end escape codes as separate fields, so
+# we can easily remove the vdev field to get what we want.
+#
+out="$(faketty TERM=xterm-256color ZFS_COLOR=1 zpool status \
+    | awk '/FAULTED/{print $1$3$4}' | sed -r $(echo -e 's/\033/ESC/g'))"
+
+log_note "Look for 'FAULTED' in red"
+log_must eval "echo \"$out\" | grep -q 'ESC\[0;31mFAULTEDESC\[0m'"
+
+log_pass "zpool status displayed colors"