atproto utils for zig
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597//! jetstream client - AT Protocol event stream via WebSocket//!//! typed, reconnecting client for the Bluesky Jetstream service.//! parses commit, identity, and account events into typed structs.//!//! see: https://github.com/bluesky-social/jetstream
const std = @import("std");const websocket = @import("websocket");const json_helpers = @import("../xrpc/json.zig");const sync = @import("sync.zig");
const mem = std.mem;const json = std.json;const posix = std.posix;const Allocator = mem.Allocator;const Io = std.Io;const log = std.log.scoped(.zat);
pub const CommitAction = sync.CommitAction;pub const AccountStatus = sync.AccountStatus;
pub const default_hosts = [_][]const u8{ "jetstream1.us-east.bsky.network", "jetstream2.us-east.bsky.network", "jetstream1.us-west.bsky.network", "jetstream2.us-west.bsky.network", "jetstream.waow.tech", "jetstream.fire.hose.cam", "jet.firehose.stream", "sfo.firehose.stream", "nyc.firehose.stream", "london.firehose.stream", "frankfurt.firehose.stream", "chennai.firehose.stream",};
pub const Options = struct { hosts: []const []const u8 = &default_hosts, wanted_collections: []const []const u8 = &.{}, wanted_dids: []const []const u8 = &.{}, cursor: ?i64 = null, max_message_size: usize = 1024 * 1024,};
pub const Event = union(enum) { commit: CommitEvent, identity: IdentityEvent, account: AccountEvent,
pub fn timeUs(self: Event) i64 { return switch (self) { inline else => |e| e.time_us, }; }};
pub const CommitEvent = struct { did: []const u8, time_us: i64, rev: ?[]const u8 = null, operation: CommitAction, collection: []const u8, rkey: []const u8, record: ?json.Value = null, cid: ?[]const u8 = null,};
pub const IdentityEvent = struct { did: []const u8, time_us: i64, handle: ?[]const u8 = null, seq: ?i64 = null, time: ?[]const u8 = null,};
pub const AccountEvent = struct { did: []const u8, time_us: i64, active: bool, status: ?AccountStatus = null, seq: ?i64 = null, time: ?[]const u8 = null,};
/// parse a raw JSON payload into a typed Event./// allocator is used for JSON structural data (ObjectMaps for record fields)./// string slices in the returned Event reference the source `payload` bytes./// keep both `payload` and allocator-owned memory alive while using the Event.pub fn parseEvent(allocator: Allocator, payload: []const u8) !Event { const parsed = try json.parseFromSlice(json.Value, allocator, payload, .{}); const root = parsed.value;
const kind_str = json_helpers.getString(root, "kind") orelse return error.MissingKind; const did = json_helpers.getString(root, "did") orelse return error.MissingDid; const time_us = json_helpers.getInt(root, "time_us") orelse return error.MissingTimeUs;
if (mem.eql(u8, kind_str, "commit")) { const op_str = json_helpers.getString(root, "commit.operation") orelse return error.MissingOperation; return .{ .commit = .{ .did = did, .time_us = time_us, .operation = CommitAction.parse(op_str) orelse return error.UnknownOperation, .collection = json_helpers.getString(root, "commit.collection") orelse return error.MissingCollection, .rkey = json_helpers.getString(root, "commit.rkey") orelse return error.MissingRkey, .rev = json_helpers.getString(root, "commit.rev"), .cid = json_helpers.getString(root, "commit.cid"), .record = json_helpers.getPath(root, "commit.record"), } }; } else if (mem.eql(u8, kind_str, "identity")) { return .{ .identity = .{ .did = did, .time_us = time_us, .handle = json_helpers.getString(root, "identity.handle"), .seq = json_helpers.getInt(root, "identity.seq"), .time = json_helpers.getString(root, "identity.time"), } }; } else if (mem.eql(u8, kind_str, "account")) { const status_str = json_helpers.getString(root, "account.status"); return .{ .account = .{ .did = did, .time_us = time_us, .active = json_helpers.getBool(root, "account.active") orelse true, .status = if (status_str) |s| AccountStatus.parse(s) else null, .seq = json_helpers.getInt(root, "account.seq"), .time = json_helpers.getString(root, "account.time"), } }; }
return error.UnknownKind;}
pub const JetstreamClient = struct { io: Io, allocator: Allocator, options: Options, last_time_us: ?i64 = null,
pub fn init(io: Io, allocator: Allocator, options: Options) JetstreamClient { return .{ .io = io, .allocator = allocator, .options = options, .last_time_us = options.cursor, }; }
pub fn deinit(_: *JetstreamClient) void {}
/// subscribe with a user-provided handler. /// handler must implement: fn onEvent(*@TypeOf(handler), Event) void /// optional: fn onError(*@TypeOf(handler), anyerror) void /// optional: fn onConnect(*@TypeOf(handler), []const u8) void — called with host on connect /// blocks forever — reconnects with exponential backoff on disconnect. /// rotates through hosts on each reconnect attempt. pub fn subscribe(self: *JetstreamClient, handler: anytype) Io.Cancelable!void { var backoff: u64 = 1; var host_index: usize = 0; const max_backoff: u64 = 60; var prev_host_index: usize = 0;
while (true) { const host = self.options.hosts[host_index % self.options.hosts.len]; const effective_index = host_index % self.options.hosts.len;
// rewind cursor by 10s on host switch (different instances may lag) if (host_index > 0 and effective_index != prev_host_index) { if (self.last_time_us) |t| { self.last_time_us = t - 10_000_000; } backoff = 1; }
log.info("connecting to host {d}/{d}: {s}", .{ effective_index + 1, self.options.hosts.len, host });
self.connectAndRead(host, handler) catch |err| { if (comptime @hasDecl(@TypeOf(handler.*), "onError")) { handler.onError(err); } else { log.err("jetstream error: {s}, reconnecting in {d}s...", .{ @errorName(err), backoff }); } };
prev_host_index = effective_index; host_index += 1; try self.io.sleep(Io.Duration.fromSeconds(@intCast(backoff)), .awake); backoff = @min(backoff * 2, max_backoff); } }
fn connectAndRead(self: *JetstreamClient, host: []const u8, handler: anytype) !void { var path_buf: [2048]u8 = undefined; const path = try self.buildSubscribePath(&path_buf);
log.info("connecting to wss://{s}{s}", .{ host, path });
var client = try websocket.Client.init(self.io, self.allocator, .{ .host = host, .port = 443, .tls = true, .max_size = self.options.max_message_size, }); defer client.deinit();
var host_header_buf: [256]u8 = undefined; const host_header = std.fmt.bufPrint(&host_header_buf, "Host: {s}\r\n", .{host}) catch host;
try client.handshake(path, .{ .headers = host_header }); configureKeepalive(&client);
log.info("jetstream connected to {s}", .{host});
if (comptime @hasDecl(@TypeOf(handler.*), "onConnect")) { handler.onConnect(host); }
var ws_handler = WsHandler(@TypeOf(handler.*)){ .allocator = self.allocator, .handler = handler, .client_state = self, }; try client.readLoop(&ws_handler); }
fn buildSubscribePath(self: *JetstreamClient, buf: *[2048]u8) ![]const u8 { var w: std.Io.Writer = .fixed(buf);
try w.writeAll("/subscribe");
var has_param = false;
for (self.options.wanted_collections) |col| { try w.writeByte(if (!has_param) '?' else '&'); try w.writeAll("wantedCollections="); try w.writeAll(col); has_param = true; }
for (self.options.wanted_dids) |did| { try w.writeByte(if (!has_param) '?' else '&'); try w.writeAll("wantedDids="); try w.writeAll(did); has_param = true; }
if (self.last_time_us) |cursor| { try w.writeByte(if (!has_param) '?' else '&'); try w.print("cursor={d}", .{cursor}); }
return w.buffered(); }};
fn WsHandler(comptime H: type) type { return struct { allocator: Allocator, handler: *H, client_state: *JetstreamClient,
const Self = @This();
pub fn serverMessage(self: *Self, data: []const u8) !void { var arena = std.heap.ArenaAllocator.init(self.allocator); defer arena.deinit();
const event = parseEvent(arena.allocator(), data) catch |err| { log.debug("message parse error: {s}", .{@errorName(err)}); return; };
self.client_state.last_time_us = event.timeUs(); self.handler.onEvent(event); }
pub fn close(_: *Self) void { log.info("jetstream connection closed", .{}); } };}
/// enable TCP keepalive so reads don't block forever when a peer/// disappears without FIN/RST (network partition, crash, power loss)./// detection time: 10s idle + 5s × 2 probes = 20s.fn configureKeepalive(client: *websocket.Client) void { const fd = client.stream.stream.socket.handle; const builtin = @import("builtin"); posix.setsockopt(fd, posix.SOL.SOCKET, posix.SO.KEEPALIVE, &std.mem.toBytes(@as(i32, 1))) catch return; const tcp: i32 = @intCast(posix.IPPROTO.TCP); if (builtin.os.tag == .linux) { posix.setsockopt(fd, tcp, posix.TCP.KEEPIDLE, &std.mem.toBytes(@as(i32, 10))) catch return; } else if (builtin.os.tag == .macos) { posix.setsockopt(fd, tcp, posix.TCP.KEEPALIVE, &std.mem.toBytes(@as(i32, 10))) catch return; } posix.setsockopt(fd, tcp, posix.TCP.KEEPINTVL, &std.mem.toBytes(@as(i32, 5))) catch return; posix.setsockopt(fd, tcp, posix.TCP.KEEPCNT, &std.mem.toBytes(@as(i32, 2))) catch return;}
// === tests ===
test "parse commit event" { const payload = \\{ \\ "did": "did:plc:abc123", \\ "time_us": 1700000000000, \\ "kind": "commit", \\ "commit": { \\ "rev": "3mbspmpaidl2a", \\ "operation": "create", \\ "collection": "app.bsky.feed.post", \\ "rkey": "xyz789", \\ "cid": "bafyreitest", \\ "record": { \\ "text": "hello world", \\ "$type": "app.bsky.feed.post" \\ } \\ } \\} ;
var arena = std.heap.ArenaAllocator.init(std.testing.allocator); defer arena.deinit();
const event = try parseEvent(arena.allocator(), payload); const commit = event.commit;
try std.testing.expectEqualStrings("did:plc:abc123", commit.did); try std.testing.expectEqual(@as(i64, 1700000000000), commit.time_us); try std.testing.expectEqualStrings("3mbspmpaidl2a", commit.rev.?); try std.testing.expectEqual(CommitAction.create, commit.operation); try std.testing.expectEqualStrings("app.bsky.feed.post", commit.collection); try std.testing.expectEqualStrings("xyz789", commit.rkey); try std.testing.expectEqualStrings("bafyreitest", commit.cid.?); try std.testing.expect(commit.record != null); try std.testing.expectEqualStrings("hello world", json_helpers.getString(commit.record.?, "text").?);}
test "parse identity event" { const payload = \\{ \\ "did": "did:plc:abc123", \\ "time_us": 1700000000000, \\ "kind": "identity", \\ "identity": { \\ "handle": "alice.bsky.social", \\ "seq": 42, \\ "time": "2024-01-01T00:00:00Z" \\ } \\} ;
var arena = std.heap.ArenaAllocator.init(std.testing.allocator); defer arena.deinit();
const event = try parseEvent(arena.allocator(), payload); const identity = event.identity;
try std.testing.expectEqualStrings("did:plc:abc123", identity.did); try std.testing.expectEqual(@as(i64, 1700000000000), identity.time_us); try std.testing.expectEqualStrings("alice.bsky.social", identity.handle.?); try std.testing.expectEqual(@as(i64, 42), identity.seq.?); try std.testing.expectEqualStrings("2024-01-01T00:00:00Z", identity.time.?);}
test "parse account event" { const payload = \\{ \\ "did": "did:plc:abc123", \\ "time_us": 1700000000000, \\ "kind": "account", \\ "account": { \\ "active": false, \\ "status": "suspended", \\ "seq": 99, \\ "time": "2024-01-01T00:00:00Z" \\ } \\} ;
var arena = std.heap.ArenaAllocator.init(std.testing.allocator); defer arena.deinit();
const event = try parseEvent(arena.allocator(), payload); const account = event.account;
try std.testing.expectEqualStrings("did:plc:abc123", account.did); try std.testing.expectEqual(@as(i64, 1700000000000), account.time_us); try std.testing.expectEqual(false, account.active); try std.testing.expectEqual(AccountStatus.suspended, account.status.?); try std.testing.expectEqual(@as(i64, 99), account.seq.?); try std.testing.expectEqualStrings("2024-01-01T00:00:00Z", account.time.?);}
test "parse unknown kind returns error" { const payload = \\{ \\ "did": "did:plc:abc123", \\ "time_us": 1700000000000, \\ "kind": "unknown_kind" \\} ;
var arena = std.heap.ArenaAllocator.init(std.testing.allocator); defer arena.deinit();
try std.testing.expectError(error.UnknownKind, parseEvent(arena.allocator(), payload));}
test "parse commit with unknown operation returns error" { const payload = \\{ \\ "did": "did:plc:abc123", \\ "time_us": 1700000000000, \\ "kind": "commit", \\ "commit": { \\ "operation": "archive", \\ "collection": "app.bsky.feed.post", \\ "rkey": "xyz789" \\ } \\} ;
var arena = std.heap.ArenaAllocator.init(std.testing.allocator); defer arena.deinit();
try std.testing.expectError(error.UnknownOperation, parseEvent(arena.allocator(), payload));}
test "cursor tracking via time_us" { const payloads = [_][]const u8{ \\{"did":"did:plc:a","time_us":100,"kind":"commit","commit":{"operation":"create","collection":"app.bsky.feed.post","rkey":"1"}} , \\{"did":"did:plc:b","time_us":200,"kind":"commit","commit":{"operation":"create","collection":"app.bsky.feed.post","rkey":"2"}} , };
var arena = std.heap.ArenaAllocator.init(std.testing.allocator); defer arena.deinit();
const e1 = try parseEvent(arena.allocator(), payloads[0]); const e2 = try parseEvent(arena.allocator(), payloads[1]);
try std.testing.expect(e1.timeUs() > 0); try std.testing.expect(e2.timeUs() > e1.timeUs());}
test "Event.timeUs works for all variants" { var arena = std.heap.ArenaAllocator.init(std.testing.allocator); defer arena.deinit();
const commit = try parseEvent(arena.allocator(), \\{"did":"did:plc:a","time_us":100,"kind":"commit","commit":{"operation":"create","collection":"x","rkey":"1"}} ); const identity = try parseEvent(arena.allocator(), \\{"did":"did:plc:a","time_us":200,"kind":"identity","identity":{}} ); const account = try parseEvent(arena.allocator(), \\{"did":"did:plc:a","time_us":300,"kind":"account","account":{"active":true}} );
try std.testing.expectEqual(@as(i64, 100), commit.timeUs()); try std.testing.expectEqual(@as(i64, 200), identity.timeUs()); try std.testing.expectEqual(@as(i64, 300), account.timeUs());}
test "build subscribe path" { var client = JetstreamClient.init(std.Options.debug_io, std.testing.allocator, .{ .wanted_collections = &.{"app.bsky.feed.post"}, });
var buf: [2048]u8 = undefined; const path = try client.buildSubscribePath(&buf); try std.testing.expectEqualStrings("/subscribe?wantedCollections=app.bsky.feed.post", path);}
test "build subscribe path with multiple params" { var client = JetstreamClient.init(std.Options.debug_io, std.testing.allocator, .{ .wanted_collections = &.{ "app.bsky.feed.post", "app.bsky.feed.like" }, .wanted_dids = &.{"did:plc:abc123"}, .cursor = 1700000000000, });
var buf: [2048]u8 = undefined; const path = try client.buildSubscribePath(&buf); try std.testing.expectEqualStrings( "/subscribe?wantedCollections=app.bsky.feed.post&wantedCollections=app.bsky.feed.like&wantedDids=did:plc:abc123&cursor=1700000000000", path, );}
test "build subscribe path no params" { var client = JetstreamClient.init(std.Options.debug_io, std.testing.allocator, .{});
var buf: [2048]u8 = undefined; const path = try client.buildSubscribePath(&buf); try std.testing.expectEqualStrings("/subscribe", path);}
test "parse commit event with delete operation" { const payload = \\{ \\ "did": "did:plc:abc123", \\ "time_us": 1700000000000, \\ "kind": "commit", \\ "commit": { \\ "operation": "delete", \\ "collection": "app.bsky.feed.post", \\ "rkey": "xyz789" \\ } \\} ;
var arena = std.heap.ArenaAllocator.init(std.testing.allocator); defer arena.deinit();
const commit = (try parseEvent(arena.allocator(), payload)).commit;
try std.testing.expectEqual(CommitAction.delete, commit.operation); try std.testing.expect(commit.record == null); try std.testing.expect(commit.rev == null); try std.testing.expect(commit.cid == null);}
test "parse identity event with minimal fields" { const payload = \\{ \\ "did": "did:plc:abc123", \\ "time_us": 1700000000000, \\ "kind": "identity", \\ "identity": {} \\} ;
var arena = std.heap.ArenaAllocator.init(std.testing.allocator); defer arena.deinit();
const identity = (try parseEvent(arena.allocator(), payload)).identity;
try std.testing.expectEqualStrings("did:plc:abc123", identity.did); try std.testing.expect(identity.handle == null); try std.testing.expect(identity.seq == null); try std.testing.expect(identity.time == null);}
test "parse missing did returns error" { const payload = \\{ \\ "time_us": 1700000000000, \\ "kind": "commit" \\} ;
var arena = std.heap.ArenaAllocator.init(std.testing.allocator); defer arena.deinit();
try std.testing.expectError(error.MissingDid, parseEvent(arena.allocator(), payload));}
test "default hosts contains known jetstream instances" { try std.testing.expectEqual(@as(usize, 12), default_hosts.len); try std.testing.expectEqualStrings("jetstream1.us-east.bsky.network", default_hosts[0]); try std.testing.expectEqualStrings("jetstream2.us-east.bsky.network", default_hosts[1]); try std.testing.expectEqualStrings("jetstream1.us-west.bsky.network", default_hosts[2]); try std.testing.expectEqualStrings("jetstream2.us-west.bsky.network", default_hosts[3]); try std.testing.expectEqualStrings("jetstream.waow.tech", default_hosts[4]); try std.testing.expectEqualStrings("jetstream.fire.hose.cam", default_hosts[5]); try std.testing.expectEqualStrings("jet.firehose.stream", default_hosts[6]); try std.testing.expectEqualStrings("chennai.firehose.stream", default_hosts[11]);}
test "round-robin cycles through hosts" { const hosts = [_][]const u8{ "host-a", "host-b", "host-c" }; // simulate the index logic from subscribe() for (0..9) |i| { const host = hosts[i % hosts.len]; const expected: []const u8 = switch (i % 3) { 0 => "host-a", 1 => "host-b", 2 => "host-c", else => unreachable, }; try std.testing.expectEqualStrings(expected, host); }}
test "options default hosts are used" { const opts = Options{}; try std.testing.expectEqual(@as(usize, 12), opts.hosts.len); try std.testing.expectEqualStrings("jetstream1.us-east.bsky.network", opts.hosts[0]);}
test "options custom single host" { const opts = Options{ .hosts = &.{"my-custom-host.example.com"} }; try std.testing.expectEqual(@as(usize, 1), opts.hosts.len); try std.testing.expectEqualStrings("my-custom-host.example.com", opts.hosts[0]);}