Add a _ioprio C extension module for exporting SYS_ioprio_* syscall numbers

The Linux kernel doesn't have a way to get the syscall numbers at runtime.

Suggested-by: Adrian Bunk

3 files changed, 66 insertions, 3 deletions

diff --git a/iotop/_ioprio.c b/iotop/_ioprio.c
new file mode 100644
index 0000000..ce8ae6c
--- /dev/null
+++ b/iotop/_ioprio.c
@@ -0,0 +1,55 @@
+#include <sys/syscall.h>
+#include <Python.h>
+
+#if PY_MAJOR_VERSION < 2 || (PY_MAJOR_VERSION == 2 && PY_MINOR_VERSION <= 3)
+#define NO_METHODS NoMethods
+static PyMethodDef NoMethods[] = {
+    {NULL, NULL, 0, NULL}
+};
+#else
+#define NO_METHODS NULL
+#endif
+
+#if PY_MAJOR_VERSION >= 3
+PyDoc_STRVAR(doc, "SYS_ioprio_get and SYS_ioprio_set for syscalls");
+static struct PyModuleDef mod = {
+    PyModuleDef_HEAD_INIT,
+    "_ioprio",
+    doc,
+    -1,
+    NO_METHODS,
+    NULL,
+    NULL,
+    NULL,
+    NULL
+};
+
+#define INIT_FUNC PyInit__ioprio
+#define INIT_MOD  PyModule_Create(&mod)
+#define INIT_ERR  NULL
+#define INIT_RET  return m
+#else
+#define INIT_FUNC init_ioprio
+#define INIT_MOD  Py_InitModule("_ioprio", NO_METHODS)
+#define INIT_ERR
+#define INIT_RET
+#endif  // PY_MAJOR_VERSION >= 3
+
+PyMODINIT_FUNC INIT_FUNC(void)
+{
+    PyObject *m = INIT_MOD;
+
+    if (m == NULL) {
+        return INIT_ERR;
+    }
+
+#ifdef SYS_ioprio_get
+    PyModule_AddIntConstant(m, "SYS_ioprio_get", SYS_ioprio_get);
+#endif
+
+#ifdef SYS_ioprio_set
+    PyModule_AddIntConstant(m, "SYS_ioprio_set", SYS_ioprio_set);
+#endif
+
+    INIT_RET;
+}
diff --git a/iotop/ioprio.py b/iotop/ioprio.py
index cc139c2..361c96a 100644
--- a/iotop/ioprio.py
+++ b/iotop/ioprio.py
@@ -76,8 +76,13 @@ class IoprioSetError(Exception):
         except TypeError:
             self.err = err
 
-__NR_ioprio_get = find_ioprio_syscall_number(IOPRIO_GET_ARCH_SYSCALL)
-__NR_ioprio_set = find_ioprio_syscall_number(IOPRIO_SET_ARCH_SYSCALL)
+try:
+    from iotop import _ioprio
+    __NR_ioprio_get = _ioprio.SYS_ioprio_get
+    __NR_ioprio_set = _ioprio.SYS_ioprio_set
+except (ImportError, AttributeError):
+    __NR_ioprio_get = find_ioprio_syscall_number(IOPRIO_GET_ARCH_SYSCALL)
+    __NR_ioprio_set = find_ioprio_syscall_number(IOPRIO_SET_ARCH_SYSCALL)
 
 try:
     ctypes_handle = ctypes.CDLL(None, use_errno=True)
diff --git a/setup.py b/setup.py
index e87bfaa..9b6022b 100755
--- a/setup.py
+++ b/setup.py
@@ -1,9 +1,11 @@
 #!/usr/bin/env python3
 
-from distutils.core import setup
+from distutils.core import setup, Extension
 from distutils.command import install as distutils_install
 from iotop.version import VERSION
 
+_ioprio = Extension('iotop._ioprio', sources = ['iotop/_ioprio.c'])
+
 # Dirty hack to make setup.py install the iotop script to sbin/ instead of bin/
 # while still honoring the choice of installing into local/ or not.
 if hasattr(distutils_install, 'INSTALL_SCHEMES'):
@@ -22,5 +24,6 @@ show of behalf of which process is the I/O going on.''',
       scripts=['sbin/iotop'],
       data_files=[('share/man/man8', ['iotop.8'])],
       packages=['iotop'],
+      ext_modules = [_ioprio],
       license='GPL'
 )