put more cache under mutex

src/cache/Cache.zig

@@ -12,6 +12,11 @@ io: std.Io,
 lru: Lru,
 /// Cache directory for L2 persistent cache
 cache_dir: ?[]const u8,
+/// Request handlers share one cache across httpz's thread pool. Exclusive
+/// rather than read/write because even reads mutate: `Lru` bumps access counts
+/// and drops expired entries, and an L2 hit promotes into L1. Held across L2
+/// file access too, so a reader never sees a file another thread is rewriting.
+mutex: std.Io.Mutex = .init,
 
 pub const Config = struct {
     max_entries: usize = 1_000,
@@ -25,6 +30,7 @@ pub const Config = struct {
 /// An entry returned by `getStale`. `expires` lets the caller say how old the
 /// data it is serving is.
 pub const Entry = struct {
+    /// Owned by the caller of `getStale`.
     value: []const u8,
     /// Expiry time, in milliseconds since the epoch. In the past for stale data.
     expires: i64,
@@ -64,21 +70,31 @@ fn evictCallback(ctx: *anyopaque, key: []const u8) void {
     self.deleteFile(key);
 }
 
-/// Returns the value only if it has not expired.
-pub fn get(self: *Cache, key: []const u8) ?[]const u8 {
+/// Returns a copy of the value, owned by the caller, only if it has not
+/// expired. A copy because the cached one can be replaced or evicted by
+/// another thread the moment the lock is released.
+pub fn get(self: *Cache, allocator: std.mem.Allocator, key: []const u8) !?[]const u8 {
+    self.mutex.lockUncancelable(self.io);
+    defer self.mutex.unlock(self.io);
+
     const now = std.Io.Timestamp.now(self.io, .real).toMilliseconds();
     const entry = self.lookup(key, now) orelse return null;
     if (now > entry.expires) return null;
-    return entry.value;
+    return try allocator.dupe(u8, entry.value);
 }
 
-/// Returns the value whether fresh or expired-but-retained. Meant as a
-/// fallback when refreshing an expired entry has failed.
-pub fn getStale(self: *Cache, key: []const u8) ?Entry {
+/// Returns the value whether fresh or expired-but-retained, copied as in
+/// `get`. Meant as a fallback when refreshing an expired entry has failed.
+pub fn getStale(self: *Cache, allocator: std.mem.Allocator, key: []const u8) !?Entry {
+    self.mutex.lockUncancelable(self.io);
+    defer self.mutex.unlock(self.io);
+
     const now = std.Io.Timestamp.now(self.io, .real).toMilliseconds();
-    return self.lookup(key, now);
+    const entry = self.lookup(key, now) orelse return null;
+    return .{ .value = try allocator.dupe(u8, entry.value), .expires = entry.expires };
 }
 
+/// Caller must hold `mutex`. The returned value is borrowed from L1.
 fn lookup(self: *Cache, key: []const u8, now_ms: i64) ?Entry {
     // Check L1 (memory) first. A retained-but-expired L1 entry is the newest
     // copy there is, so it also short-circuits the disk read.
@@ -97,6 +113,9 @@ fn lookup(self: *Cache, key: []const u8, now_ms: i64) ?Entry {
 }
 
 pub fn put(self: *Cache, key: []const u8, value: []const u8, ttl_seconds: u64) !void {
+    self.mutex.lockUncancelable(self.io);
+    defer self.mutex.unlock(self.io);
+
     const now = std.Io.Timestamp.now(self.io, .real).toMilliseconds();
     const expires = now + @as(i64, @intCast(ttl_seconds * 1000));
 
@@ -382,12 +401,16 @@ test "L1/L2 cache flow" {
 
     const cache = try Cache.init(allocator, std.testing.io, .{ .max_entries = 10, .cache_dir = cache_dir });
     defer cache.deinit();
+    // Values from get/getStale are copies the caller owns
+    var arena: std.heap.ArenaAllocator = .init(allocator);
+    defer arena.deinit();
+    const copies = arena.allocator();
 
     // Put item in cache
     try cache.put("key1", "value1", 900);
 
     // Should be in L1
-    try std.testing.expectEqualStrings("value1", cache.get("key1").?);
+    try std.testing.expectEqualStrings("value1", (try cache.get(copies, "key1")).?);
 
     // Manually remove from L1 to simulate eviction
     if (cache.lru.map.fetchRemove("key1")) |kv| {
@@ -399,7 +422,7 @@ test "L1/L2 cache flow" {
     try std.testing.expect(cache.lru.get("key1") == null);
 
     // Get should promote from L2 to L1
-    try std.testing.expectEqualStrings("value1", cache.get("key1").?);
+    try std.testing.expectEqualStrings("value1", (try cache.get(copies, "key1")).?);
 
     // Now it should be in L1
     try std.testing.expectEqualStrings("value1", cache.lru.get("key1").?);
@@ -421,28 +444,77 @@ test "stale entries are retained for getStale but never returned by get" {
         .stale_retention_s = 3600,
     });
     defer cache.deinit();
+    var arena: std.heap.ArenaAllocator = .init(allocator);
+    defer arena.deinit();
+    const copies = arena.allocator();
 
     const now = std.Io.Timestamp.now(std.testing.io, .real).toMilliseconds();
 
     // Expired a minute ago, still inside the hour of retention. Written only to
     // L2 so the disk path is what gets exercised, as it is after a restart.
     try cache.saveToFile("recent", "old-but-usable", now - 60_000);
-    try std.testing.expect(cache.get("recent") == null);
-    const stale = cache.getStale("recent").?;
+    try std.testing.expect(try cache.get(copies, "recent") == null);
+    const stale = (try cache.getStale(copies, "recent")).?;
     try std.testing.expectEqualStrings("old-but-usable", stale.value);
     try std.testing.expectEqual(now - 60_000, stale.expires);
     // Now promoted to L1, and still not fresh from there either
-    try std.testing.expect(cache.get("recent") == null);
-    try std.testing.expectEqualStrings("old-but-usable", cache.getStale("recent").?.value);
+    try std.testing.expect(try cache.get(copies, "recent") == null);
+    try std.testing.expectEqualStrings("old-but-usable", (try cache.getStale(copies, "recent")).?.value);
 
     // Expired two hours ago: past retention, unusable even as stale, and the
     // file is cleaned up.
     try cache.saveToFile("ancient", "too-old", now - 2 * 3600 * 1000);
-    try std.testing.expect(cache.getStale("ancient") == null);
+    try std.testing.expect(try cache.getStale(copies, "ancient") == null);
     try std.testing.expectError(error.FileNotFound, cache.loadFromFile("ancient"));
 
     // Fresh entries come back from both
     try cache.put("fresh", "new", 900);
-    try std.testing.expectEqualStrings("new", cache.get("fresh").?);
-    try std.testing.expectEqualStrings("new", cache.getStale("fresh").?.value);
+    try std.testing.expectEqualStrings("new", (try cache.get(copies, "fresh")).?);
+    try std.testing.expectEqualStrings("new", (try cache.getStale(copies, "fresh")).?.value);
+}
+
+test "returned values survive the entry being replaced" {
+    const allocator = std.testing.allocator;
+    const cache = try Cache.init(allocator, std.testing.io, .{ .max_entries = 10, .cache_dir = null });
+    defer cache.deinit();
+
+    try cache.put("k", "first", 900);
+    const value = (try cache.get(allocator, "k")).?;
+    defer allocator.free(value);
+
+    // Before values were copied, this freed the memory `value` pointed at
+    try cache.put("k", "second", 900);
+    try std.testing.expectEqualStrings("first", value);
+}
+
+test "concurrent puts and gets on shared keys" {
+    // Many threads hammering a handful of keys through a cache small enough to
+    // keep evicting. Without the lock this corrupts the map or frees values out
+    // from under readers; under the debug allocator that crashes.
+    const allocator = std.testing.allocator;
+    const cache = try Cache.init(allocator, std.testing.io, .{ .max_entries = 4, .cache_dir = null });
+    defer cache.deinit();
+
+    const Worker = struct {
+        fn run(c: *Cache, id: usize) void {
+            var key_buf: [16]u8 = undefined;
+            var value_buf: [32]u8 = undefined;
+            for (0..2_000) |i| {
+                const key = std.fmt.bufPrint(&key_buf, "k{d}", .{i % 8}) catch unreachable;
+                const value = std.fmt.bufPrint(&value_buf, "t{d}-i{d}", .{ id, i }) catch unreachable;
+                c.put(key, value, 900) catch @panic("put failed");
+                if (c.get(std.testing.allocator, key) catch @panic("get failed")) |got| {
+                    // Whatever thread wrote it, it must be an intact value
+                    std.debug.assert(got[0] == 't');
+                    std.testing.allocator.free(got);
+                }
+            }
+        }
+    };
+
+    var threads: [8]std.Thread = undefined;
+    for (&threads, 0..) |*t, id| t.* = try std.Thread.spawn(.{}, Worker.run, .{ cache, id });
+    for (threads) |t| t.join();
+
+    try std.testing.expect(cache.lru.map.count() <= 4);
 }

src/cache/Lru.zig

@@ -1,3 +1,11 @@
+//! Size-capped map with per-entry expiry.
+//!
+//! Not thread-safe, and every operation mutates (even `get` bumps an access
+//! count and may drop an expired entry), so owners shared between threads must
+//! serialize all access with an exclusive lock. Returned slices point into
+//! entries the next `put`, eviction, or expiry frees, so copy them before
+//! releasing that lock.
+
 const std = @import("std");
 
 const Lru = @This();
@@ -57,9 +65,11 @@ pub fn lookup(self: *Lru, key: []const u8, now_ms: i64) ?Iterator.Item {
 }
 
 pub fn put(self: *Lru, key: []const u8, value: []const u8, expires: i64) !void {
-    if (self.map.get(key)) |old_entry| {
-        self.allocator.free(old_entry.value);
-        _ = self.map.remove(key);
+    // Replacing, not evicting: the callback is skipped, since the owner may
+    // already have persisted the new value under the same key.
+    if (self.map.fetchRemove(key)) |old| {
+        self.allocator.free(old.value.value);
+        self.allocator.free(old.key);
     }
 
     if (self.map.count() >= self.max_entries) {
@@ -67,7 +77,9 @@ pub fn put(self: *Lru, key: []const u8, value: []const u8, expires: i64) !void {
     }
 
     const key_copy = try self.allocator.dupe(u8, key);
+    errdefer self.allocator.free(key_copy);
     const value_copy = try self.allocator.dupe(u8, value);
+    errdefer self.allocator.free(value_copy);
 
     try self.map.put(key_copy, .{
         .value = value_copy,
@@ -191,3 +203,14 @@ test "LRU retains expired entries for lookup within the retention window" {
     try std.testing.expect(lru.lookup("ancient", now) == null);
     try std.testing.expect(!lru.map.contains("ancient"));
 }
+
+test "LRU replacing a key frees the old key and value" {
+    // std.testing.allocator fails the test on a leak, which is the assertion
+    var lru = try Lru.init(std.testing.allocator, std.testing.io, 10);
+    defer lru.deinit();
+
+    try lru.put("key1", "value1", 9999999999999);
+    try lru.put("key1", "value2", 9999999999999);
+    try std.testing.expectEqualStrings("value2", lru.get("key1").?);
+    try std.testing.expectEqual(@as(u32, 1), lru.map.count());
+}

src/location/GeoCache.zig

@@ -19,6 +19,10 @@ last_save: i64,
 /// in memory only with a cap and a TTL, because real places are a bounded set
 /// but nonsense queries are unbounded and chosen by whoever is scanning us.
 not_found: Lru,
+/// Request handlers share one GeoCache across httpz's thread pool. Exclusive
+/// rather than read/write because `get` may rewrite the file and `not_found`
+/// lookups mutate (see `Lru`).
+mutex: std.Io.Mutex = .init,
 
 const not_found_max_entries = 1_000;
 const not_found_ttl_ms = 24 * std.time.ms_per_hour;
@@ -71,19 +75,46 @@ pub fn deinit(self: *GeoCache) void {
     if (self.cache_file) |f| self.allocator.free(f);
 }
 
-pub fn get(self: *GeoCache, query: []const u8) ?CachedLocation {
+/// Returns a copy of the cached location, whose `name` the caller owns and
+/// must free with `allocator`. A copy because another thread's `put` can free
+/// the cached name the moment the lock is released.
+pub fn get(self: *GeoCache, allocator: std.mem.Allocator, query: []const u8) !?CachedLocation {
+    self.mutex.lockUncancelable(self.io);
+    defer self.mutex.unlock(self.io);
+
     self.saveIfNeeded();
-    return self.cache.get(query);
+    const cached = self.cache.get(query) orelse return null;
+    return .{
+        .name = try allocator.dupe(u8, cached.name),
+        .coords = cached.coords,
+        .iso_country = cached.iso_country,
+    };
 }
 
 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),
+    self.mutex.lockUncancelable(self.io);
+    defer self.mutex.unlock(self.io);
+
+    const name = try self.allocator.dupe(u8, location.name);
+    errdefer self.allocator.free(name);
+
+    const gop = try self.cache.getOrPut(query);
+    if (gop.found_existing) {
+        // Keep the stored key; replacing the value used to leak both the old
+        // name and a fresh copy of the key
+        self.allocator.free(gop.value_ptr.name);
+    } else {
+        // `getOrPut` stored the caller's slice; swap in one we own
+        gop.key_ptr.* = self.allocator.dupe(u8, query) catch |err| {
+            self.cache.removeByPtr(gop.key_ptr);
+            return err;
+        };
+    }
+    gop.value_ptr.* = .{
+        .name = name,
         .coords = location.coords,
         .iso_country = location.iso_country,
     };
-    try self.cache.put(key, value);
     self.dirty = true;
 }
 
@@ -91,17 +122,24 @@ pub fn put(self: *GeoCache, query: []const u8, location: CachedLocation) !void {
 /// "no such place": a geocoder that failed or was unreachable says nothing
 /// about the query, and caching that would hide real places.
 pub fn putNotFound(self: *GeoCache, query: []const u8) !void {
+    self.mutex.lockUncancelable(self.io);
+    defer self.mutex.unlock(self.io);
+
     const now = std.Io.Timestamp.now(self.io, .real).toMilliseconds();
     try self.not_found.put(query, "", now + not_found_ttl_ms);
 }
 
 /// Whether the geocoder recently had no result for `query`.
 pub fn isKnownNotFound(self: *GeoCache, query: []const u8) bool {
+    self.mutex.lockUncancelable(self.io);
+    defer self.mutex.unlock(self.io);
+
     return self.not_found.get(query) != null;
 }
 
-/// Save cache to disk if dirty and enough time has passed (15 minutes)
-pub fn saveIfNeeded(self: *GeoCache) void {
+/// Save cache to disk if dirty and enough time has passed (15 minutes).
+/// Caller must hold `mutex`.
+fn saveIfNeeded(self: *GeoCache) void {
     if (!self.dirty) return;
 
     const cache_file = self.cache_file orelse return;
@@ -231,10 +269,31 @@ test "GeoCache basic operations" {
         },
     });
 
-    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);
+    const result = (try cache.get(allocator, "London")) orelse return error.TestUnexpectedResult;
+    defer allocator.free(result.name);
+    try std.testing.expectEqualStrings("London, UK", result.name);
+    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 replacing an entry frees the old one and keeps returned copies valid" {
+    // std.testing.allocator fails the test on a leak, which is half the assertion
+    const allocator = std.testing.allocator;
+    var cache = try GeoCache.init(allocator, std.testing.io, null);
+    defer cache.deinit();
+
+    try cache.put("Portland", .{ .name = "Portland, Maine", .coords = .{ .latitude = 43.66, .longitude = -70.26 } });
+    const first = (try cache.get(allocator, "Portland")).?;
+    defer allocator.free(first.name);
+
+    try cache.put("Portland", .{ .name = "Portland, Oregon", .coords = .{ .latitude = 45.52, .longitude = -122.68 } });
+    const second = (try cache.get(allocator, "Portland")).?;
+    defer allocator.free(second.name);
+
+    try std.testing.expectEqual(@as(u32, 1), cache.cache.count());
+    try std.testing.expectEqualStrings("Portland, Oregon", second.name);
+    // The first copy is unaffected by the entry it came from being replaced
+    try std.testing.expectEqualStrings("Portland, Maine", first.name);
 }
 
 test "GeoCache miss returns null" {
@@ -242,8 +301,7 @@ test "GeoCache miss returns null" {
     var cache = try GeoCache.init(allocator, std.testing.io, null);
     defer cache.deinit();
 
-    const result = cache.get("NonExistent");
-    try std.testing.expect(result == null);
+    try std.testing.expect(try cache.get(allocator, "NonExistent") == null);
 }
 
 test "GeoCache remembers queries with no result, separately and for a limited time" {
@@ -257,7 +315,7 @@ test "GeoCache remembers queries with no result, separately and for a limited ti
     try std.testing.expect(!cache.isKnownNotFound("Somewhereville"));
 
     // Not a place, so not in the positive cache and never written to disk
-    try std.testing.expect(cache.get("Nowhereville") == null);
+    try std.testing.expect(try cache.get(allocator, "Nowhereville") == null);
     try std.testing.expectEqual(@as(u32, 0), cache.cache.count());
     try std.testing.expect(!cache.dirty);
 

src/location/Ip2location.zig

@@ -194,6 +194,10 @@ pub const Cache = struct {
     path: []const u8,
     entries: std.AutoHashMap(u128, Location),
     file: ?std.Io.File,
+    /// Request handlers reach this from httpz's thread pool, outside GeoIp's
+    /// lock. Guards `entries` and the appends to `file`, where two unguarded
+    /// writers would read the same end offset and overwrite each other.
+    mutex: std.Io.Mutex = .init,
 
     /// On-disk shape of one entry.
     ///
@@ -308,6 +312,9 @@ pub const Cache = struct {
     }
 
     pub fn get(self: *Cache, ip: u128) ?Location {
+        self.mutex.lockUncancelable(self.io);
+        defer self.mutex.unlock(self.io);
+
         const entry = self.entries.get(ip) orelse return null;
         return .{
             .allocator = self.allocator,
@@ -322,11 +329,17 @@ pub const Cache = struct {
     /// Unit selection only needs the country, and doing it this way keeps that
     /// path free of both an allocation and any chance of a network request.
     pub fn getCountry(self: *Cache, ip: u128) ?[2]u8 {
+        self.mutex.lockUncancelable(self.io);
+        defer self.mutex.unlock(self.io);
+
         const entry = self.entries.getPtr(ip) orelse return null;
         return entry.iso_country;
     }
 
     pub fn put(self: *Cache, ip: u128, family: u8, loc: Location) !void {
+        self.mutex.lockUncancelable(self.io);
+        defer self.mutex.unlock(self.io);
+
         const name_copy = try self.allocator.dupe(u8, loc.name);
         errdefer self.allocator.free(name_copy);
 

src/location/resolver.zig

@@ -183,10 +183,10 @@ pub const Resolver = struct {
         // Don't bother passing this through
         if (std.mem.eql(u8, name, "favicon.ico")) return error.GeocodingFailed;
         // Check cache first
-        if (self.geocache.get(name)) |cached| {
+        if (try self.geocache.get(self.allocator, name)) |cached| {
             return Location{
                 .allocator = self.allocator,
-                .name = try self.allocator.dupe(u8, cached.name),
+                .name = cached.name,
                 .coords = cached.coords,
                 .iso_country = cached.iso_country,
             };

src/weather/Provider.zig

@@ -35,13 +35,16 @@ pub fn fetch(self: WeatherProvider, allocator: std.mem.Allocator, coords: Coordi
     const cache_key = try std.fmt.bufPrint(&buf, "{d:.4},{d:.4}", .{ coords.latitude, coords.longitude });
 
     // Check cache first
-    if (self.cache.get(cache_key)) |cached_raw|
+    if (try self.cache.get(allocator, cache_key)) |cached_raw| {
+        defer allocator.free(cached_raw);
         return self.vtable.parse(self.ptr, allocator, cached_raw);
+    }
 
     // Cache miss - fetch raw data from provider, falling back to an expired
     // entry if the upstream is down
     const raw = self.vtable.fetchRaw(self.ptr, allocator, coords) catch |err| {
-        if (self.cache.getStale(cache_key)) |stale| {
+        if (try self.cache.getStale(allocator, cache_key)) |stale| {
+            defer allocator.free(stale.value);
             const now_ms = std.Io.Timestamp.now(self.cache.io, .real).toMilliseconds();
             const stale_for_s = @divTrunc(now_ms - stale.expires, std.time.ms_per_s);
             log.warn("weather fetch for {s} failed ({}); serving cached data that expired {d}s ago", .{