Merge pull request #123 from utkarshgupta137/rsky-relay/fix-export

rsky-relay: handle plc export edge case; fix queue drain at start; optimize plc_directory

Author: rudyfraser

.gitignore

@@ -14,6 +14,8 @@ Rocket.toml
 **/data/
 db/
 *.db
+*.db-shm
+*.db-wal
 *.db-journal
 rsky-relay.log.*
 

rsky-relay/Cargo.toml

@@ -30,7 +30,7 @@ p256 = "0.13"
 reqwest = { version = "0.12", default-features = false, features = ["blocking", "gzip", "hickory-dns", "http2", "json", "rustls-tls-webpki-roots-no-provider"] }
 rs-car-sync = "0.4"
 rtrb = "0.3"
-rusqlite = { version = "0.35", features = ["bundled", "chrono"] }
+rusqlite = { version = "0.36", features = ["bundled", "chrono"] }
 rustls = "0.23"
 rustls-pemfile = "2"
 serde = { version = "1", features = ["derive"] }

rsky-relay/crawler.py

@@ -14,30 +14,66 @@
 def create_database():
     """Create SQLite database and table if they don't exist."""
     conn = sqlite3.connect(DB_FILE)
+    conn.execute("""PRAGMA auto_vacuum = INCREMENTAL""")
+    conn.execute("""PRAGMA journal_mode = WAL""")
     cursor = conn.cursor()
 
+    # Set PRAGMAs
+    conn.execute("""PRAGMA cache_size = -64000""")
+    conn.execute("""PRAGMA journal_size_limit = 6144000""")
+    conn.execute("""PRAGMA mmap_size = 268435456""")
+    conn.execute("""PRAGMA secure_delete = OFF""")
+    conn.execute("""PRAGMA synchronous = NORMAL""")
+    conn.execute("""PRAGMA temp_store = MEMORY""")
+
     # Create table for PLC operations
     cursor.execute("""
     CREATE TABLE IF NOT EXISTS plc_operations (
-        did TEXT,
-        created_at TEXT,
-        nullified BOOLEAN,
-        cid TEXT,
-        operation TEXT
+        cid TEXT NOT NULL PRIMARY KEY ON CONFLICT REPLACE,
+        did TEXT NOT NULL,
+        created_at TEXT NOT NULL,
+        nullified BOOLEAN NOT NULL,
+        operation BLOB NOT NULL,
+        pds_endpoint TEXT GENERATED ALWAYS AS (
+            json_extract(operation, '$.services.atproto_pds.endpoint')
+        ) STORED,
+        atproto_key TEXT GENERATED ALWAYS AS (
+            json_extract(operation, '$.verificationMethods.atproto')
+        ) STORED,
+        labeler_endpoint TEXT GENERATED ALWAYS AS (
+            json_extract(operation, '$.services.atproto_labeler.endpoint')
+        ) STORED,
+        atproto_label_key TEXT GENERATED ALWAYS AS (
+            json_extract(operation, '$.verificationMethods.atproto_label')
+        ) STORED
     )
     """)
+
+    # Drop all views
+    cursor.execute("""DROP VIEW IF EXISTS plc_labelers""")
+    cursor.execute("""DROP VIEW IF EXISTS plc_pdses""")
+    cursor.execute("""DROP VIEW IF EXISTS plc_keys""")
+    cursor.execute("""DROP VIEW IF EXISTS plc_latest""")
+
+    # Create indexes
     cursor.execute("""
-    CREATE INDEX IF NOT EXISTS did_index ON plc_operations (
-	    did
-    )
+    CREATE INDEX IF NOT EXISTS idx_plc_operations_did_created_at
+        ON plc_operations (did, created_at DESC)
     """)
     cursor.execute("""
-    CREATE INDEX IF NOT EXISTS created_at_index ON plc_operations (
-	    created_at ASC
-    )
+    CREATE INDEX IF NOT EXISTS idx_plc_operations_pds_endpoint
+        ON plc_operations (pds_endpoint, created_at)
+        WHERE pds_endpoint IS NOT NULL
     """)
     cursor.execute("""
-    CREATE VIEW IF NOT EXISTS plc_latest AS
+    CREATE INDEX IF NOT EXISTS idx_plc_operations_labeler_endpoint
+        ON plc_operations (labeler_endpoint, created_at)
+        WHERE labeler_endpoint IS NOT NULL
+    """)
+
+    # Create views
+    cursor.execute("""
+    CREATE VIEW plc_latest AS
         SELECT *
         FROM plc_operations
         WHERE created_at = (
@@ -47,16 +83,45 @@ def create_database():
         )
     """)
     cursor.execute("""
-    CREATE VIEW IF NOT EXISTS plc_keys AS
+    CREATE VIEW plc_keys AS
+        SELECT
+            did,
+            created_at,
+            pds_endpoint,
+            atproto_key AS pds_key,
+            labeler_endpoint,
+            atproto_label_key AS labeler_key
+        FROM plc_latest
+    """)
+    cursor.execute("""
+    CREATE VIEW plc_pdses AS
+        SELECT
+            MIN(created_at) AS first,
+            MAX(created_at) AS last,
+            count() AS accounts,
+            pds_endpoint
+        FROM plc_latest
+        WHERE pds_endpoint IS NOT NULL
+        GROUP BY pds_endpoint
+        ORDER BY last
+    """)
+    cursor.execute("""
+    CREATE VIEW plc_labelers AS
         SELECT
-	        did,
-	        created_at,
-	        json_extract(operation, '$.services.atproto_pds.endpoint') AS endpoint,
-	        json_extract(operation, '$.verificationMethods.atproto') AS key
+            did,
+            created_at,
+            labeler_endpoint
         FROM plc_latest
+        WHERE labeler_endpoint IS NOT NULL
+        ORDER BY created_at
     """)
 
     conn.commit()
+
+    # Vacuum & optimize
+    cursor.execute("""PRAGMA incremental_vacuum""")
+    cursor.execute("""PRAGMA optimize = 0x10002""")
+
     return conn
 
 
@@ -83,15 +148,15 @@ def insert_operations(conn, operations):
     for op in operations:
         cursor.execute(
             """
-            INSERT INTO plc_operations (did, cid, nullified, created_at, operation)
+            INSERT INTO plc_operations (cid, did, created_at, nullified, operation)
             VALUES (?, ?, ?, ?, ?)
             """,
             (
-                op.get("did"),
                 op.get("cid"),
-                op.get("nullified"),
+                op.get("did"),
                 op.get("createdAt"),
-                json.dumps(op.get("operation")),
+                op.get("nullified"),
+                json.dumps(op.get("operation"), separators=(",", ":")).encode("utf-8"),
             ),
         )
 
@@ -131,7 +196,7 @@ def main():
     after = latest_timestamp
 
     total_processed = get_count(conn) or 0
-    request_count = 0
+    request_count = total_processed // 999
 
     try:
         print("Starting PLC Directory API crawl...")
@@ -145,7 +210,8 @@ def main():
                 break
 
             insert_operations(conn, operations)
-            total_processed += len(operations)
+            prev_processed = total_processed
+            total_processed = get_count(conn) or 0
 
             # Get the last timestamp for the next request
             last_op = operations[-1]
@@ -156,6 +222,9 @@ def main():
                 f"Request #{request_count}: Fetched {len(operations)}, "
                 f"Total {total_processed}, Last timestamp: {after}"
             )
+            ignored = len(operations) - (total_processed - prev_processed)
+            if ignored != 1:
+                print(f"IGNORED: {ignored}")
 
             # Check if we got fewer records than requested (end of data)
             if len(operations) < COUNT_PER_REQUEST:

rsky-relay/src/crawler/manager.rs

@@ -178,7 +178,7 @@ impl Manager {
     fn get_cursor(&self, host: &str) -> Result<Option<Cursor>, ManagerError> {
         let mut stmt = self.conn.prepare_cached("SELECT * FROM hosts WHERE host = ?1")?;
         Ok(stmt
-            .query_row((&host,), |row| Ok(row.get_unwrap::<_, u64>("cursor")))
+            .query_one((&host,), |row| Ok(row.get_unwrap::<_, u64>("cursor")))
             .optional()?
             .map(Into::into))
     }

rsky-relay/src/validator/event.rs

@@ -13,6 +13,9 @@ use rsky_common::tid::TID;
 
 use crate::types::Cursor;
 
+pub type DidEndpoint = Option<Box<str>>;
+pub type DidKey = [u8; 35];
+
 #[derive(Debug, Error)]
 pub enum ParseError {
     #[error("header error: {0}")]

rsky-relay/src/validator/manager.rs

@@ -92,18 +92,26 @@ impl Manager {
             self.repos.insert(did, state);
             repos += 1;
         }
-        let mut queue = 0;
+
+        let mut cursor = self.firehose.last_key_value()?.map(|(k, _)| k.into()).unwrap_or_default();
+        let mut queue_drained = 0;
+        let mut queue_pending = 0;
         for res in self.queue.keys() {
             let key = res?;
             #[expect(clippy::unwrap_used)]
             let key = std::str::from_utf8(&key).unwrap();
             #[expect(clippy::unwrap_used)]
-            self.resolver.resolve(key.split('>').next().unwrap())?;
-            queue += 1;
+            let did = key.split('>').next().unwrap();
+            if self.resolver.resolve(did)?.is_some() {
+                self.scan_did(&mut cursor, did)?;
+                queue_drained += 1;
+            } else {
+                queue_pending += 1;
+            }
         }
-        let mut seq = self.firehose.last_key_value()?.map(|(k, _)| k.into()).unwrap_or_default();
-        tracing::info!(%hosts, %repos, %queue, %seq, "loaded state");
-        while self.update(&mut seq).await? {}
+
+        tracing::info!(%hosts, %repos, %queue_drained, %queue_pending, %cursor, "loaded state");
+        while self.update(&mut cursor).await? {}
         tracing::info!("shutting down validator");
         SHUTDOWN.store(true, Ordering::Relaxed);
         Ok(())
@@ -266,83 +274,87 @@ impl Manager {
             self.hosts.insert(host.clone(), (seq, time));
         }
 
-        let mut batch: Option<Batch> = None;
         for did in self.resolver.poll().await? {
-            let Some((pds, key)) = self.resolver.resolve(&did)? else {
-                continue;
-            };
+            self.scan_did(cursor, &did)?;
+        }
 
-            for res in self.queue.prefix(&did) {
-                let (k, input) = res?;
-                batch.get_or_insert_with(|| DB.batch()).remove(&self.queue, k.clone());
-
-                #[expect(clippy::unwrap_used)]
-                let host = std::str::from_utf8(&k).unwrap().split('>').nth(1).unwrap();
-                let span = tracing::debug_span!("msg_read", %host, len = %input.len());
-                let _enter = span.enter();
-
-                #[expect(clippy::unwrap_used)]
-                let event = SubscribeReposEvent::parse(&input)?.unwrap(); // already parsed
-                let type_ = event.type_();
-                let seq = event.seq();
-                let time = event.time();
-                let did = event.did();
-                let span = tracing::debug_span!("msg_data", type = %type_, %seq, %time, %did);
-                let _enter = span.enter();
-
-                #[expect(clippy::unwrap_used)]
-                let (commit, head) = event.commit()?.unwrap(); // already parsed
-                let span =
-                    tracing::debug_span!("validate", rev = %commit.rev, data = %commit.data, %head);
-                let _enter = span.enter();
-
-                if let Some(pds) = pds {
-                    if host != pds {
-                        tracing::debug!(%pds, "hostname pds mismatch");
-                        continue;
-                    }
+        Ok(true)
+    }
+
+    fn scan_did(&mut self, cursor: &mut Cursor, did: &str) -> Result<(), ManagerError> {
+        let Some((pds, key)) = self.resolver.resolve(did)? else { unreachable!() };
+
+        let mut batch: Option<Batch> = None;
+        for res in self.queue.prefix(&did) {
+            let (k, input) = res?;
+            batch.get_or_insert_with(|| DB.batch()).remove(&self.queue, k.clone());
+
+            #[expect(clippy::unwrap_used)]
+            let host = std::str::from_utf8(&k).unwrap().split('>').nth(1).unwrap();
+            let span = tracing::debug_span!("msg_read", %host, len = %input.len());
+            let _enter = span.enter();
+
+            #[expect(clippy::unwrap_used)]
+            let event = SubscribeReposEvent::parse(&input)?.unwrap(); // already parsed
+            let type_ = event.type_();
+            let seq = event.seq();
+            let time = event.time();
+            let did = event.did();
+            let span = tracing::debug_span!("msg_data", type = %type_, %seq, %time, %did);
+            let _enter = span.enter();
+
+            #[expect(clippy::unwrap_used)]
+            let (commit, head) = event.commit()?.unwrap(); // already parsed
+            let span =
+                tracing::debug_span!("validate", rev = %commit.rev, data = %commit.data, %head);
+            let _enter = span.enter();
+
+            if let Some(pds) = pds {
+                if host != pds {
+                    tracing::debug!(%pds, "hostname pds mismatch");
+                    continue;
                 }
+            }
 
-                // verify signature
-                match utils::verify_commit_sig(&commit, key) {
-                    Ok(valid) => {
-                        if !valid {
-                            tracing::debug!(?key, "signature mismatch");
-                            continue;
-                        }
-                    }
-                    Err(err) => {
-                        tracing::debug!(%err, ?key, "signature check error");
+            // verify signature
+            match utils::verify_commit_sig(&commit, key) {
+                Ok(valid) => {
+                    if !valid {
+                        tracing::debug!(?key, "signature mismatch");
                         continue;
                     }
                 }
+                Err(err) => {
+                    tracing::debug!(%err, ?key, "signature check error");
+                    continue;
+                }
+            }
 
-                // verify commit message
-                let rev = commit.rev;
-                let data = commit.data;
-                let entry = self.repos.entry(commit.did);
-                if let SubscribeReposEvent::Commit(commit) = &event {
-                    // TODO: should still validate records existing in blocks, etc
-                    if let Entry::Occupied(prev) = &entry {
-                        let prev = prev.get();
-                        let span = tracing::debug_span!("previous", rev = %prev.rev, data = %prev.data, head = %prev.head);
-                        let _enter = span.enter();
-                        if !utils::verify_commit_event(commit, data, prev) {
-                            continue;
-                        }
+            // verify commit message
+            let rev = commit.rev;
+            let data = commit.data;
+            let entry = self.repos.entry(commit.did);
+            if let SubscribeReposEvent::Commit(commit) = &event {
+                // TODO: should still validate records existing in blocks, etc
+                if let Entry::Occupied(prev) = &entry {
+                    let prev = prev.get();
+                    let span = tracing::debug_span!("previous", rev = %prev.rev, data = %prev.data, head = %prev.head);
+                    let _enter = span.enter();
+                    if !utils::verify_commit_event(commit, data, prev) {
+                        continue;
                     }
                 }
-
-                let msg = event.serialize(input.len(), cursor.next())?;
-                self.firehose.insert(*cursor, msg)?;
-                entry.insert(RepoState { rev, data, head });
             }
+
+            let msg = event.serialize(input.len(), cursor.next())?;
+            self.firehose.insert(*cursor, msg)?;
+            entry.insert(RepoState { rev, data, head });
         }
         if let Some(batch) = batch {
             batch.commit()?;
         }
 
-        Ok(true)
+        Ok(())
     }
 }