parakeet / consumer / src / backfill / downloader.rs

use super::{DL_DONE_KEY, PDS_SERVICE_ID};
use crate::db;
use chrono::prelude::*;
use deadpool_postgres::{Client as PgClient, Pool};
use did_resolver::Resolver;
use futures::TryStreamExt;
use metrics::{counter, histogram};
use parakeet_db::types::{ActorStatus, ActorSyncState};
use redis::aio::MultiplexedConnection;
use redis::AsyncTypedCommands;
use reqwest::header::HeaderMap;
use reqwest::Client as HttpClient;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::sync::watch::Receiver as WatchReceiver;
use tokio::time::{Duration, Instant};
use tokio_postgres::types::Type;
use tokio_util::io::StreamReader;
use tokio_util::task::TaskTracker;
use tracing::instrument;

const BF_RESET_KEY: &str = "bf_download_ratelimit_reset";
const BF_REM_KEY: &str = "bf_download_ratelimit_rem";
const DL_DUP_KEY: &str = "bf_downloaded";

pub async fn downloader(
    mut rc: MultiplexedConnection,
    pool: Pool,
    resolver: Arc<Resolver>,
    tmp_dir: PathBuf,
    concurrency: usize,
    buffer: usize,
    tracker: TaskTracker,
    stop: WatchReceiver<bool>,
) {
    let (tx, rx) = flume::bounded(64);
    let mut conn = pool.get().await.unwrap();

    let http = HttpClient::new();

    for _ in 0..concurrency {
        tracker.spawn(download_thread(
            rc.clone(),
            pool.clone(),
            resolver.clone(),
            http.clone(),
            rx.clone(),
            tmp_dir.clone(),
        ));
    }

    let status_stmt = conn.prepare_typed_cached(
        "INSERT INTO actors (did, sync_state, last_indexed) VALUES ($1, 'processing', NOW()) ON CONFLICT (did) DO UPDATE SET sync_state = 'processing', last_indexed=NOW()",
        &[Type::TEXT]
    ).await.unwrap();

    loop {
        if stop.has_changed().unwrap_or(true) {
            tracing::info!("stopping downloader");
            break;
        }

        if let Ok(count) = rc.llen(DL_DONE_KEY).await {
            if count > buffer {
                tracing::info!("waiting due to full buffer");
                tokio::time::sleep(Duration::from_secs(5)).await;
                continue;
            }
        }

        let did: String = match rc.lpop("backfill_queue", None).await {
            Ok(Some(did)) => did,
            Ok(None) => {
                tokio::time::sleep(Duration::from_millis(250)).await;
                continue;
            }
            Err(e) => {
                tracing::error!("failed to get item from backfill queue: {e}");
                continue;
            }
        };

        tracing::trace!("resolving repo {did}");

        // has the repo already been downloaded?
        if rc.sismember(DL_DUP_KEY, &did).await.unwrap_or_default() {
            tracing::warn!("skipping duplicate repo {did}");
            continue;
        }

        // check if they're already synced in DB too
        match db::actor_get_statuses(&mut conn, &did).await {
            Ok(Some((_, state))) => {
                if state == ActorSyncState::Synced || state == ActorSyncState::Processing {
                    tracing::warn!("skipping duplicate repo {did}");
                    continue;
                }
            }
            Ok(None) => {}
            Err(e) => {
                tracing::error!(did, "failed to check current repo status: {e}");
                db::backfill_job_write(&mut conn, &did, "failed.resolve")
                    .await
                    .unwrap();
            }
        }

        match resolver.resolve_did(&did).await {
            Ok(Some(did_doc)) => {
                let Some(service) = did_doc.find_service_by_id(PDS_SERVICE_ID) else {
                    tracing::warn!("bad DID doc for {did}");
                    db::backfill_job_write(&mut conn, &did, "failed.resolve")
                        .await
                        .unwrap();
                    continue;
                };
                let service = service.service_endpoint.clone();

                // set the repo to processing
                if let Err(e) = conn.execute(&status_stmt, &[&did]).await {
                    tracing::error!("failed to update repo status for {did}: {e}");
                    continue;
                }

                let handle = did_doc
                    .also_known_as
                    .and_then(|akas| akas.first().map(|v| v[5..].to_owned()));

                tracing::trace!("resolved repo {did} {service}");
                if let Err(e) = tx.send_async((service, did, handle)).await {
                    tracing::error!("failed to send: {e}");
                }
            }
            Ok(None) => {
                tracing::warn!(did, "bad DID doc");
                db::actor_set_sync_status(&mut conn, &did, ActorSyncState::Dirty, Utc::now())
                    .await
                    .unwrap();
                db::backfill_job_write(&mut conn, &did, "failed.resolve")
                    .await
                    .unwrap();
            }
            Err(e) => {
                tracing::error!(did, "failed to resolve DID doc: {e}");
                db::actor_set_sync_status(&mut conn, &did, ActorSyncState::Dirty, Utc::now())
                    .await
                    .unwrap();
                db::backfill_job_write(&mut conn, &did, "failed.resolve")
                    .await
                    .unwrap();
            }
        }
    }
}

async fn download_thread(
    mut rc: MultiplexedConnection,
    pool: Pool,
    resolver: Arc<Resolver>,
    http: reqwest::Client,
    rx: flume::Receiver<(String, String, Option<String>)>,
    tmp_dir: PathBuf,
) {
    tracing::debug!("spawning thread");

    // this will return Err(_) and exit when all senders (only held above) are dropped
    while let Ok((pds, did, maybe_handle)) = rx.recv_async().await {
        if let Err(e) = enforce_ratelimit(&mut rc, &pds).await {
            tracing::error!("ratelimiter error: {e}");
            continue;
        };

        {
            tracing::trace!("getting DB conn...");
            let mut conn = pool.get().await.unwrap();
            tracing::trace!("got DB conn...");
            match check_and_update_repo_status(&http, &mut conn, &pds, &did).await {
                Ok(true) => {}
                Ok(false) => continue,
                Err(e) => {
                    tracing::error!(pds, did, "failed to check repo status: {e}");
                    db::backfill_job_write(&mut conn, &did, "failed.resolve")
                        .await
                        .unwrap();
                    continue;
                }
            }

            tracing::debug!("trying to resolve handle...");
            if let Some(handle) = maybe_handle {
                if let Err(e) = resolve_and_set_handle(&conn, &resolver, &did, &handle).await {
                    tracing::error!(pds, did, "failed to resolve handle: {e}");
                    db::backfill_job_write(&mut conn, &did, "failed.resolve")
                        .await
                        .unwrap();
                }
            }
        }

        let start = Instant::now();

        tracing::trace!("downloading repo {did}");

        match download_car(&http, &tmp_dir, &pds, &did).await {
            Ok(Some((rem, reset))) => {
                let _ = rc.zadd(BF_REM_KEY, &pds, rem).await;
                let _ = rc.zadd(BF_RESET_KEY, &pds, reset).await;
            }
            Ok(_) => tracing::warn!(pds, "got response with no ratelimit headers."),
            Err(e) => {
                tracing::error!(pds, did, "failed to download repo: {e}");
                continue;
            }
        }

        histogram!("backfill_download_dur", "pds" => pds).record(start.elapsed().as_secs_f64());

        let _ = rc.sadd(DL_DUP_KEY, &did).await;
        if let Err(e) = rc.rpush(DL_DONE_KEY, &did).await {
            tracing::error!(did, "failed to mark download complete: {e}");
        } else {
            counter!("backfill_downloaded").increment(1);
        }
    }

    tracing::debug!("thread exiting");
}

async fn enforce_ratelimit(rc: &mut MultiplexedConnection, pds: &str) -> eyre::Result<()> {
    let score = rc.zscore(BF_REM_KEY, pds).await?;

    if let Some(rem) = score {
        if (rem as i32) < 100 {
            // if we've got None for some reason, just hope that the next req will contain the reset header.
            if let Some(at) = rc.zscore(BF_RESET_KEY, pds).await? {
                tracing::debug!("rate limit for {pds} resets at {at}");
                let time = chrono::DateTime::from_timestamp(at as i64, 0).unwrap();
                let delta = (time - Utc::now()).num_milliseconds().max(0);

                tokio::time::sleep(Duration::from_millis(delta as u64)).await;
            };
        }
    }

    Ok(())
}

// you wouldn't...
#[instrument(skip(http, tmp_dir, pds))]
async fn download_car(
    http: &HttpClient,
    tmp_dir: &Path,
    pds: &str,
    did: &str,
) -> eyre::Result<Option<(i32, i32)>> {
    let mut file = tokio::fs::File::create_new(tmp_dir.join(did)).await?;

    let res = http
        .get(format!("{pds}/xrpc/com.atproto.sync.getRepo?did={did}"))
        .send()
        .await?
        .error_for_status()?;

    let headers = res.headers();
    let ratelimit_rem = header_to_int(headers, "ratelimit-remaining");
    let ratelimit_reset = header_to_int(headers, "ratelimit-reset");

    let strm = res.bytes_stream().map_err(std::io::Error::other);
    let mut reader = StreamReader::new(strm);

    tokio::io::copy(&mut reader, &mut file).await?;

    Ok(ratelimit_rem.zip(ratelimit_reset))
}

// there's no ratelimit handling here because we pretty much always call download_car after.
#[instrument(skip(http, conn, pds))]
async fn check_and_update_repo_status(
    http: &HttpClient,
    conn: &mut PgClient,
    pds: &str,
    repo: &str,
) -> eyre::Result<bool> {
    match super::check_pds_repo_status(http, pds, repo).await? {
        Some(status) => {
            if !status.active {
                tracing::debug!("repo is inactive");

                let status = status
                    .status
                    .unwrap_or(crate::firehose::AtpAccountStatus::Deleted);
                conn.execute(
                    "UPDATE actors SET sync_state='dirty', status=$2 WHERE did=$1",
                    &[&repo, &ActorStatus::from(status)],
                )
                .await?;

                Ok(false)
            } else {
                Ok(true)
            }
        }
        None => {
            // this repo can't be found - set dirty and assume deleted.
            tracing::debug!("repo was deleted");
            conn.execute(
                "UPDATE actors SET sync_state='dirty', status='deleted' WHERE did=$1",
                &[&repo],
            )
            .await?;

            Ok(false)
        }
    }
}

async fn resolve_and_set_handle(
    conn: &PgClient,
    resolver: &Resolver,
    did: &str,
    handle: &str,
) -> eyre::Result<()> {
    if let Some(handle_did) = resolver.resolve_handle(handle).await? {
        if handle_did == did {
            conn.execute("UPDATE actors SET handle=$2 WHERE did=$1", &[&did, &handle])
                .await?;
        } else {
            tracing::warn!("requested DID ({did}) doesn't match handle");
        }
    }

    Ok(())
}

fn header_to_int(headers: &HeaderMap, name: &str) -> Option<i32> {
    headers
        .get(name)
        .and_then(|v| v.to_str().ok())
        .and_then(|v| v.parse().ok())
}