sd-event: add an explicit API for leaving the ratelimit state

Sometimes, it might make sense to end the ratelimit window early.

man/rules/meson.build

@@ -633,6 +633,7 @@ manpages = [
   '3',
   ['sd_event_source_get_ratelimit',
    'sd_event_source_is_ratelimited',
+   'sd_event_source_leave_ratelimit',
    'sd_event_source_set_ratelimit_expire_callback'],
   ''],
  ['sd_event_source_set_userdata', '3', ['sd_event_source_get_userdata'], ''],

man/sd_event_source_set_ratelimit.xml

@@ -20,6 +20,7 @@
     <refname>sd_event_source_get_ratelimit</refname>
     <refname>sd_event_source_is_ratelimited</refname>
     <refname>sd_event_source_set_ratelimit_expire_callback</refname>
+    <refname>sd_event_source_leave_ratelimit</refname>
 
     <refpurpose>Configure rate limiting on event sources</refpurpose>
   </refnamediv>
@@ -53,6 +54,10 @@
         <paramdef>sd_event_handler_t<parameter>callback</parameter></paramdef>
       </funcprototype>
 
+      <funcprototype>
+        <funcdef>int <function>sd_event_source_leave_ratelimit</function></funcdef>
+        <paramdef>sd_event_source *<parameter>source</parameter></paramdef>
+      </funcprototype>
     </funcsynopsis>
   </refsynopsisdiv>
 
@@ -85,10 +90,14 @@
     is currently affected by rate limiting, i.e. it has recently hit the rate limit and is currently
     temporarily disabled due to that.</para>
 
-    <para><function>sd_event_source_set_ratelimit_expire_callback</function> may be used to set a callback
+    <para><function>sd_event_source_set_ratelimit_expire_callback()</function> may be used to set a callback
     function that is invoked every time the event source leaves rate limited state. Note that function is
     called in the same event loop iteration in which state transition occurred.</para>
 
+    <para><function>sd_event_source_leave_ratelimit()</function> may be used to immediately reenable an event
+    source that was temporarily disabled due to rate limiting. This will reset the ratelimit counters for the
+    current time interval.</para>
+
     <para>Rate limiting is currently implemented for I/O, timer, signal, defer and inotify event
     sources.</para>
   </refsect1>
@@ -98,10 +107,12 @@
 
     <para>On success, <function>sd_event_source_set_ratelimit()</function>,
     <function>sd_event_source_set_ratelimit_expire_callback</function> and
-    <function>sd_event_source_get_ratelimit()</function> return a non-negative integer. On failure, they return
-    a negative errno-style error code. <function>sd_event_source_is_ratelimited</function> returns zero if rate
-    limiting is currently not in effect and greater than zero if it is in effect; it returns a negative
-    errno-style error code on failure.</para>
+    <function>sd_event_source_get_ratelimit()</function> return a non-negative integer. On failure, they
+    return a negative errno-style error code. <function>sd_event_source_is_ratelimited()</function> returns
+    zero if rate limiting is currently not in effect and greater than zero if it is in effect; it returns a
+    negative errno-style error code on failure. <function>sd_event_source_leave_ratelimit()</function>
+    returns zero if rate limiting wasn't in effect on the specified event source, and positive if it was and
+    rate limiting is now turned off again; it returns a negative errno-style error code on failure.</para>
 
     <refsect2>
       <title>Errors</title>

src/libsystemd/libsystemd.sym

@@ -824,4 +824,5 @@ global:
         sd_event_source_set_memory_pressure_period;
         sd_event_trim_memory;
         sd_pid_notify_barrier;
+        sd_event_source_leave_ratelimit;
 } LIBSYSTEMD_253;

src/libsystemd/sd-event/sd-event.c

@@ -5193,6 +5193,27 @@ _public_ int sd_event_source_is_ratelimited(sd_event_source *s) {
         return s->ratelimited;
 }
 
+_public_ int sd_event_source_leave_ratelimit(sd_event_source *s) {
+        int r;
+
+        assert_return(s, -EINVAL);
+
+        if (!EVENT_SOURCE_CAN_RATE_LIMIT(s->type))
+                return 0;
+
+        if (!ratelimit_configured(&s->rate_limit))
+                return 0;
+
+        if (!s->ratelimited)
+                return 0;
+
+        r = event_source_leave_ratelimit(s, /* run_callback */ false);
+        if (r < 0)
+                return r;
+
+        return 1; /* tell caller that we indeed just left the ratelimit state */
+}
+
 _public_ int sd_event_set_signal_exit(sd_event *e, int b) {
         bool change = false;
         int r;

src/libsystemd/sd-event/test-event.c

@@ -828,4 +828,75 @@ TEST(fork) {
         assert_se(r >= 0);
 }
 
+static int hup_callback(sd_event_source *s, int fd, uint32_t revents, void *userdata) {
+        unsigned *c = userdata;
+
+        assert_se(revents == EPOLLHUP);
+
+        (*c)++;
+        return 0;
+}
+
+TEST(leave_ratelimit) {
+        bool expect_ratelimit = false, manually_left_ratelimit = false;
+        _cleanup_(sd_event_source_unrefp) sd_event_source *s = NULL;
+        _cleanup_(sd_event_unrefp) sd_event *e = NULL;
+        _cleanup_(close_pairp) int pfd[2] = PIPE_EBADF;
+        unsigned c = 0;
+        int r;
+
+        assert_se(sd_event_default(&e) >= 0);
+
+        /* Create an event source that will continously fire by creating a pipe whose write side is closed,
+         * and which hence will only see EOF and constant EPOLLHUP */
+        assert_se(pipe2(pfd, O_CLOEXEC) >= 0);
+        assert_se(sd_event_add_io(e, &s, pfd[0], EPOLLIN, hup_callback, &c) >= 0);
+        assert_se(sd_event_source_set_io_fd_own(s, true) >= 0);
+        assert_se(sd_event_source_set_ratelimit(s, 5*USEC_PER_MINUTE, 5) >= 0);
+
+        pfd[0] = -EBADF;
+        pfd[1] = safe_close(pfd[1]); /* Trigger continous EOF */
+
+        for (;;) {
+                r = sd_event_prepare(e);
+                assert_se(r >= 0);
+
+                if (r == 0) {
+                        r = sd_event_wait(e, UINT64_MAX);
+                        assert_se(r > 0);
+                }
+
+                r = sd_event_dispatch(e);
+                assert_se(r > 0);
+
+                r = sd_event_source_is_ratelimited(s);
+                assert_se(r >= 0);
+
+                if (c < 5)
+                        /* First four dispatches should just work */
+                        assert_se(!r);
+                else if (c == 5) {
+                        /* The fifth dispatch should still work, but we now expect the ratelimit to be hit subsequently */
+                        if (!expect_ratelimit) {
+                                assert_se(!r);
+                                assert_se(sd_event_source_leave_ratelimit(s) == 0); /* this should be a NOP, and return 0 hence */
+                                expect_ratelimit = true;
+                        } else {
+                                /* We expected the ratelimit, let's leave it manually, and verify it */
+                                assert_se(r);
+                                assert_se(sd_event_source_leave_ratelimit(s) > 0); /* we are ratelimited, hence should return > 0 */
+                                assert_se(sd_event_source_is_ratelimited(s) == 0);
+
+                                manually_left_ratelimit = true;
+                        }
+
+                } else if (c == 6)
+                        /* On the sixth iteration let's just exit */
+                        break;
+        }
+
+        /* Verify we definitely hit the ratelimit and left it manually again */
+        assert_se(manually_left_ratelimit);
+}
+
 DEFINE_TEST_MAIN(LOG_DEBUG);

src/systemd/sd-event.h

@@ -173,6 +173,7 @@ int sd_event_source_set_ratelimit(sd_event_source *s, uint64_t interval_usec, un
 int sd_event_source_get_ratelimit(sd_event_source *s, uint64_t *ret_interval_usec, unsigned *ret_burst);
 int sd_event_source_is_ratelimited(sd_event_source *s);
 int sd_event_source_set_ratelimit_expire_callback(sd_event_source *s, sd_event_handler_t callback);
+int sd_event_source_leave_ratelimit(sd_event_source *s);
 
 int sd_event_trim_memory(void);