modules/python-http-handlers: better daemon handling

This commit extends the python http handler template with the following
things:

- Integration with systemd notify
- Named systemd sockets, to avoid depending on expected FD number
- Graceful and customizable handling of SIGHUP reloads
- Graceful shutdown on SIGTERM with timeout
- Multithreaded HTTP server, one thread per request
- Better error message on undefined handler

base/services/irqbalance.nix

@@ -1,22 +1,4 @@
-{ config, pkgs, lib, ... }:
-let
-  cfg = config.services.irqbalance;
-in
+{ ... }:
 {
   services.irqbalance.enable = true;
-
-  # irqbalance only has meaningful work to do on multi-socket machines, so
-  # skip starting it pointlessly everywhere else.
-  systemd.services.irqbalance.serviceConfig.ExecCondition = let
-    isMultiSocket = pkgs.writeShellApplication {
-      name = "irqbalance-is-multi-socket";
-      runtimeInputs = [ pkgs.coreutils ];
-      text = ''
-        sockets=$(cat /sys/devices/system/cpu/cpu*/topology/physical_package_id | sort -u | wc -l)
-        [ "$sockets" -gt 1 ]
-      '';
-    };
-  in lib.mkIf cfg.enable [
-    (lib.getExe isMultiSocket)
-  ];
-}
+}

hosts/bicep/services/mysql/backup.nix

@@ -41,38 +41,24 @@ in
     path = with pkgs; [
       cfg.package
       coreutils
-      diffutils
       zstd
     ];
 
-    script = ''
+    script = let
+      rotations = 2;
+    in ''
       set -euo pipefail
 
-      dump() {
-        local name="$1" out tmp
-        out="$STATE_DIRECTORY/$name.sql.zst"
-        tmp="$out.tmp"
-        shift
-        "$@" | zstd -9 --rsyncable -f -o "$tmp"
-        if cmp -s "$tmp" "$out" 2>/dev/null; then
-          rm -f "$tmp"
-        else
-          mv -f "$tmp" "$out"
-        fi
-      }
+      OUT_FILE="$STATE_DIRECTORY/mysql-dump-$(date --iso-8601).sql.zst"
 
-      declare -A keep
-      while IFS= read -r db; do
-        [ -n "$db" ] || continue
-        dump "$db" mysqldump --skip-dump-date --databases "$db"
-        keep["$db.sql.zst"]=1
-      done < <(mysql -N -e 'SHOW DATABASES' | grep -vE '^(information_schema|performance_schema)$')
+      mysqldump --all-databases | zstd --compress -9 --rsyncable -o "$OUT_FILE"
 
-      # drop dumps of databases that no longer exist
-      for f in "$STATE_DIRECTORY"/*.sql.zst; do
-        [ -e "$f" ] || continue
-        base="$(basename "$f")"
-        [ -n "''${keep[$base]:-}" ] || rm -f "$f"
+      # NOTE: this needs to be a hardlink for rrsync to allow sending it
+      rm "$STATE_DIRECTORY/mysql-dump-latest.sql.zst" ||:
+      ln -T "$OUT_FILE" "$STATE_DIRECTORY/mysql-dump-latest.sql.zst"
+
+      while [ "$(find "$STATE_DIRECTORY" -type f -printf '.' | wc -c)" -gt '${toString (rotations + 1)}' ]; do
+        rm "$(find "$STATE_DIRECTORY" -type f -printf '%T+ %p\n' | sort | head -n 1 | cut -d' ' -f2)"
       done
     '';
 

hosts/bicep/services/postgresql/backup.nix

@@ -41,42 +41,25 @@ in
 
     path = with pkgs; [
       coreutils
-      diffutils
       zstd
       cfg.package
     ];
 
-    script = ''
+    script = let
+      rotations = 2;
+    in ''
       set -euo pipefail
 
-      dump() {
-        local name="$1" out tmp
-        out="$STATE_DIRECTORY/$name.sql.zst"
-        tmp="$out.tmp"
-        shift
-        "$@" | zstd -9 --rsyncable -f -o "$tmp"
-        if cmp -s "$tmp" "$out" 2>/dev/null; then
-          rm -f "$tmp"
-        else
-          mv -f "$tmp" "$out"
-        fi
-      }
+      OUT_FILE="$STATE_DIRECTORY/postgresql-dump-$(date --iso-8601).sql.zst"
 
-      declare -A keep
-      dump globals pg_dumpall -U postgres --globals-only --restrict-key=backup
-      keep[globals.sql.zst]=1
+      pg_dumpall -U postgres | zstd --compress -9 --rsyncable -o "$OUT_FILE"
 
-      while IFS= read -r db; do
-        [ -n "$db" ] || continue
-        dump "$db" pg_dump -U postgres -C -d "$db" --restrict-key=backup
-        keep["$db.sql.zst"]=1
-      done < <(psql -U postgres -tAc "SELECT datname FROM pg_database WHERE datallowconn ORDER BY datname")
+      # NOTE: this needs to be a hardlink for rrsync to allow sending it
+      rm "$STATE_DIRECTORY/postgresql-dump-latest.sql.zst" ||:
+      ln -T "$OUT_FILE" "$STATE_DIRECTORY/postgresql-dump-latest.sql.zst"
 
-      # drop dumps of databases that no longer exist
-      for f in "$STATE_DIRECTORY"/*.sql.zst; do
-        [ -e "$f" ] || continue
-        base="$(basename "$f")"
-        [ -n "''${keep[$base]:-}" ] || rm -f "$f"
+      while [ "$(find "$STATE_DIRECTORY" -type f -printf '.' | wc -c)" -gt '${toString (rotations + 1)}' ]; do
+        rm "$(find "$STATE_DIRECTORY" -type f -printf '%T+ %p\n' | sort | head -n 1 | cut -d' ' -f2)"
       done
     '';
 

hosts/bicep/services/postgresql/default.nix

@@ -14,7 +14,6 @@ in
     extensions = ps: with ps; [ pg_repack ];
     enableTCPIP = true;
 
-    # NOTE: md5 accepts both md5 and scram-sha-256
     authentication = ''
       host all all ${values.ipv4-space} md5
       host all all ${values.ipv6-space} md5
@@ -77,15 +76,6 @@ in
       maintenance_io_concurrency = 100;
       wal_recycle = true;
 
-      # -------------------------------- #
-
-      # Authentication
-      password_encryption = "scram-sha-256";
-
-      # Logging
-      log_connections = "authorization";
-      log_disconnections = true;
-
       # SSL
       ssl = true;
       ssl_cert_file = "/run/credentials/postgresql.service/cert";

modules/python-http-handlers.nix

@@ -36,6 +36,9 @@ in
           description = ''
             For each item in this list, a `ListenStream`
             option in the `[Socket]` section will be created.
+
+            Only a single `ListenStream` is currently supported by the handler script; if
+            you need more than one, you'll have to adjust {option}`handler` accordingly.
           '';
         };
 
@@ -98,8 +101,10 @@ in
             "E303" # too many blank lines
             "E305" # expected 2 blank lines after end of function or class
             "E306" # expected 1 blank line before a nested definition
+            "E402" # module level import not at top of file
             "E501" # max line length
             "E704" # multiple statements on one line (def)
+            "F811" # redefined while unused
           ];
           description = ''
             A list of flake8 rules to ignore while linting the python code.
@@ -122,6 +127,11 @@ in
                     self.send_header("Content-Length", str(len(data)))
                     self.end_headers()
                     self.wfile.write(data)
+
+            def on_reload():
+                # This function is called when the service receives SIGHUP (e.g. via `systemctl reload`).
+                # You can use it to clear caches or re-read state. Completely optional
+                pass
           '';
           description = ''
             Python code including the HTTP handler for the server.
@@ -145,6 +155,7 @@ in
           inherit (v) listenStreams;
           socketConfig = {
             Accept = false;
+            FileDescriptorName = v.name;
           } // v.socketConfig;
         };
       }))
@@ -157,30 +168,126 @@ in
         inherit (v) name;
         value = {
           serviceConfig = {
-            Type = "simple";
+            Type = "notify-reload";
+            NotifyAccess = "main";
             DynamicUser = true;
+            TimeoutStopSec = "35s";
 
             ExecStart = let
               package = pkgs.writers.writePython3Bin "${v.name}-bin" {
                 inherit (v) libraries flakeIgnore;
               } ''
+                import os
+                import signal
                 import socket
+                import socketserver
+                import threading
+                import time
                 from http.server import HTTPServer, BaseHTTPRequestHandler
 
-                ${v.handler}
+                SOCKET_NAME = "${v.name}"
+                SHUTDOWN_TIMEOUT = 30
+
+                def sd_notify(message: str):
+                    addr = os.environ.get("NOTIFY_SOCKET")
+                    if not addr:
+                        return
+                    if addr[0] == "@":
+                        addr = "\0" + addr[1:]
+                    with socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM | socket.SOCK_CLOEXEC) as sock:
+                        sock.connect(addr)
+                        sock.sendall(message.encode())
+
+                def sd_listen_fd(name: str) -> int:
+                    if os.environ.get("LISTEN_PID") != str(os.getpid()):
+                        raise RuntimeError("No sockets were passed to this service by systemd")
+
+                    try:
+                        count = int(os.environ.get("LISTEN_FDS", "0"))
+                    except ValueError:
+                        count = 0
+
+                    raw_names = os.environ.get("LISTEN_FDNAMES")
+                    names = raw_names.split(":") if raw_names else []
+
+                    for i in range(count):
+                        if i < len(names) and names[i] == name:
+                            return 3 + i
+
+                    raise RuntimeError(
+                        "No systemd socket named %r was passed to this service; check the "
+                        "FileDescriptorName= of the corresponding .socket unit" % name
+                    )
+
+                class Server(socketserver.ThreadingMixIn, HTTPServer):
+                    daemon_threads = True
 
-                class NoBindHTTPServer(HTTPServer):
                     def server_bind(): pass
                     def server_activate(): pass
 
+                    def __init__(self, *args, **kwargs):
+                        super().__init__(*args, **kwargs)
+                        self._request_threads = []
+                        self._request_threads_lock = threading.Lock()
+
+                    def process_request(self, request, client_address):
+                        thread = threading.Thread(
+                            target=self.process_request_thread,
+                            args=(request, client_address),
+                        )
+                        thread.daemon = self.daemon_threads
+                        with self._request_threads_lock:
+                            self._request_threads.append(thread)
+                        thread.start()
+
+                    def join_request_threads(self, timeout):
+                        deadline = time.monotonic() + timeout
+                        with self._request_threads_lock:
+                            threads = list(self._request_threads)
+                        for thread in threads:
+                            thread.join(max(deadline - time.monotonic(), 0))
+
+                def handle_reload(signum, frame):
+                    monotonic_usec = time.clock_gettime_ns(time.CLOCK_MONOTONIC) // 1000
+                    sd_notify("RELOADING=1\nMONOTONIC_USEC=%d" % monotonic_usec)
+
+                    on_reload = globals().get("on_reload")
+                    if callable(on_reload):
+                        on_reload()
+
+                    sd_notify("READY=1")
+
+                shutdown_requested = threading.Event()
+
+                def handle_sigterm(signum, frame):
+                    sd_notify("STOPPING=1")
+                    shutdown_requested.set()
+
+                ${v.handler}
+
+                assert "Handler" in globals(), "You must define a class Handler(BaseHTTPRequestHandler) in the handler code"
+
                 def main():
-                    httpd = NoBindHTTPServer(
-                        ("", 0),
-                        Handler,
-                        bind_and_activate=False,
-                    )
-                    httpd.socket = socket.fromfd(3, socket.AF_INET, socket.SOCK_STREAM)
-                    httpd.serve_forever()
+                    signal.signal(signal.SIGHUP, handle_reload)
+                    signal.signal(signal.SIGTERM, handle_sigterm)
+
+                    fd = sd_listen_fd(SOCKET_NAME)
+
+                    httpd = Server(("", 0), Handler, bind_and_activate=False)
+                    httpd.socket = socket.socket(fileno=fd)
+
+                    server_thread = threading.Thread(target=httpd.serve_forever, name="http-server", daemon=True)
+                    server_thread.start()
+
+                    sd_notify("READY=1")
+                    shutdown_requested.wait()
+
+                    deadline = time.monotonic() + SHUTDOWN_TIMEOUT
+
+                    httpd.shutdown()
+                    server_thread.join(max(deadline - time.monotonic(), 0))
+                    httpd.join_request_threads(max(deadline - time.monotonic(), 0))
+                    httpd.server_close()
 
                 if __name__ == '__main__':
                     main()