From 0c5884fe008aefce92218a4a2884eb62cc687335 Mon Sep 17 00:00:00 2001 From: dawn <90008@klbr.net> Date: Sun, 7 Jun 2026 16:06:17 +0300 Subject: [PATCH] [relay] improve how host connections with stale cursors get treated, hopefully fix a nasty bug with message deduping that caused connections to stall, more debugging --- AGENTS.md | 1 + docs/api/firehose.md | 7 + scripts/compare_relay_host_status.nu | 32 +- scripts/compare_relay_repo_coverage.nu | 254 +++++++++++++++ scripts/compare_repo_statuses.nu | 220 +++++++++++++ src/control/firehose.rs | 100 +++++- src/control/mod.rs | 197 ++++++++++-- src/control/seed.rs | 421 ++++++++++++++++++++----- src/db/keys/mod.rs | 8 - src/db/migration/mod.rs | 16 + src/db/migration/v6.rs | 7 +- src/db/migration/v7.rs | 111 +++++++ src/db/mod.rs | 16 +- src/ingest/firehose.rs | 308 ++++++++++++++++-- src/ingest/relay.rs | 102 ++++-- src/ingest/stream.rs | 4 +- src/types.rs | 136 +++++++- src/util/throttle.rs | 37 ++- tests/api_firehose_sources.nu | 9 + 19 files changed, 1787 insertions(+), 199 deletions(-) create mode 100644 scripts/compare_relay_repo_coverage.nu create mode 100644 scripts/compare_repo_statuses.nu create mode 100644 src/db/migration/v7.rs diff --git a/AGENTS.md b/AGENTS.md index 854c57b..b2192d7 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -89,6 +89,7 @@ See `examples/statusphere.rs` for a usage example. - **Cursors**: Store cursors as big-endian bytes (`u64`/`i64`). - **Compression**: Configurable via `HYDRANT_DATA_COMPRESSION` (`lz4`, `zstd`, `none`). Per-keyspace zstd dictionaries can be trained via `POST /db/train` and are stored as `dict_{keyspace}.bin` in the database directory. - **Keyspaces**: Use the `keys.rs` module to maintain consistent composite key formats. +- **Schema evolution**: Treat versioned wire types in `src/types.rs` as frozen snapshots once shipped. If a stored type changes shape, add a new versioned type, export the newest version for live code, and add a forward migration that explicitly deserializes the previous version and writes the new one. Do not modify older migration input/output types in place. For example, if `RepoState` changes after `v7`, add `v8::RepoState` and a `v7 -> v8` migration instead of mutating `v7`. ## Database schema (keyspaces) diff --git a/docs/api/firehose.md b/docs/api/firehose.md index ca7ef4f..3a34b79 100644 --- a/docs/api/firehose.md +++ b/docs/api/firehose.md @@ -16,6 +16,11 @@ list all known firehose sources, including offline ones waiting for the retry lo "consecutive_failures": 3, "throttled_until": 1717240000, "retry_in_secs": 42, + "last_failure": { + "at": 1717239958, + "kind": "tcp_refused", + "detail": "connection refused" + }, "host_status": "offline", "pds": { "host": "127.0.0.1", @@ -28,6 +33,8 @@ list all known firehose sources, including offline ones waiting for the retry lo `is_pds: true` means the source is a direct PDS connection with host authority enforcement enabled. `host_status` and `pds` are only present for PDS sources. +`last_failure` is present while hydrant has recorded failure/backoff state for a source. `kind` is a compact category such as `dns`, `tcp_refused`, `tcp_timeout`, `tls`, `http_upgrade`, `websocket`, `decode`, `relay_error`, or `config`; `detail` contains the underlying error text. + ### query parameters all filters are exact-match and optional. multiple filters are combined with logical `AND`. diff --git a/scripts/compare_relay_host_status.nu b/scripts/compare_relay_host_status.nu index f84ea0c..75d1073 100644 --- a/scripts/compare_relay_host_status.nu +++ b/scripts/compare_relay_host_status.nu @@ -124,6 +124,32 @@ def fetch-host-status [relay_base: string, host: string, timeout_secs: int] { } } +def host-status-from-list [hosts: list, host: string] { + let matches = ($hosts | where hostname == $host) + + if ($matches | is-empty) { + { + http_status: 404, + hostname: $host, + status: null, + seq: null, + accountCount: null, + error: "not found in listHosts" + } + } else { + let listed = ($matches | first) + $listed + | merge { + http_status: 200, + hostname: ($listed.hostname? | default $host), + status: ($listed.status? | default null), + seq: ($listed.seq? | default null), + accountCount: ($listed.accountCount? | default null), + error: null + } + } +} + def compare-statuses [left: record, right: record] { { status_match: ($left.status == $right.status), @@ -219,9 +245,9 @@ def main [ let rows = ( $host_rows - | par-each --threads $threads --keep-order { |host| - let left = (fetch-host-status $source_relay $host $timeout_secs) - let right = (fetch-host-status $compare_relay $host $timeout_secs) + | each { |host| + let left = (host-status-from-list $source_hosts $host) + let right = (host-status-from-list $compare_hosts $host) let cmp = (compare-statuses $left $right) let comparison = (classify-comparison $left $right $cmp) let source_listed = ($source_hostnames | any {|it| $it == $host }) diff --git a/scripts/compare_relay_repo_coverage.nu b/scripts/compare_relay_repo_coverage.nu new file mode 100644 index 0000000..a0e25c7 --- /dev/null +++ b/scripts/compare_relay_repo_coverage.nu @@ -0,0 +1,254 @@ +#!/usr/bin/env nu + +def curl-json [url: string, timeout_secs: int] { + let result = ( + ^curl + --silent + --show-error + --max-time ($timeout_secs | into string) + $url + | complete + ) + + if $result.exit_code != 0 { + error make { + msg: $"request failed for ($url)" + label: { + text: ($result.stderr | str trim) + span: { start: 0, end: 0 } + } + } + } + + $result.stdout | from json +} + +def list-repos-url [relay_base: string, page_size: int, cursor: string = ""] { + if $cursor == "" { + $"($relay_base)/xrpc/com.atproto.sync.listRepos?limit=($page_size)" + } else { + let encoded_cursor = ($cursor | url encode) + $"($relay_base)/xrpc/com.atproto.sync.listRepos?limit=($page_size)&cursor=($encoded_cursor)" + } +} + +def normalize-repo [repo: record] { + [ + $repo.did + ($repo.active? | default null | into string) + ($repo.status? | default "") + ($repo.rev? | default "") + ($repo.head? | default "") + ] | str join "\t" +} + +def dump-relay-repos [ + relay_base: string, + out_path: string, + page_size: int, + timeout_secs: int, + progress_interval: int +] { + mut cursor = null + mut pages = 0 + mut total = 0 + + [] | save --force $out_path + + loop { + let page = (curl-json (list-repos-url $relay_base $page_size ($cursor | default "")) $timeout_secs) + let repos = ($page.repos? | default []) + + if ($repos | is-empty) { + break + } + + $repos + | each { |repo| (normalize-repo $repo) } + | str join "\n" + | $"($in)\n" + | save --append $out_path + + $pages = ($pages + 1) + $total = ($total + ($repos | length)) + $cursor = ($page.cursor? | default null) + + if $progress_interval > 0 and ($pages mod $progress_interval) == 0 { + print --stderr $"[repo-compare] relay=($relay_base) pages=($pages) repos=($total) cursor=($cursor | default '')" + } + + if $cursor == null { + break + } + } + + { + relay: $relay_base, + pages: $pages, + repos: $total, + out_path: $out_path + } +} + +def sort-repo-file [src: string, dst: string] { + let result = ( + ^sort + --field-separator $'\t' + --key 1,1 + $src + | save --force $dst + ) + $result +} + +def trim-lines [path: string, count: int] { + if $count <= 0 { + [] + } else { + open $path | lines | first $count + } +} + +def count-lines [path: string] { + ( + ^wc + -l + $path + | complete + ).stdout + | str trim + | split row " " + | where $it != "" + | first + | into int +} + +def compare-sorted-repo-files [ + source_sorted: string, + compare_sorted: string, + work_dir: string, + sample_size: int +] { + let source_dids = $"($work_dir)/source.dids.tsv" + let compare_dids = $"($work_dir)/compare.dids.tsv" + let shared_dids = $"($work_dir)/shared.dids.tsv" + let source_only = $"($work_dir)/source_only.tsv" + let compare_only = $"($work_dir)/compare_only.tsv" + let joined = $"($work_dir)/joined.tsv" + let field_mismatches = $"($work_dir)/field_mismatches.tsv" + + ^cut -f1 $source_sorted | save --force $source_dids + ^cut -f1 $compare_sorted | save --force $compare_dids + + ^comm -12 $source_dids $compare_dids | save --force $shared_dids + ^comm -23 $source_dids $compare_dids | save --force $source_only + ^comm -13 $source_dids $compare_dids | save --force $compare_only + + ^join -t $'\t' $source_sorted $compare_sorted | save --force $joined + ^awk -F $'\t' '($2 != $6) || ($3 != $7) || ($4 != $8) || ($5 != $9)' $joined | save --force $field_mismatches + + let source_total = (count-lines $source_sorted) + let compare_total = (count-lines $compare_sorted) + let shared_total = (count-lines $shared_dids) + let source_only_total = (count-lines $source_only) + let compare_only_total = (count-lines $compare_only) + let mismatch_total = (count-lines $field_mismatches) + + { + source_total: $source_total, + compare_total: $compare_total, + shared_total: $shared_total, + source_only_total: $source_only_total, + compare_only_total: $compare_only_total, + field_mismatch_total: $mismatch_total, + source_coverage_vs_compare: ( + if $compare_total == 0 { + null + } else { + ((($shared_total | into float) / ($compare_total | into float)) * 100.0) + } + ), + samples: { + source_only: (trim-lines $source_only $sample_size), + compare_only: (trim-lines $compare_only $sample_size), + field_mismatches: (trim-lines $field_mismatches $sample_size) + }, + paths: { + source_sorted: $source_sorted, + compare_sorted: $compare_sorted, + shared_dids: $shared_dids, + source_only: $source_only, + compare_only: $compare_only, + joined: $joined, + field_mismatches: $field_mismatches + } + } +} + +def main [ + --source-relay: string = "https://relay.klbr.net", + --compare-relay: string = "https://bsky.network", + --page-size: int = 1000, + --timeout-secs: int = 20, + --sample-size: int = 20, + --progress-interval: int = 100, + --keep-workdir, + --workdir: string, + --json +] { + let tmp = ($workdir | default (mktemp -d -t hydrant_repo_compare.XXXXXX)) + let source_raw = $"($tmp)/source.raw.tsv" + let compare_raw = $"($tmp)/compare.raw.tsv" + let source_sorted = $"($tmp)/source.sorted.tsv" + let compare_sorted = $"($tmp)/compare.sorted.tsv" + + let source_dump = (dump-relay-repos $source_relay $source_raw $page_size $timeout_secs $progress_interval) + let compare_dump = (dump-relay-repos $compare_relay $compare_raw $page_size $timeout_secs $progress_interval) + + sort-repo-file $source_raw $source_sorted + sort-repo-file $compare_raw $compare_sorted + + let comparison = (compare-sorted-repo-files $source_sorted $compare_sorted $tmp $sample_size) + let output = { + source_relay: $source_relay, + compare_relay: $compare_relay, + page_size: $page_size, + timeout_secs: $timeout_secs, + sample_size: $sample_size, + progress_interval: $progress_interval, + source_dump: $source_dump, + compare_dump: $compare_dump, + comparison: $comparison, + workdir: $tmp + } + + if (not $keep_workdir) and (not $json) { + print $"workdir: ($tmp)" + } + + if $json { + $output | to json -r | print + } else { + print $"source relay: ($source_relay)" + print $"compare relay: ($compare_relay)" + print $"source repos: ($comparison.source_total)" + print $"compare repos: ($comparison.compare_total)" + print $"shared repos: ($comparison.shared_total)" + print $"source only: ($comparison.source_only_total)" + print $"compare only: ($comparison.compare_only_total)" + print $"field mismatches on shared repos: ($comparison.field_mismatch_total)" + if $comparison.source_coverage_vs_compare != null { + print $"source coverage vs compare: (($comparison.source_coverage_vs_compare | math round -p 4))%" + } + print $"workdir: ($tmp)" + print "" + print "sample compare-only dids:" + ($comparison.samples.compare_only | each { |line| $line } | str join "\n") | print + print "" + print "sample source-only dids:" + ($comparison.samples.source_only | each { |line| $line } | str join "\n") | print + print "" + print "sample field mismatches:" + ($comparison.samples.field_mismatches | each { |line| $line } | str join "\n") | print + } +} diff --git a/scripts/compare_repo_statuses.nu b/scripts/compare_repo_statuses.nu new file mode 100644 index 0000000..bfdda9e --- /dev/null +++ b/scripts/compare_repo_statuses.nu @@ -0,0 +1,220 @@ +#!/usr/bin/env nu + +def curl-json-status [url: string, timeout_secs: int] { + let result = ( + ^curl + --silent + --show-error + --max-time ($timeout_secs | into string) + --write-out "\n%{http_code}" + $url + | complete + ) + + if $result.exit_code != 0 { + return { + http: null, + body: null, + error: ($result.stderr | str trim) + } + } + + let lines = ($result.stdout | lines) + let http = ($lines | last | into int) + let body_text = ($lines | drop nth (($lines | length) - 1) | str join "\n") + let body = (try { $body_text | from json } catch { null }) + + { + http: $http, + body: $body, + error: (if $http == 200 { null } else { ($body.error? | default ($body.message? | default ($body_text | str trim))) }) + } +} + +def fetch-repo-status [relay_base: string, did: string, timeout_secs: int] { + let encoded_did = ($did | url encode) + let res = (curl-json-status $"($relay_base)/xrpc/com.atproto.sync.getRepoStatus?did=($encoded_did)" $timeout_secs) + { + http: $res.http, + active: ($res.body.active? | default null), + rev: ($res.body.rev? | default null), + status: ($res.body.status? | default null), + error: $res.error + } +} + +def resolve-pds-host [did: string, timeout_secs: int] { + let res = (curl-json-status $"https://plc.directory/($did)" $timeout_secs) + if $res.http != 200 or $res.body == null { + return { + pds: null, + host: null, + plc_error: $res.error + } + } + + let service = ( + $res.body.service? + | default [] + | where id == "#atproto_pds" + ) + let pds = (if ($service | is-empty) { null } else { $service | get serviceEndpoint | first }) + let host = ( + if $pds == null { + null + } else { + try { $pds | url parse | get host } catch { null } + } + ) + + { + pds: $pds, + host: $host, + plc_error: null + } +} + +def classify-row [left: record, right: record] { + if $left.http == 200 and $right.http == 200 and $left.active == $right.active and $left.rev == $right.rev and $left.status == $right.status { + "same" + } else if $left.http == 404 and $right.http == 200 { + "missing_on_source" + } else if $left.http == 200 and $right.http == 404 { + "missing_on_compare" + } else if $left.http == 200 and $right.http == 200 and $left.active != $right.active { + "active_mismatch" + } else if $left.http == 200 and $right.http == 200 and $left.status != $right.status { + "status_mismatch" + } else if $left.http == 200 and $right.http == 200 and $left.rev != $right.rev { + "rev_mismatch" + } else { + "other" + } +} + +def summarize [rows: list] { + let grouped = ($rows | group-by classification | transpose classification entries) + let counts = ($grouped | each {|row| {classification: $row.classification count: ($row.entries | length)}} | sort-by classification) + { + total: ($rows | length), + counts: $counts + } +} + +def main [ + dids_file?: string, + --source-relay: string = "http://volsinii:13579", + --compare-relay: string = "https://bsky.network", + --timeout-secs: int = 15, + --threads: int = 16, + --json +] { + let default_dids = [ + "did:web:prod-sea0.stream.place" + "did:plc:v54cvudsxw5bsdb53rz3rpgw" + "did:plc:wc3uljbzjvw6xqq3b2yoirgf" + "did:plc:mdnglmdfkeyskljqgm6urpz6" + "did:plc:2yyvxfpy4s63ars26dzlpbf2" + "did:plc:p3xq5k27o4od3ttbvmd6cvrq" + "did:plc:4mnxkcsm6rvflpfujxvykald" + "did:plc:j2vrb6so7edkfy4ax6vztdke" + "did:plc:b6dl7ze2sawauxyacb6xy3cf" + "did:plc:5roprrkswuzn5cyvtddpvszh" + "did:plc:p5pcswhxiigbi4lefml2qktv" + "did:plc:xb2urvqt5f4zzccjs46hysbf" + "did:plc:vm3bwfvsdsfxiyqwepdxchmx" + "did:plc:waekfwabp6s6e546dozi5ryy" + "did:plc:bhvrrngtey3i3amzy4ndb5aq" + "did:plc:uiv7petmywo5zjcmburvucbp" + "did:plc:rajqfvoufme434pc55wqed7x" + "did:plc:oztyg3gp2jt74y6k53n2q3cu" + "did:plc:jvlk3extyrhndypkhu3372yz" + "did:plc:lkfazkg3ejp3oosadwmq6hhj" + "did:plc:4wc4qpf7gpsktpri4lnmqykn" + "did:plc:clii72ny72fpagrywk36ypig" + "did:plc:pppq65ew7bw6tyeppgmcwojh" + "did:plc:64p3c6ff3ccxc3kwftctvzx3" + "did:plc:rtawxkoomhd4c5kpvn5vp3kf" + "did:plc:z37762lrjhrqr4ghs7v2i3vj" + "did:plc:chewvd2le7fnipvfdnku3omq" + "did:plc:4wbbbi22wxvaneqlokv4rayc" + "did:plc:yiwegfgtew7hhra7f2fuhqha" + "did:plc:75kzcv3wdwajpgazpywco6d7" + ] + + let dids = ( + if $dids_file == null { + $default_dids + } else { + open $dids_file | lines | each { |line| $line | str trim } | where $it != "" + } + ) + + let rows = ( + $dids + | par-each --threads $threads --keep-order { |did| + let source = (fetch-repo-status $source_relay $did $timeout_secs) + let compare = (fetch-repo-status $compare_relay $did $timeout_secs) + let resolved = (resolve-pds-host $did $timeout_secs) + let classification = (classify-row $source $compare) + + { + did: $did, + host: $resolved.host, + pds: $resolved.pds, + classification: $classification, + source: $source, + compare: $compare, + plc_error: $resolved.plc_error + } + } + ) + + let grouped_by_host = ( + $rows + | group-by host + | transpose host entries + | each {|row| + { + host: $row.host, + count: ($row.entries | length), + classifications: ( + $row.entries + | group-by classification + | transpose classification vals + | each {|c| {classification: $c.classification count: ($c.vals | length)}} + | sort-by classification + ), + dids: ($row.entries | get did) + } + } + | sort-by count -r + ) + + if $json { + { + source_relay: $source_relay, + compare_relay: $compare_relay, + summary: (summarize $rows), + grouped_by_host: $grouped_by_host, + rows: $rows + } + | to json -r + | print + return + } + + print $"source relay: ($source_relay)" + print $"compare relay: ($compare_relay)" + print "" + print "summary:" + summarize $rows | table + print "" + print "grouped by host:" + $grouped_by_host | table -e + print "" + print "rows:" + $rows + | select did host classification source.http compare.http source.active compare.active source.status compare.status source.rev compare.rev source.error compare.error plc_error + | table -e +} diff --git a/src/control/firehose.rs b/src/control/firehose.rs index d92c994..255c131 100644 --- a/src/control/firehose.rs +++ b/src/control/firehose.rs @@ -5,7 +5,7 @@ use std::time::Duration; use miette::{IntoDiagnostic, Result}; use rand::RngExt; use tokio_util::sync::CancellationToken; -use tracing::{error, info}; +use tracing::{debug, error, info}; use url::Url; use crate::config::FirehoseSource; @@ -39,6 +39,14 @@ pub struct FirehosePdsInfo { pub status: &'static str, } +/// details for the most recent recorded firehose source failure. +#[derive(Debug, Clone, serde::Serialize)] +pub struct FirehoseFailureInfo { + pub at: i64, + pub kind: String, + pub detail: String, +} + /// a snapshot of a single firehose relay's runtime state. #[derive(Debug, Clone, serde::Serialize)] pub struct FirehoseSourceInfo { @@ -54,6 +62,8 @@ pub struct FirehoseSourceInfo { #[serde(skip_serializing_if = "Option::is_none")] pub retry_in_secs: Option, #[serde(skip_serializing_if = "Option::is_none")] + pub last_failure: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub host_status: Option<&'static str>, #[serde(skip_serializing_if = "Option::is_none")] pub pds: Option, @@ -120,12 +130,24 @@ impl FirehoseHandle { tokio::spawn({ let relay_url = source.url.clone(); + let is_pds = source.is_pds; let tasks = self.tasks.clone(); let token = cancel.clone(); async move { // jitter connection start so we dont cause thundering herd problems if delay_startup { - let jitter_ms = rand::rng().random_range(0u64..2000); + let max_jitter_ms = if is_pds && !cfg!(debug_assertions) { + 60_000 + } else { + 2_000 + }; + let jitter_ms = rand::rng().random_range(0u64..max_jitter_ms); + debug!( + relay = %relay_url, + is_pds, + jitter_ms, + "delaying firehose ingestor startup" + ); tokio::select! { _ = tokio::time::sleep(Duration::from_millis(jitter_ms)) => {} _ = token.cancelled() => { @@ -191,6 +213,15 @@ impl FirehoseHandle { let throttle = self.state.throttler.snapshot(url); let pds = is_pds.then(|| self.pds_info(url, &meta)).flatten(); let host_status = pds.as_ref().map(|pds| pds.status); + let last_failure = + throttle + .last_failure + .clone() + .map(|failure| FirehoseFailureInfo { + at: failure.at, + kind: failure.kind, + detail: failure.detail, + }); out.push(FirehoseSourceInfo { url: url.clone(), is_pds, @@ -201,6 +232,7 @@ impl FirehoseHandle { throttled_until: (throttle.throttled_until != 0) .then_some(throttle.throttled_until), retry_in_secs: throttle.retry_in_secs(now), + last_failure, host_status, pds, }); @@ -281,6 +313,70 @@ impl FirehoseHandle { Ok(()) } + /// add PDS sources discovered from a seed relay. + /// + /// seed pages can contain thousands of hosts. persist the whole page in one + /// batch and stagger startup so a fresh relay does not stampede DNS, TCP, + /// TLS, and remote PDS websocket endpoints at once. + pub(super) async fn add_seeded_sources(&self, urls: Vec) -> Result { + let shared = self + .shared + .get() + .ok_or_else(|| miette::miette!("firehose worker not started"))?; + + let mut sources = Vec::with_capacity(urls.len()); + for url in urls { + if self.is_source_known(&url) { + continue; + } + sources.push(FirehoseSource { url, is_pds: true }); + } + + if sources.is_empty() { + return Ok(0); + } + + tokio::task::spawn_blocking({ + let state = self.state.clone(); + let sources = sources.clone(); + move || { + let mut batch = state.db.inner.batch(); + for source in &sources { + let value = rmp_serde::to_vec(&db::FirehoseSourceMeta { + is_pds: source.is_pds, + }) + .map_err(|e| { + miette::miette!("failed to serialize firehose source meta: {e}") + })?; + batch.insert( + &state.db.crawler, + keys::firehose_source_key(source.url.as_str()), + &value, + ); + } + batch.commit().into_diagnostic()?; + state.db.persist() + } + }) + .await + .into_diagnostic()??; + + let mut added = 0usize; + for source in sources { + if self + .known_sources + .insert_async(source.url.clone(), true) + .await + .is_ok() + { + self.spawn_firehose_ingestor(&source, shared, true).await?; + added += 1; + } + } + + Ok(added) + } + /// remove a firehose source at runtime. /// /// returns `true` if the source was found and removed, `false` otherwise. diff --git a/src/control/mod.rs b/src/control/mod.rs index 301bf7d..7086a3f 100644 --- a/src/control/mod.rs +++ b/src/control/mod.rs @@ -30,6 +30,7 @@ pub use repos::{ListedRecord, Record, RecordList, RepoHandle, RepoInfo, ReposCon use smol_str::{SmolStr, ToSmolStr}; use std::collections::BTreeMap; +use std::collections::BTreeSet; use std::future::Future; use std::pin::Pin; use std::sync::Arc; @@ -39,7 +40,7 @@ use std::task::{Context, Poll}; use futures::{FutureExt, Stream}; use miette::{IntoDiagnostic, Result, WrapErr}; use tokio::sync::{mpsc, watch}; -use tracing::{debug, error, info}; +use tracing::{debug, error, info, warn}; #[cfg(feature = "indexer")] use crate::backfill::BackfillWorker; @@ -467,6 +468,13 @@ impl Hydrant { .expect("firehose shared already set"); let fire_shared = firehose.shared.get().unwrap(); + // refresh seed snapshots before spawning any direct PDS sources. persisted + // sources may predate cursor seeding, and an already-open websocket will not + // pick up a later cursor update until it reconnects. + if !config.seed_hosts.is_empty() { + seed::refresh_seed_snapshots(&config.seed_hosts, &state).await; + } + // add hosts from config let relay_hosts = config.relays.clone(); if !relay_hosts.is_empty() { @@ -510,9 +518,8 @@ impl Hydrant { .await?; } // we use spawn_firehose_ingestor directly here since we dont want - // to go through the whole add_source machinery and checks - // its ok since we block here before running stuff like seed_hosts - // and whatnot + // to go through the whole add_source machinery and checks. seed + // cursor snapshots were refreshed before this loop. // 10c. seed firehose PDS sources from listHosts on configured seed URLs if !config.seed_hosts.is_empty() { @@ -533,6 +540,7 @@ impl Hydrant { tokio::time::sleep(retry_interval).await; let mut to_restart: Vec<(Url, bool)> = Vec::new(); + let mut banned = 0usize; { let meta = firehose.state.pds_meta.load(); firehose @@ -543,6 +551,7 @@ impl Hydrant { } let host = url.host_str().unwrap_or(url.as_str()); if meta.is_banned(host) { + banned += 1; return true; } to_restart.push((url.clone(), is_pds)); @@ -551,8 +560,19 @@ impl Hydrant { .await; } + if !to_restart.is_empty() || banned != 0 { + info!( + restart_count = to_restart.len(), + banned, + retry_interval_secs = retry_interval.as_secs(), + "offline firehose retry tick" + ); + } + for (url, is_pds) in to_restart { - let _ = firehose.restart_source(url, is_pds).await; + if let Err(e) = firehose.restart_source(url.clone(), is_pds).await { + warn!(url = %url, is_pds, err = %e, "failed to restart firehose source"); + } } } } @@ -925,43 +945,68 @@ impl Hydrant { *end.last_mut().unwrap() += 1; end }; - let start_bound = match cursor.as_deref() { - Some(host) => std::ops::Bound::Excluded(keys::firehose_cursor_key(host)), - None => std::ops::Bound::Included(keys::FIREHOSE_CURSOR_PREFIX.to_vec()), - }; - // fetch one extra item to detect whether there is a next page - let mut hosts: Vec = Vec::with_capacity(limit + 1); - for item in state - .db - .cursors - .range((start_bound, std::ops::Bound::Excluded(prefix_end))) - .take(limit + 1) - { - let (k, v) = item.into_inner().into_diagnostic()?; + let mut all_hostnames = BTreeSet::new(); + for item in state.db.cursors.range(( + std::ops::Bound::Included(keys::FIREHOSE_CURSOR_PREFIX.to_vec()), + std::ops::Bound::Excluded(prefix_end), + )) { + let (k, _) = item.into_inner().into_diagnostic()?; let hostname = std::str::from_utf8(&k[keys::FIREHOSE_CURSOR_PREFIX.len()..]) .into_diagnostic() .wrap_err("firehose cursor key contains non-utf8 hostname")?; - let seq = i64::from_be_bytes( - v.as_ref() - .try_into() - .into_diagnostic() - .wrap_err("cursor value is not 8 bytes")?, - ); + all_hostnames.insert(SmolStr::new(hostname)); + } + + { + let meta = state.pds_meta.load(); + for hostname in meta.hosts.keys() { + all_hostnames.insert(SmolStr::new(hostname)); + } + } + + let hostnames = all_hostnames.into_iter().collect::>(); + let start_idx = cursor + .as_deref() + .map(|after| hostnames.partition_point(|host| host.as_str() <= after)) + .unwrap_or(0); + + let selected = hostnames + .iter() + .skip(start_idx) + .take(limit + 1) + .cloned() + .collect::>(); + + let mut hosts: Vec = Vec::with_capacity(selected.len().min(limit)); + for hostname in selected.iter().take(limit) { + let seq = state + .db + .cursors + .get(keys::firehose_cursor_key(hostname)) + .into_diagnostic()? + .map(|v| { + v.as_ref() + .try_into() + .into_diagnostic() + .wrap_err("cursor value is not 8 bytes") + .map(i64::from_be_bytes) + }) + .transpose()? + .unwrap_or(0); let account_count = state .db .get_count_sync(&keys::pds_account_count_key(hostname)); let status = state.pds_meta.load().status(hostname); hosts.push(Host { - name: hostname.into(), + name: hostname.clone(), seq, account_count, status, }); } - let next_cursor = if hosts.len() > limit { - hosts.pop(); + let next_cursor = if selected.len() > limit { hosts.last().map(|h| h.name.clone()) } else { None @@ -1114,3 +1159,101 @@ mod tests { Ok(()) } } + +#[cfg(test)] +mod host_listing_tests { + use super::*; + use crate::db::{keys, set_ks_count}; + use crate::pds_meta::HostStatus; + + fn test_config(path: &std::path::Path) -> Config { + Config { + database_path: path.to_path_buf(), + ..Default::default() + } + } + + #[tokio::test] + async fn list_hosts_includes_seeded_hosts_without_cursors() -> Result<()> { + let tmp = tempfile::tempdir().into_diagnostic()?; + let hydrant = Hydrant::new(test_config(tmp.path())).await?; + + { + let state = hydrant.state.clone(); + tokio::task::spawn_blocking(move || -> Result<()> { + let mut batch = state.db.inner.batch(); + crate::db::pds_meta::set_status( + &mut batch, + &state.db.filter, + "offline.example", + HostStatus::Offline, + )?; + crate::db::pds_meta::set_status( + &mut batch, + &state.db.filter, + "active.example", + HostStatus::Active, + )?; + set_ks_count( + &mut batch, + &state.db, + &keys::pds_account_count_key("offline.example"), + 7, + ); + set_ks_count( + &mut batch, + &state.db, + &keys::pds_account_count_key("active.example"), + 42, + ); + state + .db + .cursors + .insert( + keys::firehose_cursor_key("active.example"), + 123_i64.to_be_bytes(), + ) + .into_diagnostic()?; + batch.commit().into_diagnostic()?; + state.db.persist() + }) + .await + .into_diagnostic()??; + + crate::pds_meta::PdsMeta::update_host( + &hydrant.state.pds_meta, + "offline.example", + |h| { + h.status = HostStatus::Offline; + }, + ); + crate::pds_meta::PdsMeta::update_host(&hydrant.state.pds_meta, "active.example", |h| { + h.status = HostStatus::Active; + }); + hydrant.state.db.update_count("p|offline.example", 7); + hydrant.state.db.update_count("p|active.example", 42); + } + + let (hosts, next) = hydrant.list_hosts(None, 10).await?; + assert!(next.is_none()); + + let offline = hosts + .iter() + .find(|h| h.name == "offline.example") + .expect("seeded host without cursor should be listed"); + let active = hosts + .iter() + .find(|h| h.name == "active.example") + .expect("host with cursor should be listed"); + + assert_eq!(offline.seq, 0); + assert_eq!(offline.account_count, 7); + assert_eq!(offline.status, HostStatus::Offline); + + assert_eq!(active.seq, 123); + assert_eq!(active.account_count, 42); + assert_eq!(active.status, HostStatus::Active); + + Ok(()) + } +} diff --git a/src/control/seed.rs b/src/control/seed.rs index f8f67b3..d68ce67 100644 --- a/src/control/seed.rs +++ b/src/control/seed.rs @@ -1,18 +1,20 @@ use std::sync::Arc; +use std::sync::atomic::{AtomicI64, Ordering}; use std::time::Duration; use futures::StreamExt; use jacquard_api::com_atproto::sync::HostStatus; -use jacquard_api::com_atproto::sync::list_hosts::ListHostsOutput; -use miette::IntoDiagnostic; -use tracing::{info, warn}; +use jacquard_api::com_atproto::sync::list_hosts::{Host, ListHostsOutput}; +use miette::{Context, IntoDiagnostic}; +use tracing::{debug, info, warn}; use url::Url; use super::firehose::FirehoseHandle; -use crate::db::{self, keys}; +use crate::db::keys; use crate::state::AppState; const MAX_CONCURRENT_SEEDS: usize = 4; +const LIST_HOSTS_BODY_PREVIEW_BYTES: usize = 512; /// seed firehose pds sources by calling `com.atproto.sync.listHosts` on each seed URL. /// banned pds' are not added, everything else is (including offline) @@ -23,79 +25,98 @@ pub(crate) async fn seed_from_list_hosts( ) { info!("will seed urls..."); - let http = reqwest::Client::builder() - .user_agent(concat!( - env!("CARGO_PKG_NAME"), - "/", - env!("CARGO_PKG_VERSION") - )) - .timeout(Duration::from_secs(10)) - .build() - .expect("that reqwest will build"); + let http = seed_http_client(); let mut futs = futures::stream::iter(seed_urls.iter().cloned()) .map(|seed_url| { - let firehose = firehose.clone(); + let firehose = Some(firehose.clone()); let state = state.clone(); let http = http.clone(); - async move { seed_one(&seed_url, &firehose, &state, &http).await } + async move { seed_one(&seed_url, firehose, &state, &http).await } }) .buffer_unordered(MAX_CONCURRENT_SEEDS); while futs.next().await.is_some() {} } +/// refresh seed relay host status/cursor snapshots without adding sources. +/// +/// this runs before persisted sources are spawned so existing PDS tasks start +/// from the seed relay's current cursor instead of racing ahead with no cursor. +pub(crate) async fn refresh_seed_snapshots(seed_urls: &[Url], state: &Arc) { + info!("will refresh seed snapshots..."); + + let http = seed_http_client(); + + let mut futs = futures::stream::iter(seed_urls.iter().cloned()) + .map(|seed_url| { + let state = state.clone(); + let http = http.clone(); + async move { seed_one(&seed_url, None, &state, &http).await } + }) + .buffer_unordered(MAX_CONCURRENT_SEEDS); + + while futs.next().await.is_some() {} +} + +fn seed_http_client() -> reqwest::Client { + reqwest::Client::builder() + .user_agent(concat!( + env!("CARGO_PKG_NAME"), + "/", + env!("CARGO_PKG_VERSION") + )) + .timeout(Duration::from_secs(10)) + .build() + .expect("that reqwest will build") +} + #[tracing::instrument(skip_all, fields(seed_url = %seed_url))] async fn seed_one( seed_url: &Url, - firehose: &FirehoseHandle, + firehose: Option, state: &Arc, http: &reqwest::Client, ) { - let cursor_key = keys::seed_cursor_key(seed_url.as_str()); - - // resume from the last saved cursor so we don't re-page through already-seen hosts - let mut cursor: Option = { - let ks = state.db.cursors.clone(); - let key = cursor_key.clone(); - match db::Db::get(ks, key).await { - Ok(Some(b)) => rmp_serde::from_slice::(b.as_ref()).ok(), - Ok(None) => None, - Err(e) => { - warn!(err = %e, "failed to load seed cursor, starting from scratch"); - None - } - } - }; - - if cursor.is_some() { - info!(cursor = ?cursor, "resuming seed from saved cursor"); - } else { - info!("seeding firehose sources from listHosts"); - } + // always start from the beginning. `listHosts` is small, and resuming from + // an old completed-run cursor hides earlier hosts after a restart. + let mut cursor: Option = None; + info!("seeding firehose sources from listHosts"); let mut total = 0usize; let mut added = 0usize; loop { let url = list_hosts_url(seed_url, cursor.as_deref()); - let resp = match http.get(url).send().await { + let resp = match http.get(url.clone()).send().await { Ok(r) => r, Err(e) => { - warn!(err = %e, "failed to fetch listHosts, stopping"); + warn!(url = %url, err = %e, "failed to fetch listHosts, stopping"); break; } }; if !resp.status().is_success() { - warn!(status = %resp.status(), "listHosts returned error status, stopping"); + let status = resp.status(); + let body = resp + .bytes() + .await + .ok() + .map(|bytes| body_preview(&bytes)) + .unwrap_or_else(|| "".to_string()); + warn!( + url = %url, + status = %status, + body = %body, + "listHosts returned error status, stopping" + ); break; } let bytes = match resp.bytes().await { Ok(b) => b, Err(e) => { - warn!(err = %e, "failed to read listHosts response, stopping"); + warn!(url = %url, err = %e, "failed to read listHosts response, stopping"); break; } }; @@ -103,69 +124,84 @@ async fn seed_one( let body: ListHostsOutput<'_> = match serde_json::from_slice(&bytes) { Ok(b) => b, Err(e) => { - warn!(err = %e, "failed to parse listHosts response, stopping"); + warn!( + url = %url, + err = %e, + body = %body_preview(&bytes), + "failed to parse listHosts response, stopping" + ); break; } }; let next_cursor = body.cursor.as_deref().map(str::to_owned); total += body.hosts.len(); + let page_hosts = body.hosts.len(); - for host in &body.hosts { - // skip banned hosts; everything else (active, idle, offline, throttled) is included - // since the firehose ingestor handles reconnection for transiently-unavailable hosts - if matches!(host.status, Some(HostStatus::Banned)) { - continue; - } + if let Err(e) = apply_seed_snapshot(state, &body.hosts) { + warn!(err = %e, "failed to apply listHosts seed snapshot"); + } - let wss_url_str = format!("wss://{}/", host.hostname); - let wss_url = match Url::parse(&wss_url_str) { - Ok(u) => u, - Err(e) => { - warn!(hostname = %host.hostname, err = %e, "invalid hostname in listHosts response, skipping"); + let mut banned = 0usize; + let mut invalid = 0usize; + let mut already_known = 0usize; + let mut queued = 0usize; + let mut page_added = 0usize; + if let Some(firehose) = firehose.as_ref() { + let mut seed_sources = Vec::with_capacity(body.hosts.len()); + for host in &body.hosts { + // skip banned hosts; everything else (active, idle, offline, throttled) is included + // since the firehose ingestor handles reconnection for transiently-unavailable hosts + if matches!(host.status, Some(HostStatus::Banned)) { + banned += 1; continue; } - }; - // skip sources that are already running - if firehose.tasks.contains_async(&wss_url).await { - continue; - } + let wss_url_str = format!("wss://{}/", host.hostname); + let wss_url = match Url::parse(&wss_url_str) { + Ok(u) => u, + Err(e) => { + invalid += 1; + warn!(hostname = %host.hostname, err = %e, "invalid hostname in listHosts response, skipping"); + continue; + } + }; - match firehose.add_source(wss_url, true).await { - Ok(()) => added += 1, - Err(e) => { - warn!(hostname = %host.hostname, err = %e, "failed to add firehose source"); + // skip sources that are already tracked; offline retries are handled separately + if firehose.is_source_known(&wss_url) { + already_known += 1; + continue; } - } - } - cursor = next_cursor; + queued += 1; + seed_sources.push(wss_url); + } - // persist cursor after each page so a restart can resume where we left off - if let Some(ref c) = cursor { - let value = match rmp_serde::to_vec(c) { - Ok(v) => v, + match firehose.add_seeded_sources(seed_sources).await { + Ok(n) => { + page_added = n; + added += n; + } Err(e) => { - warn!(err = %e, "failed to serialize seed cursor"); - continue; + warn!(err = %e, "failed to add seeded firehose sources"); } - }; - let state = state.clone(); - let key: Vec = cursor_key.clone(); - let result = tokio::task::spawn_blocking(move || -> miette::Result<()> { - let mut batch = state.db.inner.batch(); - batch.insert(&state.db.cursors, key, &value); - batch.commit().into_diagnostic() - }) - .await - .into_diagnostic() - .flatten(); - if let Err(e) = result { - warn!(err = %e, "failed to persist seed cursor"); } } + info!( + page_hosts, + banned, + invalid, + already_known, + queued, + added = page_added, + adding_sources = firehose.is_some(), + next_cursor = ?next_cursor, + "processed listHosts page" + ); + + cursor = next_cursor; + if cursor.is_none() { break; } @@ -177,6 +213,15 @@ async fn seed_one( ); } +fn body_preview(bytes: &[u8]) -> String { + let len = bytes.len().min(LIST_HOSTS_BODY_PREVIEW_BYTES); + let mut body = String::from_utf8_lossy(&bytes[..len]).into_owned(); + if bytes.len() > len { + body.push_str("..."); + } + body +} + fn list_hosts_url(base: &Url, cursor: Option<&str>) -> Url { let mut url = base.clone(); url.set_path("/xrpc/com.atproto.sync.listHosts"); @@ -189,3 +234,211 @@ fn list_hosts_url(base: &Url, cursor: Option<&str>) -> Url { } url } + +fn apply_seed_snapshot(state: &Arc, hosts: &[Host<'_>]) -> miette::Result<()> { + let mut batch = state.db.inner.batch(); + let mut cursor_updates = Vec::with_capacity(hosts.len()); + let mut status_updates = Vec::with_capacity(hosts.len()); + + for host in hosts { + let hostname = host.hostname.as_ref(); + let status = map_seed_status(host.status.as_ref()); + + crate::db::pds_meta::set_status(&mut batch, &state.db.filter, hostname, status)?; + status_updates.push((hostname.to_string(), status)); + + let Some(seq) = host + .seq + .and_then(|seq| i64::try_from(seq).ok()) + .filter(|seq| *seq > 0) + else { + continue; + }; + + let cursor_key = keys::firehose_cursor_key(hostname); + let existing_seq = state + .db + .cursors + .get(&cursor_key) + .into_diagnostic()? + .map(|bytes| { + bytes + .as_ref() + .try_into() + .into_diagnostic() + .wrap_err("cursor value is not 8 bytes") + .map(i64::from_be_bytes) + }) + .transpose()? + .unwrap_or(0); + if seq > existing_seq { + batch.insert(&state.db.cursors, cursor_key, seq.to_be_bytes()); + cursor_updates.push((hostname.to_string(), seq)); + } + } + + batch.commit().into_diagnostic()?; + debug!( + hosts = hosts.len(), + status_updates = status_updates.len(), + cursor_updates = cursor_updates.len(), + "applied listHosts seed snapshot" + ); + + state.pds_meta.rcu(|meta| { + let mut next = (**meta).clone(); + for (hostname, status) in &status_updates { + next.update_host_entry(hostname, |entry| entry.status = *status); + } + next + }); + for (hostname, seq) in cursor_updates { + let Ok(url) = Url::parse(&format!("wss://{hostname}/")) else { + continue; + }; + let _ = state + .firehose_cursors + .insert_sync(url.clone(), AtomicI64::new(seq)); + state.firehose_cursors.peek_with(&url, |_, cursor| { + if seq > cursor.load(Ordering::SeqCst) { + cursor.store(seq, Ordering::SeqCst); + } + }); + } + + Ok(()) +} + +fn map_seed_status(status: Option<&HostStatus<'_>>) -> crate::pds_meta::HostStatus { + match status { + Some(HostStatus::Active) | None => crate::pds_meta::HostStatus::Active, + Some(HostStatus::Idle) => crate::pds_meta::HostStatus::Idle, + Some(HostStatus::Offline) => crate::pds_meta::HostStatus::Offline, + Some(HostStatus::Throttled) => crate::pds_meta::HostStatus::Throttled, + Some(HostStatus::Banned) => crate::pds_meta::HostStatus::Banned, + Some(HostStatus::Other(_)) => crate::pds_meta::HostStatus::Active, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::Config; + use jacquard_common::CowStr; + use tempfile::tempdir; + + fn persisted_cursor(state: &AppState, hostname: &str) -> miette::Result> { + state + .db + .cursors + .get(keys::firehose_cursor_key(hostname)) + .into_diagnostic()? + .map(|bytes| { + bytes + .as_ref() + .try_into() + .into_diagnostic() + .map(i64::from_be_bytes) + }) + .transpose() + } + + #[test] + fn apply_seed_snapshot_persists_statuses_and_cursors() -> miette::Result<()> { + let tmp = tempdir().into_diagnostic()?; + let cfg = Config { + database_path: tmp.path().to_path_buf(), + ..Default::default() + }; + let state = Arc::new(AppState::new(&cfg)?); + + let hosts = vec![ + Host { + hostname: CowStr::Borrowed("active.example"), + account_count: Some(42), + seq: Some(100), + status: Some(HostStatus::Active), + extra_data: None, + }, + Host { + hostname: CowStr::Borrowed("offline.example"), + account_count: Some(7), + seq: Some(5), + status: Some(HostStatus::Offline), + extra_data: None, + }, + ]; + + apply_seed_snapshot(&state, &hosts)?; + + assert_eq!( + state + .db + .get_count_sync(&keys::pds_account_count_key("active.example")), + 0 + ); + assert_eq!( + state + .db + .get_count_sync(&keys::pds_account_count_key("offline.example")), + 0 + ); + + let meta = state.pds_meta.load(); + assert_eq!( + meta.status("active.example"), + crate::pds_meta::HostStatus::Active + ); + assert_eq!( + meta.status("offline.example"), + crate::pds_meta::HostStatus::Offline + ); + assert!(meta.hosts.contains_key("active.example")); + assert!(meta.hosts.contains_key("offline.example")); + + assert_eq!(persisted_cursor(&state, "active.example")?, Some(100)); + assert_eq!(persisted_cursor(&state, "offline.example")?, Some(5)); + + let active_url = Url::parse("wss://active.example/").into_diagnostic()?; + let in_memory = state + .firehose_cursors + .peek_with(&active_url, |_, cursor| cursor.load(Ordering::SeqCst)); + assert_eq!(in_memory, Some(100)); + + Ok(()) + } + + #[test] + fn apply_seed_snapshot_does_not_lower_existing_cursor() -> miette::Result<()> { + let tmp = tempdir().into_diagnostic()?; + let cfg = Config { + database_path: tmp.path().to_path_buf(), + ..Default::default() + }; + let state = Arc::new(AppState::new(&cfg)?); + let url = Url::parse("wss://active.example/").into_diagnostic()?; + + crate::db::set_firehose_cursor(&state.db, &url, 150)?; + let _ = state + .firehose_cursors + .insert_sync(url.clone(), AtomicI64::new(150)); + + let hosts = vec![Host { + hostname: CowStr::Borrowed("active.example"), + account_count: Some(42), + seq: Some(100), + status: Some(HostStatus::Active), + extra_data: None, + }]; + + apply_seed_snapshot(&state, &hosts)?; + + assert_eq!(persisted_cursor(&state, "active.example")?, Some(150)); + let in_memory = state + .firehose_cursors + .peek_with(&url, |_, cursor| cursor.load(Ordering::SeqCst)); + assert_eq!(in_memory, Some(150)); + + Ok(()) + } +} diff --git a/src/db/keys/mod.rs b/src/db/keys/mod.rs index 1f0c502..f2db6c0 100644 --- a/src/db/keys/mod.rs +++ b/src/db/keys/mod.rs @@ -146,14 +146,6 @@ pub fn did_collection_prefix(did: &Did) -> Vec { key } -pub const SEED_CURSOR_PREFIX: &[u8] = b"seed_cursor|"; - -pub fn seed_cursor_key(url: &str) -> Vec { - let mut key = SEED_CURSOR_PREFIX.to_vec(); - key.extend_from_slice(url.as_bytes()); - key -} - pub const FIREHOSE_CURSOR_PREFIX: &[u8] = b"firehose_cursor|"; pub const FIREHOSE_SOURCE_PREFIX: &[u8] = b"firehose|"; diff --git a/src/db/migration/mod.rs b/src/db/migration/mod.rs index 61c0cb7..3ad4e58 100644 --- a/src/db/migration/mod.rs +++ b/src/db/migration/mod.rs @@ -10,9 +10,24 @@ mod v3; mod v4; mod v5; mod v6; +mod v7; type MigrationFn = fn(&Db, &mut OwnedWriteBatch) -> Result<()>; +/// ordered list of schema migrations. +/// +/// invariants: +/// - migration `vN` upgrades on-disk data from schema version `N-1` to `N` +/// - once a migration ships, its input/output wire types are frozen +/// - never "fix up" an older migration to match a newer live type; add a new +/// schema version and a new migration instead +/// - if a stored type changes shape, define a new versioned type in +/// `src/types.rs`, export the newest one for live code, and have the new +/// migration explicitly deserialize the previous version and write the new one +/// +/// example: if `RepoState` changes again after `types::v7::RepoState`, add a new +/// `types::v8::RepoState`, keep `v7` frozen, and add a migration that rewrites +/// `v7 -> v8`. /// ordered list of migrations. migration at index `i` upgrades the schema from version `i` to `i+1`. const MIGRATIONS: &[(&str, MigrationFn)] = &[ ("stable_firehose_cursors", v1::stable_firehose_cursors), @@ -21,6 +36,7 @@ const MIGRATIONS: &[(&str, MigrationFn)] = &[ ("repo_state_active", v4::repo_state_active), ("pds_meta_layout", v5::pds_meta_layout), ("rebuild_pds_account_counts", v6::rebuild_pds_account_counts), + ("repo_state_event_clocks", v7::repo_state_event_clocks), ]; fn read_version(db: &Db) -> Result { diff --git a/src/db/migration/v6.rs b/src/db/migration/v6.rs index 92ab249..21a05f7 100644 --- a/src/db/migration/v6.rs +++ b/src/db/migration/v6.rs @@ -1,7 +1,8 @@ use std::collections::BTreeMap; use crate::db::keys::{self, COUNT_KS_PREFIX}; -use crate::db::{Db, deser_repo_state, set_ks_count}; +use crate::db::{Db, set_ks_count}; +use crate::types::v4; use fjall::OwnedWriteBatch; use miette::{Context, IntoDiagnostic, Result}; use smol_str::SmolStr; @@ -12,7 +13,9 @@ pub(crate) fn rebuild_pds_account_counts(db: &Db, batch: &mut OwnedWriteBatch) - for guard in db.repos.iter() { let (_, value) = guard.into_inner().into_diagnostic()?; - let state = deser_repo_state(value.as_ref())?; + let state: v4::RepoState = rmp_serde::from_slice(value.as_ref()) + .into_diagnostic() + .wrap_err("invalid v5 repo state")?; if !state.active { continue; } diff --git a/src/db/migration/v7.rs b/src/db/migration/v7.rs new file mode 100644 index 0000000..8149950 --- /dev/null +++ b/src/db/migration/v7.rs @@ -0,0 +1,111 @@ +use fjall::OwnedWriteBatch; +use miette::{Context, IntoDiagnostic, Result}; + +use crate::db::Db; +use crate::types::{RepoState, v4}; + +pub(crate) fn repo_state_event_clocks(db: &Db, batch: &mut OwnedWriteBatch) -> Result<()> { + for guard in db.repos.iter() { + let (key, value) = guard.into_inner().into_diagnostic()?; + let old: v4::RepoState = rmp_serde::from_slice(value.as_ref()) + .into_diagnostic() + .wrap_err("invalid v6 repo state")?; + + let new_state = RepoState { + active: old.active, + status: old.status, + root: old.root, + last_message_time: old.last_message_time, + last_identity_time: None, + last_account_time: None, + last_updated_at: old.last_updated_at, + signing_key: old.signing_key, + pds: old.pds, + handle: old.handle, + }; + + batch.insert( + &db.repos, + key, + rmp_serde::to_vec(&new_state) + .into_diagnostic() + .wrap_err("cant serialize v7 repo state")?, + ); + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::Config; + use crate::db::{deser_repo_state, keys}; + use crate::types::{RepoStatus, v4}; + use jacquard_common::CowStr; + use jacquard_common::types::string::Did; + + #[test] + fn repo_state_event_clock_migration_rewrites_v6_repo_states() -> Result<()> { + let tmp = tempfile::tempdir().into_diagnostic()?; + let cfg = Config { + database_path: tmp.path().to_path_buf(), + ..Default::default() + }; + let did = Did::new("did:web:migration.test").into_diagnostic()?; + let repo_key = keys::repo_key(&did); + + let legacy_bytes = { + let db = Db::open(&cfg)?; + let legacy = v4::RepoState { + active: true, + status: RepoStatus::Synced, + root: None, + last_message_time: Some(1_234), + last_updated_at: 9_876, + signing_key: None, + pds: Some(CowStr::Borrowed("https://pds.example/")), + handle: None, + }; + let legacy_bytes = rmp_serde::to_vec(&legacy).into_diagnostic()?; + + let mut batch = db.inner.batch(); + batch.insert(&db.repos, &repo_key, &legacy_bytes); + batch.insert(&db.counts, keys::VERSIONING_KEY, 6_u64.to_be_bytes()); + batch.commit().into_diagnostic()?; + db.persist()?; + + legacy_bytes + }; + + let db = Db::open(&cfg)?; + let migrated_bytes = db + .repos + .get(&repo_key) + .into_diagnostic()? + .expect("repo state should exist after migration"); + + assert_ne!(migrated_bytes.as_ref(), legacy_bytes.as_slice()); + + let migrated_state = deser_repo_state(migrated_bytes.as_ref())?; + assert!(migrated_state.active); + assert_eq!(migrated_state.status, RepoStatus::Synced); + assert_eq!(migrated_state.last_message_time, Some(1_234)); + assert_eq!(migrated_state.last_updated_at, 9_876); + assert_eq!(migrated_state.last_identity_time, None); + assert_eq!(migrated_state.last_account_time, None); + assert_eq!(migrated_state.pds.as_deref(), Some("https://pds.example/")); + + let version_bytes = db + .counts + .get(keys::VERSIONING_KEY) + .into_diagnostic()? + .expect("db version should be set"); + assert_eq!( + u64::from_be_bytes(version_bytes.as_ref().try_into().into_diagnostic()?), + 7 + ); + + Ok(()) + } +} diff --git a/src/db/mod.rs b/src/db/mod.rs index 48db0e3..9443efb 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -1305,10 +1305,18 @@ mod tests { let mut insert_repo = |did_str: &str, pds: &'static str, active: bool| -> Result<()> { let did = Did::new(did_str).into_diagnostic()?; - let mut state = RepoState::backfilling(); - state.active = active; - state.pds = Some(CowStr::Borrowed(pds)); - batch.insert(&db.repos, keys::repo_key(&did), ser_repo_state(&state)?); + let state = crate::types::v4::RepoState { + active, + status: crate::types::v4::RepoStatus::Synced, + root: None, + last_message_time: None, + last_updated_at: 0, + signing_key: None, + pds: Some(CowStr::Borrowed(pds)), + handle: None, + }; + let bytes = rmp_serde::to_vec(&state).into_diagnostic()?; + batch.insert(&db.repos, keys::repo_key(&did), bytes); Ok(()) }; diff --git a/src/ingest/firehose.rs b/src/ingest/firehose.rs index 7047a1c..423ea0f 100644 --- a/src/ingest/firehose.rs +++ b/src/ingest/firehose.rs @@ -11,9 +11,10 @@ use miette::{IntoDiagnostic, Result}; use rand::RngExt; use rand::rngs::SmallRng; use std::borrow::Cow; +use std::io::ErrorKind; use std::sync::Arc; use std::sync::atomic::Ordering; -use std::time::Duration; +use std::time::{Duration, Instant}; use tokio::sync::watch; use tracing::{Span, debug, error, info, trace, warn}; use url::Url; @@ -21,6 +22,103 @@ use url::Url; // these match ref relay const MAX_BACKOFF: Duration = Duration::from_secs(60); +#[derive(Debug, Clone)] +struct FirehoseFailure { + kind: &'static str, + detail: String, +} + +impl FirehoseFailure { + fn new(kind: &'static str, detail: impl Into) -> Self { + Self { + kind, + detail: detail.into(), + } + } +} + +fn classify_firehose_error(err: &FirehoseError) -> FirehoseFailure { + match err { + FirehoseError::WebSocket(err) => classify_websocket_error(err), + FirehoseError::UnknownScheme(scheme) => { + FirehoseFailure::new("config", format!("unsupported URL scheme `{scheme}`")) + } + FirehoseError::InvalidUri(err) => { + FirehoseFailure::new("config", format!("invalid websocket URI: {err}")) + } + FirehoseError::Decode(err) => { + FirehoseFailure::new("decode", format!("failed to decode firehose frame: {err}")) + } + FirehoseError::EmptyFrame => FirehoseFailure::new("protocol", "received empty frame"), + FirehoseError::RelayError { error, message } => FirehoseFailure::new( + "relay_error", + message + .as_deref() + .map(|message| format!("{error}: {message}")) + .unwrap_or_else(|| error.clone()), + ), + FirehoseError::UnknownOp(op) => { + FirehoseFailure::new("protocol", format!("unknown frame op {op}")) + } + FirehoseError::MissingType => FirehoseFailure::new("protocol", "missing frame type header"), + FirehoseError::UnknownType(ty) => { + FirehoseFailure::new("protocol", format!("unknown frame type `{ty}`")) + } + FirehoseError::Cbor(err) => { + FirehoseFailure::new("decode", format!("cbor decode error: {err}")) + } + FirehoseError::StreamClosed { code, reason } => { + FirehoseFailure::new("stream_closed", format!("close code {code}: {reason}")) + } + FirehoseError::TcpDropped => FirehoseFailure::new("tcp_dropped", "tcp layer dropped"), + FirehoseError::FutureCursor => FirehoseFailure::new("cursor", "future cursor"), + } +} + +fn classify_websocket_error(err: &tokio_websockets::Error) -> FirehoseFailure { + match err { + tokio_websockets::Error::CannotResolveHost => { + FirehoseFailure::new("dns", "host could not be resolved") + } + tokio_websockets::Error::Io(err) => { + let kind = match err.kind() { + ErrorKind::ConnectionRefused => "tcp_refused", + ErrorKind::ConnectionReset => "tcp_reset", + ErrorKind::ConnectionAborted => "tcp_aborted", + ErrorKind::TimedOut => "tcp_timeout", + ErrorKind::UnexpectedEof => "tcp_eof", + ErrorKind::AddrInUse => "tcp_addr_in_use", + ErrorKind::AddrNotAvailable => "tcp_addr_not_available", + ErrorKind::PermissionDenied => "tcp_permission", + _ => "io", + }; + FirehoseFailure::new(kind, format!("{err}")) + } + tokio_websockets::Error::InvalidDNSName(_) => { + FirehoseFailure::new("tls_invalid_dns_name", format!("{err}")) + } + tokio_websockets::Error::Rustls(_) => FirehoseFailure::new("tls", format!("{err}")), + tokio_websockets::Error::Upgrade(upgrade) => { + let kind = match upgrade { + tokio_websockets::upgrade::Error::DidNotSwitchProtocols(_) => "http_upgrade", + _ => "websocket_upgrade", + }; + FirehoseFailure::new(kind, format!("{upgrade}")) + } + tokio_websockets::Error::UnsupportedScheme => { + FirehoseFailure::new("config", "unsupported websocket URL scheme") + } + tokio_websockets::Error::Protocol(err) => { + FirehoseFailure::new("websocket_protocol", format!("{err}")) + } + tokio_websockets::Error::PayloadTooLong { len, max_len } => FirehoseFailure::new( + "websocket_payload", + format!("payload length {len} > {max_len}"), + ), + _ => FirehoseFailure::new("websocket", format!("{err}")), + } +} + trait AddJitter: rand::Rng { fn add_jitter(&mut self, timeout: Duration) -> Duration { let timeout_secs = timeout.as_secs_f32(); @@ -94,23 +192,58 @@ impl FirehoseIngestor { None => info!("no cursor found, live tailing"), } - let mut stream = match FirehoseStream::connect(self.relay_host.clone(), start_cursor) - .await - { - Ok(s) => s, - Err(e) => { - let Some(secs) = self.on_failure().await else { - break Ok(()); - }; - let timeout = rng.add_jitter(Duration::from_secs(secs).min(MAX_BACKOFF)); - let fmt = humantime::format_duration(timeout); - error!(err = %e, in = %fmt, "failed to connect to firehose, retrying later"); - tokio::time::sleep(timeout).await; - continue; - } - }; + let host_status = self.is_pds.then(|| { + let meta = self.state.pds_meta.load(); + meta.status(host).as_str() + }); + debug!( + is_pds = self.is_pds, + cursor = ?start_cursor, + host_status = ?host_status, + consecutive_failures = self.throttle.consecutive_failures(), + throttled_until = self.throttle.throttled_until(), + "connecting to firehose" + ); + let connect_started = Instant::now(); + let mut stream = + match FirehoseStream::connect(self.relay_host.clone(), start_cursor).await { + Ok(s) => s, + Err(e) => { + let failure = classify_firehose_error(&e); + let secs = match self.on_failure(&failure).await { + Some(secs) => secs, + None => { + error!( + err = %e, + failure_kind = failure.kind, + failure = %failure.detail, + failures = self.throttle.consecutive_failures(), + max_failures = self.max_failures, + "failed to connect to firehose, giving up" + ); + break Ok(()); + } + }; + let timeout = rng.add_jitter(Duration::from_secs(secs).min(MAX_BACKOFF)); + let fmt = humantime::format_duration(timeout); + error!( + err = %e, + failure_kind = failure.kind, + failure = %failure.detail, + failures = self.throttle.consecutive_failures(), + max_failures = self.max_failures, + in = %fmt, + "failed to connect to firehose, retrying later" + ); + tokio::time::sleep(timeout).await; + continue; + } + }; - info!("firehose connected"); + info!( + elapsed_ms = connect_started.elapsed().as_millis(), + "firehose connected" + ); let mut marked_active = false; let active_sleep_secs = if cfg!(debug_assertions) { 1 } else { 60 }; let mut active_sleep = @@ -218,29 +351,71 @@ impl FirehoseIngestor { { error!(err = %e, "failed to update host status to idle"); } - debug!("outdated cursor, backing off"); - tokio::time::sleep(Duration::from_secs(60)).await; + if let Err(e) = self.clear_stale_cursor() { + error!(err = %e, "failed to clear outdated cursor"); + } + warn!("outdated cursor, cleared stored cursor and retrying from live tail"); + tokio::time::sleep(Duration::from_secs(1)).await; } Err(FirehoseError::RelayError { error, message }) => { let message = message .as_deref() .map_or(Cow::Borrowed(""), Cow::Borrowed); + let failure = + FirehoseFailure::new("relay_error", format!("{error}: {message}")); error!(err = %error, "relay sent error: {message}"); - let Some(secs) = self.on_failure().await else { - break Ok(()); + let secs = match self.on_failure(&failure).await { + Some(secs) => secs, + None => { + error!( + failure_kind = failure.kind, + failure = %failure.detail, + failures = self.throttle.consecutive_failures(), + max_failures = self.max_failures, + "firehose disconnected, giving up" + ); + break Ok(()); + } }; let timeout = rng.add_jitter(Duration::from_secs(secs).min(MAX_BACKOFF)); let fmt = humantime::format_duration(timeout); - error!(in = %fmt, "firehose disconnected, reconnecting later"); + error!( + failure_kind = failure.kind, + failure = %failure.detail, + failures = self.throttle.consecutive_failures(), + max_failures = self.max_failures, + in = %fmt, + "firehose disconnected, reconnecting later" + ); tokio::time::sleep(timeout).await; } Err(e) => { - let Some(secs) = self.on_failure().await else { - break Ok(()); + let failure = classify_firehose_error(&e); + let secs = match self.on_failure(&failure).await { + Some(secs) => secs, + None => { + error!( + err = %e, + failure_kind = failure.kind, + failure = %failure.detail, + failures = self.throttle.consecutive_failures(), + max_failures = self.max_failures, + "firehose stream error, giving up" + ); + break Ok(()); + } }; let timeout = rng.add_jitter(Duration::from_secs(secs).min(MAX_BACKOFF)); let fmt = humantime::format_duration(timeout); - error!(err = %e, in = %fmt, "firehose stream error, reconnecting later"); + error!( + err = %e, + failure_kind = failure.kind, + failure = %failure.detail, + failures = self.throttle.consecutive_failures(), + max_failures = self.max_failures, + in = %fmt, + "firehose stream error, reconnecting later" + ); tokio::time::sleep(timeout).await; } } @@ -249,14 +424,23 @@ impl FirehoseIngestor { /// record a failure and return the backoff duration in seconds, /// or `None` if the failure threshold was reached and the subscriber should stop - async fn on_failure(&self) -> Option { - let secs = self.throttle.record_failure().unwrap_or_else(|| { - let until = self.throttle.throttled_until(); - 0.max(until - chrono::Utc::now().timestamp()) as u64 - }); + async fn on_failure(&self, failure: &FirehoseFailure) -> Option { + let secs = self + .throttle + .record_failure_detail(failure.kind, failure.detail.clone()) + .unwrap_or_else(|| { + let until = self.throttle.throttled_until(); + 0.max(until - chrono::Utc::now().timestamp()) as u64 + }); let failures = self.throttle.consecutive_failures(); if failures >= self.max_failures { - warn!(failures, "too many consecutive failures, giving up on host"); + warn!( + failures, + max_failures = self.max_failures, + failure_kind = failure.kind, + failure = %failure.detail, + "too many consecutive failures, giving up on host" + ); if self.is_pds && let Err(e) = self.set_host_status(HostStatus::Offline) { @@ -283,6 +467,17 @@ impl FirehoseIngestor { Ok(()) } + fn clear_stale_cursor(&self) -> Result<()> { + let key = crate::db::keys::firehose_cursor_key_from_url(&self.relay_host); + self.state.db.cursors.remove(key).into_diagnostic()?; + self.state + .firehose_cursors + .peek_with(&self.relay_host, |_, cursor| { + cursor.store(0, Ordering::SeqCst) + }); + Ok(()) + } + async fn handle_message(&self, msg: SubscribeReposMessage<'_>) { let did = match &msg { SubscribeReposMessage::Commit(commit) => &commit.repo, @@ -367,3 +562,54 @@ impl FirehoseIngestor { .flatten() } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::Config; + use std::sync::atomic::AtomicI64; + + #[tokio::test] + async fn clear_stale_cursor_resets_db_and_memory() -> Result<()> { + let tmp = tempfile::tempdir().into_diagnostic()?; + let cfg = Config { + database_path: tmp.path().to_path_buf(), + ..Default::default() + }; + let state = Arc::new(AppState::new(&cfg)?); + let relay_host = Url::parse("wss://example.com/").into_diagnostic()?; + + crate::db::set_firehose_cursor(&state.db, &relay_host, 1234)?; + let _ = state + .firehose_cursors + .insert_async(relay_host.clone(), AtomicI64::new(1234)) + .await; + + let ingestor = FirehoseIngestor::new( + state.clone(), + BufferTx::channel(1).0, + relay_host.clone(), + true, + state.filter.clone(), + state.firehose_enabled.subscribe(), + false, + 1, + ) + .await; + + ingestor.clear_stale_cursor()?; + + assert!( + crate::db::get_firehose_cursor(&state.db, &relay_host) + .await? + .is_none() + ); + let in_memory = state + .firehose_cursors + .peek_with(&relay_host, |_, cursor| cursor.load(Ordering::SeqCst)) + .unwrap_or(-1); + assert_eq!(in_memory, 0); + + Ok(()) + } +} diff --git a/src/ingest/relay.rs b/src/ingest/relay.rs index 9a6e685..da3c556 100644 --- a/src/ingest/relay.rs +++ b/src/ingest/relay.rs @@ -6,7 +6,7 @@ use std::sync::atomic::Ordering; use fjall::OwnedWriteBatch; use jacquard_api::com_atproto::sync::get_repo_status::{ - GetRepoStatus, GetRepoStatusError, GetRepoStatusOutputStatus, + GetRepoStatus, GetRepoStatusError, GetRepoStatusOutput, GetRepoStatusOutputStatus, }; use jacquard_common::types::crypto::PublicKey; use jacquard_common::types::did::Did; @@ -39,6 +39,28 @@ use crate::types::{RepoState, RepoStatus}; use crate::util; use smol_str::{SmolStr, ToSmolStr}; +fn map_repo_status_probe(output: Option>) -> Option> { + let output = output?; + + let mut repo_state = RepoState::backfilling(); + repo_state.active = output.active; + repo_state.status = match output.status { + Some(GetRepoStatusOutputStatus::Takendown) => RepoStatus::Takendown, + Some(GetRepoStatusOutputStatus::Suspended) => RepoStatus::Suspended, + Some(GetRepoStatusOutputStatus::Deactivated) => RepoStatus::Deactivated, + Some(GetRepoStatusOutputStatus::Deleted) => RepoStatus::Deleted, + Some(GetRepoStatusOutputStatus::Desynchronized) => RepoStatus::Desynchronized, + Some(GetRepoStatusOutputStatus::Throttled) => RepoStatus::Throttled, + Some(GetRepoStatusOutputStatus::Other(s)) => RepoStatus::Error(s.into()), + None => output + .active + .then_some(RepoStatus::Synced) + .unwrap_or_else(|| RepoStatus::Error("unknown".into())), + }; + + Some(repo_state) +} + struct WorkerContext<'a> { verify_signatures: bool, state: &'a AppState, @@ -485,11 +507,11 @@ impl RelayWorker { is_pds: bool, ) -> Result<()> { let event_ms = identity.time.0.timestamp_millis(); - if repo_state.last_message_time.is_some_and(|t| event_ms <= t) { + if !repo_state.should_process_identity_time(event_ms) { debug!("skipping stale/duplicate identity event"); return Ok(()); } - repo_state.advance_message_time(event_ms); + repo_state.advance_identity_time(event_ms); let was_active = repo_state.active; let was_pds_host = Self::pds_host(repo_state.pds.as_deref()); @@ -577,12 +599,12 @@ impl RelayWorker { _is_pds: bool, ) -> Result<()> { let event_ms = account.time.0.timestamp_millis(); - if repo_state.last_message_time.is_some_and(|t| event_ms <= t) { + if !repo_state.should_process_account_time(event_ms) { debug!("skipping stale/duplicate account event"); return Ok(()); } - repo_state.advance_message_time(event_ms); + repo_state.advance_account_time(event_ms); // always capture was_active for count tracking, not just in indexer mode let was_active = repo_state.active; @@ -877,35 +899,17 @@ impl WorkerContext<'_> { Ok(r) => match r.into_output() { Ok(o) => o, Err(XrpcError::Xrpc(GetRepoStatusError::RepoNotFound(_))) => { - // pds explicitly says it doesn't have this repo - // we shouldnt really get here unless the pds is buggy? - // or somehow the repo gets gon right after we receive the event - let mut repo_state = RepoState::backfilling(); - repo_state.active = false; - repo_state.status = RepoStatus::Error("not_found".into()); - return Ok(Some(repo_state)); + // treat probe-time 404s like any other transient probe failure. + // we already have a live event from the authoritative host, so + // inserting an inactive placeholder here can wedge the repo until + // a later account event arrives. + return Ok(map_repo_status_probe(None)); } Err(_) => return Ok(None), }, }; - let mut repo_state = RepoState::backfilling(); - repo_state.active = output.active; - repo_state.status = match output.status { - Some(GetRepoStatusOutputStatus::Takendown) => RepoStatus::Takendown, - Some(GetRepoStatusOutputStatus::Suspended) => RepoStatus::Suspended, - Some(GetRepoStatusOutputStatus::Deactivated) => RepoStatus::Deactivated, - Some(GetRepoStatusOutputStatus::Deleted) => RepoStatus::Deleted, - Some(GetRepoStatusOutputStatus::Desynchronized) => RepoStatus::Desynchronized, - Some(GetRepoStatusOutputStatus::Throttled) => RepoStatus::Throttled, - Some(GetRepoStatusOutputStatus::Other(s)) => RepoStatus::Error(s.into()), - None => output - .active - .then_some(RepoStatus::Synced) - .unwrap_or_else(|| RepoStatus::Error("unknown".into())), - }; - - Ok(Some(repo_state)) + Ok(map_repo_status_probe(Some(output))) } fn load_repo_state(&mut self, msg: &WorkerMessage) -> Result>> { @@ -1055,3 +1059,43 @@ enum AuthorityOutcome { /// host did not match even after doc resolution. WrongHost { expected: SmolStr }, } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn missing_repo_status_probe_falls_back_to_live_discovery() { + assert!(map_repo_status_probe(None).is_none()); + } + + #[test] + fn active_repo_status_probe_maps_to_synced_repo_state() { + let repo_state = map_repo_status_probe(Some(GetRepoStatusOutput { + did: Did::new("did:plc:testrepo").expect("valid did"), + active: true, + status: None, + rev: None, + extra_data: None, + })) + .expect("probe should map"); + + assert!(repo_state.active); + assert_eq!(repo_state.status, RepoStatus::Synced); + } + + #[test] + fn throttled_repo_status_probe_preserves_host_status() { + let repo_state = map_repo_status_probe(Some(GetRepoStatusOutput { + did: Did::new("did:plc:testrepo").expect("valid did"), + active: true, + status: Some(GetRepoStatusOutputStatus::Throttled), + rev: None, + extra_data: None, + })) + .expect("probe should map"); + + assert!(repo_state.active); + assert_eq!(repo_state.status, RepoStatus::Throttled); + } +} diff --git a/src/ingest/stream.rs b/src/ingest/stream.rs index 6ec8b91..90f1142 100644 --- a/src/ingest/stream.rs +++ b/src/ingest/stream.rs @@ -26,6 +26,8 @@ pub enum FirehoseError { WebSocket(#[from] tokio_websockets::Error), #[error("unknown scheme: {0}")] UnknownScheme(String), + #[error("invalid websocket uri: {0}")] + InvalidUri(String), #[error("decode error: {0}")] Decode(#[from] DecodeError), #[error("empty frame")] @@ -77,7 +79,7 @@ impl FirehoseStream { let uri: Uri = relay .as_str() .parse() - .map_err(|e| FirehoseError::Cbor(format!("invalid uri: {e}")))?; + .map_err(|e| FirehoseError::InvalidUri(format!("{e}")))?; let (ws, _) = ClientBuilder::from_uri(uri).connect().await?; diff --git a/src/types.rs b/src/types.rs index b2e045d..1105b47 100644 --- a/src/types.rs +++ b/src/types.rs @@ -21,6 +21,18 @@ use crate::db::types::{DbAction, DbRkey}; use crate::ingest::stream::Datetime; use crate::resolver::MiniDoc; +// on-disk schema snapshots for stored types. +// +// rules: +// - versioned modules (`v2`, `v4`, `v7`, ...) are wire-format snapshots, not +// convenience namespaces +// - once a version is referenced by a shipped migration, do not change its +// serialized shape again +// - when a stored type changes, add a new versioned module and export the +// newest version for live code via `pub(crate) use vN::*` +// - migrations should explicitly read the previous version and write the new one +// - for example, if `RepoState` changes shape again, add `v8::RepoState`, +// keep `v7::RepoState` frozen, export `v8`, and migrate `v7 -> v8` pub(crate) mod v2 { use super::*; @@ -107,9 +119,6 @@ pub(crate) mod v4 { pub active: bool, pub status: RepoStatus, pub root: Option, - /// ms since epoch of the last firehose message we processed for this repo. - /// used to deduplicate identity / account events that can arrive from multiple relays at - /// different wall-clock times but represent the same underlying PDS event. pub last_message_time: Option, /// this is when we *ingested* any last updates pub last_updated_at: i64, // unix timestamp @@ -130,7 +139,45 @@ pub(crate) mod v4 { } } -pub(crate) use v4::*; +pub(crate) mod v7 { + use super::*; + pub(crate) use v4::{Commit, RepoMetadata, RepoStatus}; + + #[derive(Debug, Clone, Serialize, Deserialize)] + #[serde(bound(deserialize = "'i: 'de"))] + pub(crate) struct RepoState<'i> { + /// whether the upstream considers this account active. + /// services should use the `active` flag to control overall account visibility + pub active: bool, + pub status: RepoStatus, + pub root: Option, + /// ms since epoch of the last firehose message we processed for this repo. + /// used to deduplicate identity / account events that can arrive from multiple relays at + /// different wall-clock times but represent the same underlying PDS event. + pub last_message_time: Option, + /// high-water mark for identity events only. + /// + /// sharing a single clock with commit/sync/account events can suppress valid identity + /// updates when those event classes arrive out of order but with overlapping timestamps. + pub last_identity_time: Option, + /// high-water mark for account events only. + /// + /// sharing a single clock with commit/sync/identity events can suppress valid account + /// status transitions when those event classes arrive out of order but with overlapping + /// timestamps. + pub last_account_time: Option, + /// this is when we *ingested* any last updates + pub last_updated_at: i64, // unix timestamp + #[serde(borrow)] + pub signing_key: Option>, + #[serde(borrow)] + pub pds: Option>, + #[serde(borrow)] + pub handle: Option>, + } +} + +pub(crate) use v7::*; impl<'c> From> for Commit { fn from(value: AtpCommit<'c>) -> Self { @@ -246,6 +293,8 @@ impl<'i> RepoState<'i> { pds: None, signing_key: None, last_message_time: None, + last_identity_time: None, + last_account_time: None, } } @@ -254,6 +303,24 @@ impl<'i> RepoState<'i> { self.last_message_time = Some(event_ms.max(self.last_message_time.unwrap_or(0))); } + pub fn should_process_identity_time(&self, event_ms: i64) -> bool { + self.last_identity_time.is_none_or(|t| event_ms > t) + } + + pub fn advance_identity_time(&mut self, event_ms: i64) { + self.last_identity_time = Some(event_ms.max(self.last_identity_time.unwrap_or(0))); + self.advance_message_time(event_ms); + } + + pub fn should_process_account_time(&self, event_ms: i64) -> bool { + self.last_account_time.is_none_or(|t| event_ms > t) + } + + pub fn advance_account_time(&mut self, event_ms: i64) { + self.last_account_time = Some(event_ms.max(self.last_account_time.unwrap_or(0))); + self.advance_message_time(event_ms); + } + // updates last_updated_at to now pub fn touch(&mut self) { self.last_updated_at = chrono::Utc::now().timestamp(); @@ -284,6 +351,8 @@ impl<'i> IntoStatic for RepoState<'i> { pds: self.pds.map(IntoStatic::into_static), signing_key: self.signing_key.map(IntoStatic::into_static), last_message_time: self.last_message_time, + last_identity_time: self.last_identity_time, + last_account_time: self.last_account_time, } } } @@ -612,3 +681,62 @@ pub(crate) enum RelayBroadcast { #[allow(dead_code)] Ephemeral(u64, bytes::Bytes), } + +#[cfg(test)] +mod tests { + use super::*; + use miette::IntoDiagnostic; + + #[test] + fn identity_dedupe_does_not_depend_on_commit_clock() { + let mut state = RepoState::backfilling(); + + state.advance_message_time(2_000); + + assert!(state.should_process_identity_time(1_500)); + state.advance_identity_time(1_500); + assert_eq!(state.last_message_time, Some(2_000)); + assert_eq!(state.last_identity_time, Some(1_500)); + + assert!(!state.should_process_identity_time(1_500)); + assert!(!state.should_process_identity_time(1_499)); + assert!(state.should_process_identity_time(1_501)); + } + + #[test] + fn account_dedupe_does_not_depend_on_commit_clock() { + let mut state = RepoState::backfilling(); + + state.advance_message_time(3_000); + + assert!(state.should_process_account_time(2_500)); + state.advance_account_time(2_500); + assert_eq!(state.last_message_time, Some(3_000)); + assert_eq!(state.last_account_time, Some(2_500)); + + assert!(!state.should_process_account_time(2_500)); + assert!(!state.should_process_account_time(2_499)); + assert!(state.should_process_account_time(2_501)); + } + + #[test] + fn into_static_preserves_per_event_clocks() -> miette::Result<()> { + let mut state = RepoState::backfilling(); + state.last_message_time = Some(10); + state.last_identity_time = Some(20); + state.last_account_time = Some(30); + state.handle = Some(Handle::new("alice.test").into_diagnostic()?); + state.pds = Some(CowStr::Borrowed("https://pds.example")); + state.signing_key = Some(DidKey::from_did_key( + "did:key:zQ3shokFTS3brHcDQrn82RUDfCZESWL1ZdCEJwekUDPQiYBme", + )?); + + let static_state = state.into_static(); + + assert_eq!(static_state.last_message_time, Some(10)); + assert_eq!(static_state.last_identity_time, Some(20)); + assert_eq!(static_state.last_account_time, Some(30)); + + Ok(()) + } +} diff --git a/src/util/throttle.rs b/src/util/throttle.rs index ed86b39..60c86e5 100644 --- a/src/util/throttle.rs +++ b/src/util/throttle.rs @@ -24,27 +24,35 @@ pub struct Throttler { states: Arc>>, } -#[derive(Debug, Clone, Copy, Default)] +#[derive(Debug, Clone, Default)] pub struct ThrottleSnapshot { pub consecutive_failures: usize, pub throttled_until: i64, + pub last_failure: Option, } impl ThrottleSnapshot { - pub fn is_failing(self) -> bool { + pub fn is_failing(&self) -> bool { self.consecutive_failures != 0 || self.throttled_until != 0 } - pub fn is_throttled(self, now: i64) -> bool { + pub fn is_throttled(&self, now: i64) -> bool { self.throttled_until != 0 && now < self.throttled_until } - pub fn retry_in_secs(self, now: i64) -> Option { + pub fn retry_in_secs(&self, now: i64) -> Option { self.is_throttled(now) .then_some((self.throttled_until - now) as u64) } } +#[derive(Debug, Clone)] +pub struct FailureSnapshot { + pub at: i64, + pub kind: String, + pub detail: String, +} + impl Throttler { pub fn new() -> Self { Self { @@ -79,6 +87,7 @@ impl Throttler { .read_sync(url, |_, state| ThrottleSnapshot { consecutive_failures: state.consecutive_failures.load(Ordering::Acquire), throttled_until: state.throttled_until.load(Ordering::Acquire), + last_failure: state.last_failure.lock().clone(), }) .unwrap_or_default() } @@ -88,6 +97,7 @@ struct State { throttled_until: AtomicI64, consecutive_failures: AtomicUsize, consecutive_timeouts: AtomicUsize, + last_failure: Mutex>, /// only fires on hard failures (timeout, TLS, bad gateway, etc). /// ratelimits do NOT fire this — they just store `throttled_until` and /// let tasks exit naturally, deferring to the background retry loop. @@ -102,6 +112,7 @@ impl State { throttled_until: AtomicI64::new(0), consecutive_failures: AtomicUsize::new(0), consecutive_timeouts: AtomicUsize::new(0), + last_failure: Mutex::new(None), failure_notify: Notify::new(), semaphore: Semaphore::new(PER_PDS_CONCURRENCY), rate_limiter: RateLimiter::new(), @@ -128,6 +139,7 @@ impl ThrottleHandle { self.state.consecutive_failures.store(0, Ordering::Release); self.state.consecutive_timeouts.store(0, Ordering::Release); self.state.throttled_until.store(0, Ordering::Release); + *self.state.last_failure.lock() = None; } /// called on a 429 response. `retry_after_secs` comes from the `Retry-After` @@ -145,6 +157,23 @@ impl ThrottleHandle { /// always increments `consecutive_failures`. only sets a new `throttled_until` /// (and notifies waiters) if not already throttled. pub fn record_failure(&self) -> Option { + self.record_failure_inner() + } + + pub fn record_failure_detail( + &self, + kind: impl Into, + detail: impl Into, + ) -> Option { + *self.state.last_failure.lock() = Some(FailureSnapshot { + at: chrono::Utc::now().timestamp(), + kind: kind.into(), + detail: detail.into(), + }); + self.record_failure_inner() + } + + fn record_failure_inner(&self) -> Option { let failures = self .state .consecutive_failures diff --git a/tests/api_firehose_sources.nu b/tests/api_firehose_sources.nu index 2cc5e90..2d024df 100644 --- a/tests/api_firehose_sources.nu +++ b/tests/api_firehose_sources.nu @@ -101,6 +101,15 @@ def test-firehose-source-filters [url: string, pid: int, mock_port: int] { stop-mock-pds $mock_pds fail "expected failing source to expose retry_in_secs" $pid } + let last_failure = ($failing.last_failure? | default null) + if $last_failure == null { + stop-mock-pds $mock_pds + fail "expected failing source to expose last_failure" $pid + } + if (($last_failure.kind? | default "") == "") or (($last_failure.detail? | default "") == "") or (($last_failure.at? | default 0 | into int) <= 0) { + stop-mock-pds $mock_pds + fail "expected last_failure to include kind, detail, and timestamp" $pid + } let failing_pds = ($failing.pds? | default null) if $failing_pds == null { stop-mock-pds $mock_pds -- 2.51.2