1 
2 #include "Python.h"
3 #include "structseq.h"
4 #include <sys/resource.h>
5 #include <sys/time.h>
6 #include <string.h>
7 #include <errno.h>
8 /* for sysconf */
9 #if defined(HAVE_UNISTD_H)
10 #include <unistd.h>
11 #endif
12 
13 /* On some systems, these aren't in any header file.
14    On others they are, with inconsistent prototypes.
15    We declare the (default) return type, to shut up gcc -Wall;
16    but we can't declare the prototype, to avoid errors
17    when the header files declare it different.
18    Worse, on some Linuxes, getpagesize() returns a size_t... */
19 
20 #define doubletime(TV) ((double)(TV).tv_sec + (TV).tv_usec * 0.000001)
21 
22 static PyObject *ResourceError;
23 
24 PyDoc_STRVAR(struct_rusage__doc__,
25 "struct_rusage: Result from getrusage.\n\n"
26 "This object may be accessed either as a tuple of\n"
27 "    (utime,stime,maxrss,ixrss,idrss,isrss,minflt,majflt,\n"
28 "    nswap,inblock,oublock,msgsnd,msgrcv,nsignals,nvcsw,nivcsw)\n"
29 "or via the attributes ru_utime, ru_stime, ru_maxrss, and so on.");
30 
31 static PyStructSequence_Field struct_rusage_fields[] = {
32     {"ru_utime",        "user time used"},
33     {"ru_stime",        "system time used"},
34     {"ru_maxrss",       "max. resident set size"},
35     {"ru_ixrss",        "shared memory size"},
36     {"ru_idrss",        "unshared data size"},
37     {"ru_isrss",        "unshared stack size"},
38     {"ru_minflt",       "page faults not requiring I/O"},
39     {"ru_majflt",       "page faults requiring I/O"},
40     {"ru_nswap",        "number of swap outs"},
41     {"ru_inblock",      "block input operations"},
42     {"ru_oublock",      "block output operations"},
43     {"ru_msgsnd",       "IPC messages sent"},
44     {"ru_msgrcv",       "IPC messages received"},
45     {"ru_nsignals",     "signals received"},
46     {"ru_nvcsw",        "voluntary context switches"},
47     {"ru_nivcsw",       "involuntary context switches"},
48     {0}
49 };
50 
51 static PyStructSequence_Desc struct_rusage_desc = {
52     "resource.struct_rusage",           /* name */
53     struct_rusage__doc__,       /* doc */
54     struct_rusage_fields,       /* fields */
55     16          /* n_in_sequence */
56 };
57 
58 static int initialized;
59 static PyTypeObject StructRUsageType;
60 
61 static PyObject *
resource_getrusage(PyObject * self,PyObject * args)62 resource_getrusage(PyObject *self, PyObject *args)
63 {
64     int who;
65     struct rusage ru;
66     PyObject *result;
67 
68     if (!PyArg_ParseTuple(args, "i:getrusage", &who))
69         return NULL;
70 
71     if (getrusage(who, &ru) == -1) {
72         if (errno == EINVAL) {
73             PyErr_SetString(PyExc_ValueError,
74                             "invalid who parameter");
75             return NULL;
76         }
77         PyErr_SetFromErrno(ResourceError);
78         return NULL;
79     }
80 
81     result = PyStructSequence_New(&StructRUsageType);
82     if (!result)
83         return NULL;
84 
85     PyStructSequence_SET_ITEM(result, 0,
86                     PyFloat_FromDouble(doubletime(ru.ru_utime)));
87     PyStructSequence_SET_ITEM(result, 1,
88                     PyFloat_FromDouble(doubletime(ru.ru_stime)));
89     PyStructSequence_SET_ITEM(result, 2, PyInt_FromLong(ru.ru_maxrss));
90     PyStructSequence_SET_ITEM(result, 3, PyInt_FromLong(ru.ru_ixrss));
91     PyStructSequence_SET_ITEM(result, 4, PyInt_FromLong(ru.ru_idrss));
92     PyStructSequence_SET_ITEM(result, 5, PyInt_FromLong(ru.ru_isrss));
93     PyStructSequence_SET_ITEM(result, 6, PyInt_FromLong(ru.ru_minflt));
94     PyStructSequence_SET_ITEM(result, 7, PyInt_FromLong(ru.ru_majflt));
95     PyStructSequence_SET_ITEM(result, 8, PyInt_FromLong(ru.ru_nswap));
96     PyStructSequence_SET_ITEM(result, 9, PyInt_FromLong(ru.ru_inblock));
97     PyStructSequence_SET_ITEM(result, 10, PyInt_FromLong(ru.ru_oublock));
98     PyStructSequence_SET_ITEM(result, 11, PyInt_FromLong(ru.ru_msgsnd));
99     PyStructSequence_SET_ITEM(result, 12, PyInt_FromLong(ru.ru_msgrcv));
100     PyStructSequence_SET_ITEM(result, 13, PyInt_FromLong(ru.ru_nsignals));
101     PyStructSequence_SET_ITEM(result, 14, PyInt_FromLong(ru.ru_nvcsw));
102     PyStructSequence_SET_ITEM(result, 15, PyInt_FromLong(ru.ru_nivcsw));
103 
104     if (PyErr_Occurred()) {
105         Py_DECREF(result);
106         return NULL;
107     }
108 
109     return result;
110 }
111 
112 
113 static PyObject *
resource_getrlimit(PyObject * self,PyObject * args)114 resource_getrlimit(PyObject *self, PyObject *args)
115 {
116     struct rlimit rl;
117     int resource;
118 
119     if (!PyArg_ParseTuple(args, "i:getrlimit", &resource))
120         return NULL;
121 
122     if (resource < 0 || resource >= RLIM_NLIMITS) {
123         PyErr_SetString(PyExc_ValueError,
124                         "invalid resource specified");
125         return NULL;
126     }
127 
128     if (getrlimit(resource, &rl) == -1) {
129         PyErr_SetFromErrno(ResourceError);
130         return NULL;
131     }
132 
133 #if defined(HAVE_LONG_LONG)
134     if (sizeof(rl.rlim_cur) > sizeof(long)) {
135         return Py_BuildValue("LL",
136                              (PY_LONG_LONG) rl.rlim_cur,
137                              (PY_LONG_LONG) rl.rlim_max);
138     }
139 #endif
140     return Py_BuildValue("ll", (long) rl.rlim_cur, (long) rl.rlim_max);
141 }
142 
143 static PyObject *
resource_setrlimit(PyObject * self,PyObject * args)144 resource_setrlimit(PyObject *self, PyObject *args)
145 {
146     struct rlimit rl;
147     int resource;
148     PyObject *curobj, *maxobj;
149 
150     if (!PyArg_ParseTuple(args, "i(OO):setrlimit",
151                           &resource, &curobj, &maxobj))
152         return NULL;
153 
154     if (resource < 0 || resource >= RLIM_NLIMITS) {
155         PyErr_SetString(PyExc_ValueError,
156                         "invalid resource specified");
157         return NULL;
158     }
159 
160 #if !defined(HAVE_LARGEFILE_SUPPORT)
161     rl.rlim_cur = PyInt_AsLong(curobj);
162     if (rl.rlim_cur == (rlim_t)-1 && PyErr_Occurred())
163         return NULL;
164     rl.rlim_max = PyInt_AsLong(maxobj);
165     if (rl.rlim_max == (rlim_t)-1 && PyErr_Occurred())
166         return NULL;
167 #else
168     /* The limits are probably bigger than a long */
169     rl.rlim_cur = PyLong_Check(curobj) ?
170         PyLong_AsLongLong(curobj) : PyInt_AsLong(curobj);
171     if (rl.rlim_cur == (rlim_t)-1 && PyErr_Occurred())
172         return NULL;
173     rl.rlim_max = PyLong_Check(maxobj) ?
174         PyLong_AsLongLong(maxobj) : PyInt_AsLong(maxobj);
175     if (rl.rlim_max == (rlim_t)-1 && PyErr_Occurred())
176         return NULL;
177 #endif
178 
179     rl.rlim_cur = rl.rlim_cur & RLIM_INFINITY;
180     rl.rlim_max = rl.rlim_max & RLIM_INFINITY;
181     if (setrlimit(resource, &rl) == -1) {
182         if (errno == EINVAL)
183             PyErr_SetString(PyExc_ValueError,
184                             "current limit exceeds maximum limit");
185         else if (errno == EPERM)
186             PyErr_SetString(PyExc_ValueError,
187                             "not allowed to raise maximum limit");
188         else
189             PyErr_SetFromErrno(ResourceError);
190         return NULL;
191     }
192     Py_INCREF(Py_None);
193     return Py_None;
194 }
195 
196 static PyObject *
resource_getpagesize(PyObject * self,PyObject * unused)197 resource_getpagesize(PyObject *self, PyObject *unused)
198 {
199     long pagesize = 0;
200 #if defined(HAVE_GETPAGESIZE)
201     pagesize = getpagesize();
202 #elif defined(HAVE_SYSCONF)
203 #if defined(_SC_PAGE_SIZE)
204     pagesize = sysconf(_SC_PAGE_SIZE);
205 #else
206     /* Irix 5.3 has _SC_PAGESIZE, but not _SC_PAGE_SIZE */
207     pagesize = sysconf(_SC_PAGESIZE);
208 #endif
209 #endif
210     return Py_BuildValue("i", pagesize);
211 
212 }
213 
214 /* List of functions */
215 
216 static struct PyMethodDef
217 resource_methods[] = {
218     {"getrusage",    resource_getrusage,   METH_VARARGS},
219     {"getrlimit",    resource_getrlimit,   METH_VARARGS},
220     {"setrlimit",    resource_setrlimit,   METH_VARARGS},
221     {"getpagesize",  resource_getpagesize, METH_NOARGS},
222     {NULL, NULL}                             /* sentinel */
223 };
224 
225 
226 /* Module initialization */
227 
228 PyMODINIT_FUNC
initresource(void)229 initresource(void)
230 {
231     PyObject *m, *v;
232 
233     /* Create the module and add the functions */
234     m = Py_InitModule("resource", resource_methods);
235     if (m == NULL)
236         return;
237 
238     /* Add some symbolic constants to the module */
239     if (ResourceError == NULL) {
240         ResourceError = PyErr_NewException("resource.error",
241                                            NULL, NULL);
242     }
243     Py_INCREF(ResourceError);
244     PyModule_AddObject(m, "error", ResourceError);
245     if (!initialized)
246         PyStructSequence_InitType(&StructRUsageType,
247                                   &struct_rusage_desc);
248     Py_INCREF(&StructRUsageType);
249     PyModule_AddObject(m, "struct_rusage",
250                        (PyObject*) &StructRUsageType);
251 
252     /* insert constants */
253 #ifdef RLIMIT_CPU
254     PyModule_AddIntConstant(m, "RLIMIT_CPU", RLIMIT_CPU);
255 #endif
256 
257 #ifdef RLIMIT_FSIZE
258     PyModule_AddIntConstant(m, "RLIMIT_FSIZE", RLIMIT_FSIZE);
259 #endif
260 
261 #ifdef RLIMIT_DATA
262     PyModule_AddIntConstant(m, "RLIMIT_DATA", RLIMIT_DATA);
263 #endif
264 
265 #ifdef RLIMIT_STACK
266     PyModule_AddIntConstant(m, "RLIMIT_STACK", RLIMIT_STACK);
267 #endif
268 
269 #ifdef RLIMIT_CORE
270     PyModule_AddIntConstant(m, "RLIMIT_CORE", RLIMIT_CORE);
271 #endif
272 
273 #ifdef RLIMIT_NOFILE
274     PyModule_AddIntConstant(m, "RLIMIT_NOFILE", RLIMIT_NOFILE);
275 #endif
276 
277 #ifdef RLIMIT_OFILE
278     PyModule_AddIntConstant(m, "RLIMIT_OFILE", RLIMIT_OFILE);
279 #endif
280 
281 #ifdef RLIMIT_VMEM
282     PyModule_AddIntConstant(m, "RLIMIT_VMEM", RLIMIT_VMEM);
283 #endif
284 
285 #ifdef RLIMIT_AS
286     PyModule_AddIntConstant(m, "RLIMIT_AS", RLIMIT_AS);
287 #endif
288 
289 #ifdef RLIMIT_RSS
290     PyModule_AddIntConstant(m, "RLIMIT_RSS", RLIMIT_RSS);
291 #endif
292 
293 #ifdef RLIMIT_NPROC
294     PyModule_AddIntConstant(m, "RLIMIT_NPROC", RLIMIT_NPROC);
295 #endif
296 
297 #ifdef RLIMIT_MEMLOCK
298     PyModule_AddIntConstant(m, "RLIMIT_MEMLOCK", RLIMIT_MEMLOCK);
299 #endif
300 
301 #ifdef RLIMIT_SBSIZE
302     PyModule_AddIntConstant(m, "RLIMIT_SBSIZE", RLIMIT_SBSIZE);
303 #endif
304 
305 #ifdef RUSAGE_SELF
306     PyModule_AddIntConstant(m, "RUSAGE_SELF", RUSAGE_SELF);
307 #endif
308 
309 #ifdef RUSAGE_CHILDREN
310     PyModule_AddIntConstant(m, "RUSAGE_CHILDREN", RUSAGE_CHILDREN);
311 #endif
312 
313 #ifdef RUSAGE_BOTH
314     PyModule_AddIntConstant(m, "RUSAGE_BOTH", RUSAGE_BOTH);
315 #endif
316 
317 #if defined(HAVE_LONG_LONG)
318     if (sizeof(RLIM_INFINITY) > sizeof(long)) {
319         v = PyLong_FromLongLong((PY_LONG_LONG) RLIM_INFINITY);
320     } else
321 #endif
322     {
323         v = PyInt_FromLong((long) RLIM_INFINITY);
324     }
325     if (v) {
326         PyModule_AddObject(m, "RLIM_INFINITY", v);
327     }
328     initialized = 1;
329 }
330