tokens cleanup · microcosm.blue/lightrail@c4417d2

+14 -9

src/sync/backfill.rs

··· 15 15 url::Host, 16 16 {IntoStatic, xrpc::XrpcExt}, 17 17 }; 18 + use tokio::time::Duration; 18 19 use tokio_util::sync::CancellationToken; 19 20 use tracing::{info, trace, warn}; 20 21 ··· 27 28 backfill_progress::{BackfillProgress, get, set}, 28 29 resync_queue::ResyncItem, 29 30 }, 30 - util::unix_now, 31 + util::{TokenExt, unix_now}, 31 32 }; 32 33 33 34 const PAGE_LIMIT: i64 = 500; ··· 84 85 let now = unix_now(); 85 86 86 87 let dids = match &resolver { 87 - Some(r) => validate_dids(dids, r, &host).await, 88 + Some(r) => validate_dids(dids, r, &host, &token).await, 88 89 None => dids, 89 90 }; 90 91 ··· 137 138 /// Fetch one `listRepos` page with retry logic. 138 139 /// 139 140 /// Returns `None` if the host gives a 4xx, exceeds `MAX_PAGE_FAILURES` 140 - /// transient errors, or the token is cancelled during a retry sleep. 141 + /// transient errors, or the token is cancelled (including mid-request). 141 142 async fn fetch_page( 142 143 host: &Host, 143 144 base: &jacquard_common::url::Url, ··· 151 152 }; 152 153 let mut failures: u32 = 0; 153 154 loop { 154 - let result = match client.xrpc(base.clone()).send(&req).await { 155 + let result = match token.run(client.xrpc(base.clone()).send(&req)).await? { 155 156 Err(e) => { 156 157 let is_client_err = matches!( 157 158 e.kind(), ··· 191 192 "listRepos page failed {MAX_PAGE_FAILURES} times; giving up on this host"); 192 193 return None; 193 194 } 194 - tokio::select! { 195 - _ = token.cancelled() => return None, 196 - _ = tokio::time::sleep(tokio::time::Duration::from_secs(RETRY_DELAY_SECS)) => {} 195 + if !token.sleep(Duration::from_secs(RETRY_DELAY_SECS)).await { 196 + return None; 197 197 } 198 198 } 199 199 } ··· 209 209 /// Bypasses the identity cache to avoid populating it with entries that may 210 210 /// be rejected. Accepted DIDs are inserted into the cache explicitly so the 211 211 /// resync worker gets a warm hit when it picks them up shortly after. 212 + /// Returns early with whatever has been validated so far if the token is cancelled. 212 213 async fn validate_dids( 213 214 dids: Vec<Did<'static>>, 214 215 resolver: &Resolver, 215 216 host: &Host, 217 + token: &CancellationToken, 216 218 ) -> Vec<Did<'static>> { 217 219 let host_str = host.to_string(); 218 220 let mut valid = Vec::with_capacity(dids.len()); 219 221 for did in dids { 220 - match resolver.resolve_no_cache(&did).await { 222 + let Some(r) = token.run(resolver.resolve_no_cache(&did)).await else { 223 + break; 224 + }; 225 + match r { 221 226 Ok((pds_url, pubkey)) if pds_url.host_str() == Some(host_str.as_str()) => { 222 227 resolver.insert_resolved(did.clone().into_static(), pds_url, pubkey); 223 228 valid.push(did); ··· 234 239 trace!(did = %did, error = %e, 235 240 "DID resolution failed during backfill validation; skipping"); 236 241 } 237 - } 242 + }; 238 243 } 239 244 valid 240 245 }

+45 -50

src/sync/deep_crawl.rs

··· 9 9 10 10 use jacquard_api::com_atproto::sync::list_hosts::ListHosts; 11 11 use jacquard_common::{url::Host, xrpc::XrpcExt}; 12 - use tokio::task::JoinSet; 12 + use tokio::{task::JoinSet, time::Duration}; 13 13 use tokio_util::sync::CancellationToken; 14 14 use tracing::{info, warn}; 15 15 ··· 19 19 identity::Resolver, 20 20 storage::{DbRef, backfill_progress, list_hosts_cursor}, 21 21 sync::backfill, 22 + util::TokenExt, 22 23 }; 23 24 24 25 const PAGE_LIMIT: i64 = 500; ··· 67 68 "deep crawl pass complete; re-polling in 20 hours" 68 69 ); 69 70 70 - tokio::select! { 71 - _ = token.cancelled() => return Ok(()), 72 - _ = tokio::time::sleep(tokio::time::Duration::from_secs(REPOLL_SECS)) => {} 71 + if !token.sleep(Duration::from_secs(REPOLL_SECS)).await { 72 + return Ok(()); 73 73 } 74 74 } 75 75 } ··· 144 144 } 145 145 146 146 // If at capacity, wait for one worker to finish before spawning. 147 - if workers.len() >= max_workers 148 - && let Some(result) = workers.join_next().await 149 - { 150 - log_worker_result(result); 151 - metrics::gauge!("lightrail_deep_crawl_workers").set(workers.len() as f64); 147 + if workers.len() >= max_workers { 148 + tokio::select! { 149 + _ = token.cancelled() => { 150 + drain_workers(&mut workers).await; 151 + return Ok(()); 152 + } 153 + Some(result) = workers.join_next() => { 154 + log_worker_result(result); 155 + metrics::gauge!("lightrail_deep_crawl_workers").set(workers.len() as f64); 156 + } 157 + } 152 158 } 153 159 154 160 let db2 = db.clone(); ··· 197 203 /// Fetch one `listHosts` page with retry logic. 198 204 /// 199 205 /// Returns `None` if the host exceeds `MAX_PAGE_FAILURES` transient errors or 200 - /// the token is cancelled during a retry sleep. 206 + /// the token is cancelled (including mid-request). 201 207 async fn fetch_hosts_page( 202 208 relay_host: &Host, 203 209 base: &jacquard_common::url::Url, ··· 209 215 cursor: cursor.map(Into::into), 210 216 limit: Some(PAGE_LIMIT), 211 217 }; 212 - let mut failures: u32 = 0; 213 - loop { 214 - let result = match client.xrpc(base.clone()).send(&req).await { 215 - Err(e) => { 216 - warn!(error = %e, relay = %relay_host, "listHosts request failed"); 217 - None 218 - } 219 - Ok(resp) => match resp.parse() { 220 - Ok(out) => { 221 - let next = out.cursor.as_deref().map(str::to_owned); 222 - let hostnames = out 223 - .hosts 224 - .into_iter() 225 - .map(|h| h.hostname.to_string()) 226 - .collect::<Vec<_>>(); 227 - Some((hostnames, next)) 228 - } 229 - Err(e) => { 230 - warn!(error = %e, relay = %relay_host, "listHosts response parse failed"); 231 - None 232 - } 233 - }, 218 + for attempt in 0..MAX_PAGE_FAILURES { 219 + if attempt > 0 && !token.sleep(Duration::from_secs(RETRY_DELAY_SECS)).await { 220 + return None; 221 + } 222 + 223 + let Ok(res) = token 224 + .run(client.xrpc(base.clone()).send(&req)) 225 + .await? 226 + .inspect_err(|e| info!(error = %e, relay = %relay_host, "listHosts request failed")) 227 + else { 228 + continue; 229 + }; 230 + 231 + let Ok(out) = res.parse().inspect_err( 232 + |e| info!(error = %e, relay = %relay_host, "listHosts response parse failed"), 233 + ) else { 234 + continue; 234 235 }; 235 236 236 - match result { 237 - Some(page) => return Some(page), 238 - None => { 239 - failures += 1; 240 - if failures >= MAX_PAGE_FAILURES { 241 - warn!( 242 - relay = %relay_host, 243 - failures, 244 - "listHosts page failed {MAX_PAGE_FAILURES} times; abandoning pass" 245 - ); 246 - return None; 247 - } 248 - tokio::select! { 249 - _ = token.cancelled() => return None, 250 - _ = tokio::time::sleep(tokio::time::Duration::from_secs(RETRY_DELAY_SECS)) => {} 251 - } 252 - } 253 - } 237 + let next = out.cursor.as_deref().map(str::to_owned); 238 + let hostnames = out 239 + .hosts 240 + .into_iter() 241 + .map(|h| h.hostname.to_string()) 242 + .collect::<Vec<_>>(); 243 + return Some((hostnames, next)); 254 244 } 245 + warn!( 246 + relay = %relay_host, 247 + "listHosts page failed {MAX_PAGE_FAILURES} times; abandoning pass" 248 + ); 249 + None 255 250 } 256 251 257 252 // ---------------------------------------------------------------------------

+9 -12

src/sync/firehose/mod.rs

··· 40 40 41 41 use crate::error::Result; 42 42 use crate::storage::{self, DbRef}; 43 + use crate::util::TokenExt; 43 44 use event_dispatcher::CommitDispatcher; 44 45 45 46 /// Maximum reconnect delay. ··· 111 112 "connecting to firehose", 112 113 ); 113 114 114 - let stream = tokio::select! { 115 - _ = token.cancelled() => return Ok(()), 116 - r = client.subscription(base.clone()).subscribe(&params) => r, 115 + let sub = client.subscription(base.clone()); 116 + let Some(stream) = token.run(sub.subscribe(&params)).await else { 117 + return Ok(()); 117 118 }; 118 119 119 120 let stream = match stream { ··· 125 126 backoff_secs, 126 127 "firehose connect failed; will retry", 127 128 ); 128 - tokio::select! { 129 - _ = token.cancelled() => return Ok(()), 130 - _ = tokio::time::sleep(Duration::from_secs(backoff_secs)) => {} 129 + if !token.sleep(Duration::from_secs(backoff_secs)).await { 130 + return Ok(()); 131 131 } 132 132 backoff_secs = (backoff_secs * 2).min(MAX_BACKOFF_SECS); 133 133 continue 'reconnect; ··· 204 204 backoff_secs, 205 205 "firehose stream error; will reconnect", 206 206 ); 207 - tokio::select! { 208 - _ = token.cancelled() => return Ok(()), 209 - _ = tokio::time::sleep( 210 - Duration::from_secs(backoff_secs) 211 - ) => {} 212 - } 207 + if !token.sleep(Duration::from_secs(backoff_secs)).await { 208 + return Ok(()); 209 + }; 213 210 backoff_secs = (backoff_secs * 2).min(MAX_BACKOFF_SECS); 214 211 continue 'reconnect; 215 212 }

+6 -8

src/sync/resync/dispatcher.rs

··· 26 26 repo::{AccountStatus, RepoInfo, RepoState}, 27 27 resync_queue::ResyncItem, 28 28 }; 29 - use crate::util::unix_now; 29 + use crate::util::{TokenExt, unix_now}; 30 30 31 31 /// How long to wait between queue polls when no workers are running. 32 32 const IDLE_POLL_MS: u64 = 500; ··· 146 146 147 147 if workers.is_empty() { 148 148 // Nothing running, nothing claimable right now. Yield before polling. 149 - tokio::select! { 150 - _ = token.cancelled() => return Ok(()), 151 - _ = tokio::time::sleep(tokio::time::Duration::from_millis(IDLE_POLL_MS)) => {} 152 - } 149 + if !token.sleep(Duration::from_millis(IDLE_POLL_MS)).await { 150 + return Ok(()); 151 + }; 153 152 continue; 154 153 } 155 154 156 155 // Wait for the next worker to finish, then loop to try claiming more. 157 - let next = tokio::select! { 158 - _ = token.cancelled() => return Ok(()), 159 - r = workers.join_next_with_id() => r, 156 + let Some(next) = token.run(workers.join_next_with_id()).await else { 157 + return Ok(()); 160 158 }; 161 159 if let Some(join_result) = next { 162 160 // JoinError also carries the task ID, so we can recover the DID

+20

src/util.rs

··· 6 6 .unwrap_or_default() 7 7 .as_secs() 8 8 } 9 + 10 + pub trait TokenExt { 11 + /// alias for CancellationToken::run_until_cancelled 12 + fn run<F: Future>(&self, fut: F) -> impl Future<Output = Option<F::Output>>; 13 + /// sleep that ends early if if cancelled 14 + /// 15 + /// returns `true` when completed after the full sleep completed 16 + fn sleep(&self, d: tokio::time::Duration) -> impl Future<Output = bool>; 17 + } 18 + 19 + impl TokenExt for tokio_util::sync::CancellationToken { 20 + async fn run<F: Future>(&self, f: F) -> Option<F::Output> { 21 + self.run_until_cancelled(f).await 22 + } 23 + async fn sleep(&self, d: tokio::time::Duration) -> bool { 24 + self.run_until_cancelled(tokio::time::sleep(d)) 25 + .await 26 + .is_none() 27 + } 28 + }

Configure Feed

Configure Feed