diff --git a/examples/lookup.zig b/examples/lookup.zig index c34abc5..763e882 100644 --- a/examples/lookup.zig +++ b/examples/lookup.zig @@ -11,10 +11,12 @@ pub fn main() !void { defer _ = gpa.detectLeaks(); var db = try maxminddb.Reader.open(allocator, db_path, max_db_size); - defer db.close(); + defer db.close(allocator); + // Note, for better performance use arena allocator and reset it after calling lookup(). + // You won't need to call city.deinit() in that case. const ip = try std.net.Address.parseIp("89.160.20.128", 0); - const city = try db.lookup(maxminddb.geoip2.City, &ip); + const city = try db.lookup(allocator, maxminddb.geoip2.City, &ip); defer city.deinit(); var it = city.country.names.?.iterator(); diff --git a/examples/within.zig b/examples/within.zig index 66d0c7a..3af77c4 100644 --- a/examples/within.zig +++ b/examples/within.zig @@ -10,13 +10,13 @@ pub fn main() !void { const allocator = gpa.allocator(); defer _ = gpa.detectLeaks(); - var db = try maxminddb.Reader.open_mmap(allocator, db_path); - defer db.close(); + var db = try maxminddb.Reader.mmap(allocator, db_path); + defer db.unmap(); const network = maxminddb.Network{ .ip = try std.net.Address.parseIp("0.0.0.0", 0), }; - var it = try db.within(maxminddb.geolite2.City, network); + var it = try db.within(allocator, maxminddb.geolite2.City, network); defer it.deinit(); var n: usize = 0; diff --git a/src/maxminddb.zig b/src/maxminddb.zig index 6888d59..a45ced2 100644 --- a/src/maxminddb.zig +++ b/src/maxminddb.zig @@ -32,19 +32,20 @@ fn expectEqualMaps( } } +const allocator = std.testing.allocator; const expectEqual = std.testing.expectEqual; const expectEqualStrings = std.testing.expectEqualStrings; const expectEqualDeep = std.testing.expectEqualDeep; test "GeoLite2 Country" { - var db = try Reader.open_mmap( - std.testing.allocator, + var db = try Reader.mmap( + allocator, "test-data/test-data/GeoLite2-Country-Test.mmdb", ); - defer db.close(); + defer db.unmap(); const ip = try std.net.Address.parseIp("89.160.20.128", 0); - const got = try db.lookup(geolite2.Country, &ip); + const got = try db.lookup(allocator, geolite2.Country, &ip); defer got.deinit(); try expectEqualStrings("EU", got.continent.code); @@ -77,14 +78,14 @@ test "GeoLite2 Country" { } test "GeoLite2 City" { - var db = try Reader.open_mmap( - std.testing.allocator, + var db = try Reader.mmap( + allocator, "test-data/test-data/GeoLite2-City-Test.mmdb", ); - defer db.close(); + defer db.unmap(); const ip = try std.net.Address.parseIp("89.160.20.128", 0); - const got = try db.lookup(geolite2.City, &ip); + const got = try db.lookup(allocator, geolite2.City, &ip); defer got.deinit(); try expectEqual(2694762, got.city.geoname_id); @@ -146,14 +147,14 @@ test "GeoLite2 City" { } test "GeoLite2 ASN" { - var db = try Reader.open_mmap( - std.testing.allocator, + var db = try Reader.mmap( + allocator, "test-data/test-data/GeoLite2-ASN-Test.mmdb", ); - defer db.close(); + defer db.unmap(); const ip = try std.net.Address.parseIp("89.160.20.128", 0); - const got = try db.lookup(geolite2.ASN, &ip); + const got = try db.lookup(allocator, geolite2.ASN, &ip); const want = geolite2.ASN{ .autonomous_system_number = 29518, @@ -163,14 +164,14 @@ test "GeoLite2 ASN" { } test "GeoIP2 Country" { - var db = try Reader.open_mmap( - std.testing.allocator, + var db = try Reader.mmap( + allocator, "test-data/test-data/GeoIP2-Country-Test.mmdb", ); - defer db.close(); + defer db.unmap(); const ip = try std.net.Address.parseIp("89.160.20.128", 0); - const got = try db.lookup(geoip2.Country, &ip); + const got = try db.lookup(allocator, geoip2.Country, &ip); defer got.deinit(); try expectEqualStrings("EU", got.continent.code); @@ -210,14 +211,14 @@ test "GeoIP2 Country" { } test "GeoIP2 City" { - var db = try Reader.open_mmap( - std.testing.allocator, + var db = try Reader.mmap( + allocator, "test-data/test-data/GeoIP2-City-Test.mmdb", ); - defer db.close(); + defer db.unmap(); const ip = try std.net.Address.parseIp("89.160.20.128", 0); - const got = try db.lookup(geoip2.City, &ip); + const got = try db.lookup(allocator, geoip2.City, &ip); defer got.deinit(); try expectEqual(2694762, got.city.geoname_id); @@ -286,14 +287,14 @@ test "GeoIP2 City" { } test "GeoIP2 Enterprise" { - var db = try Reader.open_mmap( - std.testing.allocator, + var db = try Reader.mmap( + allocator, "test-data/test-data/GeoIP2-Enterprise-Test.mmdb", ); - defer db.close(); + defer db.unmap(); const ip = try std.net.Address.parseIp("74.209.24.0", 0); - const got = try db.lookup(geoip2.Enterprise, &ip); + const got = try db.lookup(allocator, geoip2.Enterprise, &ip); defer got.deinit(); try expectEqual(11, got.city.confidence); @@ -380,14 +381,14 @@ test "GeoIP2 Enterprise" { } test "GeoIP2 ISP" { - var db = try Reader.open_mmap( - std.testing.allocator, + var db = try Reader.mmap( + allocator, "test-data/test-data/GeoIP2-ISP-Test.mmdb", ); - defer db.close(); + defer db.unmap(); const ip = try std.net.Address.parseIp("89.160.20.112", 0); - const got = try db.lookup(geoip2.ISP, &ip); + const got = try db.lookup(allocator, geoip2.ISP, &ip); const want = geoip2.ISP{ .autonomous_system_number = 29518, @@ -399,14 +400,14 @@ test "GeoIP2 ISP" { } test "GeoIP2 Connection-Type" { - var db = try Reader.open_mmap( - std.testing.allocator, + var db = try Reader.mmap( + allocator, "test-data/test-data/GeoIP2-Connection-Type-Test.mmdb", ); - defer db.close(); + defer db.unmap(); const ip = try std.net.Address.parseIp("96.1.20.112", 0); - const got = try db.lookup(geoip2.ConnectionType, &ip); + const got = try db.lookup(allocator, geoip2.ConnectionType, &ip); const want = geoip2.ConnectionType{ .connection_type = "Cable/DSL", @@ -415,14 +416,14 @@ test "GeoIP2 Connection-Type" { } test "GeoIP2 Anonymous-IP" { - var db = try Reader.open_mmap( - std.testing.allocator, + var db = try Reader.mmap( + allocator, "test-data/test-data/GeoIP2-Anonymous-IP-Test.mmdb", ); - defer db.close(); + defer db.unmap(); const ip = try std.net.Address.parseIp("81.2.69.0", 0); - const got = try db.lookup(geoip2.AnonymousIP, &ip); + const got = try db.lookup(allocator, geoip2.AnonymousIP, &ip); const want = geoip2.AnonymousIP{ .is_anonymous = true, @@ -436,14 +437,14 @@ test "GeoIP2 Anonymous-IP" { } test "GeoIP2 DensityIncome" { - var db = try Reader.open_mmap( - std.testing.allocator, + var db = try Reader.mmap( + allocator, "test-data/test-data/GeoIP2-DensityIncome-Test.mmdb", ); - defer db.close(); + defer db.unmap(); const ip = try std.net.Address.parseIp("5.83.124.123", 0); - const got = try db.lookup(geoip2.DensityIncome, &ip); + const got = try db.lookup(allocator, geoip2.DensityIncome, &ip); const want = geoip2.DensityIncome{ .average_income = 32323, @@ -453,14 +454,14 @@ test "GeoIP2 DensityIncome" { } test "GeoIP2 Domain" { - var db = try Reader.open_mmap( - std.testing.allocator, + var db = try Reader.mmap( + allocator, "test-data/test-data/GeoIP2-Domain-Test.mmdb", ); - defer db.close(); + defer db.unmap(); const ip = try std.net.Address.parseIp("66.92.80.123", 0); - const got = try db.lookup(geoip2.Domain, &ip); + const got = try db.lookup(allocator, geoip2.Domain, &ip); const want = geoip2.Domain{ .domain = "speakeasy.net", diff --git a/src/reader.zig b/src/reader.zig index aeacbe2..569822b 100644 --- a/src/reader.zig +++ b/src/reader.zig @@ -1,6 +1,6 @@ const std = @import("std"); const decoder = @import("decoder.zig"); -const mmap = @import("mmap.zig"); +const memorymap = @import("mmap.zig"); const net = @import("net.zig"); pub const ReadError = error{ @@ -49,7 +49,6 @@ pub const Reader = struct { offset: usize, ipv4_start: usize, metadata: Metadata, - allocator: std.mem.Allocator, // Loads a MaxMind DB file into memory. pub fn open(allocator: std.mem.Allocator, path: []const u8, max_db_size: usize) !Reader { @@ -77,7 +76,6 @@ pub const Reader = struct { .offset = search_tree_size + data_section_separator_size, .ipv4_start = 0, .metadata = metadata, - .allocator = allocator, }; r.ipv4_start = try r.findIPv4Start(); @@ -85,13 +83,21 @@ pub const Reader = struct { return r; } + // Frees the memory occupied by the DB file. + // From this point all the DB records are unusable because their fields were backed by the same memory. + // Note, the records still have to be deinited since they might contain arrays or maps. + pub fn close(self: *Reader, allocator: std.mem.Allocator) void { + self.metadata.deinit(); + allocator.free(self.src); + } + // Maps a MaxMind DB file into memory. - pub fn open_mmap(allocator: std.mem.Allocator, path: []const u8) !Reader { + pub fn mmap(allocator: std.mem.Allocator, path: []const u8) !Reader { var f = try std.fs.cwd().openFile(path, .{}); errdefer f.close(); - const src = try mmap.map(f); - errdefer mmap.unmap(src); + const src = try memorymap.map(f); + errdefer memorymap.unmap(src); // Decode database metadata which is stored as a separate data section, // see https://maxmind.github.io/MaxMind-DB/#database-metadata. @@ -111,7 +117,6 @@ pub const Reader = struct { .offset = search_tree_size + data_section_separator_size, .ipv4_start = 0, .metadata = metadata, - .allocator = allocator, }; r.ipv4_start = try r.findIPv4Start(); @@ -119,34 +124,39 @@ pub const Reader = struct { return r; } - // Frees the memory occupied by the DB file. + // Unmaps the DB file. // From this point all the DB records are unusable because their fields were backed by the same memory. // Note, the records still have to be deinited since they might contain arrays or maps. - pub fn close(self: *Reader) void { + pub fn unmap(self: *Reader) void { self.metadata.deinit(); - if (self.mapped_file == null) { - self.allocator.free(self.src); - return; - } - - mmap.unmap(self.src); + memorymap.unmap(self.src); self.mapped_file.?.close(); } // Looks up a record by an IP address. - pub fn lookup(self: *Reader, comptime T: type, address: *const std.net.Address) !T { + pub fn lookup( + self: *Reader, + allocator: std.mem.Allocator, + comptime T: type, + address: *const std.net.Address, + ) !T { const ip_bytes = net.ipToBytes(address); const pointer, _ = try self.findAddressInTree(ip_bytes); if (pointer == 0) { return ReadError.AddressNotFound; } - return try self.resolveDataPointerAndDecode(T, pointer); + return try self.resolveDataPointerAndDecode(allocator, T, pointer); } // Iterates over blocks of IP networks. - pub fn within(self: *Reader, comptime T: type, network: net.Network) !Iterator(T) { + pub fn within( + self: *Reader, + allocator: std.mem.Allocator, + comptime T: type, + network: net.Network, + ) !Iterator(T) { const ip_bytes = net.IP.init(network.ip); const prefix_len: usize = network.prefix_len; const bit_count: usize = ip_bytes.bitCount(); @@ -154,7 +164,7 @@ pub const Reader = struct { var node = self.startNode(bit_count); const node_count = self.metadata.node_count; - var stack = try std.ArrayList(WithinNode).initCapacity(self.allocator, bit_count - prefix_len); + var stack = try std.ArrayList(WithinNode).initCapacity(allocator, bit_count - prefix_len); errdefer stack.deinit(); // Traverse down the tree to the level that matches the CIDR mark. @@ -186,10 +196,16 @@ pub const Reader = struct { .reader = self, .node_count = node_count, .stack = stack, + .allocator = allocator, }; } - fn resolveDataPointerAndDecode(self: *Reader, comptime T: type, pointer: usize) !T { + fn resolveDataPointerAndDecode( + self: *Reader, + allocator: std.mem.Allocator, + comptime T: type, + pointer: usize, + ) !T { const record_offset = try self.resolveDataPointer(pointer); var d = decoder.Decoder{ @@ -197,7 +213,7 @@ pub const Reader = struct { .offset = record_offset, }; - return try d.decodeRecord(self.allocator, T); + return try d.decodeRecord(allocator, T); } fn resolveDataPointer(self: *Reader, pointer: usize) !usize { @@ -315,6 +331,7 @@ fn Iterator(comptime T: type) type { reader: *Reader, node_count: usize, stack: std.ArrayList(WithinNode), + allocator: std.mem.Allocator, const Self = @This(); @@ -341,7 +358,11 @@ fn Iterator(comptime T: type) type { if (current.node > self.node_count) { const ip_net = current.ip_bytes.network(current.prefix_len); - const record = try reader.resolveDataPointerAndDecode(T, current.node); + const record = try reader.resolveDataPointerAndDecode( + self.allocator, + T, + current.node, + ); return Item{ .net = ip_net,