Handle graceful server termination

Author: 0xShug0

app/server/http.cpp

@@ -6,6 +6,7 @@
 #include <array>
 #include <cctype>
 #include <cstdint>
+#include <cerrno>
 #include <iostream>
 #include <limits>
 #include <optional>
@@ -23,6 +24,7 @@ constexpr SocketHandle kInvalidSocket = INVALID_SOCKET;
 #else
 #include <arpa/inet.h>
 #include <netinet/in.h>
+#include <sys/select.h>
 #include <sys/socket.h>
 #include <unistd.h>
 using SocketHandle = int;
@@ -273,6 +275,33 @@ void handle_client(SocketHandle client, IHttpHandler & handler) {
     }
 }
 
+bool wait_for_client(SocketHandle socket, int timeout_ms) {
+    fd_set read_set;
+    FD_ZERO(&read_set);
+    FD_SET(socket, &read_set);
+
+    timeval timeout{};
+    timeout.tv_sec = timeout_ms / 1000;
+    timeout.tv_usec = (timeout_ms % 1000) * 1000;
+
+#ifdef _WIN32
+    const int ready = select(0, &read_set, nullptr, nullptr, &timeout);
+#else
+    const int ready = select(socket + 1, &read_set, nullptr, nullptr, &timeout);
+#endif
+    if (ready < 0) {
+#ifdef _WIN32
+        if (WSAGetLastError() == WSAEINTR) {
+#else
+        if (errno == EINTR) {
+#endif
+            return false;
+        }
+        throw std::runtime_error("server select failed");
+    }
+    return ready > 0 && FD_ISSET(socket, &read_set);
+}
+
 }  // namespace
 
 HttpResponse json_response(std::string body, int status) {
@@ -285,11 +314,14 @@ HttpResponse error_response(int status, const std::string & message, const std::
     return json_response(body, status);
 }
 
-void serve_http(const std::string & host, int port, IHttpHandler & handler) {
+void serve_http(const std::string & host, int port, IHttpHandler & handler, ShutdownRequested shutdown_requested) {
     SocketRuntime sockets;
     auto listen_socket = bind_listen_socket(host, port);
     std::cout << "audiocpp_server listening on http://" << host << ":" << port << "\n";
-    while (true) {
+    while (!shutdown_requested()) {
+        if (!wait_for_client(listen_socket.get(), 250)) {
+            continue;
+        }
         sockaddr_in client_addr{};
 #ifdef _WIN32
         int client_len = sizeof(client_addr);
@@ -301,10 +333,20 @@ void serve_http(const std::string & host, int port, IHttpHandler & handler) {
             reinterpret_cast<sockaddr *>(&client_addr),
             &client_len);
         if (client == kInvalidSocket) {
+#ifdef _WIN32
+            const int error = WSAGetLastError();
+            const bool transient = error == WSAEINTR || error == WSAEWOULDBLOCK;
+#else
+            const bool transient = errno == EINTR || errno == EAGAIN || errno == EWOULDBLOCK;
+#endif
+            if (shutdown_requested() || transient) {
+                continue;
+            }
             throw std::runtime_error("accept failed");
         }
         std::thread(handle_client, client, std::ref(handler)).detach();
     }
+    std::cout << "audiocpp_server stopped\n";
 }
 
 }  // namespace minitts::server

app/server/http.h

@@ -26,8 +26,10 @@ class IHttpHandler {
     virtual HttpResponse handle(const HttpRequest & request) = 0;
 };
 
+using ShutdownRequested = bool (*)();
+
 HttpResponse json_response(std::string body, int status = 200);
 HttpResponse error_response(int status, const std::string & message, const std::string & type);
-void serve_http(const std::string & host, int port, IHttpHandler & handler);
+void serve_http(const std::string & host, int port, IHttpHandler & handler, ShutdownRequested shutdown_requested);
 
 }  // namespace minitts::server

app/server/main.cpp

@@ -4,6 +4,7 @@
 
 #include "engine/framework/debug/trace.h"
 
+#include <csignal>
 #include <filesystem>
 #include <iostream>
 #include <optional>
@@ -12,6 +13,16 @@
 
 namespace {
 
+volatile std::sig_atomic_t g_shutdown_requested = 0;
+
+void request_shutdown(int) {
+    g_shutdown_requested = 1;
+}
+
+bool shutdown_requested() {
+    return g_shutdown_requested != 0;
+}
+
 std::optional<std::string> arg_value(int argc, char ** argv, const std::string & name) {
     for (int i = 1; i + 1 < argc; ++i) {
         if (argv[i] == name) {
@@ -63,6 +74,8 @@ int main(int argc, char ** argv) {
             has_arg(argc, argv, "--log") || log_file.has_value(),
             log_file,
         });
+        std::signal(SIGINT, request_shutdown);
+        std::signal(SIGTERM, request_shutdown);
 
         auto config = minitts::server::load_server_config(*config_path);
         if (const auto host = arg_value(argc, argv, "--host")) {
@@ -85,7 +98,7 @@ int main(int argc, char ** argv) {
         }
 
         minitts::server::ServerState state(config, std::filesystem::current_path());
-        minitts::server::serve_http(config.host, config.port, state);
+        minitts::server::serve_http(config.host, config.port, state, shutdown_requested);
         return 0;
     } catch (const std::exception & ex) {
         std::cerr << "audiocpp_server failed: " << ex.what() << "\n";

scripts/build_linux.sh

@@ -11,6 +11,7 @@ WITH_TESTS="OFF"
 WITH_EXAMPLES="OFF"
 WITH_WARMBENCH="OFF"
 NATIVE_CPU="ON"
+LLAMAFILE="ON"
 TARGETS=()
 JOBS=""
 
@@ -79,6 +80,21 @@ while [[ $# -gt 0 ]]; do
             esac
             shift 2
             ;;
+        --llamafile)
+            case "$2" in
+                ON|on|On|1|true|TRUE|yes|YES)
+                    LLAMAFILE="ON"
+                    ;;
+                OFF|off|Off|0|false|FALSE|no|NO)
+                    LLAMAFILE="OFF"
+                    ;;
+                *)
+                    echo "--llamafile must be ON or OFF" >&2
+                    exit 1
+                    ;;
+            esac
+            shift 2
+            ;;
         --target)
             TARGETS+=("$2")
             shift 2
@@ -178,6 +194,7 @@ echo "Using build dir: $BUILD_DIR"
 echo "Including CUDA backend: $ENGINE_ENABLE_CUDA"
 echo "Including Vulkan backend: $ENGINE_ENABLE_VULKAN"
 echo "Native CPU optimization: $NATIVE_CPU"
+echo "llamafile SGEMM: $LLAMAFILE"
 echo "Building examples: $WITH_EXAMPLES"
 echo "Building tests: $WITH_TESTS"
 echo "Building warmbench: $WITH_WARMBENCH"
@@ -190,6 +207,7 @@ echo "Building warmbench: $WITH_WARMBENCH"
     -DENGINE_ENABLE_CUDA="$ENGINE_ENABLE_CUDA" \
     -DENGINE_ENABLE_VULKAN="$ENGINE_ENABLE_VULKAN" \
     -DENGINE_ENABLE_NATIVE_CPU="$NATIVE_CPU" \
+    -DENGINE_ENABLE_LLAMAFILE="$LLAMAFILE" \
     -DENGINE_BUILD_EXAMPLES="$WITH_EXAMPLES" \
     -DENGINE_BUILD_TESTS="$WITH_TESTS" \
     -DENGINE_BUILD_WARMBENCH="$WITH_WARMBENCH"

scripts/build_windows.ps1

@@ -8,6 +8,8 @@ param(
     [string]$CudaArchitectures = "auto",
     [ValidateSet("ON", "OFF")]
     [string]$NativeCpu = $null,
+    [ValidateSet("ON", "OFF")]
+    [string]$Llamafile = $null,
     [string]$VsInstall = ""
 )
 
@@ -261,6 +263,7 @@ function Get-PresetSettings {
                 BuildType = "Release"
                 BuildTests = "OFF"
                 Native = "ON"
+                Llamafile = "ON"
                 EnableCuda = "OFF"
                 EnableCudaGraphs = "OFF"
                 CFlagsDebug = ""
@@ -272,6 +275,7 @@ function Get-PresetSettings {
                 BuildType = "Debug"
                 BuildTests = "ON"
                 Native = "ON"
+                Llamafile = "ON"
                 EnableCuda = "ON"
                 EnableCudaGraphs = "ON"
                 CFlagsDebug = "/O2 /Zi"
@@ -283,6 +287,7 @@ function Get-PresetSettings {
                 BuildType = "Release"
                 BuildTests = "OFF"
                 Native = "ON"
+                Llamafile = "ON"
                 EnableCuda = "ON"
                 EnableCudaGraphs = "ON"
                 CFlagsDebug = ""
@@ -294,6 +299,7 @@ function Get-PresetSettings {
                 BuildType = "Debug"
                 BuildTests = "ON"
                 Native = "ON"
+                Llamafile = "ON"
                 EnableCuda = "ON"
                 EnableCudaGraphs = "ON"
                 CFlagsDebug = "/O2 /Zi"
@@ -310,6 +316,9 @@ $settings = Get-PresetSettings $Preset
 if ($null -ne $NativeCpu) {
     $settings.Native = $NativeCpu
 }
+if ($null -ne $Llamafile) {
+    $settings.Llamafile = $Llamafile
+}
 $isCudaPreset = $settings.EnableCuda -eq "ON"
 
 if ($isCudaPreset) {
@@ -356,6 +365,7 @@ if ($arch -ne "") {
     Write-Host "CUDA architectures: $arch"
 }
 Write-Host "Native CPU optimization: $($settings.Native)"
+Write-Host "llamafile SGEMM: $($settings.Llamafile)"
 
 if ($Clean) {
     $buildDirForClean = Join-Path (Join-Path (Split-Path $PSScriptRoot -Parent) "build") $Preset
@@ -387,6 +397,7 @@ $configureArgs = @(
     "-DENGINE_ENABLE_METAL=OFF",
     "-DGGML_OPENMP=ON",
     "-DENGINE_ENABLE_NATIVE_CPU=$($settings.Native)",
+    "-DENGINE_ENABLE_LLAMAFILE=$($settings.Llamafile)",
     "-DENGINE_BUILD_TESTS=$($settings.BuildTests)"
 )
 if ($settings.CFlagsDebug -ne "") {