diff --git a/src/maxminddb.zig b/src/maxminddb.zig index f66ca76..8e125a3 100644 --- a/src/maxminddb.zig +++ b/src/maxminddb.zig @@ -796,3 +796,19 @@ test "lookup with custom record" { try expectEqual(2694762, got.value.city.geoname_id); try expectEqualStrings("Linköping", got.value.city.names.en); } + +test "within returns all networks" { + var db = try Reader.mmap( + allocator, + "test-data/test-data/GeoLite2-City-Test.mmdb", + ); + defer db.unmap(); + + var it = try db.within(allocator, geolite2.City, net.Network.all_ipv6, .{}); + defer it.deinit(); + + var n: usize = 0; + while (try it.next()) |_| : (n += 1) {} + + try expectEqual(242, n); +} diff --git a/src/reader.zig b/src/reader.zig index a98a766..c473548 100644 --- a/src/reader.zig +++ b/src/reader.zig @@ -214,7 +214,7 @@ pub const Reader = struct { .node_count = node_count, .stack = stack, .allocator = allocator, - .arena = std.heap.ArenaAllocator.init(allocator), + .cache = .{}, .fields = options.only, }; } @@ -381,8 +381,74 @@ pub fn Iterator(T: type) type { node_count: usize, stack: std.ArrayList(WithinNode), allocator: std.mem.Allocator, - arena: std.heap.ArenaAllocator, fields: ?decoder.Fields, + cache: Cache, + + // Ring buffer cache of recently decoded records. + // Many adjacent networks in the tree share the same data pointer, + // so caching avoids re-decoding the same record repeatedly. + // Once full, new entries overwrite the oldest slot in a circular fashion. + // Each entry owns an arena that backs the decoded value's allocations; + // the arena is freed on eviction. + const Cache = struct { + const Entry = struct { + pointer: usize, + value: T, + arena: std.heap.ArenaAllocator, + }; + + // 16 showed a good tradeoff in DuckDB table scan, + // see https://github.com/marselester/duckdb-maxmind. + const cache_size = 16; + entries: [cache_size]Entry = undefined, + // Indicates number of entries in the cache. + len: usize = 0, + // It's an index in the entries array where a new item will be written at. + write_pos: usize = 0, + + fn lookup(self: *Cache, pointer: usize) ?T { + for (self.entries[0..self.len]) |e| { + if (e.pointer == pointer) { + return e.value; + } + } + + return null; + } + + fn insert( + self: *Cache, + pointer: usize, + value: T, + arena: std.heap.ArenaAllocator, + ) void { + if (self.len < cache_size) { + self.entries[self.len] = .{ + .pointer = pointer, + .value = value, + .arena = arena, + }; + self.len += 1; + + return; + } + + // Evict oldest entry. + self.entries[self.write_pos].arena.deinit(); + self.entries[self.write_pos] = .{ + .pointer = pointer, + .value = value, + .arena = arena, + }; + self.write_pos = (self.write_pos + 1) % cache_size; + } + + fn deinit(self: *Cache) void { + for (self.entries[0..self.len]) |*e| { + e.arena.deinit(); + } + } + }; const Self = @This(); @@ -392,7 +458,7 @@ pub fn Iterator(T: type) type { }; /// Returns the next network and its value. - /// The iterator owns the value; each call invalidates the previous Item. + /// The iterator owns the value; each call eventually invalidates the previous Item. pub fn next(self: *Self) !?Item { while (self.stack.pop()) |current| { const reader = self.reader; @@ -411,15 +477,27 @@ pub fn Iterator(T: type) type { if (current.node > self.node_count) { const ip_net = current.ip_bytes.network(current.prefix_len); - _ = self.arena.reset(.retain_capacity); + // Check the ring buffer cache. + // Recently decoded records are reused. + if (self.cache.lookup(current.node)) |cached_value| { + return Item{ + .network = ip_net, + .value = cached_value, + }; + } + + var entry_arena = std.heap.ArenaAllocator.init(self.allocator); + errdefer entry_arena.deinit(); const value = try reader.resolveDataPointerAndDecode( - self.arena.allocator(), + entry_arena.allocator(), T, current.node, self.fields, ); + self.cache.insert(current.node, value, entry_arena); + return Item{ .network = ip_net, .value = value, @@ -457,7 +535,7 @@ pub fn Iterator(T: type) type { } pub fn deinit(self: *Self) void { - self.arena.deinit(); + self.cache.deinit(); self.stack.deinit(self.allocator); } };