diff --git a/src/Config.zig b/src/Config.zig index 4f66aec..78bcc94 100644 --- a/src/Config.zig +++ b/src/Config.zig @@ -17,6 +17,20 @@ cache_dir: []const u8, /// fallback provider is used (ipwho.is by default, or IP2Location) geolite_path: []const u8, +/// How old the GeoLite2 database may get before the background refresher +/// replaces it. Upstream rebuilds it continuously and IP allocations move +/// between cities, so a stale database quietly resolves clients to the wrong +/// place. `WTTR_GEOLITE_MAX_AGE_DAYS=0` refreshes on every check. +geolite_max_age_seconds: u64, + +/// How often the refresher wakes up to compare the database's age against +/// `geolite_max_age_seconds`. `WTTR_GEOLITE_CHECK_INTERVAL_HOURS=0` disables +/// background refreshing entirely. +/// +/// "Disable" lives on the interval rather than on the max age so that neither +/// value has an ambiguous zero. +geolite_check_interval_seconds: u64, + /// Geocache file stores location lookups /// (e.g. "Portland -> 45.52345°N, -122.67621° W). When not found in cache, /// a web service from Nominatum (https://nominatim.org/) is used @@ -67,6 +81,20 @@ pub fn load(allocator: std.mem.Allocator, env: *const std.process.Environ.Map) ! env.get("WTTR_CACHE_DIR") orelse default_cache_dir, }); }, + .geolite_max_age_seconds = blk: { + const days = if (env.get("WTTR_GEOLITE_MAX_AGE_DAYS")) |v| + try std.fmt.parseInt(u64, v, 10) + else + 30; + break :blk days * std.time.s_per_day; + }, + .geolite_check_interval_seconds = blk: { + const hours = if (env.get("WTTR_GEOLITE_CHECK_INTERVAL_HOURS")) |v| + try std.fmt.parseInt(u64, v, 10) + else + 24; + break :blk hours * std.time.s_per_hour; + }, .geocache_file = if (env.get("WTTR_GEOCACHE_FILE")) |v| try allocator.dupe(u8, v) else try std.fs.path.join(allocator, &[_][]const u8{ default_cache_dir, "geocache.json" }), .geoip_fallback = blk: { if (env.get("WTTR_GEOIP_FALLBACK")) |v| { diff --git a/src/location/GeoIp.zig b/src/location/GeoIp.zig index 348b106..1fbd889 100644 --- a/src/location/GeoIp.zig +++ b/src/location/GeoIp.zig @@ -57,17 +57,23 @@ const FallbackClient = union(enum) { mmdb: *c.MMDB_s, fallback_client: FallbackClient, allocator: std.mem.Allocator, +io: std.Io, +/// Retained so `reload` can reopen the database after it is replaced on disk. +db_path: []const u8, +/// Guards `mmdb`. Lookups hold this shared; `reload` holds it exclusively to +/// swap the handle. Request threads and the background refresher both touch +/// `mmdb`, so the swap cannot be unsynchronized. +lock: std.Io.RwLock, pub fn init(allocator: std.mem.Allocator, io: std.Io, db_path: []const u8, config: Config) !GeoIP { - const path_z = try std.heap.c_allocator.dupeZ(u8, db_path); - defer std.heap.c_allocator.free(path_z); + const mmdb = try openDatabase(allocator, db_path); + errdefer { + c.MMDB_close(mmdb); + allocator.destroy(mmdb); + } - const mmdb = try allocator.create(c.MMDB_s); - errdefer allocator.destroy(mmdb); - - const status = c.MMDB_open(path_z.ptr, c.MMDB_MODE_MMAP, mmdb); - if (status != c.MMDB_SUCCESS) - return error.CannotOpenDatabase; + const db_path_copy = try allocator.dupe(u8, db_path); + errdefer allocator.free(db_path_copy); const fallback_client: FallbackClient = switch (config.geoip_fallback) { .ip2location => blk: { @@ -93,25 +99,76 @@ pub fn init(allocator: std.mem.Allocator, io: std.Io, db_path: []const u8, confi .mmdb = mmdb, .fallback_client = fallback_client, .allocator = allocator, + .io = io, + .db_path = db_path_copy, + .lock = .init, }; } +/// Opens an mmdb file, returning an owned handle. +fn openDatabase(allocator: std.mem.Allocator, db_path: []const u8) !*c.MMDB_s { + const path_z = try std.heap.c_allocator.dupeZ(u8, db_path); + defer std.heap.c_allocator.free(path_z); + + const mmdb = try allocator.create(c.MMDB_s); + errdefer allocator.destroy(mmdb); + + const status = c.MMDB_open(path_z.ptr, c.MMDB_MODE_MMAP, mmdb); + if (status != c.MMDB_SUCCESS) + return error.CannotOpenDatabase; + + return mmdb; +} + +/// Reopens the database from `db_path`, replacing the in-use handle. +/// +/// The file is opened *before* the exclusive lock is taken so that lookups only +/// block for the pointer swap rather than for the open. The previous handle is +/// closed after the lock is released: lookups never retain the pointer beyond +/// their own critical section, so once the swap is published no thread can still +/// reach the old handle. +/// +/// Reopening is required after the file is replaced on disk. The database is +/// mapped with `MMDB_MODE_MMAP`, so an atomic rename leaves this process mapped +/// to the old inode until it opens the new one. +pub fn reload(self: *GeoIP) !void { + const new_mmdb = try openDatabase(self.allocator, self.db_path); + + self.lock.lockUncancelable(self.io); + const old_mmdb = self.mmdb; + self.mmdb = new_mmdb; + self.lock.unlock(self.io); + + c.MMDB_close(old_mmdb); + self.allocator.destroy(old_mmdb); + log.info("GeoLite2 database reloaded from {s}", .{self.db_path}); +} + pub fn deinit(self: *GeoIP) void { c.MMDB_close(self.mmdb); self.allocator.destroy(self.mmdb); + self.allocator.free(self.db_path); self.fallback_client.deinit(self.allocator); } pub fn lookup(self: *GeoIP, ip: []const u8) ?Location { - // Try MaxMind first - const result = lookupInternal(self.mmdb, ip) catch return null; + // Try MaxMind first. The shared lock has to cover `extractCoordinates` as + // well as the lookup itself: `MMDB_lookup_result_s.entry` holds a pointer + // back to the `MMDB_s`, which `extractCoordinates` dereferences. + const from_db: ?Location = blk: { + self.lock.lockSharedUncancelable(self.io); + defer self.lock.unlockShared(self.io); - log.debug("lookup geoip db for ip {s}. Found: {}", .{ ip, result.found_entry }); - if (result.found_entry) - if (self.extractCoordinates(ip, result)) |coords| - return coords; + const result = lookupInternal(self.mmdb, ip) catch break :blk null; - // Fallback to configured online provider + log.debug("lookup geoip db for ip {s}. Found: {}", .{ ip, result.found_entry }); + if (!result.found_entry) break :blk null; + break :blk self.extractCoordinates(ip, result); + }; + if (from_db) |coords| return coords; + + // Fallback to configured online provider. Deliberately outside the lock: + // this performs a network request and must not block a reload. return self.fallback_client.lookup(ip); } @@ -131,6 +188,11 @@ fn lookupInternal(mmdb: *c.MMDB_s, ip: []const u8) !c.MMDB_lookup_result_s { } pub fn isUSIp(self: *GeoIP, ip: []const u8) bool { + // `country_data` borrows from the mapped database via `result.entry`, so the + // whole read stays inside the shared lock. + self.lock.lockSharedUncancelable(self.io); + defer self.lock.unlockShared(self.io); + var result = lookupInternal(self.mmdb, ip) catch return false; if (!result.found_entry) return false; @@ -305,3 +367,46 @@ test "lookup works" { try std.testing.expect(result.coords.longitude < -121.0 and result.coords.longitude > -123.0); try std.testing.expect(result.name.len > 0); } + +test "reload swaps the database while lookups keep working" { + const allocator = std.testing.allocator; + const config = try Config.loadForTest(allocator); + defer config.deinit(allocator); + const build_options = @import("build_options"); + const db_path = config.geolite_path; + + if (build_options.download_geoip) { + const GeoLite2 = @import("GeoLite2.zig"); + try GeoLite2.ensureDatabase(allocator, std.testing.io, db_path); + } + + var geoip = GeoIP.init(allocator, std.testing.io, db_path, config) catch + return error.SkipZigTest; + defer geoip.deinit(); + + // A residential IP that resolves from the database rather than the online + // fallback, so the assertions below exercise the mmap and not the network. + const test_ip = "73.158.64.1"; + + const before = geoip.lookup(test_ip) orelse return error.SkipZigTest; + defer before.deinit(); + const handle_before = geoip.mmdb; + + try geoip.reload(); + + // A fresh handle must be installed, and the old one must not be reused. + try std.testing.expect(geoip.mmdb != handle_before); + + // Same query, same answer: the swap installed an equivalent database and + // left the lookup path intact. + const after = geoip.lookup(test_ip) orelse return error.TestUnexpectedResult; + defer after.deinit(); + + try std.testing.expectEqualStrings(before.name, after.name); + try std.testing.expectEqual(before.coords.latitude, after.coords.latitude); + try std.testing.expectEqual(before.coords.longitude, after.coords.longitude); + + // Lookups that go through the shared lock a second time still work, which + // catches a reload that left the lock in a bad state. + try std.testing.expect(geoip.isUSIp(test_ip)); +} diff --git a/src/location/GeoLite2.zig b/src/location/GeoLite2.zig index 87d7a49..baca3de 100644 --- a/src/location/GeoLite2.zig +++ b/src/location/GeoLite2.zig @@ -1,65 +1,50 @@ const std = @import("std"); const log = std.log.scoped(.geolite2); -/// Ensures a GeoLite2 database exists at `path`, downloading it when missing -/// and refreshing it once it is older than `max_age_seconds`. +/// Ensures a GeoLite2 database exists at `path`, downloading it if absent. /// -/// Refresh matters: the upstream release is rebuilt continuously and IP -/// allocations move between cities, so a stale database silently resolves -/// clients to the wrong place. Before this, the database was only fetched when -/// absent, so a deployment could -- and did -- run for months on the copy it -/// happened to download first. -/// -/// A missing database is fatal (there is nothing to fall back on). A failed -/// *refresh* is not: an out-of-date database still answers lookups, so a -/// network blip must not stop the service from starting. -pub fn ensureDatabase( - allocator: std.mem.Allocator, - io: std.Io, - path: []const u8, - max_age_seconds: u64, -) !void { - const stat = std.Io.Dir.cwd().statFile(io, path, .{}) catch |err| switch (err) { - error.FileNotFound => { - log.info("GeoLite2 database not found at {s}, will download", .{path}); - try downloadDatabase(allocator, io, path); - log.info("GeoLite2 database downloaded successfully", .{}); - return; - }, - else => return err, - }; - - const age_seconds = ageInSeconds(io, stat.mtime); - if (age_seconds < max_age_seconds) { - log.info("GeoLite2 database is {d} day(s) old, no refresh needed", .{@divFloor(age_seconds, std.time.s_per_day)}); - return; - } - - log.info( - "GeoLite2 database is {d} day(s) old (limit {d}), refreshing", - .{ @divFloor(age_seconds, std.time.s_per_day), @divFloor(max_age_seconds, std.time.s_per_day) }, - ); - downloadDatabase(allocator, io, path) catch |err| { - // Keep serving from the stale database rather than refusing to start. - log.warn( - "GeoLite2 refresh failed ({t}); continuing with the existing database at {s}", - .{ err, path }, - ); +/// A missing database is fatal: `GeoIp.init` cannot open anything, so there is +/// nothing to serve lookups with. Keeping the database *fresh* is a separate +/// concern handled off-request by `Refresher`, because a stale database still +/// answers lookups and a network problem must not stop the service from +/// starting. +pub fn ensureDatabase(allocator: std.mem.Allocator, io: std.Io, path: []const u8) !void { + std.Io.Dir.cwd().access(io, path, .{}) catch { + log.info("GeoLite2 database not found at {s}, downloading", .{path}); + try download(allocator, io, path); + log.info("GeoLite2 database downloaded successfully", .{}); return; }; - log.info("GeoLite2 database refreshed successfully", .{}); } -/// Age of `mtime` in seconds, clamped at 0 so a clock skew into the future -/// cannot read as an enormous age and trigger a pointless download. -fn ageInSeconds(io: std.Io, mtime: std.Io.Timestamp) u64 { - const now_ns = std.Io.Timestamp.now(io, .real).nanoseconds; - const age_ns = now_ns - mtime.nanoseconds; +/// Age of the database at `path`, in seconds, or null if it cannot be stat'd. +pub fn ageInSeconds(io: std.Io, path: []const u8) ?u64 { + const stat = std.Io.Dir.cwd().statFile(io, path, .{}) catch |err| { + log.warn("could not stat GeoLite2 database {s}: {t}", .{ path, err }); + return null; + }; + return ageOf(std.Io.Timestamp.now(io, .real), stat.mtime); +} + +/// Seconds between `mtime` and `now`, clamped at 0. +/// +/// Split out from `ageInSeconds` so the arithmetic is testable without a file. +/// The clamp matters: an mtime in the future (clock skew, or a file copied with +/// its timestamps from a machine that is ahead) would otherwise wrap the +/// unsigned result into a huge age and trigger a pointless download. +pub fn ageOf(now: std.Io.Timestamp, mtime: std.Io.Timestamp) u64 { + const age_ns = now.nanoseconds - mtime.nanoseconds; if (age_ns <= 0) return 0; return @intCast(@divFloor(age_ns, std.time.ns_per_s)); } -fn downloadDatabase(allocator: std.mem.Allocator, io: std.Io, path: []const u8) !void { +/// Downloads the current database and atomically replaces `path`. +/// +/// The download lands on a sibling temp file and is renamed into place, so a +/// download that dies partway cannot leave a truncated database behind, and any +/// process with the old file mapped keeps reading a consistent inode until it +/// reopens. +pub fn download(allocator: std.mem.Allocator, io: std.Io, path: []const u8) !void { const latest_url = try getLatestReleaseUrl(allocator, io); defer allocator.free(latest_url); @@ -73,23 +58,19 @@ fn downloadDatabase(allocator: std.mem.Allocator, io: std.Io, path: []const u8) try std.Io.Dir.cwd().createDirPath(io, dir); } - // Download to a sibling temp file and rename into place. The rename is - // atomic, so a download that dies partway cannot leave a truncated - // database behind, and a database already mmap'd by this process keeps - // pointing at the old inode until it is reopened. const tmp_path = try std.fmt.allocPrint(allocator, "{s}.download", .{path}); defer allocator.free(tmp_path); { const file = try std.Io.Dir.cwd().createFile(io, tmp_path, .{}); - // The temp file is only useful if everything below succeeds. + // A partial download is worse than no download; do not leave it around. errdefer std.Io.Dir.cwd().deleteFile(io, tmp_path) catch |err| log.warn("could not remove partial download {s}: {t}", .{ tmp_path, err }); defer file.close(io); - // Stream the response straight to disk. Buffering the whole database in - // memory needed an allocation sized by guesswork, and the guess (64 MiB) - // was already within ~1 MB of the real database size. + // Stream straight to disk. Buffering the whole database in memory needed + // an allocation sized by guesswork, and the guess (64 MiB) was already + // within about 1 MB of the real database size. var buf: [64 * 1024]u8 = undefined; var file_writer = file.writer(io, &buf); const result = try client.fetch(.{ diff --git a/src/location/Refresher.zig b/src/location/Refresher.zig new file mode 100644 index 0000000..3423c2c --- /dev/null +++ b/src/location/Refresher.zig @@ -0,0 +1,105 @@ +const std = @import("std"); +const GeoLite2 = @import("GeoLite2.zig"); +const GeoIp = @import("GeoIp.zig"); + +const Refresher = @This(); +const log = std.log.scoped(.geolite2_refresh); + +allocator: std.mem.Allocator, +io: std.Io, +geoip: *GeoIp, +db_path: []const u8, +max_age_seconds: u64, +check_interval_seconds: u64, + +/// Background refresh of the GeoLite2 database. +/// +/// Runs off-request on its own unit of concurrency: the download is tens of +/// megabytes, so doing it inside a request handler would stall that request and +/// occupy a worker thread. Doing it at startup would either delay boot or, worse, +/// make a transient network failure fatal. +/// +/// Replacing the file is not enough on its own. The database is mapped with +/// `MMDB_MODE_MMAP`, so this process keeps reading the old inode until it +/// reopens; `GeoIp.reload` performs that reopen under its lock. +pub fn run(self: Refresher) void { + // Check once up front so a database that is already stale at boot is + // replaced promptly rather than after a full interval. + while (true) { + self.checkOnce(); + + std.Io.sleep( + self.io, + .fromNanoseconds(@intCast(self.check_interval_seconds * std.time.ns_per_s)), + .real, + ) catch { + // Cancelled, which is how shutdown stops this loop. + log.debug("refresher stopping", .{}); + return; + }; + } +} + +fn checkOnce(self: Refresher) void { + const age = GeoLite2.ageInSeconds(self.io, self.db_path) orelse return; + if (age < self.max_age_seconds) { + log.debug( + "GeoLite2 database is {d} day(s) old, under the {d} day limit", + .{ @divFloor(age, std.time.s_per_day), @divFloor(self.max_age_seconds, std.time.s_per_day) }, + ); + return; + } + + log.info( + "GeoLite2 database is {d} day(s) old (limit {d} day(s)), refreshing", + .{ @divFloor(age, std.time.s_per_day), @divFloor(self.max_age_seconds, std.time.s_per_day) }, + ); + + GeoLite2.download(self.allocator, self.io, self.db_path) catch |err| { + // A stale database still answers lookups. Keep serving from it. + log.warn( + "GeoLite2 download failed ({t}); continuing with the existing database", + .{err}, + ); + return; + }; + + self.geoip.reload() catch |err| { + // The new file is on disk but this process could not map it. The old + // handle is still installed and usable, so keep going and retry on the + // next interval. + log.err( + "GeoLite2 downloaded but reload failed ({t}); still serving the previous database", + .{err}, + ); + return; + }; + + log.info("GeoLite2 database refreshed", .{}); +} + +test "ageOf clamps a future mtime to zero" { + const now: std.Io.Timestamp = .fromNanoseconds(1_000 * std.time.ns_per_s); + const future: std.Io.Timestamp = .fromNanoseconds(2_000 * std.time.ns_per_s); + try std.testing.expectEqual(@as(u64, 0), GeoLite2.ageOf(now, future)); +} + +test "ageOf reports whole seconds elapsed" { + const mtime: std.Io.Timestamp = .fromNanoseconds(1_000 * std.time.ns_per_s); + const now: std.Io.Timestamp = .fromNanoseconds(1_090 * std.time.ns_per_s); + try std.testing.expectEqual(@as(u64, 90), GeoLite2.ageOf(now, mtime)); +} + +test "ageOf truncates sub-second remainders rather than rounding up" { + const mtime: std.Io.Timestamp = .fromNanoseconds(0); + const now: std.Io.Timestamp = .fromNanoseconds(std.time.ns_per_s - 1); + try std.testing.expectEqual(@as(u64, 0), GeoLite2.ageOf(now, mtime)); +} + +test "ageOf spans days for the staleness comparison" { + const mtime: std.Io.Timestamp = .fromNanoseconds(0); + const now: std.Io.Timestamp = .fromNanoseconds(31 * std.time.s_per_day * std.time.ns_per_s); + const age = GeoLite2.ageOf(now, mtime); + try std.testing.expectEqual(@as(u64, 31), @divFloor(age, std.time.s_per_day)); + try std.testing.expect(age >= 30 * std.time.s_per_day); +} diff --git a/src/main.zig b/src/main.zig index 186699b..e1a2155 100644 --- a/src/main.zig +++ b/src/main.zig @@ -9,6 +9,7 @@ const GeoCache = @import("location/GeoCache.zig"); const Airports = @import("location/Airports.zig"); const Resolver = @import("location/resolver.zig").Resolver; const GeoLite2 = @import("location/GeoLite2.zig"); +const Refresher = @import("location/Refresher.zig"); const version = @import("build_options").version; /// Zig 0.16 entry point: the runtime supplies allocators, the `Io` @@ -36,7 +37,7 @@ pub fn main(init: std.process.Init) !u8 { }; defer metno.deinit(); - // Ensure GeoLite2 database exists + // A missing database is fatal; keeping it fresh is handled below. try GeoLite2.ensureDatabase(allocator, io, cfg.geolite_path); // Initialize GeoIP database with configured fallback @@ -85,6 +86,39 @@ pub fn main(init: std.process.Init) !u8 { // Only set up the server instance in debug mode if (@import("builtin").mode == .Debug) @import("http/handler.zig").server_instance = &server.httpz_server; + // Refresh the GeoLite2 database in the background. `concurrent` rather than + // `async` because this must get a real unit of concurrency: `async` is + // allowed to run the function inline, which would block startup on a + // multi-megabyte download. + var refresher_future = if (cfg.geolite_check_interval_seconds == 0) blk: { + std.log.info("GeoLite2 background refresh disabled", .{}); + break :blk null; + } else blk: { + const refresher: Refresher = .{ + .allocator = allocator, + .io = io, + .geoip = &geoip, + .db_path = cfg.geolite_path, + .max_age_seconds = cfg.geolite_max_age_seconds, + .check_interval_seconds = cfg.geolite_check_interval_seconds, + }; + std.log.info( + "GeoLite2 refresh: checking every {d}h, refreshing when older than {d}d", + .{ + @divFloor(cfg.geolite_check_interval_seconds, std.time.s_per_hour), + @divFloor(cfg.geolite_max_age_seconds, std.time.s_per_day), + }, + ); + break :blk io.concurrent(Refresher.run, .{refresher}) catch |err| { + // Not fatal: without a refresher the database simply ages, which is + // the behavior that existed before. + std.log.warn("could not start GeoLite2 refresher ({t}); database will not auto-refresh", .{err}); + break :blk null; + }; + }; + // `Refresher.run` returns void, so cancelling yields nothing to discard. + defer if (refresher_future) |*f| f.cancel(io); + try server.listen(); std.debug.print("shutting down\n", .{}); return 0; @@ -108,4 +142,5 @@ test { _ = @import("location/Airports.zig"); _ = @import("location/resolver.zig"); _ = @import("location/IpWhoIs.zig"); + _ = @import("location/Refresher.zig"); }