58 lines
2.2 KiB
C
58 lines
2.2 KiB
C
#include "resource.h"
|
|
#include "err.h"
|
|
#include "utils/log.h"
|
|
|
|
#include <stdlib.h>
|
|
#include <sys/resource.h>
|
|
|
|
#define UNSET_ENV(name) \
|
|
do { \
|
|
if (unsetenv(name)) { \
|
|
LOG_ERR("Failed to unset environment variable %s", name); \
|
|
exit(ERR_UNSETENV); \
|
|
} \
|
|
} while (0)
|
|
|
|
#define SET_LIMIT(resource, limit, name, unit) \
|
|
do { \
|
|
if (limit) { \
|
|
LOG_INFO("Setting " name " limit to %ld " unit, limit); \
|
|
if (setrlimit(resource, &(struct rlimit){limit, limit})) { \
|
|
LOG_ERR("Failed to set " name " limit"); \
|
|
exit(ERR_RLIMIT_SET); \
|
|
} \
|
|
} \
|
|
} while (0)
|
|
|
|
void setup_rlimit(void) {
|
|
LOG_INFO("Setting resource limit");
|
|
|
|
char *mem_limit_str = getenv(LIMIT_MEMORY); // in mb
|
|
char *nproc_limit_str = getenv(LIMIT_NPROC);
|
|
char *time_limit_str = getenv(LIMIT_TIME); // in ms
|
|
|
|
long mem_limit = 0, nproc_limit = 0, time_limit = 0;
|
|
|
|
if (mem_limit_str) {
|
|
// convert to bytes and double the memory limit
|
|
mem_limit = strtol(mem_limit_str, NULL, 10) * 1024 * 1024 * 2;
|
|
}
|
|
|
|
if (nproc_limit_str) {
|
|
nproc_limit = strtol(nproc_limit_str, NULL, 10);
|
|
}
|
|
|
|
if (time_limit_str) {
|
|
long limit = strtol(time_limit_str, NULL, 10);
|
|
time_limit = limit ? (limit + 1000) / 1000 : 0;
|
|
}
|
|
|
|
SET_LIMIT(RLIMIT_AS, mem_limit, "memory", "bytes");
|
|
SET_LIMIT(RLIMIT_NPROC, nproc_limit, "nproc", "processes"); // per user
|
|
SET_LIMIT(RLIMIT_CPU, time_limit, "time", "seconds"); // except blocked time
|
|
|
|
UNSET_ENV(LIMIT_MEMORY);
|
|
UNSET_ENV(LIMIT_NPROC);
|
|
UNSET_ENV(LIMIT_TIME);
|
|
}
|