+76

src/storage/list_hosts_cursor.rs

reviewed

··· 1 1 + //! listHosts pagination cursor for the upstream relay. 2 2 + //! 3 3 + //! Tracks progress through `com.atproto.sync.listHosts` for deep-crawl mode. 4 4 + //! 5 5 + //! Key: `"lhs"<upstream_host>`, value: plain string cursor 6 6 + 7 7 + use jacquard_common::url::Host; 8 8 + 9 9 + use crate::storage::{ 10 10 + DbRef, PREFIX_LIST_HOSTS, 11 11 + error::{StorageError, StorageResult}, 12 12 + }; 13 13 + 14 14 + fn key(host: &Host) -> Vec<u8> { 15 15 + let host = host.to_string(); 16 16 + let mut k = Vec::with_capacity(PREFIX_LIST_HOSTS.len() + host.len()); 17 17 + k.extend_from_slice(&PREFIX_LIST_HOSTS); 18 18 + k.extend_from_slice(host.as_bytes()); 19 19 + k 20 20 + } 21 21 + 22 22 + /// Read the `listHosts` cursor for `host`. 23 23 + /// 24 24 + /// Returns `None` if no walk has been started yet. 25 25 + pub fn get(db: &DbRef, host: &Host) -> StorageResult<Option<String>> { 26 26 + let k = key(host); 27 27 + match db.ks.get(k)? { 28 28 + None => Ok(None), 29 29 + Some(v) => { 30 30 + let s = std::str::from_utf8(v.as_ref()).map_err(|_| StorageError::Corrupt { 31 31 + key: host.to_string(), 32 32 + reason: "listHosts cursor not valid UTF-8", 33 33 + })?; 34 34 + Ok(Some(s.to_owned())) 35 35 + } 36 36 + } 37 37 + } 38 38 + 39 39 + /// Persist the `listHosts` cursor for `host`. 40 40 + pub fn set(db: &DbRef, host: &Host, cursor: &str) -> StorageResult<()> { 41 41 + db.ks.insert(key(host), cursor.as_bytes())?; 42 42 + Ok(()) 43 43 + } 44 44 + 45 45 + #[cfg(test)] 46 46 + mod tests { 47 47 + use super::*; 48 48 + use crate::storage::open_temporary; 49 49 + 50 50 + fn host(s: &str) -> Host { 51 51 + Host::parse(s).unwrap() 52 52 + } 53 53 + 54 54 + #[test] 55 55 + fn none_when_missing() { 56 56 + let db = open_temporary().unwrap(); 57 57 + assert_eq!(get(&db, &host("relay.example.com")).unwrap(), None); 58 58 + } 59 59 + 60 60 + #[test] 61 61 + fn set_and_get() { 62 62 + let db = open_temporary().unwrap(); 63 63 + let h = host("relay.example.com"); 64 64 + set(&db, &h, "page2cursor").unwrap(); 65 65 + assert_eq!(get(&db, &h).unwrap(), Some("page2cursor".to_owned())); 66 66 + } 67 67 + 68 68 + #[test] 69 69 + fn overwrite_cursor() { 70 70 + let db = open_temporary().unwrap(); 71 71 + let h = host("relay.example.com"); 72 72 + set(&db, &h, "first").unwrap(); 73 73 + set(&db, &h, "second").unwrap(); 74 74 + assert_eq!(get(&db, &h).unwrap(), Some("second".to_owned())); 75 75 + } 76 76 + }

+6

src/storage/mod.rs

reviewed

··· 2 2 pub mod collection_index; 3 3 pub mod error; 4 4 pub mod firehose_cursor; 5 5 + pub mod list_hosts_cursor; 6 6 + pub mod pds_host; 5 7 pub mod repo; 6 8 pub mod resync_buffer; 7 9 pub mod resync_queue; ··· 35 37 pub(crate) const PREFIX_RESYNC_QUEUE: KeyPrefix = *b"rsq"; 36 38 /// Per-repo buffered firehose events during resync. See [`resync_buffer`]. 37 39 pub(crate) const PREFIX_RESYNC_BUFFER: KeyPrefix = *b"rsb"; 40 40 + /// Per-PDS host state (sync1.1 mode, trust, listRepos cursor/done). See [`pds_host`]. 41 41 + pub(crate) const PREFIX_PDS_HOST: KeyPrefix = *b"pdh"; 42 42 + /// listHosts walk cursor (per upstream relay host). See [`list_hosts_cursor`]. 43 43 + pub(crate) const PREFIX_LIST_HOSTS: KeyPrefix = *b"lhs"; 38 44 39 45 use std::path::Path; 40 46 use std::sync::Arc;

+259

src/storage/pds_host.rs

reviewed

··· 1 1 + //! Per-PDS-host state storage 2 2 + //! 3 3 + //! Key: `"pdh"<hostname>`, value: `<sync11_mode>\0<trust>\0<list_repos_done>\0<cursor_or_empty>` 4 4 + 5 5 + use jacquard_common::url::Host; 6 6 + 7 7 + use crate::storage::{ 8 8 + DbRef, PREFIX_PDS_HOST, 9 9 + error::{StorageError, StorageResult}, 10 10 + }; 11 11 + 12 12 + // --------------------------------------------------------------------------- 13 13 + // Key encoding 14 14 + // --------------------------------------------------------------------------- 15 15 + 16 16 + fn key(host: &Host) -> Vec<u8> { 17 17 + let host = host.to_string(); 18 18 + let mut k = Vec::with_capacity(PREFIX_PDS_HOST.len() + host.len()); 19 19 + k.extend_from_slice(&PREFIX_PDS_HOST); 20 20 + k.extend_from_slice(host.as_bytes()); 21 21 + k 22 22 + } 23 23 + 24 24 + // --------------------------------------------------------------------------- 25 25 + // Types 26 26 + // --------------------------------------------------------------------------- 27 27 + 28 28 + /// Whether this host has been seen sending sync1.1-style events. 29 29 + /// 30 30 + /// Hosts start in `Lenient` mode and flip to `Strict` once any event with 31 31 + /// `prevData` is observed. 32 32 + #[derive(Debug, Clone, PartialEq, Eq)] 33 33 + pub enum Sync11Mode { 34 34 + Lenient, 35 35 + Strict, 36 36 + } 37 37 + 38 38 + impl Sync11Mode { 39 39 + fn as_str(&self) -> &'static str { 40 40 + match self { 41 41 + Sync11Mode::Lenient => "lenient", 42 42 + Sync11Mode::Strict => "strict", 43 43 + } 44 44 + } 45 45 + 46 46 + fn from_str(s: &str) -> Option<Self> { 47 47 + Some(match s { 48 48 + "lenient" => Sync11Mode::Lenient, 49 49 + "strict" => Sync11Mode::Strict, 50 50 + _ => return None, 51 51 + }) 52 52 + } 53 53 + } 54 54 + 55 55 + /// Trust level for a PDS host. 56 56 + #[derive(Debug, Clone, PartialEq, Eq)] 57 57 + pub enum PdsTrust { 58 58 + Ok, 59 59 + Cautious, 60 60 + Blocked, 61 61 + } 62 62 + 63 63 + impl PdsTrust { 64 64 + fn as_str(&self) -> &'static str { 65 65 + match self { 66 66 + PdsTrust::Ok => "ok", 67 67 + PdsTrust::Cautious => "cautious", 68 68 + PdsTrust::Blocked => "blocked", 69 69 + } 70 70 + } 71 71 + 72 72 + fn from_str(s: &str) -> Option<Self> { 73 73 + Some(match s { 74 74 + "ok" => PdsTrust::Ok, 75 75 + "cautious" => PdsTrust::Cautious, 76 76 + "blocked" => PdsTrust::Blocked, 77 77 + _ => return None, 78 78 + }) 79 79 + } 80 80 + } 81 81 + 82 82 + /// Stored state for a PDS host. 83 83 + #[derive(Debug, Clone)] 84 84 + pub struct PdsHostInfo { 85 85 + pub sync11_mode: Sync11Mode, 86 86 + pub trust: PdsTrust, 87 87 + /// Whether the `listRepos` walk for this PDS has completed. 88 88 + pub list_repos_done: bool, 89 89 + /// Current `listRepos` pagination cursor; empty when not yet started. 90 90 + pub list_repos_cursor: String, 91 91 + } 92 92 + 93 93 + impl Default for PdsHostInfo { 94 94 + fn default() -> Self { 95 95 + PdsHostInfo { 96 96 + sync11_mode: Sync11Mode::Lenient, 97 97 + trust: PdsTrust::Ok, 98 98 + list_repos_done: false, 99 99 + list_repos_cursor: String::new(), 100 100 + } 101 101 + } 102 102 + } 103 103 + 104 104 + // --------------------------------------------------------------------------- 105 105 + // Encode / decode 106 106 + // --------------------------------------------------------------------------- 107 107 + 108 108 + /// Wire format: `<sync11_mode>\0<trust>\0<list_repos_done>\0<cursor_or_empty>` 109 109 + fn encode(info: &PdsHostInfo) -> Vec<u8> { 110 110 + let done = if info.list_repos_done { "1" } else { "0" }; 111 111 + let s = format!( 112 112 + "{}\0{}\0{}\0{}", 113 113 + info.sync11_mode.as_str(), 114 114 + info.trust.as_str(), 115 115 + done, 116 116 + info.list_repos_cursor, 117 117 + ); 118 118 + s.into_bytes() 119 119 + } 120 120 + 121 121 + fn decode(bytes: &[u8], key: &str) -> StorageResult<PdsHostInfo> { 122 122 + let s = std::str::from_utf8(bytes).map_err(|_| StorageError::Corrupt { 123 123 + key: key.to_owned(), 124 124 + reason: "pds host info not valid UTF-8", 125 125 + })?; 126 126 + let mut parts = s.splitn(4, '\0'); 127 127 + let sync11_mode = parts 128 128 + .next() 129 129 + .and_then(Sync11Mode::from_str) 130 130 + .ok_or(StorageError::Corrupt { 131 131 + key: key.to_owned(), 132 132 + reason: "invalid sync11_mode", 133 133 + })?; 134 134 + let trust = parts 135 135 + .next() 136 136 + .and_then(PdsTrust::from_str) 137 137 + .ok_or(StorageError::Corrupt { 138 138 + key: key.to_owned(), 139 139 + reason: "invalid trust", 140 140 + })?; 141 141 + let list_repos_done = match parts.next() { 142 142 + Some("0") => false, 143 143 + Some("1") => true, 144 144 + _ => { 145 145 + return Err(StorageError::Corrupt { 146 146 + key: key.to_owned(), 147 147 + reason: "invalid list_repos_done", 148 148 + }); 149 149 + } 150 150 + }; 151 151 + let list_repos_cursor = parts.next().unwrap_or("").to_owned(); 152 152 + Ok(PdsHostInfo { 153 153 + sync11_mode, 154 154 + trust, 155 155 + list_repos_done, 156 156 + list_repos_cursor, 157 157 + }) 158 158 + } 159 159 + 160 160 + // --------------------------------------------------------------------------- 161 161 + // Public API 162 162 + // --------------------------------------------------------------------------- 163 163 + 164 164 + /// Read the [`PdsHostInfo`] for `host`. 165 165 + /// 166 166 + /// Returns the default (lenient, ok, not done, empty cursor) when the key is 167 167 + /// absent — matching the implicit state for any host we haven't seen yet. 168 168 + pub fn get(db: &DbRef, host: &Host) -> StorageResult<PdsHostInfo> { 169 169 + let k = key(host); 170 170 + match db.ks.get(&k)? { 171 171 + None => Ok(PdsHostInfo::default()), 172 172 + Some(v) => { 173 173 + let key_str = String::from_utf8_lossy(&k).into_owned(); 174 174 + decode(v.as_ref(), &key_str) 175 175 + } 176 176 + } 177 177 + } 178 178 + 179 179 + /// Add a [`PdsHostInfo`] write to an existing batch. 180 180 + pub fn put_into(batch: &mut fjall::OwnedWriteBatch, db: &DbRef, host: &Host, info: &PdsHostInfo) { 181 181 + batch.insert(&db.ks, key(host), encode(info)); 182 182 + } 183 183 + 184 184 + /// Write a [`PdsHostInfo`] for `host`. 185 185 + pub fn put(db: &DbRef, host: &Host, info: &PdsHostInfo) -> StorageResult<()> { 186 186 + let mut batch = db.database.batch(); 187 187 + put_into(&mut batch, db, host, info); 188 188 + batch.commit()?; 189 189 + Ok(()) 190 190 + } 191 191 + 192 192 + // --------------------------------------------------------------------------- 193 193 + // Tests 194 194 + // --------------------------------------------------------------------------- 195 195 + 196 196 + #[cfg(test)] 197 197 + mod tests { 198 198 + use super::*; 199 199 + use crate::storage::open_temporary; 200 200 + 201 201 + fn host(s: &str) -> Host { 202 202 + Host::parse(s).unwrap() 203 203 + } 204 204 + 205 205 + #[test] 206 206 + fn default_when_missing() { 207 207 + let db = open_temporary().unwrap(); 208 208 + let info = get(&db, &host("pds.example.com")).unwrap(); 209 209 + assert_eq!(info.sync11_mode, Sync11Mode::Lenient); 210 210 + assert_eq!(info.trust, PdsTrust::Ok); 211 211 + assert!(!info.list_repos_done); 212 212 + assert!(info.list_repos_cursor.is_empty()); 213 213 + } 214 214 + 215 215 + #[test] 216 216 + fn roundtrip_default() { 217 217 + let info = PdsHostInfo::default(); 218 218 + let encoded = encode(&info); 219 219 + let decoded = decode(&encoded, "test").unwrap(); 220 220 + assert_eq!(decoded.sync11_mode, Sync11Mode::Lenient); 221 221 + assert_eq!(decoded.trust, PdsTrust::Ok); 222 222 + assert!(!decoded.list_repos_done); 223 223 + assert!(decoded.list_repos_cursor.is_empty()); 224 224 + } 225 225 + 226 226 + #[test] 227 227 + fn roundtrip_strict_blocked_done_with_cursor() { 228 228 + let info = PdsHostInfo { 229 229 + sync11_mode: Sync11Mode::Strict, 230 230 + trust: PdsTrust::Blocked, 231 231 + list_repos_done: true, 232 232 + list_repos_cursor: "abc123".to_owned(), 233 233 + }; 234 234 + let encoded = encode(&info); 235 235 + let decoded = decode(&encoded, "test").unwrap(); 236 236 + assert_eq!(decoded.sync11_mode, Sync11Mode::Strict); 237 237 + assert_eq!(decoded.trust, PdsTrust::Blocked); 238 238 + assert!(decoded.list_repos_done); 239 239 + assert_eq!(decoded.list_repos_cursor, "abc123"); 240 240 + } 241 241 + 242 242 + #[test] 243 243 + fn put_and_get() { 244 244 + let db = open_temporary().unwrap(); 245 245 + let h = host("pds.example.com"); 246 246 + let info = PdsHostInfo { 247 247 + sync11_mode: Sync11Mode::Strict, 248 248 + trust: PdsTrust::Cautious, 249 249 + list_repos_done: false, 250 250 + list_repos_cursor: "cursor42".to_owned(), 251 251 + }; 252 252 + put(&db, &h, &info).unwrap(); 253 253 + let retrieved = get(&db, &h).unwrap(); 254 254 + assert_eq!(retrieved.sync11_mode, Sync11Mode::Strict); 255 255 + assert_eq!(retrieved.trust, PdsTrust::Cautious); 256 256 + assert!(!retrieved.list_repos_done); 257 257 + assert_eq!(retrieved.list_repos_cursor, "cursor42"); 258 258 + } 259 259 + }