use std::collections::BTreeMap;
use std::time::Duration;
use anyhow::Context as _;
use bytes::Bytes;
use knot_runtime::{HttpRequest, HttpTransport, ReqwestHttp};
use tracing::{debug, warn};
use crate::config::ValidatedConfig;
const SERVED_TIMEOUT: Duration = Duration::from_secs(10);
const SERVED_POLL: Duration = Duration::from_millis(250);
fn expected_revs(config: &ValidatedConfig, changed_repos: &[String]) -> BTreeMap<String, String> {
let Some(generation) = crate::publish::current_generation(&config.output_root) else {
return BTreeMap::new();
};
let lists = config
.output_root
.join("generations")
.join(generation)
.join(crate::xrpc::ATPROTO_DIR)
.join("list");
let mut expected = BTreeMap::new();
let mut cursor = "0".to_owned();
loop {
let Some(page) = std::fs::read(lists.join(format!("1000-{cursor}.json")))
.ok()
.and_then(|bytes| serde_json::from_slice::<serde_json::Value>(&bytes).ok())
else {
return expected;
};
expected.extend(repo_revs(&page).filter(|(did, _)| changed_repos.contains(did)));
match page["cursor"].as_str() {
Some(next) => cursor = next.to_owned(),
None => return expected,
}
}
}
fn repo_revs(page: &serde_json::Value) -> impl Iterator<Item = (String, String)> + '_ {
page["repos"]
.as_array()
.into_iter()
.flatten()
.filter_map(|repo| {
Some((
repo["did"].as_str()?.to_owned(),
repo["rev"].as_str()?.to_owned(),
))
})
}
/// a reload returns before the old workers exit. a crawl now would list the old generation.
async fn wait_until_served(
config: &ValidatedConfig,
base: &url::Url,
expected: &BTreeMap<String, String>,
timeout: Duration,
) -> bool {
if expected.is_empty() {
return true;
}
let deadline = tokio::time::Instant::now() + timeout;
loop {
match served_revs(config, base).await {
Ok(served)
if expected
.iter()
.all(|(did, rev)| served.get(did) == Some(rev)) =>
{
return true;
}
Ok(_) => {}
Err(error) => debug!(target: "knot_static::crawl", "listRepos poll failed: {error:#}"),
}
if tokio::time::Instant::now() >= deadline {
return false;
}
tokio::time::sleep(SERVED_POLL).await;
}
}
async fn served_revs(
config: &ValidatedConfig,
base: &url::Url,
) -> anyhow::Result<BTreeMap<String, String>> {
// a keepalive connection could stay pinned to an old worker. each poll gets its own transport.
let transport = ReqwestHttp::new(config.http_limits()).context("init http transport")?;
let mut served = BTreeMap::new();
let mut cursor: Option<String> = None;
loop {
let mut url = base
.join("/xrpc/com.atproto.sync.listRepos")
.context("listRepos url")?;
{
let mut query = url.query_pairs_mut();
query.append_pair("limit", "1000");
if let Some(cursor) = &cursor {
query.append_pair("cursor", cursor);
}
}
let response = transport
.execute(HttpRequest::get(url))
.await
.context("listRepos request")?;
anyhow::ensure!(
response.status.is_success(),
"listRepos returned {}",
response.status
);
let page: serde_json::Value = serde_json::from_slice(&response.body)?;
served.extend(repo_revs(&page));
match page["cursor"].as_str() {
Some(next) => cursor = Some(next.to_owned()),
None => return Ok(served),
}
}
}
pub(crate) async fn wait_served(
config: &ValidatedConfig,
changed_repos: &[String],
timeout: Duration,
) -> bool {
let expected = expected_revs(config, changed_repos);
if expected.is_empty() {
return true;
}
match url::Url::parse(&format!("https://{}", config.hostname)) {
Ok(base) => wait_until_served(config, &base, &expected, timeout).await,
Err(error) => {
warn!(target: "knot_static::crawl", "cannot poll own listRepos: {error}");
false
}
}
}
/// with a firehose, crawlers follow it, and a request per push would make them relist every repo.
pub async fn announce(config: &ValidatedConfig, changed_repos: &[String]) {
if config.firehose_socket.is_some() {
return;
}
if !wait_served(config, changed_repos, SERVED_TIMEOUT).await {
warn!(
target: "knot_static::crawl",
"the new generation was not served within {SERVED_TIMEOUT:?}; requesting a crawl anyway"
);
}
request_all(config).await;
}
pub(crate) async fn request_all(config: &ValidatedConfig) {
let requests = config.crawlers.iter().map(|crawler| async move {
match request(config, crawler).await {
Ok(status) if status.is_success() => {}
Ok(status) => {
warn!(target: "knot_static::crawl", "requestCrawl at {crawler} returned {status}")
}
Err(error) => {
warn!(target: "knot_static::crawl", "requestCrawl at {crawler} failed: {error:#}")
}
}
});
futures::future::join_all(requests).await;
}
async fn request(config: &ValidatedConfig, crawler: &url::Url) -> anyhow::Result<http::StatusCode> {
let transport = ReqwestHttp::new(config.http_limits()).context("init http transport")?;
let url = crawler
.join("/xrpc/com.atproto.sync.requestCrawl")
.context("requestCrawl url")?;
let body = serde_json::to_vec(&serde_json::json!({ "hostname": config.hostname }))?;
let mut request = HttpRequest::post(url, Bytes::from(body));
request.headers.insert(
http::header::CONTENT_TYPE,
http::header::HeaderValue::from_static("application/json"),
);
transport
.execute(request)
.await
.map(|response| response.status)
.context("requestCrawl request")
}
#[cfg(test)]
mod tests {
use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
use super::*;
fn config(allow_private: bool, crawler: &str) -> ValidatedConfig {
let raw = format!(
"[knot]\nhostname = \"knot.test\"\nowner = \"did:plc:owner\"\n[repo]\nscan_path = \"r\"\n[secrets]\nsealed_key_file = \"s\"\n[output]\nroot = \"o\"\n[atproto]\nplc_directory = \"https://127.0.0.1:9\"\nallow_private_addresses = {allow_private}\n[crawl]\nurls = [\"{crawler}\"]\n"
);
ValidatedConfig::from_toml(&raw).unwrap()
}
async fn one_shot_server() -> (String, tokio::task::JoinHandle<String>) {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let url = format!("http://{}", listener.local_addr().unwrap());
let task = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.unwrap();
let mut buf = vec![0u8; 4096];
let n = stream.read(&mut buf).await.unwrap();
stream
.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")
.await
.unwrap();
String::from_utf8_lossy(&buf[..n]).into_owned()
});
(url, task)
}
#[tokio::test]
async fn request_crawl_posts_the_hostname_when_private_addresses_are_allowed() {
let (url, server) = one_shot_server().await;
let config = config(true, &url);
request_all(&config).await;
let seen = server.await.unwrap();
assert!(
seen.starts_with("POST /xrpc/com.atproto.sync.requestCrawl "),
"{seen}"
);
assert!(seen.ends_with(r#"{"hostname":"knot.test"}"#), "{seen}");
}
#[tokio::test]
async fn a_knot_with_a_firehose_asks_no_crawl_per_push() {
let (url, server) = one_shot_server().await;
let mut config = config(true, &url);
config.firehose_socket = Some("/run/f.sock".into());
announce(&config, &["did:plc:a".to_owned()]).await;
tokio::time::sleep(Duration::from_millis(200)).await;
assert!(!server.is_finished(), "a crawl was requested");
server.abort();
}
#[tokio::test]
async fn loopback_crawlers_are_refused_by_default() {
let (url, server) = one_shot_server().await;
let config = config(false, &url);
let outcome = request(&config, &config.crawlers[0]).await;
assert!(outcome.is_err(), "{outcome:?}");
server.abort();
}
async fn listing_server(
stale: usize,
) -> (url::Url, std::sync::Arc<std::sync::atomic::AtomicUsize>) {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let base = url::Url::parse(&format!("http://{}", listener.local_addr().unwrap())).unwrap();
let hits = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
let counter = hits.clone();
tokio::spawn(async move {
loop {
let (mut stream, _) = listener.accept().await.unwrap();
let seen = counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
let rev = if seen < stale { "old" } else { "new" };
let body = format!(r#"{{"repos":[{{"did":"did:plc:a","rev":"{rev}"}}]}}"#);
let mut buf = [0u8; 2048];
let _ = stream.read(&mut buf).await;
let reply = format!(
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
body.len()
);
let _ = stream.write_all(reply.as_bytes()).await;
}
});
(base, hits)
}
fn expect_new() -> BTreeMap<String, String> {
BTreeMap::from([("did:plc:a".to_owned(), "new".to_owned())])
}
#[tokio::test]
async fn waits_past_a_stale_listing_until_the_new_rev_is_served() {
let (base, hits) = listing_server(3).await;
let config = config(true, "http://crawler.test");
assert!(wait_until_served(&config, &base, &expect_new(), Duration::from_secs(10)).await);
assert_eq!(hits.load(std::sync::atomic::Ordering::SeqCst), 4);
}
#[tokio::test]
async fn gives_up_after_the_timeout_when_the_old_rev_stays() {
let (base, _) = listing_server(usize::MAX).await;
let config = config(true, "http://crawler.test");
assert!(
!wait_until_served(&config, &base, &expect_new(), Duration::from_millis(600)).await
);
}
#[tokio::test]
async fn nothing_expected_means_no_polling() {
let config = config(true, "http://crawler.test");
let unreachable = url::Url::parse("http://127.0.0.1:9").unwrap();
assert!(
wait_until_served(
&config,
&unreachable,
&BTreeMap::new(),
Duration::from_secs(1)
)
.await
);
}
#[test]
fn expected_revs_follow_the_active_generations_list_pages_for_changed_repos_only() {
let root = tempfile::tempdir().unwrap();
let lists = root.path().join("generations/g/www/_atproto/list");
std::fs::create_dir_all(&lists).unwrap();
std::fs::write(
lists.join("1000-0.json"),
r#"{"repos":[{"did":"did:plc:a","rev":"r1"},{"did":"did:plc:b","rev":"r2"}],"cursor":"1000"}"#,
)
.unwrap();
std::fs::write(
lists.join("1000-1000.json"),
r#"{"repos":[{"did":"did:plc:c","rev":"r3"}]}"#,
)
.unwrap();
std::os::unix::fs::symlink("generations/g", root.path().join("current")).unwrap();
let raw = format!(
"[knot]\nhostname = \"knot.test\"\nowner = \"did:plc:owner\"\n[repo]\nscan_path = \"r\"\n[secrets]\nsealed_key_file = \"s\"\n[output]\nroot = \"{}\"\n[atproto]\nplc_directory = \"https://127.0.0.1:9\"\n",
root.path().display()
);
let config = ValidatedConfig::from_toml(&raw).unwrap();
let changed = vec!["did:plc:a".to_owned(), "did:plc:c".to_owned()];
assert_eq!(
expected_revs(&config, &changed),
BTreeMap::from([
("did:plc:a".into(), "r1".into()),
("did:plc:c".into(), "r3".into())
])
);
}
}