From 945cfab3eb412cc4285aa16eff022a9943854149 Mon Sep 17 00:00:00 2001
From: apio <blobs.trading@gmail.com>
Date: Sat, 7 Oct 2023 14:26:35 +0200
Subject: [PATCH] libos: Add signal handling to event loops

---
 libos/include/os/EventLoop.h | 14 ++++++++++
 libos/src/EventLoop.cpp      | 50 +++++++++++++++++++++++++++++++++---
 2 files changed, 61 insertions(+), 3 deletions(-)

diff --git a/libos/include/os/EventLoop.h b/libos/include/os/EventLoop.h
index 7d1bb392..e48b9ff8 100644
--- a/libos/include/os/EventLoop.h
+++ b/libos/include/os/EventLoop.h
@@ -47,6 +47,14 @@ namespace os
          */
         Result<void> register_fd_listener(int fd, void (*listener)(int, int));
 
+        /**
+         * @brief Register a new POSIX signal handler.
+         *
+         * Unlike standard POSIX signal handling, this handler will be run synchronously as part of the event loop,
+         * eliminating many problems with asynchronous signal handling.
+         */
+        Result<void> register_signal_handler(int sig, void (*handler)(int));
+
         /**
          * @brief Run the event loop until it is asked to quit.
          *
@@ -63,9 +71,15 @@ namespace os
 
       private:
         HashMap<int, void (*)(int, int)> m_fd_listeners;
+        HashMap<int, void (*)(int)> m_signal_handlers;
         Vector<struct pollfd> m_pollfds;
 
+        static void handle_signal(int sig);
+        static void handle_quit_signal(int);
+
         bool m_should_quit { false };
         int m_quit_status { 0 };
+
+        int m_signal_pipe[2];
     };
 }
diff --git a/libos/src/EventLoop.cpp b/libos/src/EventLoop.cpp
index e3e4be4c..43655f8c 100644
--- a/libos/src/EventLoop.cpp
+++ b/libos/src/EventLoop.cpp
@@ -10,7 +10,9 @@
 #include <errno.h>
 #include <os/EventLoop.h>
 #include <os/File.h>
+#include <signal.h>
 #include <string.h>
+#include <unistd.h>
 
 static os::EventLoop* s_the = nullptr;
 
@@ -19,6 +21,13 @@ namespace os
     EventLoop::EventLoop()
     {
         s_the = this;
+
+        // Set up synchronous signal handling.
+        check(pipe(m_signal_pipe) == 0);
+        m_pollfds.try_append({ .fd = m_signal_pipe[0], .events = POLLIN, .revents = 0 }).release_value();
+
+        register_signal_handler(SIGTERM, EventLoop::handle_quit_signal).release_value();
+        register_signal_handler(SIGINT, EventLoop::handle_quit_signal).release_value();
     }
 
     EventLoop::~EventLoop()
@@ -39,6 +48,18 @@ namespace os
         return {};
     }
 
+    Result<void> EventLoop::register_signal_handler(int sig, void (*handler)(int))
+    {
+        struct sigaction sa;
+        sa.sa_handler = EventLoop::handle_signal;
+        sa.sa_flags = 0;
+        sigemptyset(&sa.sa_mask);
+        check(sigaction(sig, &sa, nullptr) == 0);
+
+        TRY(m_signal_handlers.try_set(sig, handler));
+        return {};
+    }
+
     int EventLoop::run()
     {
         while (!m_should_quit)
@@ -53,11 +74,23 @@ namespace os
                 return 1;
             }
 
+            if (m_pollfds[0].revents & POLLIN)
+            {
+                int sig;
+                read(m_signal_pipe[0], &sig, sizeof(int));
+
+                auto handler = m_signal_handlers.try_get(sig);
+                if (handler.has_value()) { handler.value()(sig); }
+            }
+
             for (usize i = 0; i < m_fd_listeners.size(); i++)
             {
-                if (m_pollfds[i].revents)
-                    m_fd_listeners.try_get(m_pollfds[i].fd)
-                        .value()(m_pollfds[i].fd, m_pollfds[i].revents); // Invoke the callback
+                auto& pfd = m_pollfds[i + 1];
+                if (pfd.revents)
+                {
+                    auto handler = m_fd_listeners.try_get(pfd.fd);
+                    if (handler.has_value()) handler.value()(pfd.fd, pfd.revents);
+                }
             }
         }
 
@@ -69,4 +102,15 @@ namespace os
         m_quit_status = status;
         m_should_quit = true;
     }
+
+    void EventLoop::handle_signal(int sig)
+    {
+        write(the().m_signal_pipe[1], &sig, sizeof(int));
+    }
+
+    void EventLoop::handle_quit_signal(int sig)
+    {
+        os::println("EventLoop: quit requested by signal %d", sig);
+        the().quit(sig);
+    }
 }