milestone 6: dns cache, rate limiting, query logging, disk monitoring and retention
This commit is contained in:
@@ -0,0 +1,468 @@
|
||||
//! `query_log` and its `domains` dimension table in `querylog.db`.
|
||||
//!
|
||||
//! Two shapes live here. The free functions follow the milestone-4 repository
|
||||
//! idiom — prepare, use, finalize — because retention runs them a handful of
|
||||
//! times per day. The flush loop is the one hot path in the program, so it gets
|
||||
//! `BatchWriter`, which owns its three statements for its whole life
|
||||
//! (`db.zig:360` names this file as the reason `db.zig` carries no statement
|
||||
//! cache).
|
||||
//!
|
||||
//! Every string in a `Row` is borrowed for the duration of the call only:
|
||||
//! `Stmt.bindText` binds with `SQLITE_TRANSIENT`, so SQLite copies before
|
||||
//! `writeBatch` returns.
|
||||
//!
|
||||
//! The rows are expendable log data. Nothing here retries, and the caller
|
||||
//! decides what a failed batch means.
|
||||
|
||||
const std = @import("std");
|
||||
|
||||
const db = @import("../db.zig");
|
||||
|
||||
/// One `query_log` row. The logger applies the privacy transforms of PLAN
|
||||
/// §11.4 before it builds this, so `domain` and `client_ip` are already
|
||||
/// whatever the operator agreed to store.
|
||||
pub const Row = struct {
|
||||
timestamp: i64,
|
||||
domain: []const u8,
|
||||
client_ip: []const u8,
|
||||
qtype: ?u16,
|
||||
blocked: bool,
|
||||
block_reason: ?[]const u8,
|
||||
response_time_us: ?i64,
|
||||
cache_hit: ?bool,
|
||||
upstream: ?[]const u8,
|
||||
};
|
||||
|
||||
const insert_domain_sql = "INSERT OR IGNORE INTO domains (domain) VALUES (?1)";
|
||||
|
||||
const select_domain_sql = "SELECT id FROM domains WHERE domain = ?1";
|
||||
|
||||
const insert_row_sql =
|
||||
\\INSERT INTO query_log
|
||||
\\ (timestamp, domain_id, client_ip, qtype, blocked, block_reason,
|
||||
\\ response_time_us, cache_hit, upstream)
|
||||
\\VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)
|
||||
;
|
||||
|
||||
/// Owns the prepared statements of the flush loop. Init once, reuse per batch.
|
||||
///
|
||||
/// `database` must outlive the writer and must not move: every `Stmt` holds a
|
||||
/// `*Db`. Neither `Db` nor `Stmt` is thread-safe, so one writer belongs to one
|
||||
/// task.
|
||||
pub const BatchWriter = struct {
|
||||
database: *db.Db,
|
||||
insert_domain: db.Stmt,
|
||||
select_domain: db.Stmt,
|
||||
insert_row: db.Stmt,
|
||||
|
||||
pub fn init(database: *db.Db) db.Error!BatchWriter {
|
||||
var insert_domain = try database.prepare(insert_domain_sql);
|
||||
errdefer insert_domain.deinit();
|
||||
var select_domain = try database.prepare(select_domain_sql);
|
||||
errdefer select_domain.deinit();
|
||||
const insert_row = try database.prepare(insert_row_sql);
|
||||
return .{
|
||||
.database = database,
|
||||
.insert_domain = insert_domain,
|
||||
.select_domain = select_domain,
|
||||
.insert_row = insert_row,
|
||||
};
|
||||
}
|
||||
|
||||
pub fn deinit(self: *BatchWriter) void {
|
||||
self.insert_row.deinit();
|
||||
self.select_domain.deinit();
|
||||
self.insert_domain.deinit();
|
||||
}
|
||||
|
||||
/// One transaction for the whole batch. Domains are interned through
|
||||
/// `INSERT OR IGNORE` followed by `SELECT id`.
|
||||
///
|
||||
/// On any failure the transaction rolls back, so a batch is all or
|
||||
/// nothing, and the writer stays usable for the next batch.
|
||||
pub fn writeBatch(self: *BatchWriter, rows: []const Row) db.Error!void {
|
||||
if (rows.len == 0) return;
|
||||
|
||||
var tx = try db.Tx.begin(self.database);
|
||||
// `errdefer`s run in reverse: the statements are released before the
|
||||
// ROLLBACK, so no read cursor is still open when it runs.
|
||||
errdefer tx.rollback();
|
||||
errdefer self.resetAll();
|
||||
|
||||
for (rows) |row| {
|
||||
const domain_id = try self.internDomain(row.domain);
|
||||
try self.write(row, domain_id);
|
||||
}
|
||||
try tx.commit();
|
||||
}
|
||||
|
||||
fn internDomain(self: *BatchWriter, domain: []const u8) db.Error!i64 {
|
||||
try self.insert_domain.reset();
|
||||
try self.insert_domain.bindText(1, domain);
|
||||
try self.insert_domain.exec();
|
||||
|
||||
try self.select_domain.reset();
|
||||
try self.select_domain.bindText(1, domain);
|
||||
// The insert above either created the row or found it already there,
|
||||
// so a miss means the table changed under this connection.
|
||||
if (!try self.select_domain.step()) return error.NotFound;
|
||||
const id = self.select_domain.columnInt(0);
|
||||
// A statement stopped on a row keeps its cursor open until it is
|
||||
// reset; the transaction must not carry that to the next row.
|
||||
try self.select_domain.reset();
|
||||
return id;
|
||||
}
|
||||
|
||||
fn write(self: *BatchWriter, row: Row, domain_id: i64) db.Error!void {
|
||||
var stmt = &self.insert_row;
|
||||
try stmt.reset();
|
||||
try stmt.bindInt(1, row.timestamp);
|
||||
try stmt.bindInt(2, domain_id);
|
||||
try stmt.bindText(3, row.client_ip);
|
||||
try bindIntOrNull(stmt, 4, if (row.qtype) |v| @as(i64, v) else null);
|
||||
try stmt.bindBool(5, row.blocked);
|
||||
try stmt.bindTextOrNull(6, row.block_reason);
|
||||
try bindIntOrNull(stmt, 7, row.response_time_us);
|
||||
try bindIntOrNull(stmt, 8, if (row.cache_hit) |v| @as(i64, @intFromBool(v)) else null);
|
||||
try stmt.bindTextOrNull(9, row.upstream);
|
||||
try stmt.exec();
|
||||
}
|
||||
|
||||
/// Best effort: this runs on the failure path, where the error that
|
||||
/// matters is the one already on its way to the caller.
|
||||
fn resetAll(self: *BatchWriter) void {
|
||||
self.insert_row.reset() catch {};
|
||||
self.select_domain.reset() catch {};
|
||||
self.insert_domain.reset() catch {};
|
||||
}
|
||||
};
|
||||
|
||||
fn bindIntOrNull(stmt: *db.Stmt, idx: c_int, value: ?i64) db.Error!void {
|
||||
if (value) |v| return stmt.bindInt(idx, v);
|
||||
return stmt.bindNull(idx);
|
||||
}
|
||||
|
||||
/// Deletes every `query_log` row strictly older than `cutoff_ts` and returns
|
||||
/// how many went.
|
||||
///
|
||||
/// Orphaned `domains` rows stay: it is a dimension table, re-interning a name
|
||||
/// costs one indexed insert, and §11.3 asks for no collection.
|
||||
pub fn pruneOlderThan(database: *db.Db, cutoff_ts: i64) db.Error!i64 {
|
||||
var stmt = try database.prepare("DELETE FROM query_log WHERE timestamp < ?1");
|
||||
defer stmt.deinit();
|
||||
try stmt.bindInt(1, cutoff_ts);
|
||||
try stmt.exec();
|
||||
return database.changes();
|
||||
}
|
||||
|
||||
/// `PRAGMA wal_checkpoint(TRUNCATE)`: moves the WAL into the database and
|
||||
/// truncates it to zero bytes, which is what keeps a day of log writes from
|
||||
/// growing the WAL past the free space the disk monitor watches.
|
||||
///
|
||||
/// SQLite reports a checkpoint blocked by a concurrent reader in the row it
|
||||
/// returns, not as an error code, so a blocked checkpoint is not an error
|
||||
/// here. Retention checkpoints after every prune, so the next pass retries.
|
||||
/// On a database that is not in WAL mode the pragma is a no-op.
|
||||
pub fn checkpointTruncate(database: *db.Db) db.Error!void {
|
||||
return database.exec("PRAGMA wal_checkpoint(TRUNCATE);");
|
||||
}
|
||||
|
||||
/// Rewrites the whole file. Retention runs this rarely by design — on an SD
|
||||
/// card a full rewrite is the most expensive thing this program does.
|
||||
pub fn vacuum(database: *db.Db) db.Error!void {
|
||||
return database.exec("VACUUM;");
|
||||
}
|
||||
|
||||
pub fn countRows(database: *db.Db) db.Error!i64 {
|
||||
return database.queryInt("SELECT count(*) FROM query_log");
|
||||
}
|
||||
|
||||
pub fn countDomains(database: *db.Db) db.Error!i64 {
|
||||
return database.queryInt("SELECT count(*) FROM domains");
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
const querylog_schema = @import("../querylog_schema.zig");
|
||||
|
||||
const testing = std.testing;
|
||||
|
||||
fn openLog() !db.Db {
|
||||
var database = try db.Db.open(":memory:", .{ .mode = .memory });
|
||||
errdefer database.close();
|
||||
try db.applyPragmas(&database, .{});
|
||||
try database.exec(querylog_schema.ddl);
|
||||
return database;
|
||||
}
|
||||
|
||||
fn plainRow(timestamp: i64, domain: []const u8) Row {
|
||||
return .{
|
||||
.timestamp = timestamp,
|
||||
.domain = domain,
|
||||
.client_ip = "192.0.2.10",
|
||||
.qtype = 1,
|
||||
.blocked = false,
|
||||
.block_reason = null,
|
||||
.response_time_us = 1200,
|
||||
.cache_hit = false,
|
||||
.upstream = "9.9.9.9",
|
||||
};
|
||||
}
|
||||
|
||||
fn domainIdOf(database: *db.Db, domain: []const u8) !i64 {
|
||||
var stmt = try database.prepare(select_domain_sql);
|
||||
defer stmt.deinit();
|
||||
try stmt.bindText(1, domain);
|
||||
try testing.expect(try stmt.step());
|
||||
return stmt.columnInt(0);
|
||||
}
|
||||
|
||||
test "writeBatch inserts every row and interns each domain once" {
|
||||
var database = try openLog();
|
||||
defer database.close();
|
||||
var writer = try BatchWriter.init(&database);
|
||||
defer writer.deinit();
|
||||
|
||||
try writer.writeBatch(&.{
|
||||
plainRow(100, "example.com"),
|
||||
plainRow(101, "example.com"),
|
||||
plainRow(102, "ads.example.net"),
|
||||
});
|
||||
|
||||
try testing.expectEqual(@as(i64, 3), try countRows(&database));
|
||||
try testing.expectEqual(@as(i64, 2), try countDomains(&database));
|
||||
|
||||
const first = try domainIdOf(&database, "example.com");
|
||||
try testing.expectEqual(
|
||||
@as(i64, 2),
|
||||
try database.queryInt("SELECT count(*) FROM query_log WHERE domain_id = 1"),
|
||||
);
|
||||
try testing.expectEqual(@as(i64, 1), first);
|
||||
}
|
||||
|
||||
test "a second batch reuses the interned domain id" {
|
||||
var database = try openLog();
|
||||
defer database.close();
|
||||
var writer = try BatchWriter.init(&database);
|
||||
defer writer.deinit();
|
||||
|
||||
try writer.writeBatch(&.{plainRow(100, "example.com")});
|
||||
const before = try domainIdOf(&database, "example.com");
|
||||
|
||||
try writer.writeBatch(&.{ plainRow(200, "example.com"), plainRow(201, "other.example") });
|
||||
const after = try domainIdOf(&database, "example.com");
|
||||
|
||||
try testing.expectEqual(before, after);
|
||||
try testing.expectEqual(@as(i64, 3), try countRows(&database));
|
||||
try testing.expectEqual(@as(i64, 2), try countDomains(&database));
|
||||
try testing.expectEqual(
|
||||
@as(i64, 2),
|
||||
try database.queryInt("SELECT count(*) FROM query_log WHERE domain_id = 1"),
|
||||
);
|
||||
}
|
||||
|
||||
test "nullable columns round-trip a value and a null" {
|
||||
var database = try openLog();
|
||||
defer database.close();
|
||||
var writer = try BatchWriter.init(&database);
|
||||
defer writer.deinit();
|
||||
|
||||
try writer.writeBatch(&.{
|
||||
.{
|
||||
.timestamp = 10,
|
||||
.domain = "blocked.example",
|
||||
.client_ip = "2001:db8::1",
|
||||
.qtype = 28,
|
||||
.blocked = true,
|
||||
.block_reason = "blocklist",
|
||||
.response_time_us = 42,
|
||||
.cache_hit = true,
|
||||
.upstream = "dns.example",
|
||||
},
|
||||
.{
|
||||
.timestamp = 11,
|
||||
.domain = "quiet.example",
|
||||
.client_ip = "hidden",
|
||||
.qtype = null,
|
||||
.blocked = false,
|
||||
.block_reason = null,
|
||||
.response_time_us = null,
|
||||
.cache_hit = null,
|
||||
.upstream = null,
|
||||
},
|
||||
});
|
||||
|
||||
var stmt = try database.prepare(
|
||||
\\SELECT d.domain, q.client_ip, q.qtype, q.blocked, q.block_reason,
|
||||
\\ q.response_time_us, q.cache_hit, q.upstream
|
||||
\\ FROM query_log q JOIN domains d ON d.id = q.domain_id
|
||||
\\ ORDER BY q.timestamp
|
||||
);
|
||||
defer stmt.deinit();
|
||||
|
||||
try testing.expect(try stmt.step());
|
||||
try testing.expectEqualStrings("blocked.example", stmt.columnText(0));
|
||||
try testing.expectEqualStrings("2001:db8::1", stmt.columnText(1));
|
||||
try testing.expectEqual(@as(i64, 28), stmt.columnInt(2));
|
||||
try testing.expect(stmt.columnBool(3));
|
||||
try testing.expectEqualStrings("blocklist", stmt.columnText(4));
|
||||
try testing.expectEqual(@as(i64, 42), stmt.columnInt(5));
|
||||
try testing.expect(stmt.columnBool(6));
|
||||
try testing.expectEqualStrings("dns.example", stmt.columnText(7));
|
||||
|
||||
try testing.expect(try stmt.step());
|
||||
try testing.expectEqualStrings("quiet.example", stmt.columnText(0));
|
||||
try testing.expectEqualStrings("hidden", stmt.columnText(1));
|
||||
try testing.expect(stmt.isNull(2));
|
||||
try testing.expect(!stmt.columnBool(3));
|
||||
try testing.expect(stmt.isNull(4));
|
||||
try testing.expect(stmt.isNull(5));
|
||||
try testing.expect(stmt.isNull(6));
|
||||
try testing.expect(stmt.isNull(7));
|
||||
|
||||
try testing.expect(!try stmt.step());
|
||||
}
|
||||
|
||||
test "an empty batch writes nothing and opens no transaction" {
|
||||
var database = try openLog();
|
||||
defer database.close();
|
||||
var writer = try BatchWriter.init(&database);
|
||||
defer writer.deinit();
|
||||
|
||||
// A transaction is already open, so a `BEGIN IMMEDIATE` from `writeBatch`
|
||||
// would fail: this is what proves the empty batch returns before it.
|
||||
var tx = try db.Tx.begin(&database);
|
||||
try writer.writeBatch(&.{});
|
||||
tx.rollback();
|
||||
|
||||
try testing.expectEqual(@as(i64, 0), try countRows(&database));
|
||||
try testing.expectEqual(@as(i64, 0), try countDomains(&database));
|
||||
}
|
||||
|
||||
test "pruneOlderThan deletes strictly older rows and returns the count" {
|
||||
var database = try openLog();
|
||||
defer database.close();
|
||||
var writer = try BatchWriter.init(&database);
|
||||
defer writer.deinit();
|
||||
|
||||
try writer.writeBatch(&.{
|
||||
plainRow(100, "old.example"),
|
||||
plainRow(199, "old.example"),
|
||||
plainRow(200, "edge.example"),
|
||||
plainRow(300, "fresh.example"),
|
||||
});
|
||||
|
||||
try testing.expectEqual(@as(i64, 2), try pruneOlderThan(&database, 200));
|
||||
try testing.expectEqual(@as(i64, 2), try countRows(&database));
|
||||
// The row exactly at the cutoff stays.
|
||||
try testing.expectEqual(
|
||||
@as(i64, 1),
|
||||
try database.queryInt("SELECT count(*) FROM query_log WHERE timestamp = 200"),
|
||||
);
|
||||
// A second pass over the same cutoff finds nothing left to do.
|
||||
try testing.expectEqual(@as(i64, 0), try pruneOlderThan(&database, 200));
|
||||
}
|
||||
|
||||
test "pruneOlderThan leaves the domains dimension table intact" {
|
||||
var database = try openLog();
|
||||
defer database.close();
|
||||
var writer = try BatchWriter.init(&database);
|
||||
defer writer.deinit();
|
||||
|
||||
try writer.writeBatch(&.{ plainRow(10, "a.example"), plainRow(11, "b.example") });
|
||||
try testing.expectEqual(@as(i64, 2), try pruneOlderThan(&database, 1000));
|
||||
|
||||
try testing.expectEqual(@as(i64, 0), try countRows(&database));
|
||||
try testing.expectEqual(@as(i64, 2), try countDomains(&database));
|
||||
}
|
||||
|
||||
test "a failing row rolls the whole batch back and the writer survives it" {
|
||||
var database = try openLog();
|
||||
defer database.close();
|
||||
try database.exec(
|
||||
\\CREATE TRIGGER refuse_boom BEFORE INSERT ON query_log
|
||||
\\WHEN new.client_ip = 'boom'
|
||||
\\BEGIN SELECT RAISE(ABORT, 'refused'); END;
|
||||
);
|
||||
|
||||
var writer = try BatchWriter.init(&database);
|
||||
defer writer.deinit();
|
||||
|
||||
var doomed = plainRow(20, "second.example");
|
||||
doomed.client_ip = "boom";
|
||||
try testing.expectError(error.Constraint, writer.writeBatch(&.{
|
||||
plainRow(10, "first.example"),
|
||||
doomed,
|
||||
}));
|
||||
|
||||
// The interned domain of the row that did insert is gone with it.
|
||||
try testing.expectEqual(@as(i64, 0), try countRows(&database));
|
||||
try testing.expectEqual(@as(i64, 0), try countDomains(&database));
|
||||
|
||||
try writer.writeBatch(&.{plainRow(30, "third.example")});
|
||||
try testing.expectEqual(@as(i64, 1), try countRows(&database));
|
||||
try testing.expectEqual(@as(i64, 1), try countDomains(&database));
|
||||
}
|
||||
|
||||
test "countRows and countDomains agree with what the batches wrote" {
|
||||
var database = try openLog();
|
||||
defer database.close();
|
||||
var writer = try BatchWriter.init(&database);
|
||||
defer writer.deinit();
|
||||
|
||||
try testing.expectEqual(@as(i64, 0), try countRows(&database));
|
||||
try testing.expectEqual(@as(i64, 0), try countDomains(&database));
|
||||
|
||||
var rows: [50]Row = undefined;
|
||||
var names: [50][16]u8 = undefined;
|
||||
for (&rows, &names, 0..) |*row, *name, i| {
|
||||
const written = std.fmt.bufPrint(name, "d{d}.example", .{i % 7}) catch unreachable;
|
||||
row.* = plainRow(@intCast(i), written);
|
||||
}
|
||||
try writer.writeBatch(&rows);
|
||||
|
||||
try testing.expectEqual(@as(i64, 50), try countRows(&database));
|
||||
try testing.expectEqual(@as(i64, 7), try countDomains(&database));
|
||||
}
|
||||
|
||||
// `PRAGMA wal_checkpoint` needs a real WAL, which an in-memory database cannot
|
||||
// have. `std.testing.tmpDir` creates its directory under `.zig-cache/tmp/`
|
||||
// relative to the process working directory, which is also how SQLite's VFS
|
||||
// resolves the filename it is handed (`storage_integration_test.zig:44`).
|
||||
const tmp_prefix = ".zig-cache/tmp/";
|
||||
const sub_path_len = @typeInfo(@FieldType(testing.TmpDir, "sub_path")).array.len;
|
||||
|
||||
test "checkpointTruncate and vacuum run against a WAL file database" {
|
||||
var tmp = testing.tmpDir(.{});
|
||||
defer tmp.cleanup();
|
||||
|
||||
var path_buf: [tmp_prefix.len + sub_path_len + 32]u8 = undefined;
|
||||
const path = try std.fmt.bufPrintZ(&path_buf, "{s}{s}/querylog.db", .{ tmp_prefix, &tmp.sub_path });
|
||||
|
||||
var database = try db.Db.open(path, .{ .mode = .read_write_create });
|
||||
defer database.close();
|
||||
try db.applyPragmas(&database, .{});
|
||||
{
|
||||
var stmt = try database.prepare("PRAGMA journal_mode");
|
||||
defer stmt.deinit();
|
||||
try testing.expect(try stmt.step());
|
||||
// `columnText` is borrowed until the next call on the statement, so it
|
||||
// is compared here rather than carried out of this block.
|
||||
try testing.expectEqualStrings("wal", stmt.columnText(0));
|
||||
}
|
||||
try database.exec(querylog_schema.ddl);
|
||||
|
||||
var writer = try BatchWriter.init(&database);
|
||||
defer writer.deinit();
|
||||
try writer.writeBatch(&.{ plainRow(10, "a.example"), plainRow(20, "b.example") });
|
||||
|
||||
try checkpointTruncate(&database);
|
||||
try testing.expectEqual(@as(i64, 1), try pruneOlderThan(&database, 20));
|
||||
try checkpointTruncate(&database);
|
||||
try vacuum(&database);
|
||||
|
||||
try testing.expectEqual(@as(i64, 1), try countRows(&database));
|
||||
try testing.expectEqual(@as(i64, 2), try countDomains(&database));
|
||||
}
|
||||
Reference in New Issue
Block a user