Avoid counting disk io twice for NVMe SSD block device

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 = {}

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