1 
2 #include "Python.h"
3 #include <sys/resource.h>
4 #include <sys/time.h>
5 #include <string.h>
6 #include <errno.h>
7 #include <unistd.h>
8 
9 /* On some systems, these aren't in any header file.
10    On others they are, with inconsistent prototypes.
11    We declare the (default) return type, to shut up gcc -Wall;
12    but we can't declare the prototype, to avoid errors
13    when the header files declare it different.
14    Worse, on some Linuxes, getpagesize() returns a size_t... */
15 
16 #define doubletime(TV) ((double)(TV).tv_sec + (TV).tv_usec * 0.000001)
17 
18 /*[clinic input]
19 module resource
20 [clinic start generated code]*/
21 /*[clinic end generated code: output=da39a3ee5e6b4b0d input=e89d38ed52609d7c]*/
22 
23 /*[python input]
24 class pid_t_converter(CConverter):
25     type = 'pid_t'
26     format_unit = '" _Py_PARSE_PID "'
27 [python start generated code]*/
28 /*[python end generated code: output=da39a3ee5e6b4b0d input=0c1d19f640d57e48]*/
29 
30 #include "clinic/resource.c.h"
31 
32 PyDoc_STRVAR(struct_rusage__doc__,
33 "struct_rusage: Result from getrusage.\n\n"
34 "This object may be accessed either as a tuple of\n"
35 "    (utime,stime,maxrss,ixrss,idrss,isrss,minflt,majflt,\n"
36 "    nswap,inblock,oublock,msgsnd,msgrcv,nsignals,nvcsw,nivcsw)\n"
37 "or via the attributes ru_utime, ru_stime, ru_maxrss, and so on.");
38 
39 static PyStructSequence_Field struct_rusage_fields[] = {
40     {"ru_utime",        "user time used"},
41     {"ru_stime",        "system time used"},
42     {"ru_maxrss",       "max. resident set size"},
43     {"ru_ixrss",        "shared memory size"},
44     {"ru_idrss",        "unshared data size"},
45     {"ru_isrss",        "unshared stack size"},
46     {"ru_minflt",       "page faults not requiring I/O"},
47     {"ru_majflt",       "page faults requiring I/O"},
48     {"ru_nswap",        "number of swap outs"},
49     {"ru_inblock",      "block input operations"},
50     {"ru_oublock",      "block output operations"},
51     {"ru_msgsnd",       "IPC messages sent"},
52     {"ru_msgrcv",       "IPC messages received"},
53     {"ru_nsignals",     "signals received"},
54     {"ru_nvcsw",        "voluntary context switches"},
55     {"ru_nivcsw",       "involuntary context switches"},
56     {0}
57 };
58 
59 static PyStructSequence_Desc struct_rusage_desc = {
60     "resource.struct_rusage",           /* name */
61     struct_rusage__doc__,       /* doc */
62     struct_rusage_fields,       /* fields */
63     16          /* n_in_sequence */
64 };
65 
66 typedef struct {
67   PyTypeObject *StructRUsageType;
68 } resourcemodulestate;
69 
70 
71 static inline resourcemodulestate*
get_resource_state(PyObject * module)72 get_resource_state(PyObject *module)
73 {
74     void *state = PyModule_GetState(module);
75     assert(state != NULL);
76     return (resourcemodulestate *)state;
77 }
78 
79 static struct PyModuleDef resourcemodule;
80 
81 #ifdef HAVE_GETRUSAGE
82 /*[clinic input]
83 resource.getrusage
84 
85     who: int
86     /
87 
88 [clinic start generated code]*/
89 
90 static PyObject *
resource_getrusage_impl(PyObject * module,int who)91 resource_getrusage_impl(PyObject *module, int who)
92 /*[clinic end generated code: output=8fad2880ba6a9843 input=5c857bcc5b9ccb1b]*/
93 {
94     struct rusage ru;
95     PyObject *result;
96 
97     if (getrusage(who, &ru) == -1) {
98         if (errno == EINVAL) {
99             PyErr_SetString(PyExc_ValueError,
100                             "invalid who parameter");
101             return NULL;
102         }
103         PyErr_SetFromErrno(PyExc_OSError);
104         return NULL;
105     }
106 
107     result = PyStructSequence_New(
108         get_resource_state(module)->StructRUsageType);
109     if (!result)
110         return NULL;
111 
112     PyStructSequence_SET_ITEM(result, 0,
113                     PyFloat_FromDouble(doubletime(ru.ru_utime)));
114     PyStructSequence_SET_ITEM(result, 1,
115                     PyFloat_FromDouble(doubletime(ru.ru_stime)));
116     PyStructSequence_SET_ITEM(result, 2, PyLong_FromLong(ru.ru_maxrss));
117     PyStructSequence_SET_ITEM(result, 3, PyLong_FromLong(ru.ru_ixrss));
118     PyStructSequence_SET_ITEM(result, 4, PyLong_FromLong(ru.ru_idrss));
119     PyStructSequence_SET_ITEM(result, 5, PyLong_FromLong(ru.ru_isrss));
120     PyStructSequence_SET_ITEM(result, 6, PyLong_FromLong(ru.ru_minflt));
121     PyStructSequence_SET_ITEM(result, 7, PyLong_FromLong(ru.ru_majflt));
122     PyStructSequence_SET_ITEM(result, 8, PyLong_FromLong(ru.ru_nswap));
123     PyStructSequence_SET_ITEM(result, 9, PyLong_FromLong(ru.ru_inblock));
124     PyStructSequence_SET_ITEM(result, 10, PyLong_FromLong(ru.ru_oublock));
125     PyStructSequence_SET_ITEM(result, 11, PyLong_FromLong(ru.ru_msgsnd));
126     PyStructSequence_SET_ITEM(result, 12, PyLong_FromLong(ru.ru_msgrcv));
127     PyStructSequence_SET_ITEM(result, 13, PyLong_FromLong(ru.ru_nsignals));
128     PyStructSequence_SET_ITEM(result, 14, PyLong_FromLong(ru.ru_nvcsw));
129     PyStructSequence_SET_ITEM(result, 15, PyLong_FromLong(ru.ru_nivcsw));
130 
131     if (PyErr_Occurred()) {
132         Py_DECREF(result);
133         return NULL;
134     }
135 
136     return result;
137 }
138 #endif
139 
140 static int
py2rlimit(PyObject * limits,struct rlimit * rl_out)141 py2rlimit(PyObject *limits, struct rlimit *rl_out)
142 {
143     PyObject *curobj, *maxobj;
144     limits = PySequence_Tuple(limits);
145     if (!limits)
146         /* Here limits is a borrowed reference */
147         return -1;
148 
149     if (PyTuple_GET_SIZE(limits) != 2) {
150         PyErr_SetString(PyExc_ValueError,
151                         "expected a tuple of 2 integers");
152         goto error;
153     }
154     curobj = PyTuple_GET_ITEM(limits, 0);
155     maxobj = PyTuple_GET_ITEM(limits, 1);
156 #if !defined(HAVE_LARGEFILE_SUPPORT)
157     rl_out->rlim_cur = PyLong_AsLong(curobj);
158     if (rl_out->rlim_cur == (rlim_t)-1 && PyErr_Occurred())
159         goto error;
160     rl_out->rlim_max = PyLong_AsLong(maxobj);
161     if (rl_out->rlim_max == (rlim_t)-1 && PyErr_Occurred())
162         goto error;
163 #else
164     /* The limits are probably bigger than a long */
165     rl_out->rlim_cur = PyLong_AsLongLong(curobj);
166     if (rl_out->rlim_cur == (rlim_t)-1 && PyErr_Occurred())
167         goto error;
168     rl_out->rlim_max = PyLong_AsLongLong(maxobj);
169     if (rl_out->rlim_max == (rlim_t)-1 && PyErr_Occurred())
170         goto error;
171 #endif
172 
173     Py_DECREF(limits);
174     rl_out->rlim_cur = rl_out->rlim_cur & RLIM_INFINITY;
175     rl_out->rlim_max = rl_out->rlim_max & RLIM_INFINITY;
176     return 0;
177 
178 error:
179     Py_DECREF(limits);
180     return -1;
181 }
182 
183 static PyObject*
rlimit2py(struct rlimit rl)184 rlimit2py(struct rlimit rl)
185 {
186     if (sizeof(rl.rlim_cur) > sizeof(long)) {
187         return Py_BuildValue("LL",
188                              (long long) rl.rlim_cur,
189                              (long long) rl.rlim_max);
190     }
191     return Py_BuildValue("ll", (long) rl.rlim_cur, (long) rl.rlim_max);
192 }
193 
194 /*[clinic input]
195 resource.getrlimit
196 
197     resource: int
198     /
199 
200 [clinic start generated code]*/
201 
202 static PyObject *
resource_getrlimit_impl(PyObject * module,int resource)203 resource_getrlimit_impl(PyObject *module, int resource)
204 /*[clinic end generated code: output=98327b25061ffe39 input=a697cb0004cb3c36]*/
205 {
206     struct rlimit rl;
207 
208     if (resource < 0 || resource >= RLIM_NLIMITS) {
209         PyErr_SetString(PyExc_ValueError,
210                         "invalid resource specified");
211         return NULL;
212     }
213 
214     if (getrlimit(resource, &rl) == -1) {
215         PyErr_SetFromErrno(PyExc_OSError);
216         return NULL;
217     }
218     return rlimit2py(rl);
219 }
220 
221 /*[clinic input]
222 resource.setrlimit
223 
224     resource: int
225     limits: object
226     /
227 
228 [clinic start generated code]*/
229 
230 static PyObject *
resource_setrlimit_impl(PyObject * module,int resource,PyObject * limits)231 resource_setrlimit_impl(PyObject *module, int resource, PyObject *limits)
232 /*[clinic end generated code: output=4e82ec3f34d013d1 input=6235a6ce23b4ca75]*/
233 {
234     struct rlimit rl;
235 
236     if (resource < 0 || resource >= RLIM_NLIMITS) {
237         PyErr_SetString(PyExc_ValueError,
238                         "invalid resource specified");
239         return NULL;
240     }
241 
242     if (PySys_Audit("resource.setrlimit", "iO", resource,
243                     limits ? limits : Py_None) < 0) {
244         return NULL;
245     }
246 
247     if (py2rlimit(limits, &rl) < 0) {
248         return NULL;
249     }
250 
251     if (setrlimit(resource, &rl) == -1) {
252         if (errno == EINVAL)
253             PyErr_SetString(PyExc_ValueError,
254                             "current limit exceeds maximum limit");
255         else if (errno == EPERM)
256             PyErr_SetString(PyExc_ValueError,
257                             "not allowed to raise maximum limit");
258         else
259             PyErr_SetFromErrno(PyExc_OSError);
260         return NULL;
261     }
262     Py_RETURN_NONE;
263 }
264 
265 #ifdef HAVE_PRLIMIT
266 /*[clinic input]
267 resource.prlimit
268 
269     pid: pid_t
270     resource: int
271     [
272     limits: object
273     ]
274     /
275 
276 [clinic start generated code]*/
277 
278 static PyObject *
resource_prlimit_impl(PyObject * module,pid_t pid,int resource,int group_right_1,PyObject * limits)279 resource_prlimit_impl(PyObject *module, pid_t pid, int resource,
280                       int group_right_1, PyObject *limits)
281 /*[clinic end generated code: output=ee976b393187a7a3 input=b77743bdccc83564]*/
282 {
283     struct rlimit old_limit, new_limit;
284     int retval;
285 
286     if (resource < 0 || resource >= RLIM_NLIMITS) {
287         PyErr_SetString(PyExc_ValueError,
288                         "invalid resource specified");
289         return NULL;
290     }
291 
292     if (PySys_Audit("resource.prlimit", "iiO", pid, resource,
293                     limits ? limits : Py_None) < 0) {
294         return NULL;
295     }
296 
297     if (group_right_1) {
298         if (py2rlimit(limits, &new_limit) < 0) {
299             return NULL;
300         }
301         retval = prlimit(pid, resource, &new_limit, &old_limit);
302     }
303     else {
304         retval = prlimit(pid, resource, NULL, &old_limit);
305     }
306 
307     if (retval == -1) {
308         if (errno == EINVAL) {
309             PyErr_SetString(PyExc_ValueError,
310                             "current limit exceeds maximum limit");
311         } else {
312             PyErr_SetFromErrno(PyExc_OSError);
313         }
314         return NULL;
315     }
316     return rlimit2py(old_limit);
317 }
318 #endif /* HAVE_PRLIMIT */
319 
320 /*[clinic input]
321 resource.getpagesize -> int
322 [clinic start generated code]*/
323 
324 static int
resource_getpagesize_impl(PyObject * module)325 resource_getpagesize_impl(PyObject *module)
326 /*[clinic end generated code: output=9ba93eb0f3d6c3a9 input=546545e8c1f42085]*/
327 {
328     long pagesize = 0;
329 #if defined(HAVE_GETPAGESIZE)
330     pagesize = getpagesize();
331 #elif defined(HAVE_SYSCONF) && defined(_SC_PAGE_SIZE)
332     pagesize = sysconf(_SC_PAGE_SIZE);
333 #else
334 #   error "unsupported platform: resource.getpagesize()"
335 #endif
336     return pagesize;
337 }
338 
339 /* List of functions */
340 
341 static struct PyMethodDef
342 resource_methods[] = {
343     RESOURCE_GETRUSAGE_METHODDEF
344     RESOURCE_GETRLIMIT_METHODDEF
345     RESOURCE_PRLIMIT_METHODDEF
346     RESOURCE_SETRLIMIT_METHODDEF
347     RESOURCE_GETPAGESIZE_METHODDEF
348     {NULL, NULL}                             /* sentinel */
349 };
350 
351 
352 /* Module initialization */
353 
354 static int
resource_exec(PyObject * module)355 resource_exec(PyObject *module)
356 {
357     resourcemodulestate *state = get_resource_state(module);
358 #define ADD_INT(module, value)                                    \
359     do {                                                          \
360         if (PyModule_AddIntConstant(module, #value, value) < 0) { \
361             return -1;                                            \
362         }                                                         \
363     } while (0)
364 
365     /* Add some symbolic constants to the module */
366     Py_INCREF(PyExc_OSError);
367     if (PyModule_AddObject(module, "error", PyExc_OSError) < 0) {
368         Py_DECREF(PyExc_OSError);
369         return -1;
370     }
371 
372     state->StructRUsageType = PyStructSequence_NewType(&struct_rusage_desc);
373     if (state->StructRUsageType == NULL) {
374         return -1;
375     }
376     if (PyModule_AddType(module, state->StructRUsageType) < 0) {
377         return -1;
378     }
379 
380     /* insert constants */
381 #ifdef RLIMIT_CPU
382     ADD_INT(module, RLIMIT_CPU);
383 #endif
384 
385 #ifdef RLIMIT_FSIZE
386     ADD_INT(module, RLIMIT_FSIZE);
387 #endif
388 
389 #ifdef RLIMIT_DATA
390     ADD_INT(module, RLIMIT_DATA);
391 #endif
392 
393 #ifdef RLIMIT_STACK
394     ADD_INT(module, RLIMIT_STACK);
395 #endif
396 
397 #ifdef RLIMIT_CORE
398     ADD_INT(module, RLIMIT_CORE);
399 #endif
400 
401 #ifdef RLIMIT_NOFILE
402     ADD_INT(module, RLIMIT_NOFILE);
403 #endif
404 
405 #ifdef RLIMIT_OFILE
406     ADD_INT(module, RLIMIT_OFILE);
407 #endif
408 
409 #ifdef RLIMIT_VMEM
410     ADD_INT(module, RLIMIT_VMEM);
411 #endif
412 
413 #ifdef RLIMIT_AS
414     ADD_INT(module, RLIMIT_AS);
415 #endif
416 
417 #ifdef RLIMIT_RSS
418     ADD_INT(module, RLIMIT_RSS);
419 #endif
420 
421 #ifdef RLIMIT_NPROC
422     ADD_INT(module, RLIMIT_NPROC);
423 #endif
424 
425 #ifdef RLIMIT_MEMLOCK
426     ADD_INT(module, RLIMIT_MEMLOCK);
427 #endif
428 
429 #ifdef RLIMIT_SBSIZE
430     ADD_INT(module, RLIMIT_SBSIZE);
431 #endif
432 
433 /* Linux specific */
434 #ifdef RLIMIT_MSGQUEUE
435     ADD_INT(module, RLIMIT_MSGQUEUE);
436 #endif
437 
438 #ifdef RLIMIT_NICE
439     ADD_INT(module, RLIMIT_NICE);
440 #endif
441 
442 #ifdef RLIMIT_RTPRIO
443     ADD_INT(module, RLIMIT_RTPRIO);
444 #endif
445 
446 #ifdef RLIMIT_RTTIME
447     ADD_INT(module, RLIMIT_RTTIME);
448 #endif
449 
450 #ifdef RLIMIT_SIGPENDING
451     ADD_INT(module, RLIMIT_SIGPENDING);
452 #endif
453 
454 /* target */
455 #ifdef RUSAGE_SELF
456     ADD_INT(module, RUSAGE_SELF);
457 #endif
458 
459 #ifdef RUSAGE_CHILDREN
460     ADD_INT(module, RUSAGE_CHILDREN);
461 #endif
462 
463 #ifdef RUSAGE_BOTH
464     ADD_INT(module, RUSAGE_BOTH);
465 #endif
466 
467 #ifdef RUSAGE_THREAD
468     ADD_INT(module, RUSAGE_THREAD);
469 #endif
470 
471 /* FreeBSD specific */
472 
473 #ifdef RLIMIT_SWAP
474     ADD_INT(module, RLIMIT_SWAP);
475 #endif
476 
477 #ifdef RLIMIT_SBSIZE
478     ADD_INT(module, RLIMIT_SBSIZE);
479 #endif
480 
481 #ifdef RLIMIT_NPTS
482     ADD_INT(module, RLIMIT_NPTS);
483 #endif
484 
485 #ifdef RLIMIT_KQUEUES
486     ADD_INT(module, RLIMIT_KQUEUES);
487 #endif
488 
489     PyObject *v;
490     if (sizeof(RLIM_INFINITY) > sizeof(long)) {
491         v = PyLong_FromLongLong((long long) RLIM_INFINITY);
492     } else
493     {
494         v = PyLong_FromLong((long) RLIM_INFINITY);
495     }
496     if (!v) {
497         return -1;
498     }
499 
500     if (PyModule_AddObject(module, "RLIM_INFINITY", v) < 0) {
501         Py_DECREF(v);
502         return -1;
503     }
504     return 0;
505 
506 #undef ADD_INT
507 }
508 
509 static struct PyModuleDef_Slot resource_slots[] = {
510     {Py_mod_exec, resource_exec},
511     {0, NULL}
512 };
513 
514 static int
resourcemodule_traverse(PyObject * m,visitproc visit,void * arg)515 resourcemodule_traverse(PyObject *m, visitproc visit, void *arg) {
516     Py_VISIT(get_resource_state(m)->StructRUsageType);
517     return 0;
518 }
519 
520 static int
resourcemodule_clear(PyObject * m)521 resourcemodule_clear(PyObject *m) {
522     Py_CLEAR(get_resource_state(m)->StructRUsageType);
523     return 0;
524 }
525 
526 static void
resourcemodule_free(void * m)527 resourcemodule_free(void *m) {
528     resourcemodule_clear((PyObject *)m);
529 }
530 
531 static struct PyModuleDef resourcemodule = {
532     PyModuleDef_HEAD_INIT,
533     .m_name = "resource",
534     .m_size = sizeof(resourcemodulestate),
535     .m_methods = resource_methods,
536     .m_slots = resource_slots,
537     .m_traverse = resourcemodule_traverse,
538     .m_clear = resourcemodule_clear,
539     .m_free = resourcemodule_free,
540 };
541 
542 PyMODINIT_FUNC
PyInit_resource(void)543 PyInit_resource(void)
544 {
545     return PyModuleDef_Init(&resourcemodule);
546 }
547