move disk io stuff

4 files changed, 103 insertions, 30 deletions

diff --git a/psutil/_psutil_bsd.h b/psutil/_psutil_bsd.h
index d6733d78..41d548be 100644
--- a/psutil/_psutil_bsd.h
+++ b/psutil/_psutil_bsd.h
@@ -48,7 +48,6 @@ static PyObject* psutil_boot_time(PyObject* self, PyObject* args);
 static PyObject* psutil_cpu_count_logical(PyObject* self, PyObject* args);
 static PyObject* psutil_cpu_count_phys(PyObject* self, PyObject* args);
 static PyObject* psutil_cpu_times(PyObject* self, PyObject* args);
-static PyObject* psutil_disk_io_counters(PyObject* self, PyObject* args);
 static PyObject* psutil_disk_partitions(PyObject* self, PyObject* args);
 static PyObject* psutil_net_io_counters(PyObject* self, PyObject* args);
 static PyObject* psutil_pids(PyObject* self, PyObject* args);
@@ -56,6 +55,7 @@ static PyObject* psutil_users(PyObject* self, PyObject* args);
 #ifdef __FreeBSD__
 static PyObject* psutil_virtual_mem(PyObject* self, PyObject* args);
 static PyObject* psutil_swap_mem(PyObject* self, PyObject* args);
+static PyObject* psutil_disk_io_counters(PyObject* self, PyObject* args);
 #endif
 
 #if defined(__FreeBSD_version) && __FreeBSD_version >= 800000
diff --git a/psutil/_psutil_openbsd.c b/psutil/_psutil_openbsd.c
index d453d24f..b76cb995 100644
--- a/psutil/_psutil_openbsd.c
+++ b/psutil/_psutil_openbsd.c
@@ -85,7 +85,6 @@
     #define _KERNEL  // for DTYPE_VNODE
     #include <sys/file.h>
     #undef _KERNEL
-    #include <sys/disk.h>  // struct diskstats
     #include <sys/sched.h>  // for CPUSTATES & CP_*
     #include <sys/swap.h>
     #include <kvm.h>
@@ -1644,68 +1643,76 @@ error:
 }
 
 
+#ifdef __FreeBSD__
 /*
  * Return a Python dict of tuples for disk I/O information
  */
 static PyObject *
 psutil_disk_io_counters(PyObject *self, PyObject *args) {
-    int i, dk_ndrive, mib[3];
-    size_t len;
-    struct diskstats *stats;
+    int i;
+    struct statinfo stats;
 
     PyObject *py_retdict = PyDict_New();
     PyObject *py_disk_info = NULL;
+
     if (py_retdict == NULL)
         return NULL;
-
-    mib[0] = CTL_HW;
-    mib[1] = HW_DISKSTATS;
-    len = 0;
-    if (sysctl(mib, 2, NULL, &len, NULL, 0) < 0) {
-        warn("can't get hw.diskstats size");
-        PyErr_SetFromErrno(PyExc_OSError);
+    if (devstat_checkversion(NULL) < 0) {
+        PyErr_Format(PyExc_RuntimeError, "devstat_checkversion() failed");
         goto error;
     }
-    dk_ndrive = (int)(len / sizeof(struct diskstats));
 
-    stats = malloc(len);
-    if (stats == NULL) {
-        warn("can't malloc");
+    stats.dinfo = (struct devinfo *)malloc(sizeof(struct devinfo));
+    if (stats.dinfo == NULL) {
         PyErr_NoMemory();
         goto error;
     }
-    if (sysctl(mib, 2, stats, &len, NULL, 0) < 0 ) {
-        warn("could not read hw.diskstats");
-        PyErr_SetFromErrno(PyExc_OSError);
+    bzero(stats.dinfo, sizeof(struct devinfo));
+
+    if (devstat_getdevs(NULL, &stats) == -1) {
+        PyErr_Format(PyExc_RuntimeError, "devstat_getdevs() failed");
         goto error;
     }
 
-    for (i = 0; i < dk_ndrive; i++) {
+    for (i = 0; i < stats.dinfo->numdevs; i++) {
+        py_disk_info = NULL;
+        struct devstat current;
+        char disk_name[128];
+        current = stats.dinfo->devices[i];
+        snprintf(disk_name, sizeof(disk_name), "%s%d",
+                 current.device_name,
+                 current.unit_number);
+
         py_disk_info = Py_BuildValue(
             "(KKKKLL)",
-            stats[i].ds_rxfer,
-            stats[i].ds_wxfer,
-            stats[i].ds_rbytes,
-            stats[i].ds_wbytes,
-            (long long) TV2DOUBLE(stats[i].ds_time) / 2, /* assume half read - half writes.. */
-            (long long) TV2DOUBLE(stats[i].ds_time) / 2);
+            current.operations[DEVSTAT_READ],   // no reads
+            current.operations[DEVSTAT_WRITE],  // no writes
+            current.bytes[DEVSTAT_READ],        // bytes read
+            current.bytes[DEVSTAT_WRITE],       // bytes written
+            (long long) BT2MSEC(current.duration[DEVSTAT_READ]),  // r time
+            (long long) BT2MSEC(current.duration[DEVSTAT_WRITE])  // w time
+        );      // finished transactions
         if (!py_disk_info)
             goto error;
-        if (PyDict_SetItemString(py_retdict, stats[i].ds_name, py_disk_info))
+        if (PyDict_SetItemString(py_retdict, disk_name, py_disk_info))
             goto error;
         Py_DECREF(py_disk_info);
     }
 
-    free(stats);
+    if (stats.dinfo->mem_ptr)
+        free(stats.dinfo->mem_ptr);
+    free(stats.dinfo);
     return py_retdict;
 
 error:
     Py_XDECREF(py_disk_info);
     Py_DECREF(py_retdict);
-    if (stats != NULL)
-        free(stats);
+    if (stats.dinfo != NULL)
+        free(stats.dinfo);
     return NULL;
 }
+#endif
+
 
 /*
  * Return currently connected users as a list of tuples.
diff --git a/psutil/arch/bsd/openbsd.c b/psutil/arch/bsd/openbsd.c
index d90796d7..5e949241 100644
--- a/psutil/arch/bsd/openbsd.c
+++ b/psutil/arch/bsd/openbsd.c
@@ -29,12 +29,16 @@
 #define _KERNEL  // for DTYPE_*
 #include <sys/file.h>
 #undef _KERNEL
+#include <sys/disk.h>  // struct diskstats
+
 
 #include "openbsd.h"
 
 // a signaler for connections without an actual status
 int PSUTIL_CONN_NONE = 128;
+
 #define KPT2DOUBLE(t)   (t ## _sec + t ## _usec / 1000000.0)
+#define TV2DOUBLE(t)    ((t).tv_sec + (t).tv_usec / 1000000.0)
 
 
 // ============================================================================
@@ -694,3 +698,64 @@ error:
     Py_DECREF(py_retlist);
     return NULL;
 }
+
+
+PyObject *
+psutil_disk_io_counters(PyObject *self, PyObject *args) {
+    int i, dk_ndrive, mib[3];
+    size_t len;
+    struct diskstats *stats;
+
+    PyObject *py_retdict = PyDict_New();
+    PyObject *py_disk_info = NULL;
+    if (py_retdict == NULL)
+        return NULL;
+
+    mib[0] = CTL_HW;
+    mib[1] = HW_DISKSTATS;
+    len = 0;
+    if (sysctl(mib, 2, NULL, &len, NULL, 0) < 0) {
+        warn("can't get hw.diskstats size");
+        PyErr_SetFromErrno(PyExc_OSError);
+        goto error;
+    }
+    dk_ndrive = (int)(len / sizeof(struct diskstats));
+
+    stats = malloc(len);
+    if (stats == NULL) {
+        warn("can't malloc");
+        PyErr_NoMemory();
+        goto error;
+    }
+    if (sysctl(mib, 2, stats, &len, NULL, 0) < 0 ) {
+        warn("could not read hw.diskstats");
+        PyErr_SetFromErrno(PyExc_OSError);
+        goto error;
+    }
+
+    for (i = 0; i < dk_ndrive; i++) {
+        py_disk_info = Py_BuildValue(
+            "(KKKKLL)",
+            stats[i].ds_rxfer,
+            stats[i].ds_wxfer,
+            stats[i].ds_rbytes,
+            stats[i].ds_wbytes,
+            (long long) TV2DOUBLE(stats[i].ds_time) / 2, /* assume half read - half writes.. */
+            (long long) TV2DOUBLE(stats[i].ds_time) / 2);
+        if (!py_disk_info)
+            goto error;
+        if (PyDict_SetItemString(py_retdict, stats[i].ds_name, py_disk_info))
+            goto error;
+        Py_DECREF(py_disk_info);
+    }
+
+    free(stats);
+    return py_retdict;
+
+error:
+    Py_XDECREF(py_disk_info);
+    Py_DECREF(py_retdict);
+    if (stats != NULL)
+        free(stats);
+    return NULL;
+}
diff --git a/psutil/arch/bsd/openbsd.h b/psutil/arch/bsd/openbsd.h
index 31cd8312..e859278e 100644
--- a/psutil/arch/bsd/openbsd.h
+++ b/psutil/arch/bsd/openbsd.h
@@ -25,3 +25,4 @@ PyObject *psutil_proc_num_fds(PyObject *self, PyObject *args);
 PyObject *psutil_proc_cwd(PyObject *self, PyObject *args);
 PyObject *psutil_proc_connections(PyObject *self, PyObject *args);
 PyObject *psutil_per_cpu_times(PyObject *self, PyObject *args);
+PyObject* psutil_disk_io_counters(PyObject* self, PyObject* args);