summaryrefslogtreecommitdiff
path: root/iotop/ioprio.py
blob: c76e665c10878f1612f378655ae06a3088a2d50c (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Library General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301 USA
#
# See the COPYING file for license information.
#
# Copyright (c) 2007 Guillaume Chazarain <guichaz@gmail.com>

import ctypes
import fnmatch
import os
import platform

# From https://git.kernel.org/pub/scm/utils/util-linux/util-linux.git/tree/configure.ac#n2289
# i386 bit userspace under an x86_64 kernel will have its uname() appear as
# 'x86_64' but it will use the i386 syscall number, that's why we consider both
# the architecture name and the word size.
IOPRIO_GET_ARCH_SYSCALL = [
    ('alpha',       '*',  443),
    ('arm*',        '*',  315),
    ('i*86',        '*',  290),
    ('ia64*',       '*', 1275),
    ('mips*',   '32bit', 4315),
    ('mips*',   '64bit', 5274),
    ('parisc*',     '*',  268),
    ('powerpc*',    '*',  274),
    ('s390*',       '*',  283),
    ('sparc*',      '*',  218),
    ('sh*',         '*',  289),
    ('x86_64*', '32bit',  290),
    ('x86_64*', '64bit',  252),
]

IOPRIO_SET_ARCH_SYSCALL = [
    ('alpha',       '*',  442),
    ('arm*',        '*',  314),
    ('i*86',        '*',  289),
    ('ia64*',       '*', 1274),
    ('mips*',   '32bit', 4314),
    ('mips*',   '64bit', 5273),
    ('parisc*',     '*',  267),
    ('powerpc*',    '*',  273),
    ('s390*',       '*',  282),
    ('sparc*',      '*',  196),
    ('sh*',         '*',  288),
    ('x86_64*',  '32bit', 289),
    ('x86_64*',  '64bit', 251),
]


def find_ioprio_syscall_number(syscall_list):
    arch = os.uname()[4]
    bits = platform.architecture()[0]

    for candidate_arch, candidate_bits, syscall_nr in syscall_list:
        if fnmatch.fnmatch(arch, candidate_arch) and \
           fnmatch.fnmatch(bits, candidate_bits):
            return syscall_nr


class IoprioSetError(Exception):
    def __init__(self, err):
        try:
            self.err = os.strerror(err)
        except TypeError:
            self.err = err

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)
except TypeError:
    ctypes_handle = ctypes.CDLL(None)

syscall = ctypes_handle.syscall

PRIORITY_CLASSES = [None, 'rt', 'be', 'idle']

IOPRIO_WHO_PROCESS = 1
IOPRIO_CLASS_SHIFT = 13
IOPRIO_PRIO_MASK = (1 << IOPRIO_CLASS_SHIFT) - 1


def ioprio_value(ioprio_class, ioprio_data):
    try:
        ioprio_class = PRIORITY_CLASSES.index(ioprio_class)
    except ValueError:
        ioprio_class = PRIORITY_CLASSES.index(None)
    return (ioprio_class << IOPRIO_CLASS_SHIFT) | ioprio_data


def ioprio_class(ioprio):
    return PRIORITY_CLASSES[ioprio >> IOPRIO_CLASS_SHIFT]


def ioprio_data(ioprio):
    return ioprio & IOPRIO_PRIO_MASK

sched_getscheduler = ctypes_handle.sched_getscheduler
SCHED_OTHER, SCHED_FIFO, SCHED_RR, SCHED_BATCH, SCHED_ISO, SCHED_IDLE = \
    range(6)

getpriority = ctypes_handle.getpriority
PRIO_PROCESS = 0


def get_ioprio_from_sched(pid):
    scheduler = sched_getscheduler(pid)
    nice = getpriority(PRIO_PROCESS, pid)
    ioprio_nice = (nice + 20) / 5

    if scheduler in (SCHED_FIFO, SCHED_RR):
        return 'rt/%d' % ioprio_nice
    elif scheduler == SCHED_IDLE:
        return 'idle'
    else:
        return 'be/%d' % ioprio_nice


def get(pid):
    if __NR_ioprio_get is None:
        return '?sys'

    ioprio = syscall(__NR_ioprio_get, IOPRIO_WHO_PROCESS, pid)
    if ioprio < 0:
        return '?err'

    prio_class = ioprio_class(ioprio)
    if not prio_class:
        return get_ioprio_from_sched(pid)
    if prio_class == 'idle':
        return prio_class
    return '%s/%d' % (prio_class, ioprio_data(ioprio))


def set_ioprio(which, who, ioprio_class, ioprio_data):
    if __NR_ioprio_set is None:
        raise IoprioSetError('No ioprio_set syscall found')

    ioprio_val = ioprio_value(ioprio_class, ioprio_data)
    ret = syscall(__NR_ioprio_set, which, who, ioprio_val, use_errno=True)
    if ret < 0:
        try:
            err = ctypes.get_errno()
        except AttributeError:
            err = \
                'Unknown error (errno support not available before Python2.6)'
        raise IoprioSetError(err)


def sort_key(key):
    if key[0] == '?':
        return -ord(key[1])

    if '/' in key:
        if key.startswith('rt/'):
            shift = 0
        elif key.startswith('be/'):
            shift = 1
        prio = int(key.split('/')[1])
    elif key == 'idle':
        shift = 2
        prio = 0

    return (1 << (shift * IOPRIO_CLASS_SHIFT)) + prio


def to_class_and_data(ioprio_str):
    if '/' in ioprio_str:
        split = ioprio_str.split('/')
        return (split[0], int(split[1]))
    elif ioprio_str == 'idle':
        return ('idle', 0)
    return (None, None)

if __name__ == '__main__':
    import sys
    if len(sys.argv) == 2:
        pid = int(sys.argv[1])
    else:
        pid = os.getpid()
    print('pid:', pid)
    print('ioprio:', get(pid))