automatically refresh the geoip db

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| {

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);
+    const mmdb = try openDatabase(allocator, db_path);
-    defer std.heap.c_allocator.free(path_z);
+    errdefer {
+        c.MMDB_close(mmdb);
+        allocator.destroy(mmdb);
+    }
 
-    const mmdb = try allocator.create(c.MMDB_s);
+    const db_path_copy = try allocator.dupe(u8, db_path);
-    errdefer allocator.destroy(mmdb);
+    errdefer allocator.free(db_path_copy);
-
-    const status = c.MMDB_open(path_z.ptr, c.MMDB_MODE_MMAP, mmdb);
-    if (status != c.MMDB_SUCCESS)
-        return error.CannotOpenDatabase;
 
     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
+    // Try MaxMind first. The shared lock has to cover `extractCoordinates` as
-    const result = lookupInternal(self.mmdb, ip) catch return null;
+    // 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 });
+        const result = lookupInternal(self.mmdb, ip) catch break :blk null;
-    if (result.found_entry)
-        if (self.extractCoordinates(ip, result)) |coords|
-            return coords;
 
-    // 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));
+}

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
+/// Ensures a GeoLite2 database exists at `path`, downloading it if absent.
-/// and refreshing it once it is older than `max_age_seconds`.
 ///
-/// Refresh matters: the upstream release is rebuilt continuously and IP
+/// A missing database is fatal: `GeoIp.init` cannot open anything, so there is
-/// allocations move between cities, so a stale database silently resolves
+/// nothing to serve lookups with. Keeping the database *fresh* is a separate
-/// clients to the wrong place. Before this, the database was only fetched when
+/// concern handled off-request by `Refresher`, because a stale database still
-/// absent, so a deployment could -- and did -- run for months on the copy it
+/// answers lookups and a network problem must not stop the service from
-/// happened to download first.
+/// starting.
-///
+pub fn ensureDatabase(allocator: std.mem.Allocator, io: std.Io, path: []const u8) !void {
-/// A missing database is fatal (there is nothing to fall back on). A failed
+    std.Io.Dir.cwd().access(io, path, .{}) catch {
-/// *refresh* is not: an out-of-date database still answers lookups, so a
+        log.info("GeoLite2 database not found at {s}, downloading", .{path});
-/// network blip must not stop the service from starting.
+        try download(allocator, io, path);
-pub fn ensureDatabase(
+        log.info("GeoLite2 database downloaded successfully", .{});
-    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 },
-        );
         return;
     };
-    log.info("GeoLite2 database refreshed successfully", .{});
 }
 
-/// Age of `mtime` in seconds, clamped at 0 so a clock skew into the future
+/// Age of the database at `path`, in seconds, or null if it cannot be stat'd.
-/// cannot read as an enormous age and trigger a pointless download.
+pub fn ageInSeconds(io: std.Io, path: []const u8) ?u64 {
-fn ageInSeconds(io: std.Io, mtime: std.Io.Timestamp) u64 {
+    const stat = std.Io.Dir.cwd().statFile(io, path, .{}) catch |err| {
-    const now_ns = std.Io.Timestamp.now(io, .real).nanoseconds;
+        log.warn("could not stat GeoLite2 database {s}: {t}", .{ path, err });
-    const age_ns = now_ns - mtime.nanoseconds;
+        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
+        // Stream straight to disk. Buffering the whole database in memory needed
-        // memory needed an allocation sized by guesswork, and the guess (64 MiB)
+        // an allocation sized by guesswork, and the guess (64 MiB) was already
-        // was already within ~1 MB of the real database size.
+        // 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(.{

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);
+}

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");
 }