From ecebf770c1d7db987b8eafe371d72feb19e5870d Mon Sep 17 00:00:00 2001
From: Rich Ercolani <214141+rincebrain@users.noreply.github.com>
Date: Tue, 25 May 2021 22:02:01 -0400
Subject: [PATCH] Correct flaws in arc_summary[23] and their test.

The change correctly handles BrokenPipeError and improves the
associated tests.

Reviewed-by: Brian Behlendorf <behlendorf1@llnl.gov>
Reviewed-by: John Kennedy <john.kennedy@delphix.com>
Signed-off-by: Rich Ercolani <rincebrain@gmail.com>
Closes #12037
Closes #12036
---
 cmd/arc_summary/arc_summary2                  | 97 +++++++++----------
 cmd/arc_summary/arc_summary3                  | 24 +++++
 .../cli_user/misc/arc_summary_001_pos.ksh     |  3 +
 3 files changed, 73 insertions(+), 51 deletions(-)

diff --git a/cmd/arc_summary/arc_summary2 b/cmd/arc_summary/arc_summary2
index 75b5697526..3302a802d1 100755
--- a/cmd/arc_summary/arc_summary2
+++ b/cmd/arc_summary/arc_summary2
@@ -102,18 +102,6 @@ show_tunable_descriptions = False
 alternate_tunable_layout = False
 
 
-def handle_Exception(ex_cls, ex, tb):
-    if ex is IOError:
-        if ex.errno == errno.EPIPE:
-            sys.exit()
-
-    if ex is KeyboardInterrupt:
-        sys.exit()
-
-
-sys.excepthook = handle_Exception
-
-
 def get_Kstat():
     """Collect information on the ZFS subsystem from the /proc virtual
     file system. The name "kstat" is a holdover from the Solaris utility
@@ -1137,48 +1125,55 @@ def main():
     global alternate_tunable_layout
 
     try:
-        opts, args = getopt.getopt(
-            sys.argv[1:],
-            "adp:h", ["alternate", "description", "page=", "help"]
-        )
-    except getopt.error as e:
-        sys.stderr.write("Error: %s\n" % e.msg)
-        usage()
-        sys.exit(1)
-
-    args = {}
-    for opt, arg in opts:
-        if opt in ('-a', '--alternate'):
-            args['a'] = True
-        if opt in ('-d', '--description'):
-            args['d'] = True
-        if opt in ('-p', '--page'):
-            args['p'] = arg
-        if opt in ('-h', '--help'):
-            usage()
-            sys.exit(0)
-
-    Kstat = get_Kstat()
-
-    alternate_tunable_layout = 'a' in args
-    show_tunable_descriptions = 'd' in args
-
-    pages = []
-
-    if 'p' in args:
         try:
-            pages.append(unSub[int(args['p']) - 1])
-        except IndexError:
-            sys.stderr.write('the argument to -p must be between 1 and ' +
-                             str(len(unSub)) + '\n')
+            opts, args = getopt.getopt(
+                sys.argv[1:],
+                "adp:h", ["alternate", "description", "page=", "help"]
+            )
+        except getopt.error as e:
+            sys.stderr.write("Error: %s\n" % e.msg)
+            usage()
             sys.exit(1)
-    else:
-        pages = unSub
 
-    zfs_header()
-    for page in pages:
-        page(Kstat)
-        sys.stdout.write("\n")
+        args = {}
+        for opt, arg in opts:
+            if opt in ('-a', '--alternate'):
+                args['a'] = True
+            if opt in ('-d', '--description'):
+                args['d'] = True
+            if opt in ('-p', '--page'):
+                args['p'] = arg
+            if opt in ('-h', '--help'):
+                usage()
+                sys.exit(0)
+
+        Kstat = get_Kstat()
+
+        alternate_tunable_layout = 'a' in args
+        show_tunable_descriptions = 'd' in args
+
+        pages = []
+
+        if 'p' in args:
+            try:
+                pages.append(unSub[int(args['p']) - 1])
+            except IndexError:
+                sys.stderr.write('the argument to -p must be between 1 and ' +
+                                 str(len(unSub)) + '\n')
+                sys.exit(1)
+        else:
+            pages = unSub
+
+        zfs_header()
+        for page in pages:
+            page(Kstat)
+            sys.stdout.write("\n")
+    except IOError as ex:
+        if (ex.errno == errno.EPIPE):
+            sys.exit(0)
+        raise
+    except KeyboardInterrupt:
+        sys.exit(0)
 
 
 if __name__ == '__main__':
diff --git a/cmd/arc_summary/arc_summary3 b/cmd/arc_summary/arc_summary3
index d6ddd06e71..d5b5489992 100755
--- a/cmd/arc_summary/arc_summary3
+++ b/cmd/arc_summary/arc_summary3
@@ -43,6 +43,12 @@ import subprocess
 import sys
 import time
 
+# We can't use env -S portably, and we need python3 -u to handle pipes in
+# the shell abruptly closing the way we want to, so...
+import io
+if isinstance(sys.__stderr__.buffer, io.BufferedWriter):
+    os.execv(sys.executable, [sys.executable, "-u"] + sys.argv)
+
 DESCRIPTION = 'Print ARC and other statistics for OpenZFS'
 INDENT = ' '*8
 LINE_LENGTH = 72
@@ -221,6 +227,24 @@ elif sys.platform.startswith('linux'):
 
         return descs
 
+def handle_unraisableException(exc_type, exc_value=None, exc_traceback=None,
+                               err_msg=None, object=None):
+   handle_Exception(exc_type, object, exc_traceback)
+
+def handle_Exception(ex_cls, ex, tb):
+    if ex_cls is KeyboardInterrupt:
+        sys.exit()
+
+    if ex_cls is BrokenPipeError:
+        # It turns out that while sys.exit() triggers an exception
+        # not handled message on Python 3.8+, os._exit() does not.
+        os._exit(0)
+    raise ex
+
+if hasattr(sys,'unraisablehook'): # Python 3.8+
+    sys.unraisablehook = handle_unraisableException
+sys.excepthook = handle_Exception
+
 
 def cleanup_line(single_line):
     """Format a raw line of data from /proc and isolate the name value
diff --git a/tests/zfs-tests/tests/functional/cli_user/misc/arc_summary_001_pos.ksh b/tests/zfs-tests/tests/functional/cli_user/misc/arc_summary_001_pos.ksh
index a445fbb48c..befbea986e 100755
--- a/tests/zfs-tests/tests/functional/cli_user/misc/arc_summary_001_pos.ksh
+++ b/tests/zfs-tests/tests/functional/cli_user/misc/arc_summary_001_pos.ksh
@@ -48,6 +48,9 @@ else
 	set -A args  "" "-a" "-d" "-p 1"
 fi
 
+# Without this, the below checks aren't going to work the way we hope...
+set -o pipefail
+
 typeset -i i=0
 while [[ $i -lt ${#args[*]} ]]; do
 	log_must eval "arc_summary ${args[i]} > /dev/null"