From 74eee75e4730fdc6279acef273d1045348f2dda3 Mon Sep 17 00:00:00 2001 From: Nikolay Govorov Date: Wed, 3 Dec 2025 14:47:06 +0000 Subject: [PATCH] Add a synchronous signal handler for graceful shutdown --- src/app.zig | 6 ++++ src/main.zig | 68 +++++++++++------------------------ src/server.zig | 43 +++++++++++++++++----- src/sighandler.zig | 88 ++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 149 insertions(+), 56 deletions(-) create mode 100644 src/sighandler.zig diff --git a/src/app.zig b/src/app.zig index 719dd9b7..076c14fe 100644 --- a/src/app.zig +++ b/src/app.zig @@ -19,6 +19,7 @@ pub const App = struct { telemetry: Telemetry, app_dir_path: ?[]const u8, notification: *Notification, + shutdown: bool = false, pub const RunMode = enum { help, @@ -82,9 +83,14 @@ pub const App = struct { } pub fn deinit(self: *App) void { + if (@atomicRmw(bool, &self.shutdown, .Xchg, true, .monotonic)) { + return; + } + const allocator = self.allocator; if (self.app_dir_path) |app_dir_path| { allocator.free(app_dir_path); + self.app_dir_path = null; } self.telemetry.deinit(); self.notification.deinit(); diff --git a/src/main.zig b/src/main.zig index 829c7a2a..cf594edc 100644 --- a/src/main.zig +++ b/src/main.zig @@ -23,67 +23,42 @@ const Allocator = std.mem.Allocator; const log = @import("log.zig"); const App = @import("app.zig").App; const Server = @import("server.zig").Server; +const SigHandler = @import("sighandler.zig").SigHandler; const Browser = @import("browser/browser.zig").Browser; const DumpStripMode = @import("browser/dump.zig").Opts.StripMode; const build_config = @import("build_config"); -var _app: ?*App = null; -var _server: ?Server = null; - pub fn main() !void { // allocator // - in Debug mode we use the General Purpose Allocator to detect memory leaks // - in Release mode we use the c allocator - var gpa: std.heap.DebugAllocator(.{}) = .init; - const alloc = if (builtin.mode == .Debug) gpa.allocator() else std.heap.c_allocator; + var gpa_instance: std.heap.DebugAllocator(.{}) = .init; + const gpa = if (builtin.mode == .Debug) gpa_instance.allocator() else std.heap.c_allocator; defer if (builtin.mode == .Debug) { - if (gpa.detectLeaks()) std.posix.exit(1); + if (gpa_instance.detectLeaks()) std.posix.exit(1); }; - run(alloc) catch |err| { + var arena_instance = std.heap.ArenaAllocator.init(gpa); + const arena = arena_instance.allocator(); + defer arena_instance.deinit(); + + var sighandler = SigHandler{ .arena = arena }; + try sighandler.install(); + + run(gpa, arena, &sighandler) catch |err| { // If explicit filters were set, they won't be valid anymore because - // the args_arena is gone. We need to set it to something that's not - // invalid. (We should just move the args_arena up to main) + // the arena is gone. We need to set it to something that's not + // invalid. (We should just move the arena up to main) log.opts.filter_scopes = &.{}; log.fatal(.app, "exit", .{ .err = err }); std.posix.exit(1); }; } -// Handle app shutdown gracefuly on signals. -fn shutdown() void { - const sigaction: std.posix.Sigaction = .{ - .handler = .{ - .handler = struct { - pub fn handler(_: c_int) callconv(.c) void { - // Shutdown service gracefuly. - if (_server) |server| { - server.deinit(); - } - if (_app) |app| { - app.deinit(); - } - std.posix.exit(0); - } - }.handler, - }, - .mask = std.posix.empty_sigset, - .flags = 0, - }; - // Exit the program on SIGINT signal. When running the browser in a Docker - // container, sending a CTRL-C (SIGINT) signal is catched but doesn't exit - // the program. Here we force exiting on SIGINT. - std.posix.sigaction(std.posix.SIG.INT, &sigaction, null); - std.posix.sigaction(std.posix.SIG.TERM, &sigaction, null); - std.posix.sigaction(std.posix.SIG.QUIT, &sigaction, null); -} - -fn run(alloc: Allocator) !void { - var args_arena = std.heap.ArenaAllocator.init(alloc); - defer args_arena.deinit(); - const args = try parseArgs(args_arena.allocator()); +fn run(gpa: Allocator, arena: Allocator, sighandler: *SigHandler) !void { + const args = try parseArgs(arena); switch (args.mode) { .help => { @@ -110,13 +85,13 @@ fn run(alloc: Allocator) !void { const user_agent = blk: { const USER_AGENT = "User-Agent: Lightpanda/1.0"; if (args.userAgentSuffix()) |suffix| { - break :blk try std.fmt.allocPrintSentinel(args_arena.allocator(), "{s} {s}", .{ USER_AGENT, suffix }, 0); + break :blk try std.fmt.allocPrintSentinel(arena, "{s} {s}", .{ USER_AGENT, suffix }, 0); } break :blk USER_AGENT; }; // _app is global to handle graceful shutdown. - _app = try App.init(alloc, .{ + var app = try App.init(gpa, .{ .run_mode = args.mode, .http_proxy = args.httpProxy(), .proxy_bearer_token = args.proxyBearerToken(), @@ -127,8 +102,6 @@ fn run(alloc: Allocator) !void { .http_max_concurrent = args.httpMaxConcurrent(), .user_agent = user_agent, }); - - const app = _app.?; defer app.deinit(); app.telemetry.record(.{ .run = {} }); @@ -141,10 +114,11 @@ fn run(alloc: Allocator) !void { }; // _server is global to handle graceful shutdown. - _server = try Server.init(app, address); - const server = &_server.?; + var server = try Server.init(app, address); defer server.deinit(); + try sighandler.on(Server.stop, .{&server}); + // max timeout of 1 week. const timeout = if (opts.timeout > 604_800) 604_800_000 else @as(i32, opts.timeout) * 1000; server.run(address, timeout) catch |err| { diff --git a/src/server.zig b/src/server.zig index 1719e11d..da7f16fe 100644 --- a/src/server.zig +++ b/src/server.zig @@ -38,7 +38,7 @@ const MAX_MESSAGE_SIZE = 512 * 1024 + 14 + 140; pub const Server = struct { app: *App, - shutdown: bool, + shutdown: bool = false, allocator: Allocator, client: ?posix.socket_t, listener: ?posix.socket_t, @@ -53,16 +53,36 @@ pub const Server = struct { .app = app, .client = null, .listener = null, - .shutdown = false, .allocator = allocator, .json_version_response = json_version_response, }; } + /// Interrupts the server so that main can complete normally and call all defer handlers. + pub fn stop(self: *Server) void { + if (@atomicRmw(bool, &self.shutdown, .Xchg, true, .monotonic)) { + return; + } + + // Linux and BSD/macOS handle canceling a socket blocked on accept differently. + // For Linux, we use std.shutdown, which will cause accept to return error.SocketNotListening (EINVAL). + // For BSD, shutdown will return an error. Instead we call posix.close, which will result with error.ConnectionAborted (BADF). + if (self.listener) |listener| switch (builtin.target.os.tag) { + .linux => posix.shutdown(listener, .recv) catch |err| { + log.warn(.app, "listener shutdown", .{ .err = err }); + }, + .macos, .freebsd, .netbsd, .openbsd => { + self.listener = null; + posix.close(listener); + }, + else => unreachable, + }; + } + pub fn deinit(self: *Server) void { - self.shutdown = true; if (self.listener) |listener| { posix.close(listener); + self.listener = null; } // *if* server.run is running, we should really wait for it to return // before existing from here. @@ -83,14 +103,19 @@ pub const Server = struct { try posix.listen(listener, 1); log.info(.app, "server running", .{ .address = address }); - while (true) { + while (!@atomicLoad(bool, &self.shutdown, .monotonic)) { const socket = posix.accept(listener, null, null, posix.SOCK.NONBLOCK) catch |err| { - if (self.shutdown) { - return; + switch (err) { + error.SocketNotListening, error.ConnectionAborted => { + log.info(.app, "server stopped", .{}); + break; + }, + else => { + log.err(.app, "CDP accept", .{ .err = err }); + std.Thread.sleep(std.time.ns_per_s); + continue; + }, } - log.err(.app, "CDP accept", .{ .err = err }); - std.Thread.sleep(std.time.ns_per_s); - continue; }; self.client = socket; diff --git a/src/sighandler.zig b/src/sighandler.zig new file mode 100644 index 00000000..877fe7b1 --- /dev/null +++ b/src/sighandler.zig @@ -0,0 +1,88 @@ +//! This structure processes operating system signals (SIGINT, SIGTERM) +//! and runs callbacks to clean up the system gracefully. +//! +//! The structure does not clear the memory allocated in the arena, +//! clear the entire arena when exiting the program. +const std = @import("std"); +const assert = std.debug.assert; +const Allocator = std.mem.Allocator; + +const log = @import("log.zig"); + +pub const SigHandler = struct { + arena: Allocator, + + sigset: std.posix.sigset_t = undefined, + handle_thread: ?std.Thread = null, + + attempt: u32 = 0, + listeners: std.ArrayList(Listener) = .empty, + + pub const Listener = struct { + args: []const u8, + start: *const fn (context: *const anyopaque) void, + }; + + pub fn install(self: *SigHandler) !void { + // Block SIGINT and SIGTERM for the current thread and all created from it + self.sigset = std.posix.sigemptyset(); + std.posix.sigaddset(&self.sigset, std.posix.SIG.INT); + std.posix.sigaddset(&self.sigset, std.posix.SIG.TERM); + std.posix.sigaddset(&self.sigset, std.posix.SIG.QUIT); + std.posix.sigprocmask(std.posix.SIG.BLOCK, &self.sigset, null); + + self.handle_thread = try std.Thread.spawn(.{ .allocator = self.arena }, SigHandler.sighandle, .{self}); + self.handle_thread.?.detach(); + } + + pub fn on(self: *SigHandler, func: anytype, args: std.meta.ArgsTuple(@TypeOf(func))) !void { + assert(@typeInfo(@TypeOf(func)).@"fn".return_type.? == void); + + const Args = @TypeOf(args); + const TypeErased = struct { + fn start(context: *const anyopaque) void { + const args_casted: *const Args = @ptrCast(@alignCast(context)); + @call(.auto, func, args_casted.*); + } + }; + + const buffer = try self.arena.alignedAlloc(u8, .of(Args), @sizeOf(Args)); + errdefer self.arena.free(buffer); + + const bytes: []const u8 = @ptrCast((&args)[0..1]); + @memcpy(buffer, bytes); + + try self.listeners.append(self.arena, .{ + .args = buffer, + .start = TypeErased.start, + }); + } + + fn sighandle(self: *SigHandler) noreturn { + while (true) { + var sig: c_int = 0; + + const rc = std.c.sigwait(&self.sigset, &sig); + if (rc != 0) { + log.err(.app, "Unable to process signal {}", .{rc}); + std.process.exit(1); + } + + switch (sig) { + std.posix.SIG.INT, std.posix.SIG.TERM => { + if (self.attempt > 1) { + std.process.exit(1); + } + self.attempt += 1; + + log.info(.app, "Received termination signal...", .{}); + for (self.listeners.items) |*item| { + item.start(item.args.ptr); + } + continue; + }, + else => continue, + } + } + } +};