diff --git a/Cargo.lock b/Cargo.lock index 3875ce5..8878429 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -41,6 +41,7 @@ dependencies = [ "reqwest", "serde", "serde_json", + "thiserror", "tokio", "tokio-postgres", "url", @@ -1760,6 +1761,26 @@ dependencies = [ "windows-sys 0.61.0", ] +[[package]] +name = "thiserror" +version = "2.0.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3467d614147380f2e4e374161426ff399c91084acd2363eaf549172b3d5e60c0" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c5e1be1c48b9172ee610da68fd9cd2770e7a4056cb3fc98710ee6906f0c7960" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "tinystr" version = "0.8.1" diff --git a/Cargo.toml b/Cargo.toml index 7456345..8a10031 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,6 +16,7 @@ log = "0.4.28" reqwest = { version = "0.12.23", features = ["stream"] } serde = "1.0.219" serde_json = { version = "1.0.143", features = ["raw_value"] } +thiserror = "2.0.16" tokio = { version = "1.47.1", features = ["full"] } tokio-postgres = { version = "0.7.13", features = ["with-chrono-0_4", "with-serde_json-1"] } url = "2.5.7" diff --git a/src/bin/blah.rs b/src/bin/get_backfill_chunk_adsf.rs similarity index 100% rename from src/bin/blah.rs rename to src/bin/get_backfill_chunk_adsf.rs diff --git a/src/bin/main.rs b/src/bin/main.rs index 22e8706..480fc39 100644 --- a/src/bin/main.rs +++ b/src/bin/main.rs @@ -4,10 +4,9 @@ use std::time::Duration; use tokio_postgres::NoTls; use url::Url; -use allegedly::{ExportPage, week_to_pages}; +use allegedly::{ExportPage, poll_upstream, week_to_pages}; const EXPORT_PAGE_QUEUE_SIZE: usize = 0; // rendezvous for now -const UPSTREAM_REQUEST_INTERVAL: Duration = Duration::from_millis(500); const WEEK_IN_SECONDS: u64 = 7 * 86400; #[derive(Parser)] @@ -43,12 +42,6 @@ struct Args { postgres: String, } -#[derive(Deserialize)] -#[serde(rename_all = "camelCase")] -struct OpPeek { - pub created_at: chrono::DateTime, -} - #[derive(Deserialize)] #[serde(rename_all = "camelCase")] struct Op<'a> { @@ -85,8 +78,9 @@ async fn export_upstream( upstream: Url, bulk: (Url, u64), tx: flume::Sender, - latest: Option>, + pg_client: tokio_postgres::Client, ) { + let latest = get_latest(&pg_client).await; let client = reqwest::Client::builder() .user_agent(concat!( "allegedly v", @@ -99,44 +93,9 @@ async fn export_upstream( if latest.is_none() { bulk_backfill(client.clone(), bulk, tx.clone()).await; } - let mut upstream = upstream; upstream.set_path("/export"); - let mut after = latest; - let mut tick = tokio::time::interval(UPSTREAM_REQUEST_INTERVAL); - - loop { - tick.tick().await; - let mut url = upstream.clone(); - if let Some(ref after) = after { - url.query_pairs_mut() - .append_pair("after", &after.to_rfc3339()); - } - let ops = client - .get(url) - .send() - .await - .unwrap() - .error_for_status() - .unwrap() - .text() - .await - .unwrap() - .trim() - .to_string(); - - let Some((_, last_line)) = ops.rsplit_once('\n') else { - log::trace!("no ops in response page, nothing to do"); - continue; - }; - - let op: OpPeek = serde_json::from_str(last_line).unwrap(); - after = Some(op.created_at); - - log::trace!("got some ops until {after:?}, sending them..."); - let ops = ops.split('\n').map(Into::into).collect(); - tx.send_async(ExportPage { ops }).await.unwrap(); - } + poll_upstream(&client, latest, upstream, tx).await.unwrap(); } async fn write_pages( @@ -221,6 +180,18 @@ async fn write_pages( Ok(()) } +async fn get_latest(pg_client: &tokio_postgres::Client) -> Option> { + pg_client + .query_opt( + r#"SELECT "createdAt" FROM operations + ORDER BY "createdAt" DESC LIMIT 1"#, + &[], + ) + .await + .unwrap() + .map(|r| r.get(0)) +} + #[tokio::main] async fn main() { env_logger::init(); @@ -241,17 +212,18 @@ async fn main() { } }); - let latest = pg_client - .query_opt( - r#"SELECT "createdAt" FROM operations - ORDER BY "createdAt" DESC LIMIT 1"#, - &[], - ) + log::trace!("connecting postgres 2..."); + let (pg_client2, connection2) = tokio_postgres::connect(&args.postgres, NoTls) .await - .unwrap() - .map(|r| r.get(0)); + .unwrap(); - log::info!("connected! latest: {latest:?}"); + // send the connection away to do the actual communication work + // TODO: error and shutdown handling + let conn_task2 = tokio::task::spawn(async move { + if let Err(e) = connection2.await { + eprintln!("connection error: {e}"); + } + }); let (tx, rx) = flume::bounded(EXPORT_PAGE_QUEUE_SIZE); @@ -259,12 +231,13 @@ async fn main() { args.upstream, (args.upstream_bulk, args.bulk_epoch), tx, - latest, + pg_client2, )); let writer_task = tokio::task::spawn(write_pages(rx, pg_client)); tokio::select! { z = conn_task => log::warn!("connection task ended: {z:?}"), + z = conn_task2 => log::warn!("connection task ended: {z:?}"), z = export_task => log::warn!("export task ended: {z:?}"), z = writer_task => log::warn!("writer task ended: {z:?}"), }; diff --git a/src/bin/tail_export.rs b/src/bin/tail_export.rs new file mode 100644 index 0000000..dd32b6c --- /dev/null +++ b/src/bin/tail_export.rs @@ -0,0 +1,44 @@ +use allegedly::OpPeek; +use url::Url; + +async fn get_page(client: &reqwest::Client, url: Url) -> Vec { + client + .get(url) + .send() + .await + .unwrap() + .error_for_status() + .unwrap() + .text() + .await + .unwrap() + .trim() + .split('\n') + .map(Into::into) + .collect() +} + +#[tokio::main] +async fn main() { + let client = reqwest::Client::builder() + .user_agent(concat!( + "allegedly (export) v", + env!("CARGO_PKG_VERSION"), + " (from @microcosm.blue; contact @bad-example.com)" + )) + .build() + .unwrap(); + + let mut url = Url::parse("https://plc.directory/export").unwrap(); + let ops = get_page(&client, url.clone()).await; + + println!("first: {:?}", ops.first()); + + if let Some(last_line) = ops.last() { + let x: OpPeek = serde_json::from_str(last_line).unwrap(); + url.query_pairs_mut() + .append_pair("after", &x.created_at.to_rfc3339()); + let ops2 = get_page(&client, url).await; + println!("2nd: {:?}", ops2.first()); + } +} diff --git a/src/lib.rs b/src/lib.rs index d5b8a88..8f8dd71 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,10 @@ +use serde::Deserialize; + mod backfill; +mod poll; pub use backfill::week_to_pages; +pub use poll::poll_upstream; /// One page of PLC export /// @@ -8,3 +12,9 @@ pub use backfill::week_to_pages; pub struct ExportPage { pub ops: Vec, } + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct OpPeek { + pub created_at: chrono::DateTime, +} diff --git a/src/poll.rs b/src/poll.rs new file mode 100644 index 0000000..5bd7aa1 --- /dev/null +++ b/src/poll.rs @@ -0,0 +1,60 @@ +use crate::{ExportPage, OpPeek}; +use chrono::{DateTime, Utc}; +use std::time::Duration; +use thiserror::Error; +use url::Url; + +const UPSTREAM_REQUEST_INTERVAL: Duration = Duration::from_millis(500); + +#[derive(Debug, Error)] +pub enum GetPageError { + #[error(transparent)] + ReqwestError(#[from] reqwest::Error), + #[error(transparent)] + SerdeError(#[from] serde_json::Error), +} + +pub async fn get_page( + client: &reqwest::Client, + url: Url, +) -> Result<(ExportPage, Option>), GetPageError> { + let ops: Vec = client + .get(url) + .send() + .await? + .error_for_status()? + .text() + .await? + .trim() + .split('\n') + .map(Into::into) + .collect(); + + let last_at = ops + .last() + .map(|s| serde_json::from_str::(s)) + .transpose()? + .map(|o| o.created_at); + + Ok((ExportPage { ops }, last_at)) +} + +pub async fn poll_upstream( + client: &reqwest::Client, + after: Option>, + base: Url, + dest: flume::Sender, +) -> anyhow::Result<()> { + let mut tick = tokio::time::interval(UPSTREAM_REQUEST_INTERVAL); + let mut after = after; + loop { + tick.tick().await; + let mut url = base.clone(); + if let Some(a) = after { + url.query_pairs_mut().append_pair("after", &a.to_rfc3339()); + }; + let (page, next_after) = get_page(client, url).await?; + dest.send_async(page).await?; + after = next_after; + } +}