From 3af5e95ef76e866118294c7f775393242de439ac Mon Sep 17 00:00:00 2001
From: spacewander <spacewanderlzx@gmail.com>
Date: Tue, 13 Mar 2018 11:05:27 +0800
Subject: [PATCH] Avoid counting disk io twice for NVMe SSD block device

The first NVMe block device will be named as nvme0n1,
and its partitions are nvme0n1p1, nvme0n1p2, etc.
Therefore we could not use the last letter to identify if the name is a
partition or not.
---
 psutil/_pslinux.py         | 21 +++++++++++----------
 psutil/tests/test_linux.py | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/psutil/_pslinux.py b/psutil/_pslinux.py
index b57adb34e..ec951317c 100644
--- a/psutil/_pslinux.py
+++ b/psutil/_pslinux.py
@@ -1017,23 +1017,24 @@ def disk_io_counters():
     system as a dict of raw tuples.
     """
     # determine partitions we want to look for
+    partition_suffix = re.compile(r'^p*\d+$')
+
     def get_partitions():
         partitions = []
         with open_text("%s/partitions" % get_procfs_path()) as f:
             lines = f.readlines()[2:]
         for line in reversed(lines):
             _, _, _, name = line.split()
-            if name[-1].isdigit():
-                # we're dealing with a partition (e.g. 'sda1'); 'sda' will
-                # also be around but we want to omit it
-                partitions.append(name)
-            else:
-                if not partitions or not partitions[-1].startswith(name):
-                    # we're dealing with a disk entity for which no
-                    # partitions have been defined (e.g. 'sda' but
-                    # 'sda1' was not around), see:
-                    # https://github.com/giampaolo/psutil/issues/338
+            if partitions and partitions[-1].startswith(name):
+                suffix = partitions[-1][len(name):]
+                if not partition_suffix.match(suffix):
+                    # If a disk entity (e.g. 'sda' or 'nvme0n1') has
+                    # partition(e.g. 'sda1' or 'nvme0n1p1'),
+                    # we will deal with the partition and ignore the
+                    # disk entiy.
                     partitions.append(name)
+            else:
+                partitions.append(name)
         return partitions
 
     retdict = {}
diff --git a/psutil/tests/test_linux.py b/psutil/tests/test_linux.py
index 6ba17b254..c100e2808 100755
--- a/psutil/tests/test_linux.py
+++ b/psutil/tests/test_linux.py
@@ -1014,6 +1014,39 @@ def open_mock(name, *args, **kwargs):
             self.assertEqual(ret.write_time, 0)
             self.assertEqual(ret.busy_time, 0)
 
+    def test_disk_io_counters_with_nvme(self):
+        def open_mock(name, *args, **kwargs):
+            if name == '/proc/partitions':
+                return io.StringIO(textwrap.dedent(u"""\
+                    major minor  #blocks  name
+
+                    259        0  250059096 nvme0n1
+                    259        1     266240 nvme0n1p1
+                    259        2      16384 nvme0n1p2
+                    """))
+            elif name == '/proc/diskstats':
+                return io.StringIO(textwrap.dedent(u"""\
+                        259       0 nvme0n1  3 6 9 12 15 18 21 24 27 30 33
+                        259       1 nvme0n1p1 1 2 3 4 5 6 7 8 9 10 11
+                        259       2 nvme0n1p2 2 4 6 8 10 12 14 16 18 20 22
+                    """))
+            else:
+                return orig_open(name, *args, **kwargs)
+
+        orig_open = open
+        patch_point = 'builtins.open' if PY3 else '__builtin__.open'
+        with mock.patch(patch_point, side_effect=open_mock) as m:
+            ret = psutil.disk_io_counters(nowrap=False)
+            assert m.called
+            self.assertEqual(ret.read_count, 3)
+            self.assertEqual(ret.read_merged_count, 6)
+            self.assertEqual(ret.read_bytes, 9 * SECTOR_SIZE)
+            self.assertEqual(ret.read_time, 12)
+            self.assertEqual(ret.write_count, 15)
+            self.assertEqual(ret.write_merged_count, 18)
+            self.assertEqual(ret.write_bytes, 21 * SECTOR_SIZE)
+            self.assertEqual(ret.write_time, 24)
+            self.assertEqual(ret.busy_time, 30)
 
 # =====================================================================
 # --- misc