From ee1201ee0347909ce37d943a507ad1d69d5cd22d Mon Sep 17 00:00:00 2001 From: Chris van der Loo Date: Sat, 22 Aug 2026 17:28:37 -0400 Subject: [PATCH] feat: bound repository stream sends --- README.md | 4 ++ src/config.rs | 28 +++++++++ src/config/tests.rs | 10 +++ src/http/subscription.rs | 47 +++++++++----- tests/backup_restore.rs | 1 + tests/process_lifecycle.rs | 3 +- tests/subscribe_repos_flow.rs | 115 +++++++++++++++++++++++++++++++++- tests/support/mod.rs | 1 + 8 files changed, 191 insertions(+), 18 deletions(-) diff --git a/README.md b/README.md index eea64f8..6509b5b 100644 --- a/README.md +++ b/README.md @@ -99,6 +99,7 @@ export CPDS_BLOB_DIRECTORY=/var/lib/cpds/blobs export CPDS_BLOB_ORPHAN_TTL_SECONDS=86400 export CPDS_BLOB_GC_INTERVAL_SECONDS=3600 export CPDS_REPO_EVENT_RETENTION=100000 +export CPDS_REPO_STREAM_SEND_TIMEOUT_MS=10000 export CPDS_ACCOUNT_DID=did:web:pds.example.com export CPDS_ACCOUNT_HANDLE=alice.example.com export CPDS_ACCOUNT_PASSWORD='use-a-password-manager' @@ -141,6 +142,9 @@ All configuration variables are parsed in [`src/config.rs`](src/config.rs). `CPDS_REPO_EVENT_RETENTION` bounds the durable firehose window between 1 and 10,000,000 events. Consumers behind that window receive `#info` with `OutdatedCursor` and must resynchronize from repository exports. +`CPDS_REPO_STREAM_SEND_TIMEOUT_MS` disconnects a firehose consumer when a +single frame cannot be sent within 100 milliseconds to one minute (10 seconds +by default). The consumer can reconnect from its last processed durable cursor. `CPDS_BLOB_ORPHAN_TTL_SECONDS` retains unattached uploads for 60 seconds to one year (24 hours by default). `CPDS_BLOB_GC_INTERVAL_SECONDS` runs bounded cleanup every second to 24 hours (hourly by default). Cleanup durably claims each diff --git a/src/config.rs b/src/config.rs index 13427a9..5482b6b 100644 --- a/src/config.rs +++ b/src/config.rs @@ -60,6 +60,7 @@ pub struct Config { pub blob_orphan_ttl: Duration, pub blob_gc_interval: Duration, pub repo_event_retention: usize, + pub repo_stream_send_timeout: Duration, pub account: AccountConfig, pub oauth_clients: Vec, } @@ -161,6 +162,14 @@ impl Config { "CPDS_REPO_EVENT_RETENTION must be between 1 and 10000000".to_owned(), ) })?; + let repo_stream_send_timeout = millisecond_duration_value( + &mut source, + deployment_mode, + "CPDS_REPO_STREAM_SEND_TIMEOUT_MS", + 10_000, + 100, + 60_000, + )?; let account = account_config(&mut source, deployment_mode)?; let oauth_clients = oauth_clients(&mut source, deployment_mode)?; @@ -182,6 +191,7 @@ impl Config { blob_orphan_ttl, blob_gc_interval, repo_event_retention, + repo_stream_send_timeout, account, oauth_clients, }) @@ -206,6 +216,24 @@ fn duration_value( Ok(Duration::from_secs(seconds)) } +fn millisecond_duration_value( + source: &mut impl FnMut(&str) -> Option, + deployment_mode: DeploymentMode, + name: &str, + default: u64, + minimum: u64, + maximum: u64, +) -> Result { + let milliseconds = value(source, deployment_mode, name, &default.to_string(), false)? + .parse::() + .ok() + .filter(|milliseconds| (minimum..=maximum).contains(milliseconds)) + .ok_or_else(|| { + Error::Configuration(format!("{name} must be between {minimum} and {maximum}")) + })?; + Ok(Duration::from_millis(milliseconds)) +} + fn account_config( source: &mut impl FnMut(&str) -> Option, deployment_mode: DeploymentMode, diff --git a/src/config/tests.rs b/src/config/tests.rs index 704300f..82b08a0 100644 --- a/src/config/tests.rs +++ b/src/config/tests.rs @@ -17,6 +17,7 @@ fn local_mode_uses_safe_loopback_defaults() { assert_eq!(config.blob_orphan_ttl, Duration::from_secs(86_400)); assert_eq!(config.blob_gc_interval, Duration::from_secs(3_600)); assert_eq!(config.repo_event_retention, 100_000); + assert_eq!(config.repo_stream_send_timeout, Duration::from_secs(10)); assert_eq!(config.account.did, "did:web:localhost"); assert_eq!(config.account.handle, "alice.test"); assert_eq!(config.account.password, DEVELOPMENT_PASSWORD); @@ -116,6 +117,15 @@ fn repository_event_retention_is_bounded() { } } +#[test] +fn repository_stream_send_timeout_is_bounded() { + for value in ["99", "60001", "not-a-number"] { + let values = BTreeMap::from([("CPDS_REPO_STREAM_SEND_TIMEOUT_MS", value.to_owned())]); + let error = parse(&values).err().unwrap().to_string(); + assert!(error.contains("CPDS_REPO_STREAM_SEND_TIMEOUT_MS")); + } +} + #[test] fn blob_lifecycle_intervals_are_bounded() { for (name, value) in [ diff --git a/src/http/subscription.rs b/src/http/subscription.rs index c58de2f..b0f7815 100644 --- a/src/http/subscription.rs +++ b/src/http/subscription.rs @@ -19,7 +19,6 @@ use super::{ApiError, ApiResult}; const EVENT_PAGE_SIZE: usize = 100; const POLL_INTERVAL: Duration = Duration::from_millis(250); -const SEND_TIMEOUT: Duration = Duration::from_secs(10); const MAX_CLIENT_MESSAGE_BYTES: usize = 64 * 1024; pub async fn subscribe_repos( @@ -51,9 +50,10 @@ async fn stream_repositories( _permit: SubscriptionPermit, ) { let mut shutdown = state.stream_shutdown(); + let send_timeout = state.config.repo_stream_send_timeout; loop { if *shutdown.borrow() { - close_for_shutdown(&mut socket, cursor).await; + close_for_shutdown(&mut socket, cursor, send_timeout).await; return; } let messages = match state @@ -63,7 +63,7 @@ async fn stream_repositories( { Ok(messages) => messages, Err(error) => { - send_stream_error(&mut socket, &error).await; + send_stream_error(&mut socket, &error, send_timeout).await; return; } }; @@ -78,7 +78,7 @@ async fn stream_repositories( () = sleep(POLL_INTERVAL) => {} changed = shutdown.changed() => { if changed.is_ok() { - close_for_shutdown(&mut socket, cursor).await; + close_for_shutdown(&mut socket, cursor, send_timeout).await; } return; } @@ -91,16 +91,16 @@ async fn stream_repositories( Ok(frame) => frame, Err(error) => { tracing::error!(%error, "could not encode repository stream event"); - send_stream_error(&mut socket, &error).await; + send_stream_error(&mut socket, &error, send_timeout).await; return; } }; - let send = timeout(SEND_TIMEOUT, socket.send(ws::Message::Binary(frame.into()))); + let send = timeout(send_timeout, socket.send(ws::Message::Binary(frame.into()))); let result = tokio::select! { result = send => result, changed = shutdown.changed() => { if changed.is_ok() { - close_for_shutdown(&mut socket, cursor).await; + close_for_shutdown(&mut socket, cursor, send_timeout).await; } return; } @@ -115,10 +115,16 @@ async fn stream_repositories( return; } Err(_) => { + tracing::warn!( + cursor, + ?send_timeout, + "disconnecting stalled repository stream consumer" + ); send_error_frame( &mut socket, "ConsumerTooSlow", "repository stream client could not keep up", + send_timeout, ) .await; return; @@ -128,12 +134,12 @@ async fn stream_repositories( } } -async fn close_for_shutdown(socket: &mut ws::WebSocket, cursor: i64) { +async fn close_for_shutdown(socket: &mut ws::WebSocket, cursor: i64, send_timeout: Duration) { let frame = CloseFrame { code: 1001, reason: format!("server shutdown; cursor={cursor}").into(), }; - let _ = timeout(SEND_TIMEOUT, socket.send(ws::Message::Close(Some(frame)))).await; + let _ = timeout(send_timeout, socket.send(ws::Message::Close(Some(frame)))).await; } fn message_sequence(message: &subscribe_repos::Message) -> Option { @@ -166,20 +172,33 @@ fn encode_frame(kind: &'static str, payload: &impl Serialize) -> Result, Ok(frame) } -async fn send_stream_error(socket: &mut ws::WebSocket, error: &Error) { +async fn send_stream_error(socket: &mut ws::WebSocket, error: &Error, send_timeout: Duration) { match error { - Error::Xrpc { name, message } => send_error_frame(socket, name, message).await, + Error::Xrpc { name, message } => { + send_error_frame(socket, name, message, send_timeout).await; + } error => { tracing::error!(%error, "repository stream failed"); - send_error_frame(socket, "InternalServerError", "repository stream failed").await; + send_error_frame( + socket, + "InternalServerError", + "repository stream failed", + send_timeout, + ) + .await; } } } -async fn send_error_frame(socket: &mut ws::WebSocket, name: &'static str, message: &str) { +async fn send_error_frame( + socket: &mut ws::WebSocket, + name: &'static str, + message: &str, + send_timeout: Duration, +) { let frame = encode_error(name, message); if let Ok(frame) = frame { - let _ = timeout(SEND_TIMEOUT, socket.send(ws::Message::Binary(frame.into()))).await; + let _ = timeout(send_timeout, socket.send(ws::Message::Binary(frame.into()))).await; } } diff --git a/tests/backup_restore.rs b/tests/backup_restore.rs index bf074e1..32e8822 100644 --- a/tests/backup_restore.rs +++ b/tests/backup_restore.rs @@ -132,6 +132,7 @@ fn test_config(database_path: std::path::PathBuf, blob_directory: std::path::Pat blob_orphan_ttl: std::time::Duration::from_secs(86_400), blob_gc_interval: std::time::Duration::from_secs(3_600), repo_event_retention: 100_000, + repo_stream_send_timeout: std::time::Duration::from_secs(10), account: AccountConfig { did: DID.to_owned(), handle: "alice.test".to_owned(), diff --git a/tests/process_lifecycle.rs b/tests/process_lifecycle.rs index cf436f0..176ca26 100644 --- a/tests/process_lifecycle.rs +++ b/tests/process_lifecycle.rs @@ -8,7 +8,7 @@ use std::{ time::{Duration, Instant}, }; -const CPDS_VARIABLES: [&str; 15] = [ +const CPDS_VARIABLES: [&str; 16] = [ "CPDS_ACCOUNT_DID", "CPDS_ACCOUNT_HANDLE", "CPDS_ACCOUNT_PASSWORD", @@ -24,6 +24,7 @@ const CPDS_VARIABLES: [&str; 15] = [ "CPDS_OAUTH_REDIRECT_URI", "CPDS_PUBLIC_URL", "CPDS_REPO_EVENT_RETENTION", + "CPDS_REPO_STREAM_SEND_TIMEOUT_MS", ]; #[tokio::test] diff --git a/tests/subscribe_repos_flow.rs b/tests/subscribe_repos_flow.rs index 7f49440..bce4005 100644 --- a/tests/subscribe_repos_flow.rs +++ b/tests/subscribe_repos_flow.rs @@ -19,9 +19,9 @@ use serde::Deserialize; use serde_json::json; use tokio::{ io::{AsyncReadExt as _, AsyncWriteExt as _}, - net::{TcpListener, TcpStream}, + net::{TcpListener, TcpSocket, TcpStream}, sync::oneshot, - time::timeout, + time::{sleep, timeout}, }; const DID: &str = "did:web:alice.test"; @@ -186,6 +186,63 @@ async fn subscribe_repos_reports_cursors_outside_the_retained_window() { .unwrap(); } +#[tokio::test] +async fn subscribe_repos_disconnects_a_stalled_consumer_and_allows_replay() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let state = state_with_options( + address, + Database::open_in_memory().unwrap(), + 100_000, + Duration::from_millis(100), + ); + let actor = Account { + did: DID.to_owned(), + handle: "alice.test".to_owned(), + }; + let (shutdown, stopped) = oneshot::channel(); + let server = tokio::spawn(http::serve(listener, state.clone(), async move { + let _ = stopped.await; + })); + let (stalled, _) = open_websocket_with_receive_buffer(address, "", Some(1_024)).await; + + let write = state + .repository + .create_record( + &actor, + DID, + "app.bsky.feed.post", + Some("large"), + json!({ + "$type": "app.bsky.feed.post", + "text": "x".repeat(900_000), + "createdAt": "2026-08-22T00:00:00Z" + }), + None, + ) + .await + .unwrap(); + + sleep(Duration::from_secs(1)).await; + wait_for_disconnect(stalled).await; + + let (mut replay, mut buffered) = open_websocket(address, "?cursor=0").await; + assert_initial_lifecycle_events(&mut replay, &mut buffered).await; + let frame = read_binary_frame(&mut replay, &mut buffered).await; + let (header, commit) = decode_message::(&frame); + assert_eq!(header.t.as_deref(), Some("#commit")); + assert_eq!(commit.seq, 3); + assert_eq!(commit.commit.0.to_string(), write.commit.cid); + drop(replay); + + shutdown.send(()).unwrap(); + timeout(Duration::from_secs(2), server) + .await + .unwrap() + .unwrap() + .unwrap(); +} + #[tokio::test] async fn subscribe_repos_shutdown_reports_the_last_durable_cursor() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -249,6 +306,20 @@ fn state_with_retention( address: SocketAddr, database: Database, repo_event_retention: usize, +) -> AppState { + state_with_options( + address, + database, + repo_event_retention, + Duration::from_secs(10), + ) +} + +fn state_with_options( + address: SocketAddr, + database: Database, + repo_event_retention: usize, + repo_stream_send_timeout: Duration, ) -> AppState { let config = Config { deployment_mode: DeploymentMode::Local, @@ -260,6 +331,7 @@ fn state_with_retention( blob_orphan_ttl: std::time::Duration::from_secs(86_400), blob_gc_interval: std::time::Duration::from_secs(3_600), repo_event_retention, + repo_stream_send_timeout, account: AccountConfig { did: DID.to_owned(), handle: "alice.test".to_owned(), @@ -301,7 +373,19 @@ async fn create_post( } async fn open_websocket(address: SocketAddr, query: &str) -> (TcpStream, Vec) { - let mut stream = TcpStream::connect(address).await.unwrap(); + open_websocket_with_receive_buffer(address, query, None).await +} + +async fn open_websocket_with_receive_buffer( + address: SocketAddr, + query: &str, + receive_buffer: Option, +) -> (TcpStream, Vec) { + let socket = TcpSocket::new_v4().unwrap(); + if let Some(receive_buffer) = receive_buffer { + socket.set_recv_buffer_size(receive_buffer).unwrap(); + } + let mut stream = socket.connect(address).await.unwrap(); let request = format!( "GET /xrpc/com.atproto.sync.subscribeRepos{query} HTTP/1.1\r\n\ Host: {address}\r\n\ @@ -325,6 +409,31 @@ async fn open_websocket(address: SocketAddr, query: &str) -> (TcpStream, Vec (stream, frame) } +async fn wait_for_disconnect(mut stream: TcpStream) { + timeout(Duration::from_secs(2), async move { + let mut bytes = [0; 8 * 1024]; + loop { + match stream.read(&mut bytes).await { + Ok(0) => return, + Ok(_) => {} + Err(error) + if matches!( + error.kind(), + std::io::ErrorKind::ConnectionReset + | std::io::ErrorKind::ConnectionAborted + | std::io::ErrorKind::BrokenPipe + ) => + { + return; + } + Err(error) => panic!("unexpected stalled-consumer read error: {error}"), + } + } + }) + .await + .expect("stalled repository consumer remained connected"); +} + async fn read_binary_frame(stream: &mut TcpStream, buffered: &mut Vec) -> Vec { let (opcode, payload) = read_frame(stream, buffered).await; assert_eq!(opcode, 2, "expected a binary WebSocket frame"); diff --git a/tests/support/mod.rs b/tests/support/mod.rs index db62c6d..a1e9e04 100644 --- a/tests/support/mod.rs +++ b/tests/support/mod.rs @@ -47,6 +47,7 @@ impl TestApp { blob_orphan_ttl: std::time::Duration::from_secs(86_400), blob_gc_interval: std::time::Duration::from_secs(3_600), repo_event_retention: 100_000, + repo_stream_send_timeout: std::time::Duration::from_secs(10), account: AccountConfig { did: "did:web:alice.test".to_owned(), handle: "alice.test".to_owned(), -- knotmirror 0.0.1