diff options
Diffstat (limited to 'src/lib/cpu-limit.c')
| -rw-r--r-- | src/lib/cpu-limit.c | 200 |
1 files changed, 200 insertions, 0 deletions
diff --git a/src/lib/cpu-limit.c b/src/lib/cpu-limit.c new file mode 100644 index 0000000..754717a --- /dev/null +++ b/src/lib/cpu-limit.c @@ -0,0 +1,200 @@ +/* Copyright (c) 2020 Dovecot authors, see the included COPYING file */ + +#include "lib.h" +#include "lib-signals.h" +#include "time-util.h" +#include "cpu-limit.h" + +#include <sys/time.h> +#include <sys/resource.h> + +struct cpu_limit { + struct cpu_limit *parent; + + enum cpu_limit_type type; + unsigned int cpu_limit_secs; + struct rusage initial_usage; + + bool limit_reached; +}; + +static struct cpu_limit *cpu_limit; +static struct rlimit orig_limit, last_set_rlimit; +static volatile sig_atomic_t xcpu_signal_counter; +static sig_atomic_t checked_signal_counter; +static unsigned int rlim_cur_adjust_secs; + +static void +cpu_limit_handler(const siginfo_t *si ATTR_UNUSED, void *context ATTR_UNUSED) +{ + xcpu_signal_counter++; +} + +static unsigned int +cpu_limit_get_usage_msecs_with(struct cpu_limit *climit, + enum cpu_limit_type type, + const struct rusage *rusage) +{ + struct timeval cpu_usage = { 0, 0 }; + int usage_diff; + + if ((type & CPU_LIMIT_TYPE_USER) != 0) + timeval_add(&cpu_usage, &rusage->ru_utime); + if ((type & CPU_LIMIT_TYPE_SYSTEM) != 0) + timeval_add(&cpu_usage, &rusage->ru_stime); + + struct timeval initial_total = { 0, 0 }; + if ((type & CPU_LIMIT_TYPE_USER) != 0) + timeval_add(&initial_total, &climit->initial_usage.ru_utime); + if ((type & CPU_LIMIT_TYPE_SYSTEM) != 0) + timeval_add(&initial_total, &climit->initial_usage.ru_stime); + usage_diff = timeval_diff_msecs(&cpu_usage, &initial_total); + i_assert(usage_diff >= 0); + + return (unsigned int)usage_diff; +} + +unsigned int +cpu_limit_get_usage_msecs(struct cpu_limit *climit, enum cpu_limit_type type) +{ + struct rusage rusage; + + /* Query cpu usage so far */ + if (getrusage(RUSAGE_SELF, &rusage) < 0) + i_fatal("getrusage() failed: %m"); + + return cpu_limit_get_usage_msecs_with(climit, type, &rusage); +} + +static bool +cpu_limit_update_recursive(struct cpu_limit *climit, + const struct rusage *rusage, + unsigned int *max_wait_secs) +{ + if (climit == NULL) + return FALSE; + if (cpu_limit_update_recursive(climit->parent, rusage, max_wait_secs)) { + /* parent's limit reached */ + climit->limit_reached = TRUE; + return TRUE; + } + unsigned int secs_used = + cpu_limit_get_usage_msecs_with(climit, climit->type, rusage)/1000; + if (secs_used >= climit->cpu_limit_secs) { + climit->limit_reached = TRUE; + return TRUE; + } + unsigned int secs_left = climit->cpu_limit_secs - secs_used; + if (*max_wait_secs > secs_left) + *max_wait_secs = secs_left; + return FALSE; +} + +static void cpu_limit_update_rlimit(void) +{ + struct rusage rusage; + struct rlimit rlimit; + unsigned int max_wait_secs = UINT_MAX; + + if (getrusage(RUSAGE_SELF, &rusage) < 0) + i_fatal("getrusage() failed: %m"); + + (void)cpu_limit_update_recursive(cpu_limit, &rusage, &max_wait_secs); + if (max_wait_secs == UINT_MAX) { + /* All the limits have reached now. Restore the original + limits. */ + rlimit = orig_limit; + } else { + struct timeval tv_limit = rusage.ru_utime; + timeval_add(&tv_limit, &rusage.ru_stime); + + i_zero(&rlimit); + /* Add +1 second to round up. */ + rlimit.rlim_cur = tv_limit.tv_sec + + max_wait_secs + 1 + rlim_cur_adjust_secs; + rlimit.rlim_max = orig_limit.rlim_max; + } + if (last_set_rlimit.rlim_cur != rlimit.rlim_cur) { + last_set_rlimit = rlimit; + if (setrlimit(RLIMIT_CPU, &rlimit) < 0) + i_fatal("setrlimit() failed: %m"); + } +} + +struct cpu_limit * +cpu_limit_init(unsigned int cpu_limit_secs, enum cpu_limit_type type) +{ + struct cpu_limit *climit; + struct rusage rusage; + + i_assert(cpu_limit_secs > 0); + i_assert(type != 0); + + climit = i_new(struct cpu_limit, 1); + climit->parent = cpu_limit; + climit->type = type; + climit->cpu_limit_secs = cpu_limit_secs; + + /* Query current limit */ + if (climit->parent == NULL) { + if (getrlimit(RLIMIT_CPU, &orig_limit) < 0) + i_fatal("getrlimit() failed: %m"); + } + + /* Query cpu usage so far */ + if (getrusage(RUSAGE_SELF, &rusage) < 0) + i_fatal("getrusage() failed: %m"); + climit->initial_usage = rusage; + + if (climit->parent == NULL) { + lib_signals_set_handler(SIGXCPU, LIBSIG_FLAG_RESTART, + cpu_limit_handler, NULL); + } + + cpu_limit = climit; + cpu_limit_update_rlimit(); + return climit; +} + +void cpu_limit_deinit(struct cpu_limit **_climit) +{ + struct cpu_limit *climit = *_climit; + + *_climit = NULL; + if (climit == NULL) + return; + + i_assert(climit == cpu_limit); + + cpu_limit = climit->parent; + cpu_limit_update_rlimit(); + if (climit->parent == NULL) + lib_signals_unset_handler(SIGXCPU, cpu_limit_handler, NULL); + i_free(climit); +} + +bool cpu_limit_exceeded(struct cpu_limit *climit) +{ + static struct timeval tv_last = { 0, 0 }; + struct timeval tv_now; + + if (checked_signal_counter != xcpu_signal_counter) { + i_gettimeofday(&tv_now); + if (tv_last.tv_sec != 0 && + timeval_diff_msecs(&tv_now, &tv_last) < 1000) { + /* Additional sanity check: We're getting here more + rapidly than once per second. This isn't expected + to happen, but at least in theory it could happen + because rlim_cur isn't clearly calculated from just + the user+system CPU usage. So in case rlim_cur is + too low and keeps firing XCPU signal, try to + increase rlim_cur by 1 second. Eventually it should + become large enough. */ + rlim_cur_adjust_secs++; + } + + checked_signal_counter = xcpu_signal_counter; + cpu_limit_update_rlimit(); + } + return climit->limit_reached; +} |