#1395 - disk_io_counters() - linux: mimick iostat behaviorlinux-disk-io-partitions

3 files changed, 64 insertions, 58 deletions

diff --git a/HISTORY.rst b/HISTORY.rst
index a5382fbc..6b754706 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -19,6 +19,9 @@ XXXX-XX-XX
 - 1278_: [macOS] Process.threads() incorrectly return microseconds instead of
   seconds. (patch by Nikhil Marathe)
 - 1279_: [Linux, macOS, BSD] net_if_stats() may return ENODEV.
+- 1395_: [Linux] disk_io_counters() can report inflated IO counters due to
+  counting disk device and disk partition(s) twice. This was fixed and now it
+  mirrors "iostat" cmdline tool behavior.
 
 5.4.6
 =====
diff --git a/psutil/_pslinux.py b/psutil/_pslinux.py
index 85ffe66a..1942d662 100644
--- a/psutil/_pslinux.py
+++ b/psutil/_pslinux.py
@@ -247,12 +247,12 @@ def file_flags_to_mode(flags):
     return mode
 
 
-def get_sector_size(partition):
-    """Return the sector size of a partition.
-    Used by disk_io_counters().
-    """
+def get_sector_size(name):
+    """Return the sector size of a device or partition name."""
+    # Some devices may have a slash in their name (e.g. cciss/c0d0...).
+    name = name.replace('/', '!')
     try:
-        with open("/sys/block/%s/queue/hw_sector_size" % partition, "rt") as f:
+        with open("/sys/block/%s/queue/hw_sector_size" % name, "rt") as f:
             return int(f.read())
     except (IOError, ValueError):
         # man iostat states that sectors are equivalent with blocks and
@@ -260,6 +260,23 @@ def get_sector_size(partition):
         return SECTOR_SIZE_FALLBACK
 
 
+def is_storage_device(name):
+    """Return True if the given name is a device or a partition.
+    This also includes virtual devices.
+    """
+    # Readapted from iostat source code, see:
+    # https://github.com/sysstat/sysstat/blob/
+    #     97912938cd476645b267280069e83b1c8dc0e1c7/common.c#L208
+    # Some devices may have a slash in their name (e.g. cciss/c0d0...).
+    name = name.replace('/', '!')
+    including_virtual = True
+    if including_virtual:
+        path = "/sys/block/%s" % name
+    else:
+        path = "/sys/block/%s/device" % name
+    return os.access(path, os.F_OK)
+
+
 @memoize
 def set_scputimes_ntuple(procfs_path):
     """Set a namedtuple of variable fields depending on the CPU times
@@ -1027,28 +1044,7 @@ def disk_io_counters():
     """Return disk I/O statistics for every disk installed on the
     system as a dict of raw tuples.
     """
-    # determine partitions we want to look for
-    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
-                    partitions.append(name)
-        return partitions
-
     retdict = {}
-    partitions = get_partitions()
     with open_text("%s/diskstats" % get_procfs_path()) as f:
         lines = f.readlines()
     for line in lines:
@@ -1086,12 +1082,13 @@ def disk_io_counters():
         else:
             raise ValueError("not sure how to interpret line %r" % line)
 
-        if name in partitions:
+        if is_storage_device(name):
             ssize = get_sector_size(name)
             rbytes *= ssize
             wbytes *= ssize
             retdict[name] = (reads, writes, rbytes, wbytes, rtime, wtime,
                              reads_merged, writes_merged, busy_time)
+
     return retdict
 
 
diff --git a/psutil/tests/test_linux.py b/psutil/tests/test_linux.py
index 9ea59b61..4d8ec66c 100755
--- a/psutil/tests/test_linux.py
+++ b/psutil/tests/test_linux.py
@@ -995,16 +995,18 @@ class TestSystemDisks(unittest.TestCase):
             with mock_open_content(
                     '/proc/diskstats',
                     "   3     0   1 hda 2 3 4 5 6 7 8 9 10 11 12"):
-                ret = psutil.disk_io_counters(nowrap=False)
-                self.assertEqual(ret.read_count, 1)
-                self.assertEqual(ret.read_merged_count, 2)
-                self.assertEqual(ret.read_bytes, 3 * SECTOR_SIZE)
-                self.assertEqual(ret.read_time, 4)
-                self.assertEqual(ret.write_count, 5)
-                self.assertEqual(ret.write_merged_count, 6)
-                self.assertEqual(ret.write_bytes, 7 * SECTOR_SIZE)
-                self.assertEqual(ret.write_time, 8)
-                self.assertEqual(ret.busy_time, 10)
+                with mock.patch('psutil._pslinux.is_storage_device',
+                                return_value=True):
+                    ret = psutil.disk_io_counters(nowrap=False)
+                    self.assertEqual(ret.read_count, 1)
+                    self.assertEqual(ret.read_merged_count, 2)
+                    self.assertEqual(ret.read_bytes, 3 * SECTOR_SIZE)
+                    self.assertEqual(ret.read_time, 4)
+                    self.assertEqual(ret.write_count, 5)
+                    self.assertEqual(ret.write_merged_count, 6)
+                    self.assertEqual(ret.write_bytes, 7 * SECTOR_SIZE)
+                    self.assertEqual(ret.write_time, 8)
+                    self.assertEqual(ret.busy_time, 10)
 
     def test_disk_io_counters_kernel_2_6_full_mocked(self):
         # Tests /proc/diskstats parsing format for 2.6 kernels,
@@ -1020,16 +1022,18 @@ class TestSystemDisks(unittest.TestCase):
             with mock_open_content(
                     '/proc/diskstats',
                     "   3    0   hda 1 2 3 4 5 6 7 8 9 10 11"):
-                ret = psutil.disk_io_counters(nowrap=False)
-                self.assertEqual(ret.read_count, 1)
-                self.assertEqual(ret.read_merged_count, 2)
-                self.assertEqual(ret.read_bytes, 3 * SECTOR_SIZE)
-                self.assertEqual(ret.read_time, 4)
-                self.assertEqual(ret.write_count, 5)
-                self.assertEqual(ret.write_merged_count, 6)
-                self.assertEqual(ret.write_bytes, 7 * SECTOR_SIZE)
-                self.assertEqual(ret.write_time, 8)
-                self.assertEqual(ret.busy_time, 10)
+                with mock.patch('psutil._pslinux.is_storage_device',
+                                return_value=True):
+                    ret = psutil.disk_io_counters(nowrap=False)
+                    self.assertEqual(ret.read_count, 1)
+                    self.assertEqual(ret.read_merged_count, 2)
+                    self.assertEqual(ret.read_bytes, 3 * SECTOR_SIZE)
+                    self.assertEqual(ret.read_time, 4)
+                    self.assertEqual(ret.write_count, 5)
+                    self.assertEqual(ret.write_merged_count, 6)
+                    self.assertEqual(ret.write_bytes, 7 * SECTOR_SIZE)
+                    self.assertEqual(ret.write_time, 8)
+                    self.assertEqual(ret.busy_time, 10)
 
     def test_disk_io_counters_kernel_2_6_limited_mocked(self):
         # Tests /proc/diskstats parsing format for 2.6 kernels,
@@ -1047,17 +1051,19 @@ class TestSystemDisks(unittest.TestCase):
             with mock_open_content(
                     '/proc/diskstats',
                     "   3    1   hda 1 2 3 4"):
-                ret = psutil.disk_io_counters(nowrap=False)
-                self.assertEqual(ret.read_count, 1)
-                self.assertEqual(ret.read_bytes, 2 * SECTOR_SIZE)
-                self.assertEqual(ret.write_count, 3)
-                self.assertEqual(ret.write_bytes, 4 * SECTOR_SIZE)
-
-                self.assertEqual(ret.read_merged_count, 0)
-                self.assertEqual(ret.read_time, 0)
-                self.assertEqual(ret.write_merged_count, 0)
-                self.assertEqual(ret.write_time, 0)
-                self.assertEqual(ret.busy_time, 0)
+                with mock.patch('psutil._pslinux.is_storage_device',
+                                return_value=True):
+                    ret = psutil.disk_io_counters(nowrap=False)
+                    self.assertEqual(ret.read_count, 1)
+                    self.assertEqual(ret.read_bytes, 2 * SECTOR_SIZE)
+                    self.assertEqual(ret.write_count, 3)
+                    self.assertEqual(ret.write_bytes, 4 * SECTOR_SIZE)
+
+                    self.assertEqual(ret.read_merged_count, 0)
+                    self.assertEqual(ret.read_time, 0)
+                    self.assertEqual(ret.write_merged_count, 0)
+                    self.assertEqual(ret.write_time, 0)
+                    self.assertEqual(ret.busy_time, 0)
 
 
 # =====================================================================