wttr/src/location/GeoCache.zig

419 lines
14 KiB
Zig

const std = @import("std");
const Coordinates = @import("../Coordinates.zig");
const Location = @import("resolver.zig").Location;
const srf = @import("srf");
const GeoCache = @This();
const log = std.log.scoped(.geocache);
allocator: std.mem.Allocator,
/// Zig 0.16 requires an explicit `Io` for filesystem access.
io: std.Io,
cache: std.StringHashMap(CachedLocation),
cache_file: ?[]const u8,
dirty: bool,
last_save: i64,
pub const CachedLocation = struct {
name: []const u8,
coords: Coordinates,
/// ISO 3166-1 alpha-2 country code, upper case, when the geocoder reported
/// one. Cached so that `wttr pin <cidr> <place>` records the country even
/// when the place name is answered from here rather than from Nominatim.
iso_country: ?[2]u8 = null,
};
pub fn init(allocator: std.mem.Allocator, io: std.Io, cache_file: ?[]const u8) !GeoCache {
var cache = std.StringHashMap(CachedLocation).init(allocator);
// Load from file if specified
if (cache_file) |file_path| {
loadFromFile(allocator, io, &cache, file_path) catch |err| {
log.warn("Failed to load geocoding cache from {s}: {}", .{ file_path, err });
};
}
return GeoCache{
.allocator = allocator,
.io = io,
.cache = cache,
.cache_file = if (cache_file) |f| try allocator.dupe(u8, f) else null,
.dirty = false,
.last_save = std.Io.Timestamp.now(io, .real).toMilliseconds(),
};
}
pub fn deinit(self: *GeoCache) void {
// Save to file if specified
if (self.cache_file) |file_path| {
self.saveToFile(file_path) catch |err| {
log.warn("Failed to save geocoding cache to {s}: {}", .{ file_path, err });
};
}
var it = self.cache.iterator();
while (it.next()) |entry| {
self.allocator.free(entry.key_ptr.*);
self.allocator.free(entry.value_ptr.name);
}
self.cache.deinit();
if (self.cache_file) |f| self.allocator.free(f);
}
pub fn get(self: *GeoCache, query: []const u8) ?CachedLocation {
self.saveIfNeeded();
return self.cache.get(query);
}
pub fn put(self: *GeoCache, query: []const u8, location: CachedLocation) !void {
const key = try self.allocator.dupe(u8, query);
const value = CachedLocation{
.name = try self.allocator.dupe(u8, location.name),
.coords = location.coords,
.iso_country = location.iso_country,
};
try self.cache.put(key, value);
self.dirty = true;
}
/// Save cache to disk if dirty and enough time has passed (15 minutes)
pub fn saveIfNeeded(self: *GeoCache) void {
if (!self.dirty) return;
const cache_file = self.cache_file orelse return;
const now = std.Io.Timestamp.now(self.io, .real).toMilliseconds();
const elapsed_ms = now - self.last_save;
const fifteen_minutes_ms = 15 * std.time.ms_per_min;
if (elapsed_ms < fifteen_minutes_ms) return;
self.saveToFile(cache_file) catch |err| {
log.warn("Failed to save geocoding cache to {s}: {}", .{ cache_file, err });
return;
};
self.dirty = false;
self.last_save = now;
}
/// On-disk shape of one geocoded place.
///
/// The query is stored alongside the result rather than used as an object key,
/// because SRF records are flat field lists rather than a nested map.
///
/// `iso` defaults to null so files written before it existed still load: SRF
/// fills a missing field from the Zig default and only errors when there is
/// none. It is also omitted on write when null, so unknown countries add
/// nothing to the file.
const Record = struct {
query: []const u8,
name: []const u8,
lat: f64,
lon: f64,
iso: ?[]const u8 = null,
};
fn load(allocator: std.mem.Allocator, cache: *std.StringHashMap(CachedLocation), content: []const u8) !void {
if (std.mem.trim(u8, content, " \r\n\t").len == 0) return;
var reader = std.Io.Reader.fixed(content);
var records = try srf.iterator(&reader, allocator, .{});
defer records.deinit();
var index: usize = 0;
while (records.next() catch |err| {
log.warn("stopped reading geocache after {d} entr(ies): {t}", .{ index, err });
return;
}) |fields| {
index += 1;
const record = fields.to(Record, .{}) catch continue;
if (record.query.len == 0) continue;
const cache_key = try allocator.dupe(u8, record.query);
errdefer allocator.free(cache_key);
const name_copy = try allocator.dupe(u8, record.name);
errdefer allocator.free(name_copy);
// A duplicate query in the file would otherwise leak the key and name
// already stored under it.
const existing = try cache.fetchPut(cache_key, .{
.name = name_copy,
.coords = .{ .latitude = record.lat, .longitude = record.lon },
.iso_country = if (record.iso) |iso| Location.isoFrom(iso) else null,
});
if (existing) |old| {
allocator.free(old.key);
allocator.free(old.value.name);
}
}
}
fn loadFromFile(allocator: std.mem.Allocator, io: std.Io, cache: *std.StringHashMap(CachedLocation), file_path: []const u8) !void {
const content = try std.Io.Dir.cwd().readFileAlloc(io, file_path, allocator, .limited(10 * 1024 * 1024)); // 10MB max
defer allocator.free(content);
try load(allocator, cache, content);
}
fn save(self: *GeoCache, writer: *std.Io.Writer) !void {
// Rewritten in full rather than appended to, because entries are replaced in
// place; SRF long format keeps the result legible for a file an operator may
// want to inspect or prune.
var records = try self.allocator.alloc(Record, self.cache.count());
defer self.allocator.free(records);
var it = self.cache.iterator();
var i: usize = 0;
while (it.next()) |entry| : (i += 1) {
records[i] = .{
.query = entry.key_ptr.*,
.name = entry.value_ptr.name,
.lat = entry.value_ptr.coords.latitude,
.lon = entry.value_ptr.coords.longitude,
// Borrowed from the map rather than copied: the map is not touched
// between here and the `print` below, and a by-value capture would
// leave the slice pointing at a dead loop temporary.
.iso = if (entry.value_ptr.iso_country) |*iso| iso[0..] else null,
};
}
try writer.print("{f}", .{srf.fmt(Record, records, .{ .long_format = true })});
}
fn saveToFile(self: *GeoCache, file_path: []const u8) !void {
const file = try std.Io.Dir.cwd().createFile(self.io, file_path, .{});
defer file.close(self.io);
var buffer: [4096]u8 = undefined;
var file_writer = file.writer(self.io, &buffer);
const writer = &file_writer.interface;
try self.save(writer);
try writer.flush();
}
test "GeoCache basic operations" {
const allocator = std.testing.allocator;
var cache = try GeoCache.init(allocator, std.testing.io, null);
defer cache.deinit();
// Test put and get
try cache.put("London", .{
.name = "London, UK",
.coords = .{
.latitude = 51.5074,
.longitude = -0.1278,
},
});
const result = cache.get("London");
try std.testing.expect(result != null);
try std.testing.expectApproxEqAbs(@as(f64, 51.5074), result.?.coords.latitude, 0.0001);
try std.testing.expectApproxEqAbs(@as(f64, -0.1278), result.?.coords.longitude, 0.0001);
}
test "GeoCache miss returns null" {
const allocator = std.testing.allocator;
var cache = try GeoCache.init(allocator, std.testing.io, null);
defer cache.deinit();
const result = cache.get("NonExistent");
try std.testing.expect(result == null);
}
test "save produces SRF an operator can read" {
const allocator = std.testing.allocator;
var cache = try GeoCache.init(allocator, std.testing.io, null);
defer cache.deinit();
try cache.put("London", .{
.name = "London, UK",
.coords = .{ .latitude = 51.5074, .longitude = -0.1278 },
});
try cache.put("Paris", .{
.name = "Paris, France",
.coords = .{ .latitude = 48.8566, .longitude = 2.3522 },
});
var buffer: [1024]u8 = undefined;
var writer = std.Io.Writer.fixed(&buffer);
try cache.save(&writer);
const output = buffer[0..writer.end];
// Assert the shape, not just that the values appear: a silent switch to
// compact output would still round-trip while making the file harder to read.
try std.testing.expect(std.mem.indexOf(u8, output, "#!srfv1") != null);
try std.testing.expect(std.mem.indexOf(u8, output, "#!long") != null);
try std.testing.expect(std.mem.indexOf(u8, output, "query::London") != null);
try std.testing.expect(std.mem.indexOf(u8, output, "query::Paris") != null);
try std.testing.expect(std.mem.indexOf(u8, output, "51.5074") != null);
}
test "load parses SRF records" {
const allocator = std.testing.allocator;
var cache_map = std.StringHashMap(CachedLocation).init(allocator);
defer {
var it = cache_map.iterator();
while (it.next()) |entry| {
allocator.free(entry.key_ptr.*);
allocator.free(entry.value_ptr.name);
}
cache_map.deinit();
}
const content =
\\#!srfv1
\\#!long
\\query::London
\\name::London, UK
\\lat:num:51.5074
\\lon:num:-0.1278
\\
\\query::Paris
\\name::Paris, France
\\lat:num:48.8566
\\lon:num:2.3522
\\
;
try load(allocator, &cache_map, content);
try std.testing.expectEqual(@as(usize, 2), cache_map.count());
const london = cache_map.get("London") orelse return error.TestUnexpectedResult;
try std.testing.expectEqualStrings("London, UK", london.name);
try std.testing.expectApproxEqAbs(@as(f64, 51.5074), london.coords.latitude, 0.0001);
const paris = cache_map.get("Paris") orelse return error.TestUnexpectedResult;
try std.testing.expectApproxEqAbs(@as(f64, 2.3522), paris.coords.longitude, 0.0001);
}
test "load ignores a legacy JSON file rather than failing" {
const allocator = std.testing.allocator;
var cache_map = std.StringHashMap(CachedLocation).init(allocator);
defer {
var it = cache_map.iterator();
while (it.next()) |entry| {
allocator.free(entry.key_ptr.*);
allocator.free(entry.value_ptr.name);
}
cache_map.deinit();
}
// The previous format. Entries are re-geocodable, so the file is dropped
// rather than migrated; what matters is that this is not fatal.
const legacy =
\\{
\\ "London": {"name": "London, UK", "latitude": 51.5074, "longitude": -0.1278}
\\}
;
load(allocator, &cache_map, legacy) catch {};
try std.testing.expectEqual(@as(usize, 0), cache_map.count());
}
test "save and load round-trip" {
const allocator = std.testing.allocator;
var cache1 = try GeoCache.init(allocator, std.testing.io, null);
defer cache1.deinit();
try cache1.put("Berlin", .{
.name = "Berlin, Germany",
.coords = .{ .latitude = 52.5200, .longitude = 13.4050 },
});
var buffer: [1024]u8 = undefined;
var writer = std.Io.Writer.fixed(&buffer);
try cache1.save(&writer);
var cache2 = std.StringHashMap(CachedLocation).init(allocator);
defer {
var it = cache2.iterator();
while (it.next()) |entry| {
allocator.free(entry.key_ptr.*);
allocator.free(entry.value_ptr.name);
}
cache2.deinit();
}
try load(allocator, &cache2, buffer[0..writer.end]);
const berlin = cache2.get("Berlin");
try std.testing.expect(berlin != null);
try std.testing.expectEqualStrings("Berlin, Germany", berlin.?.name);
try std.testing.expectApproxEqAbs(@as(f64, 52.5200), berlin.?.coords.latitude, 0.0001);
try std.testing.expectApproxEqAbs(@as(f64, 13.4050), berlin.?.coords.longitude, 0.0001);
}
test "save and load round-trip the country" {
const allocator = std.testing.allocator;
var cache1 = try GeoCache.init(allocator, std.testing.io, null);
defer cache1.deinit();
try cache1.put("Seattle,+Washington", .{
.name = "Seattle, King County, Washington, United States",
.coords = .{ .latitude = 47.6038321, .longitude = -122.330062 },
.iso_country = Location.isoFrom("us"),
});
// Kept as unknown rather than acquiring a country on the round trip.
try cache1.put("Nowhere", .{
.name = "Nowhere",
.coords = .{ .latitude = 1, .longitude = 2 },
});
var buffer: [2048]u8 = undefined;
var writer = std.Io.Writer.fixed(&buffer);
try cache1.save(&writer);
var cache2 = std.StringHashMap(CachedLocation).init(allocator);
defer {
var it = cache2.iterator();
while (it.next()) |entry| {
allocator.free(entry.key_ptr.*);
allocator.free(entry.value_ptr.name);
}
cache2.deinit();
}
try load(allocator, &cache2, buffer[0..writer.end]);
const seattle = cache2.get("Seattle,+Washington") orelse return error.TestUnexpectedResult;
// Normalized on the way in, so it reads back upper case even though
// Nominatim reports "us".
try std.testing.expectEqualStrings("US", &(seattle.iso_country.?));
const nowhere = cache2.get("Nowhere") orelse return error.TestUnexpectedResult;
try std.testing.expect(nowhere.iso_country == null);
}
test "load: a geocache written before the country field still loads" {
const allocator = std.testing.allocator;
var cache_map = std.StringHashMap(CachedLocation).init(allocator);
defer {
var it = cache_map.iterator();
while (it.next()) |entry| {
allocator.free(entry.key_ptr.*);
allocator.free(entry.value_ptr.name);
}
cache_map.deinit();
}
// The four-field shape. These entries are re-geocodable, but discarding them
// would mean a burst of Nominatim traffic on the first restart after upgrade.
const content =
\\#!srfv1
\\#!long
\\query::London
\\name::London, UK
\\lat:num:51.5074
\\lon:num:-0.1278
\\
;
try load(allocator, &cache_map, content);
try std.testing.expectEqual(@as(usize, 1), cache_map.count());
const london = cache_map.get("London") orelse return error.TestUnexpectedResult;
try std.testing.expectEqualStrings("London, UK", london.name);
try std.testing.expect(london.iso_country == null);
}