proc.py 8.92 KB
Newer Older
Nicolas Chauvat's avatar
Nicolas Chauvat committed
1
# copyright 2003-2011 LOGILAB S.A. (Paris, FRANCE), all rights reserved.
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# contact http://www.logilab.fr/ -- mailto:contact@logilab.fr
#
# This file is part of logilab-common.
#
# logilab-common is free software: you can redistribute it and/or modify it under
# the terms of the GNU Lesser General Public License as published by the Free
# Software Foundation, either version 2.1 of the License, or (at your option) any
# later version.
#
# logilab-common 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 Lesser General Public License for more
# details.
#
# You should have received a copy of the GNU Lesser General Public License along
17
# with logilab-common.  If not, see <http://www.gnu.org/licenses/>.
18
19
"""module providing:
* process information (linux specific: rely on /proc)
Sylvain Thenault's avatar
clenaup    
Sylvain Thenault committed
20
* a class for resource control (memory / time / cpu time)
21

22
23
This module doesn't work on windows platforms (only tested on linux)

24
:organization: Logilab
25
26
27



28
29
30
31
"""
__docformat__ = "restructuredtext en"

import os
32
import stat
33
34
35
36
from resource import getrlimit, setrlimit, RLIMIT_CPU, RLIMIT_AS
from signal import signal, SIGXCPU, SIGKILL, SIGUSR2, SIGUSR1
from threading import Timer, currentThread, Thread, Event
from time import time
37
38
39

from logilab.common.tree import Node

Laurent Peuch's avatar
Laurent Peuch committed
40
41
42
43

class NoSuchProcess(Exception):
    pass

44
45
46
47
48

def proc_exists(pid):
    """check the a pid is registered in /proc
    raise NoSuchProcess exception if not
    """
Laurent Peuch's avatar
Laurent Peuch committed
49
    if not os.path.exists("/proc/%s" % pid):
50
51
        raise NoSuchProcess()

Laurent Peuch's avatar
Laurent Peuch committed
52

53
54
55
56
57
58
59
PPID = 3
UTIME = 13
STIME = 14
CUTIME = 15
CSTIME = 16
VSIZE = 22

Laurent Peuch's avatar
Laurent Peuch committed
60

61
62
class ProcInfo(Node):
    """provide access to process information found in /proc"""
63

64
    def __init__(self, pid):
65
66
67
        self.pid = int(pid)
        Node.__init__(self, self.pid)
        proc_exists(self.pid)
Laurent Peuch's avatar
Laurent Peuch committed
68
        self.file = "/proc/%s/stat" % self.pid
69
        self.ppid = int(self.status()[PPID])
70

71
72
    def memory_usage(self):
        """return the memory usage of the process in Ko"""
Laurent Peuch's avatar
Laurent Peuch committed
73
        try:
74
75
76
            return int(self.status()[VSIZE])
        except IOError:
            return 0
77

78
    def lineage_memory_usage(self):
Laurent Peuch's avatar
Laurent Peuch committed
79
        return self.memory_usage() + sum([child.lineage_memory_usage() for child in self.children])
80
81
82
83
84
85
86
87
88
89
90
91
92
93

    def time(self, children=0):
        """return the number of jiffies that this process has been scheduled
        in user and kernel mode"""
        status = self.status()
        time = int(status[UTIME]) + int(status[STIME])
        if children:
            time += int(status[CUTIME]) + int(status[CSTIME])
        return time

    def status(self):
        """return the list of fields found in /proc/<pid>/stat"""
        return open(self.file).read().split()

94
95
96
    def name(self):
        """return the process name found in /proc/<pid>/stat
        """
Laurent Peuch's avatar
Laurent Peuch committed
97
        return self.status()[1].strip("()")
98
99
100
101
102
103

    def age(self):
        """return the age of the process
        """
        return os.stat(self.file)[stat.ST_MTIME]

Laurent Peuch's avatar
Laurent Peuch committed
104

105
106
class ProcInfoLoader:
    """manage process information"""
107

108
109
110
111
    def __init__(self):
        self._loaded = {}

    def list_pids(self):
112
        """return a list of existent process ids"""
Laurent Peuch's avatar
Laurent Peuch committed
113
        for subdir in os.listdir("/proc"):
114
115
116
117
118
119
120
121
122
123
124
125
126
            if subdir.isdigit():
                yield int(subdir)

    def load(self, pid):
        """get a ProcInfo object for a given pid"""
        pid = int(pid)
        try:
            return self._loaded[pid]
        except KeyError:
            procinfo = ProcInfo(pid)
            procinfo.manager = self
            self._loaded[pid] = procinfo
            return procinfo
127

128
129
130
131
132
133
134
135
136
137
138
139
    def load_all(self):
        """load all processes information"""
        for pid in self.list_pids():
            try:
                procinfo = self.load(pid)
                if procinfo.parent is None and procinfo.ppid:
                    pprocinfo = self.load(procinfo.ppid)
                    pprocinfo.append(procinfo)
            except NoSuchProcess:
                pass


140
141
class ResourceError(Exception):
    """Error raise when resource limit is reached"""
Laurent Peuch's avatar
Laurent Peuch committed
142

143
    limit = "Unknown Resource Limit"
144
145
146


class XCPUError(ResourceError):
147
    """Error raised when CPU Time limit is reached"""
Laurent Peuch's avatar
Laurent Peuch committed
148

149
150
    limit = "CPU Time"

Laurent Peuch's avatar
Laurent Peuch committed
151

152
153
154
class LineageMemoryError(ResourceError):
    """Error raised when the total amount of memory used by a process and
    it's child is reached"""
Laurent Peuch's avatar
Laurent Peuch committed
155

156
157
    limit = "Lineage total Memory"

Laurent Peuch's avatar
Laurent Peuch committed
158

159
160
class TimeoutError(ResourceError):
    """Error raised when the process is running for to much time"""
Laurent Peuch's avatar
Laurent Peuch committed
161

162
163
    limit = "Real Time"

Laurent Peuch's avatar
Laurent Peuch committed
164

165
166
167
168
169
170
171
172
# Can't use subclass because the StandardError MemoryError raised
RESOURCE_LIMIT_EXCEPTION = (ResourceError, MemoryError)


class MemorySentinel(Thread):
    """A class checking a process don't use too much memory in a separated
    daemonic thread
    """
Laurent Peuch's avatar
Laurent Peuch committed
173

Chris Lamb's avatar
Chris Lamb committed
174
    def __init__(self, interval, memory_limit, gpid=None):
175
176
177
178
179
        Thread.__init__(self, target=self._run, name="Test.Sentinel")
        self.memory_limit = memory_limit
        self._stop = Event()
        self.interval = interval
        self.setDaemon(True)
Chris Lamb's avatar
Chris Lamb committed
180
        self.gpid = gpid if gpid is not None else os.getpid()
181

182
183
184
    def stop(self):
        """stop ap"""
        self._stop.set()
185

186
187
188
189
    def _run(self):
        pil = ProcInfoLoader()
        while not self._stop.isSet():
            if self.memory_limit <= pil.load(self.gpid).lineage_memory_usage():
190
                os.killpg(self.gpid, SIGUSR1)
191
192
193
194
            self._stop.wait(self.interval)


class ResourceController:
Laurent Peuch's avatar
Laurent Peuch committed
195
    def __init__(self, max_cpu_time=None, max_time=None, max_memory=None, max_reprieve=60):
Alexandre Fayolle's avatar
Alexandre Fayolle committed
196
197
        if SIGXCPU == -1:
            raise RuntimeError("Unsupported platform")
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
        self.max_time = max_time
        self.max_memory = max_memory
        self.max_cpu_time = max_cpu_time
        self._reprieve = max_reprieve
        self._timer = None
        self._msentinel = None
        self._old_max_memory = None
        self._old_usr1_hdlr = None
        self._old_max_cpu_time = None
        self._old_usr2_hdlr = None
        self._old_sigxcpu_hdlr = None
        self._limit_set = 0
        self._abort_try = 0
        self._start_time = None
        self._elapse_time = 0

    def _hangle_sig_timeout(self, sig, frame):
        raise TimeoutError()
216

217
218
219
220
221
    def _hangle_sig_memory(self, sig, frame):
        if self._abort_try < self._reprieve:
            self._abort_try += 1
            raise LineageMemoryError("Memory limit reached")
        else:
222
            os.killpg(os.getpid(), SIGKILL)
223
224
225
226
227
228

    def _handle_sigxcpu(self, sig, frame):
        if self._abort_try < self._reprieve:
            self._abort_try += 1
            raise XCPUError("Soft CPU time limit reached")
        else:
229
            os.killpg(os.getpid(), SIGKILL)
230
231
232
233

    def _time_out(self):
        if self._abort_try < self._reprieve:
            self._abort_try += 1
234
            os.killpg(os.getpid(), SIGUSR2)
235
236
237
238
            if self._limit_set > 0:
                self._timer = Timer(1, self._time_out)
                self._timer.start()
        else:
239
            os.killpg(os.getpid(), SIGKILL)
240

241
242
    def setup_limit(self):
        """set up the process limit"""
Laurent Peuch's avatar
Laurent Peuch committed
243
        assert currentThread().getName() == "MainThread"
244
        os.setpgrp()
245
246
247
        if self._limit_set <= 0:
            if self.max_time is not None:
                self._old_usr2_hdlr = signal(SIGUSR2, self._hangle_sig_timeout)
Laurent Peuch's avatar
Laurent Peuch committed
248
                self._timer = Timer(max(1, int(self.max_time) - self._elapse_time), self._time_out)
249
250
251
252
253
254
255
256
                self._start_time = int(time())
                self._timer.start()
            if self.max_cpu_time is not None:
                self._old_max_cpu_time = getrlimit(RLIMIT_CPU)
                cpu_limit = (int(self.max_cpu_time), self._old_max_cpu_time[1])
                self._old_sigxcpu_hdlr = signal(SIGXCPU, self._handle_sigxcpu)
                setrlimit(RLIMIT_CPU, cpu_limit)
            if self.max_memory is not None:
Laurent Peuch's avatar
Laurent Peuch committed
257
                self._msentinel = MemorySentinel(1, int(self.max_memory))
258
259
260
261
262
263
264
265
266
267
268
269
                self._old_max_memory = getrlimit(RLIMIT_AS)
                self._old_usr1_hdlr = signal(SIGUSR1, self._hangle_sig_memory)
                as_limit = (int(self.max_memory), self._old_max_memory[1])
                setrlimit(RLIMIT_AS, as_limit)
                self._msentinel.start()
        self._limit_set += 1

    def clean_limit(self):
        """reinstall the old process limit"""
        if self._limit_set > 0:
            if self.max_time is not None:
                self._timer.cancel()
Laurent Peuch's avatar
Laurent Peuch committed
270
                self._elapse_time += int(time()) - self._start_time
271
272
273
274
275
276
277
278
279
280
281
                self._timer = None
                signal(SIGUSR2, self._old_usr2_hdlr)
            if self.max_cpu_time is not None:
                setrlimit(RLIMIT_CPU, self._old_max_cpu_time)
                signal(SIGXCPU, self._old_sigxcpu_hdlr)
            if self.max_memory is not None:
                self._msentinel.stop()
                self._msentinel = None
                setrlimit(RLIMIT_AS, self._old_max_memory)
                signal(SIGUSR1, self._old_usr1_hdlr)
        self._limit_set -= 1