retry upstream getRecord without CID if NotFound

Author: uniphil

slingshot/src/error.rs

@@ -1,3 +1,4 @@
+use crate::ErrorResponseObject;
 use thiserror::Error;
 
 #[derive(Debug, Error)]
@@ -75,4 +76,8 @@ pub enum RecordError {
     MissingUpstreamCid,
     #[error("upstream CID was not valid: {0}")]
     BadUpstreamCid(String),
+    #[error("upstream atproto-looking bad request")]
+    UpstreamBadRequest(ErrorResponseObject),
+    #[error("upstream non-atproto bad request")]
+    UpstreamBadBadNotGoodRequest(reqwest::Error),
 }

slingshot/src/lib.rs

@@ -8,5 +8,5 @@ mod server;
 pub use consumer::consume;
 pub use firehose_cache::firehose_cache;
 pub use identity::Identity;
-pub use record::{CachedRecord, Repo};
+pub use record::{CachedRecord, ErrorResponseObject, Repo};
 pub use server::serve;

slingshot/src/record.rs

@@ -2,7 +2,7 @@
 
 use crate::{Identity, error::RecordError};
 use atrium_api::types::string::{Cid, Did, Nsid, RecordKey};
-use reqwest::Client;
+use reqwest::{Client, StatusCode};
 use serde::{Deserialize, Serialize};
 use serde_json::value::RawValue;
 use std::str::FromStr;
@@ -56,6 +56,13 @@ struct RecordResponseObject {
     value: Box<RawValue>,
 }
 
+#[derive(Debug, Deserialize)]
+pub struct ErrorResponseObject {
+    error: String,
+    #[allow(dead_code)]
+    message: String,
+}
+
 #[derive(Clone)]
 pub struct Repo {
     identity: Identity,
@@ -87,39 +94,63 @@ impl Repo {
             return Err(RecordError::NotFound("could not get pds for DID"));
         };
 
-        // TODO: throttle by host probably, generally guard against outgoing requests
+        // cid gets set to None for a retry, if it's Some and we got NotFound
+        let mut cid = cid;
 
-        let mut params = vec![
-            ("repo", did.to_string()),
-            ("collection", collection.to_string()),
-            ("rkey", rkey.to_string()),
-        ];
-        if let Some(cid) = cid {
-            params.push(("cid", cid.as_ref().to_string()));
-        }
-        let mut url = Url::parse_with_params(&pds, &params)?;
-        url.set_path("/xrpc/com.atproto.repo.getRecord");
+        let res = loop {
+            // TODO: throttle outgoing requests by host probably, generally guard against outgoing requests
+            let mut params = vec![
+                ("repo", did.to_string()),
+                ("collection", collection.to_string()),
+                ("rkey", rkey.to_string()),
+            ];
+            if let Some(cid) = cid {
+                params.push(("cid", cid.as_ref().to_string()));
+            }
+            let mut url = Url::parse_with_params(&pds, &params)?;
+            url.set_path("/xrpc/com.atproto.repo.getRecord");
 
-        let res = self
-            .client
-            .get(url)
-            .send()
-            .await
-            .map_err(RecordError::SendError)?
+            let res = self
+                .client
+                .get(url.clone())
+                .send()
+                .await
+                .map_err(RecordError::SendError)?;
+
+            if res.status() == StatusCode::BAD_REQUEST {
+                // 1. if we're not able to parse json, it's not something we can handle
+                let err = res
+                    .json::<ErrorResponseObject>()
+                    .await
+                    .map_err(RecordError::UpstreamBadBadNotGoodRequest)?;
+                // 2. if we are, is it a NotFound? and if so, did we try with a CID?
+                // if so, retry with no CID (api handler will reject for mismatch but
+                // with a nice error + warm cache)
+                if err.error == "NotFound" && cid.is_some() {
+                    cid = &None;
+                    continue;
+                } else {
+                    return Err(RecordError::UpstreamBadRequest(err));
+                }
+            }
+            break res;
+        };
+
+        let data = res
             .error_for_status()
             .map_err(RecordError::StatusError)? // TODO atproto error handling (think about handling not found)
             .json::<RecordResponseObject>()
             .await
             .map_err(RecordError::ParseJsonError)?; // todo...
 
-        let Some(cid) = res.cid else {
+        let Some(cid) = data.cid else {
             return Err(RecordError::MissingUpstreamCid);
         };
         let cid = Cid::from_str(&cid).map_err(|e| RecordError::BadUpstreamCid(e.to_string()))?;
 
         Ok(CachedRecord::Found(RawRecord {
             cid,
-            record: res.value.to_string(),
+            record: data.value.to_string(),
         }))
     }
 }