feat: add WorkerCtx and worker to server
This commit is contained in:
dualfroz committed
1 parent
44c5e9b9cb
commit
48ff5cf38e
1 file changed
+159
+159
@@ -2171,3 +2171,162 @@ test "server batch applies puts and deletes atomically" {
|
|||||||
try std.testing.expect(kept == null);
|
try std.testing.expect(kept == null);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
test "server rejects a malformed frame with bad_request without crashing" {
|
||||||
|
const allocator = std.testing.allocator;
|
||||||
|
var tmp = std.testing.tmpDir(.{});
|
||||||
|
defer tmp.cleanup();
|
||||||
|
const path = try tmp.dir.realpathAlloc(allocator, ".");
|
||||||
|
defer allocator.free(path);
|
||||||
|
|
||||||
|
var db = try testDb(allocator, path);
|
||||||
|
defer db.close();
|
||||||
|
|
||||||
|
const address = try std.net.Address.parseIp("127.0.0.1", 0);
|
||||||
|
var server = try Server.init(allocator, &db, address);
|
||||||
|
try server.start();
|
||||||
|
defer server.stop();
|
||||||
|
|
||||||
|
const stream = try std.net.tcpConnectToAddress(server.boundAddress());
|
||||||
|
defer stream.close();
|
||||||
|
|
||||||
|
try protocol.writeFrame(stream.writer(), @intFromEnum(protocol.Op.set), "not a valid set payload with no length prefix boundary");
|
||||||
|
var frame = try protocol.readFrame(allocator, stream.reader());
|
||||||
|
defer frame.deinit();
|
||||||
|
try std.testing.expectEqual(@intFromEnum(protocol.Status.bad_request), frame.tag);
|
||||||
|
}
|
||||||
|
|
||||||
|
test "server serves multiple concurrent clients without corrupting state" {
|
||||||
|
const allocator = std.testing.allocator;
|
||||||
|
var tmp = std.testing.tmpDir(.{});
|
||||||
|
defer tmp.cleanup();
|
||||||
|
const path = try tmp.dir.realpathAlloc(allocator, ".");
|
||||||
|
defer allocator.free(path);
|
||||||
|
|
||||||
|
var db = try testDb(allocator, path);
|
||||||
|
defer db.close();
|
||||||
|
|
||||||
|
const address = try std.net.Address.parseIp("127.0.0.1", 0);
|
||||||
|
var server = try Server.init(allocator, &db, address);
|
||||||
|
try server.start();
|
||||||
|
defer server.stop();
|
||||||
|
|
||||||
|
const client_mod = @import("client.zig");
|
||||||
|
const WorkerCtx = struct {
|
||||||
|
allocator: Allocator,
|
||||||
|
addr: std.net.Address,
|
||||||
|
index: usize,
|
||||||
|
ok: bool = false,
|
||||||
|
};
|
||||||
|
|
||||||
|
const worker = struct {
|
||||||
|
fn run(ctx: *WorkerCtx) void {
|
||||||
|
var client = client_mod.Client.connect(ctx.allocator, ctx.addr) catch return;
|
||||||
|
defer client.close();
|
||||||
|
var buf: [32]u8 = undefined;
|
||||||
|
const key = std.fmt.bufPrint(&buf, "k{d}", .{ctx.index}) catch return;
|
||||||
|
client.set(key, "v") catch return;
|
||||||
|
const got = client.get(ctx.allocator, key) catch return;
|
||||||
|
defer if (got) |g| ctx.allocator.free(g);
|
||||||
|
ctx.ok = got != null and std.mem.eql(u8, got.?, "v");
|
||||||
|
}
|
||||||
|
}.run;
|
||||||
|
|
||||||
|
var contexts: [8]WorkerCtx = undefined;
|
||||||
|
var threads: [8]std.Thread = undefined;
|
||||||
|
for (0..8) |i| {
|
||||||
|
contexts[i] = .{ .allocator = allocator, .addr = server.boundAddress(), .index = i };
|
||||||
|
threads[i] = try std.Thread.spawn(.{}, worker, .{&contexts[i]});
|
||||||
|
}
|
||||||
|
for (0..8) |i| threads[i].join();
|
||||||
|
for (0..8) |i| try std.testing.expect(contexts[i].ok);
|
||||||
|
}
|
||||||
|
|
||||||
|
test "server enforces a token-bucket rate limit and rejects the overflow" {
|
||||||
|
const allocator = std.testing.allocator;
|
||||||
|
var tmp = std.testing.tmpDir(.{});
|
||||||
|
defer tmp.cleanup();
|
||||||
|
const path = try tmp.dir.realpathAlloc(allocator, ".");
|
||||||
|
defer allocator.free(path);
|
||||||
|
|
||||||
|
var db = try testDb(allocator, path);
|
||||||
|
defer db.close();
|
||||||
|
|
||||||
|
const address = try std.net.Address.parseIp("127.0.0.1", 0);
|
||||||
|
var server = try Server.init(allocator, &db, address);
|
||||||
|
// Slow refill (1 token/sec): the whole burst below completes in well
|
||||||
|
// under a second of wall-clock time, so refill contributes nowhere
|
||||||
|
// near a full token and the outcome is deterministic.
|
||||||
|
server.enableRateLimit(3, 1);
|
||||||
|
try server.start();
|
||||||
|
defer server.stop();
|
||||||
|
|
||||||
|
const client_mod = @import("client.zig");
|
||||||
|
var client = try client_mod.Client.connect(allocator, server.boundAddress());
|
||||||
|
defer client.close();
|
||||||
|
|
||||||
|
var allowed: usize = 0;
|
||||||
|
var rejected: usize = 0;
|
||||||
|
var i: usize = 0;
|
||||||
|
while (i < 6) : (i += 1) {
|
||||||
|
client.ping() catch |err| {
|
||||||
|
try std.testing.expectEqual(client_mod.ClientError.ServerRateLimited, err);
|
||||||
|
rejected += 1;
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
allowed += 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
try std.testing.expectEqual(@as(usize, 3), allowed);
|
||||||
|
try std.testing.expectEqual(@as(usize, 3), rejected);
|
||||||
|
try std.testing.expectEqual(@as(u64, 6), server.requestsTotal.get());
|
||||||
|
try std.testing.expectEqual(@as(u64, 3), server.rateLimitedTotal.get());
|
||||||
|
}
|
||||||
|
|
||||||
|
test "server without an enabled rate limit never rejects" {
|
||||||
|
const allocator = std.testing.allocator;
|
||||||
|
var tmp = std.testing.tmpDir(.{});
|
||||||
|
defer tmp.cleanup();
|
||||||
|
const path = try tmp.dir.realpathAlloc(allocator, ".");
|
||||||
|
defer allocator.free(path);
|
||||||
|
|
||||||
|
var db = try testDb(allocator, path);
|
||||||
|
defer db.close();
|
||||||
|
|
||||||
|
const address = try std.net.Address.parseIp("127.0.0.1", 0);
|
||||||
|
var server = try Server.init(allocator, &db, address);
|
||||||
|
try server.start();
|
||||||
|
defer server.stop();
|
||||||
|
|
||||||
|
const client_mod = @import("client.zig");
|
||||||
|
var client = try client_mod.Client.connect(allocator, server.boundAddress());
|
||||||
|
defer client.close();
|
||||||
|
|
||||||
|
var i: usize = 0;
|
||||||
|
while (i < 20) : (i += 1) try client.ping();
|
||||||
|
try std.testing.expectEqual(@as(u64, 20), server.requestsTotal.get());
|
||||||
|
try std.testing.expectEqual(@as(u64, 0), server.rateLimitedTotal.get());
|
||||||
|
}
|
||||||
|
|
||||||
|
test "per-connection rate limit gives each connection its own budget" {
|
||||||
|
const allocator = std.testing.allocator;
|
||||||
|
var tmp = std.testing.tmpDir(.{});
|
||||||
|
defer tmp.cleanup();
|
||||||
|
const path = try tmp.dir.realpathAlloc(allocator, ".");
|
||||||
|
defer allocator.free(path);
|
||||||
|
|
||||||
|
var db = try testDb(allocator, path);
|
||||||
|
defer db.close();
|
||||||
|
|
||||||
|
const address = try std.net.Address.parseIp("127.0.0.1", 0);
|
||||||
|
var server = try Server.init(allocator, &db, address);
|
||||||
|
server.enablePerConnectionRateLimit(2, 1);
|
||||||
|
try server.start();
|
||||||
|
defer server.stop();
|
||||||
|
|
||||||
|
const client_mod = @import("client.zig");
|
||||||
|
|
||||||
|
var clientA = try client_mod.Client.connect(allocator, server.boundAddress());
|
||||||
|
defer clientA.close();
|
||||||
|
var clientB = try client_mod.Client.connect(allocator, server.boundAddress());
|
||||||
|
defer clientB.close();
|
||||||
|
|
||||||
Reference in new issue
Block a user