1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
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;
}
|
