jetstream client
atproto jetstream client
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130const std = @import("std");const sdk = @import("jetstream");const LoadHandler = struct { io: std.Io, samples: []i64, count: usize = 0, failed: bool = false, delay_ns: i64, rate: i64, epoch: i64 = 0, pub fn onEvent(self: *@This(), event: sdk.Event) bool { if (event.seq != self.count + 1) { self.failed = true; return false; } var scheduled = std.fmt.parseInt(i64, event.payload.commit.rkey, 10) catch { self.failed = true; return false; }; if (self.count == 0) self.epoch = scheduled; scheduled = self.epoch + @divTrunc(@as(i64, @intCast(self.count)) * std.time.ns_per_s, self.rate); self.samples[self.count] = @intCast(std.Io.Clock.real.now(self.io).toNanoseconds() - scheduled); self.count += 1; if (self.delay_ns > 0) self.io.sleep(std.Io.Duration.fromNanoseconds(self.delay_ns), .awake) catch { self.failed = true; return false; }; return self.count < self.samples.len; } pub fn onError(self: *@This(), _: anyerror) bool { self.failed = true; return false; }};const CatchupHandler = struct { target: usize, count: usize = 0, digest: u64 = 14695981039346656037, failed: bool = false, pub fn onEvent(self: *@This(), event: sdk.Event) bool { self.count += 1; if (event.seq != self.count) { self.failed = true; return false; } for (event.payload.commit.record_cbor) |byte| self.digest = (self.digest ^ byte) *% 1099511628211; return self.count < self.target; } pub fn onError(self: *@This(), _: anyerror) bool { self.failed = true; return false; }};const Handler = struct { target: usize, count: usize = 0, stop_on_error: bool, error_count: usize = 0, max_errors: usize = 0, pub fn onEvent(self: *@This(), event: sdk.Event) bool { const c = event.payload.commit; std.debug.print("EVENT\t{d}\t{s}\t{s}\t{s}\t{s}\n", .{ event.seq, event.did, c.collection, c.rkey, @tagName(c.operation) }); self.count += 1; return self.count < self.target; } pub fn onError(self: *@This(), err: anyerror) bool { std.debug.print("ERROR\nDETAIL {s}\n", .{@errorName(err)}); self.error_count += 1; return !self.stop_on_error and (self.max_errors == 0 or self.error_count < self.max_errors); }};pub fn main(init: std.process.Init) !void { const args = try init.minimal.args.toSlice(init.arena.allocator()); if (args.len == 4 and std.mem.eql(u8, args[2], "catchup")) { const count = try std.fmt.parseInt(usize, args[3], 10); if (count == 0 or count > 1_000_000) return error.InvalidCount; var handler = CatchupHandler{ .target = count + 1 }; std.debug.print("READY\n", .{}); const start = std.Io.Clock.awake.now(init.io); try sdk.subscribe(init.io, init.gpa, .{ .hosts = &.{args[1]}, .after_seq = 0, .zstd_compression = false, .concurrency = 4 }, &handler); if (handler.failed or handler.count != count + 1) return error.InvalidDelivery; const elapsed = start.durationTo(std.Io.Clock.awake.now(init.io)).nanoseconds; std.debug.print("CATCHUP {d} {d} {d}\n", .{ handler.count, handler.digest, elapsed }); return; } if (args.len == 7 and std.mem.eql(u8, args[2], "load")) { const count = try std.fmt.parseInt(usize, args[3], 10); if (count == 0 or count > 1_000_000) return error.InvalidCount; const samples = try init.gpa.alloc(i64, count); defer init.gpa.free(samples); std.debug.print("READY\n", .{}); var handler = LoadHandler{ .io = init.io, .samples = samples, .delay_ns = try std.fmt.parseInt(i64, args[5], 10), .rate = try std.fmt.parseInt(i64, args[6], 10) }; try sdk.subscribe(init.io, init.gpa, .{ .hosts = &.{args[1]}, .zstd_compression = std.mem.eql(u8, args[4], "compressed") }, &handler); if (handler.failed or handler.count != count) return error.InvalidDelivery; std.mem.sort(i64, samples, {}, std.sort.asc(i64)); std.debug.print("LOAD {d} {d} {d} {d}\n", .{ count, samples[count / 2], samples[@min(count - 1, count * 99 / 100)], samples[count - 1] }); return; } if (args.len != 3) return error.ExpectedHostAndScenario; const mode = args[2]; const exact = std.mem.eql(u8, mode, "exact"); const wildcard = std.mem.eql(u8, mode, "wildcard"); const resuming = std.mem.eql(u8, mode, "resume"); const kinds = std.mem.startsWith(u8, mode, "archive-kinds"); const many_dids = std.mem.eql(u8, mode, "many-dids"); var did_storage: [120][32]u8 = undefined; var dids: [120][]const u8 = undefined; const alphabet = "abcdefghijklmnopqrstuvwxyz234567"; for (&did_storage, 0..) |*did, i| { @memcpy(did[0..8], "did:plc:"); @memset(did[8..30], 'a'); did[30] = alphabet[i / 32]; did[31] = alphabet[i % 32]; dids[i] = did; } var handler = Handler{ .target = if (exact or resuming or kinds) 2 else 3, .stop_on_error = std.mem.endsWith(u8, mode, "-stop"), .max_errors = if (many_dids) 3 else 0 }; try sdk.subscribe(init.io, init.gpa, .{ .hosts = &.{args[1]}, .collections = if (exact) &.{"app.bsky.feed.post"} else if (wildcard) &.{"app.bsky.feed.*"} else &.{}, .dids = if (many_dids) &dids else &.{}, .kinds = if (kinds) &.{.commit} else &.{}, .backoff_min_ns = std.time.ns_per_ms, .backoff_max_ns = std.time.ns_per_ms, .live_cursor = if (resuming) 1 else null, .after_seq = if (std.mem.startsWith(u8, mode, "archive-")) 0 else null, .snapshot_only = std.mem.eql(u8, mode, "archive-pages") or kinds, .zstd_compression = !std.mem.eql(u8, mode, "plain") and !std.mem.startsWith(u8, mode, "archive-") and !std.mem.startsWith(u8, mode, "terminal-") and !many_dids, }, &handler);}