lib: Add API for limiting CPU usage of subroutines.

diff --git a/src/lib/Makefile.am b/src/lib/Makefile.am
index 912f5582f8ac7a31ef4ac616d4b1b97bae509863..d9f4345ce0f95364d324f62fd94e4240bafd86e2 100644 (file)
--- a/src/lib/Makefile.am
+++ b/src/lib/Makefile.am
@@ -51,6 +51,7 @@ liblib_la_SOURCES = \
        child-wait.c \
        compat.c \
        connection.c \
+       cpu-limit.c \
        crc32.c \
        data-stack.c \
        eacces-error.c \
@@ -211,6 +212,7 @@ headers = \
        child-wait.h \
        compat.h \
        connection.h \
+       cpu-limit.h \
        crc32.h \
        data-stack.h \
        eacces-error.h \
@@ -376,6 +378,7 @@ test_lib_SOURCES = \
        test-byteorder.c \
        test-connection.c \
        test-crc32.c \
+       test-cpu-limit.c \
        test-data-stack.c \
        test-event-category-register.c \
        test-event-filter.c \

diff --git a/src/lib/cpu-limit.c b/src/lib/cpu-limit.c
new file mode 100644 (file)
index 0000000..f436933
--- /dev/null
+++ b/src/lib/cpu-limit.c
@@ -0,0 +1,141 @@
+/* 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>
+
+static struct cpu_limit *volatile cpu_limit = NULL;
+
+struct cpu_limit {
+       struct cpu_limit *parent;
+
+       struct timeval initial_usage;
+       struct rlimit old_limit, limit;
+
+       void (*callback)(void *context);
+       void *context;
+};
+
+static void
+cpu_limit_handler(const siginfo_t *si ATTR_UNUSED, void *context ATTR_UNUSED)
+{
+       struct cpu_limit *climit = cpu_limit;
+
+       while (climit != NULL) {
+               if (climit->callback != NULL)
+                       climit->callback(climit->context);
+               climit->callback = NULL;
+               climit->context = NULL;
+
+               if (climit->parent == NULL ||
+                   climit->limit.rlim_cur < climit->parent->limit.rlim_cur)
+                       break;
+               climit = climit->parent;
+       }
+}
+
+#undef cpu_limit_init
+struct cpu_limit *
+cpu_limit_init(unsigned int cpu_limit_sec,
+              void (*limit_callback)(void *context), void *context)
+{
+       struct cpu_limit *climit;
+       struct rusage rusage;
+
+       i_assert(cpu_limit_sec > 0);
+
+       climit = i_new(struct cpu_limit, 1);
+       climit->parent = cpu_limit;
+       climit->callback = limit_callback;
+       climit->context = context;
+
+       /* Query current limit */
+       if (climit->parent == NULL) {
+               if (getrlimit(RLIMIT_CPU, &climit->old_limit) < 0)
+                       i_fatal("getrlimit() failed: %m");
+       } else {
+               climit->old_limit = climit->parent->limit;
+       }
+
+       /* Query cpu usage so far */
+       if (getrusage(RUSAGE_SELF, &rusage) < 0)
+               i_fatal("getrusage() failed: %m");
+       climit->initial_usage = rusage.ru_utime;
+       timeval_add(&climit->initial_usage, &rusage.ru_stime);
+
+       climit->limit = climit->old_limit;
+       climit->limit.rlim_cur = timeval_round(&climit->initial_usage);
+
+       if (climit->limit.rlim_max != RLIM_INFINITY &&
+           climit->limit.rlim_cur > climit->limit.rlim_max) {
+               i_fatal("CPU resource limit already exceeded (%ld > %ld)",
+                       (long)climit->limit.rlim_cur,
+                       (long)climit->limit.rlim_max);
+       }
+       climit->limit.rlim_cur += cpu_limit_sec;
+
+       if (climit->parent == NULL) {
+               lib_signals_set_handler(SIGXCPU, LIBSIG_FLAG_RESTART,
+                                       cpu_limit_handler, climit);
+       } else {
+               if (climit->limit.rlim_cur > climit->parent->limit.rlim_cur)
+                       climit->limit.rlim_cur = climit->parent->limit.rlim_cur;
+       }
+
+       if (climit->limit.rlim_max != RLIM_INFINITY &&
+           climit->limit.rlim_cur > climit->limit.rlim_max)
+               climit->limit.rlim_cur = climit->limit.rlim_max;
+
+       if (setrlimit(RLIMIT_CPU, &climit->limit) < 0)
+               i_fatal("setrlimit() failed: %m");
+
+       cpu_limit = climit;
+       if (climit->parent != NULL && climit->callback != NULL &&
+           climit->parent->callback == NULL) {
+               /* Resolve race condition: parent hit limit before we fully
+                  initialized. */
+               raise(SIGXCPU);
+       }
+
+       return climit;
+}
+
+unsigned int cpu_limit_get_usage_msecs(struct cpu_limit *climit)
+{
+       struct rusage rusage;
+       struct timeval cpu_usage;
+       int usage_diff;
+
+       /* Query cpu usage so far */
+       if (getrusage(RUSAGE_SELF, &rusage) < 0)
+               i_fatal("getrusage() failed: %m");
+       cpu_usage = rusage.ru_utime;
+       timeval_add(&cpu_usage, &rusage.ru_stime);
+
+       usage_diff = timeval_diff_msecs(&cpu_usage, &climit->initial_usage);
+       i_assert(usage_diff >= 0);
+
+       return (unsigned int)usage_diff;
+}
+
+void cpu_limit_deinit(struct cpu_limit **_climit)
+{
+       struct cpu_limit *climit = *_climit;
+
+       *_climit = NULL;
+       if (climit == NULL)
+               return;
+
+       cpu_limit = climit->parent;
+       if (setrlimit(RLIMIT_CPU, &climit->old_limit) < 0)
+               i_fatal("setrlimit() failed: %m");
+
+       if (climit->parent == NULL)
+               lib_signals_unset_handler(SIGXCPU, cpu_limit_handler, climit);
+
+       i_free(climit);
+}

diff --git a/src/lib/cpu-limit.h b/src/lib/cpu-limit.h
new file mode 100644 (file)
index 0000000..27fb016
--- /dev/null
+++ b/src/lib/cpu-limit.h
@@ -0,0 +1,36 @@
+#ifndef CPU_LIMIT
+#define CPU_LIMIT
+
+struct cpu_limit;
+
+typedef void cpu_limit_callback_t(void *context);
+
+/* Call the callback when the CPU time limit is exceeded. The callback is called
+   in signal handler context, so be careful. The limit is enforced until
+   cpu_limit_deinit() is called. CPU time limits can be nested, upon which the
+   outer time limit applies for all when it is shorter, while an inner limit
+   will trigger alone (along with its children) when it is shorter. So, if e.g.
+   both inner and outer limits are 5s, both will trigger at 5s. If the outer
+   limit is 5s and the inner one is 10s, both with trigger at 5s. If the outer
+   limit is 10s and the inner is 5, only the inner limit with trigger at 5s.
+   Once all limits created by this API are released, the original CPU resource
+   limits are restored (if any). This uses setrlimit() with RLIMIT_CPU
+   internally, which counts both user and system CPU time.
+ */
+struct cpu_limit *
+cpu_limit_init(unsigned int cpu_limit_sec,
+              cpu_limit_callback_t *callback, void *context);
+#define cpu_limit_init(cpu_limit_sec, callback, context) \
+       cpu_limit_init(cpu_limit_sec - \
+               CALLBACK_TYPECHECK(callback, void (*)(typeof(context))), \
+               (cpu_limit_callback_t *)callback, context)
+void cpu_limit_deinit(struct cpu_limit **_climit);
+
+unsigned int cpu_limit_get_usage_msecs(struct cpu_limit *climit);
+
+static inline unsigned int cpu_limit_get_usage_secs(struct cpu_limit *climit)
+{
+       return cpu_limit_get_usage_msecs(climit) / 1000;
+}
+
+#endif

diff --git a/src/lib/test-cpu-limit.c b/src/lib/test-cpu-limit.c
new file mode 100644 (file)
index 0000000..5bf0b97
--- /dev/null
+++ b/src/lib/test-cpu-limit.c
@@ -0,0 +1,145 @@
+#include "test-lib.h"
+#include "lib-signals.h"
+#include "guid.h"
+#include "time-util.h"
+#include "cpu-limit.h"
+
+#include <unistd.h>
+#include <fcntl.h>
+#include <sys/time.h>
+#include <sys/resource.h>
+
+static bool limit_exceeded1, limit_exceeded2;
+static const char *const test_path = ".test.cpulimit";
+
+static void cpu_limit_callback1(void *context ATTR_UNUSED)
+{
+       limit_exceeded1 = TRUE;
+}
+
+static void cpu_limit_callback2(void *context ATTR_UNUSED)
+{
+       limit_exceeded2 = TRUE;
+}
+
+static time_t get_cpu_time(void)
+{
+       struct rusage rusage;
+       struct timeval cpu_usage;
+
+       /* Query cpu usage so far */
+       if (getrusage(RUSAGE_SELF, &rusage) < 0)
+               i_fatal("getrusage() failed: %m");
+       cpu_usage = rusage.ru_utime;
+       timeval_add(&cpu_usage, &rusage.ru_stime);
+
+       return timeval_round(&cpu_usage);
+}
+
+static void test_cpu_loop_once(void)
+{
+       guid_128_t guid;
+
+       /* consume some user CPU */
+       for (unsigned int i = 0; i < 10000; i++)
+               guid_128_generate(guid);
+       /* consume some system CPU */
+       int fd = creat(test_path, 0600);
+       if (fd == -1)
+               i_fatal("creat(%s) failed: %m", test_path);
+       if (write(fd, guid, sizeof(guid)) < 0)
+               i_fatal("write(%s) failed: %m", test_path);
+       i_close_fd(&fd);
+}
+
+static void test_cpu_limit_simple(void)
+{
+       struct cpu_limit *climit;
+       time_t usage;
+
+       test_begin("cpu limit - simple");
+
+       lib_signals_init();
+       climit = cpu_limit_init(2, cpu_limit_callback1, NULL);
+       usage = get_cpu_time();
+
+       limit_exceeded1 = FALSE;
+       while (!limit_exceeded1)
+               test_cpu_loop_once();
+
+       cpu_limit_deinit(&climit);
+       test_assert((get_cpu_time() - usage) == 2);
+
+       lib_signals_deinit();
+       test_end();
+}
+
+static void test_cpu_limit_nested(void)
+{
+       struct cpu_limit *climit1, *climit2;
+       time_t usage1, usage2;
+       unsigned int n;
+
+       test_begin("cpu limit - nested");
+
+       lib_signals_init();
+       climit1 = cpu_limit_init(3, cpu_limit_callback1, NULL);
+       usage1 = get_cpu_time();
+
+       limit_exceeded1 = FALSE;
+       while (!limit_exceeded1 && !test_has_failed()) {
+               climit2 = cpu_limit_init(1, cpu_limit_callback2, NULL);
+               usage2 = get_cpu_time();
+
+               limit_exceeded2 = FALSE;
+               while (!limit_exceeded2 && !test_has_failed())
+                       test_cpu_loop_once();
+
+               cpu_limit_deinit(&climit2);
+               test_assert((get_cpu_time() - usage2) == 1);
+       }
+
+       cpu_limit_deinit(&climit1);
+       test_assert((get_cpu_time() - usage1) == 3);
+
+       lib_signals_deinit();
+       test_end();
+
+       test_begin("cpu limit - nested2");
+
+       lib_signals_init();
+       climit1 = cpu_limit_init(3, cpu_limit_callback1, NULL);
+       usage1 = get_cpu_time();
+
+       limit_exceeded1 = FALSE;
+       n = 0;
+       while (!limit_exceeded1 && !test_has_failed()) {
+               if (++n >= 3) {
+                       /* Consume last second in top cpu limit */
+                       test_cpu_loop_once();
+                       continue;
+               }
+               climit2 = cpu_limit_init(1, cpu_limit_callback2, NULL);
+               usage2 = get_cpu_time();
+
+               limit_exceeded2 = FALSE;
+               while (!limit_exceeded2 && !test_has_failed())
+                       test_cpu_loop_once();
+
+               cpu_limit_deinit(&climit2);
+               test_assert((get_cpu_time() - usage2) == 1);
+       }
+
+       cpu_limit_deinit(&climit1);
+       test_assert((get_cpu_time() - usage1) == 3);
+
+       i_unlink_if_exists(test_path);
+       lib_signals_deinit();
+       test_end();
+}
+
+void test_cpu_limit(void)
+{
+       test_cpu_limit_simple();
+       test_cpu_limit_nested();
+}

diff --git a/src/lib/test-lib.inc b/src/lib/test-lib.inc
index 76df782757bebdfd821e736bd8abc13e56ee9a9d..b85edcd081e39662d757bdf195acee5bd8d70d58 100644 (file)
--- a/src/lib/test-lib.inc
+++ b/src/lib/test-lib.inc
@@ -16,6 +16,7 @@ FATAL(fatal_buffer)
 TEST(test_byteorder)
 TEST(test_connection)
 TEST(test_crc32)
+TEST(test_cpu_limit)
 TEST(test_data_stack)
 FATAL(fatal_data_stack)
 TEST(test_event_category_register)