dogfood
Jump to…
snowcorepublic2 bookmarks
snowinitial commitqoxwzsukwmkx1mo
mainroot / crates / df-worker / src / main.rsHide BlameRaw
Matt W1//! `dogfood-worker` — the job runner (spec §5).
Matt W2//!
Matt W3//! Jobs are claimed with `SELECT … FOR UPDATE SKIP LOCKED`. A `NOTIFY` on
Matt W4//! insert wakes the worker so it does not poll on the hot path; a 5-second poll
Matt W5//! remains as a fallback for missed notifications.
Matt W6
Matt W7use std::sync::Arc;
Matt W8use std::time::Duration;
Matt W9
Matt W10use anyhow::{Context, Result};
Matt W11use df_store::RepoStore;
Matt W12use sqlx::postgres::PgListener;
Matt W13use sqlx::PgPool;
Matt W14use uuid::Uuid;
Matt W15
Matt W16mod anchors;
Matt W17mod index_push;
Matt W18mod patch_id;
Matt W19
Matt W20/// Identifies this worker in `jobs.locked_by`, so a stuck job can be traced
Matt W21/// back to the process that took it.
Matt W22fn worker_id() -> String {
Matt W23 format!(
Matt W24 "{}-{}",
Matt W25 std::env::var("HOSTNAME").unwrap_or_else(|_| "worker".into()),
Matt W26 std::process::id()
Matt W27 )
Matt W28}
Matt W29
Matt W30#[tokio::main]
Matt W31async fn main() -> Result<()> {
Matt W32 let _ = dotenvy::dotenv();
Matt W33 init_tracing();
Matt W34
Matt W35 let database_url = std::env::var("DATABASE_URL").context("DATABASE_URL must be set")?;
Matt W36 let repo_root = std::env::var("REPO_ROOT").unwrap_or_else(|_| "/srv/repos".into());
Matt W37 let max_conn = std::env::var("DATABASE_MAX_CONNECTIONS")
Matt W38 .ok()
Matt W39 .and_then(|v| v.parse().ok())
Matt W40 .unwrap_or(5);
Matt W41
Matt W42 let db = df_db::connect(&database_url, max_conn)
Matt W43 .await
Matt W44 .context("connecting to the database")?;
Matt W45
Matt W46 // The worker does NOT run migrations — `web` owns that (spec §10), and two
Matt W47 // processes racing to migrate is exactly what the advisory lock exists to
Matt W48 // prevent. If the schema is behind, the queries below fail loudly.
Matt W49 let store: Arc<dyn RepoStore> = Arc::new(df_store::GitStore::new(&repo_root));
Matt W50
Matt W51 let me = worker_id();
Matt W52 tracing::info!(worker = %me, root = %repo_root, "dogfood-worker started");
Matt W53
Matt W54 // LISTEN for wake-ups. If the listener cannot be established the worker
Matt W55 // still functions on the poll fallback, so this is a warning, not fatal.
Matt W56 let mut listener = match PgListener::connect(&database_url).await {
Matt W57 Ok(mut l) => match l.listen("dogfood_jobs").await {
Matt W58 Ok(()) => Some(l),
Matt W59 Err(e) => {
Matt W60 tracing::warn!("LISTEN failed, falling back to polling: {e}");
Matt W61 None
Matt W62 }
Matt W63 },
Matt W64 Err(e) => {
Matt W65 tracing::warn!("could not open a listener, falling back to polling: {e}");
Matt W66 None
Matt W67 }
Matt W68 };
Matt W69
Matt W70 let mut shutdown = std::pin::pin!(shutdown_signal());
Matt W71
Matt W72 loop {
Matt W73 // Drain everything currently runnable before sleeping.
Matt W74 loop {
Matt W75 match claim_and_run(&db, store.as_ref(), &me).await {
Matt W76 Ok(true) => continue,
Matt W77 Ok(false) => break,
Matt W78 Err(e) => {
Matt W79 tracing::error!("job loop error: {e:#}");
Matt W80 break;
Matt W81 }
Matt W82 }
Matt W83 }
Matt W84
Matt W85 // Wait for a notification, the poll interval, or shutdown.
Matt W86 let wake = async {
Matt W87 match listener.as_mut() {
Matt W88 Some(l) => {
Matt W89 let _ = l.recv().await;
Matt W90 }
Matt W91 None => std::future::pending::<()>().await,
Matt W92 }
Matt W93 };
Matt W94
Matt W95 tokio::select! {
Matt W96 _ = wake => {}
Matt W97 _ = tokio::time::sleep(Duration::from_secs(5)) => {}
Matt W98 _ = &mut shutdown => {
Matt W99 tracing::info!("shutting down");
Matt W100 return Ok(());
Matt W101 }
Matt W102 }
Matt W103 }
Matt W104}
Matt W105
Matt W106/// Claim one job and run it. Returns whether a job was found.
Matt W107async fn claim_and_run(db: &PgPool, store: &dyn RepoStore, me: &str) -> Result<bool> {
Matt W108 let mut tx = db.begin().await?;
Matt W109
Matt W110 // SKIP LOCKED lets several workers share the queue without blocking.
Matt W111 let job: Option<(Uuid, String, serde_json::Value, i32, i32)> = sqlx::query_as(
Matt W112 "SELECT id, kind, payload, attempts, max_attempts
Matt W113 FROM jobs
Matt W114 WHERE locked_at IS NULL AND run_at <= now()
Matt W115 ORDER BY run_at
Matt W116 FOR UPDATE SKIP LOCKED
Matt W117 LIMIT 1",
Matt W118 )
Matt W119 .fetch_optional(&mut *tx)
Matt W120 .await?;
Matt W121
Matt W122 let Some((id, kind, payload, attempts, max_attempts)) = job else {
Matt W123 tx.rollback().await?;
Matt W124 return Ok(false);
Matt W125 };
Matt W126
Matt W127 sqlx::query("UPDATE jobs SET locked_at = now(), locked_by = $2, attempts = attempts + 1 WHERE id = $1")
Matt W128 .bind(id)
Matt W129 .bind(me)
Matt W130 .execute(&mut *tx)
Matt W131 .await?;
Matt W132 tx.commit().await?;
Matt W133
Matt W134 tracing::info!(job = %id, %kind, attempt = attempts + 1, "running job");
Matt W135 let started = std::time::Instant::now();
Matt W136
Matt W137 let result = match kind.as_str() {
Matt W138 "index_push" => run_index_push(db, store, &payload).await,
Matt W139 other => Err(anyhow::anyhow!("unknown job kind: {other}")),
Matt W140 };
Matt W141
Matt W142 match result {
Matt W143 Ok(()) => {
Matt W144 sqlx::query("DELETE FROM jobs WHERE id = $1")
Matt W145 .bind(id)
Matt W146 .execute(db)
Matt W147 .await?;
Matt W148 tracing::info!(job = %id, ms = started.elapsed().as_millis(), "job complete");
Matt W149 }
Matt W150 Err(e) => {
Matt W151 let attempts_now = attempts + 1;
Matt W152 let give_up = attempts_now >= max_attempts;
Matt W153 tracing::error!(job = %id, attempt = attempts_now, "job failed: {e:#}");
Matt W154
Matt W155 if give_up {
Matt W156 // Leave the row, unlocked and exhausted, so it is visible for
Matt W157 // inspection rather than silently vanishing.
Matt W158 sqlx::query(
Matt W159 "UPDATE jobs SET locked_at = NULL, locked_by = NULL, last_error = $2
Matt W160 WHERE id = $1",
Matt W161 )
Matt W162 .bind(id)
Matt W163 .bind(format!("{e:#}"))
Matt W164 .execute(db)
Matt W165 .await?;
Matt W166 tracing::error!(job = %id, "job exhausted its attempts");
Matt W167 } else {
Matt W168 // Exponential backoff, capped.
Matt W169 let delay = Duration::from_secs(2u64.saturating_pow(attempts_now as u32).min(300));
Matt W170 sqlx::query(
Matt W171 "UPDATE jobs SET locked_at = NULL, locked_by = NULL, last_error = $2,
Matt W172 run_at = now() + $3
Matt W173 WHERE id = $1",
Matt W174 )
Matt W175 .bind(id)
Matt W176 .bind(format!("{e:#}"))
Matt W177 .bind(delay)
Matt W178 .execute(db)
Matt W179 .await?;
Matt W180 }
Matt W181 }
Matt W182 }
Matt W183
Matt W184 Ok(true)
Matt W185}
Matt W186
Matt W187async fn run_index_push(
Matt W188 db: &PgPool,
Matt W189 store: &dyn RepoStore,
Matt W190 payload: &serde_json::Value,
Matt W191) -> Result<()> {
Matt W192 let repo_id: Uuid = payload
Matt W193 .get("repo_id")
Matt W194 .and_then(|v| v.as_str())
Matt W195 .and_then(|s| s.parse().ok())
Matt W196 .context("index_push payload is missing a valid repo_id")?;
Matt W197
Matt W198 let pushed_by: Option<Uuid> = payload
Matt W199 .get("pushed_by")
Matt W200 .and_then(|v| v.as_str())
Matt W201 .and_then(|s| s.parse().ok());
Matt W202
Matt W203 let outcome = index_push::run(db, store, repo_id, pushed_by).await?;
Matt W204 tracing::info!(
Matt W205 repo = %repo_id,
Matt W206 changes = outcome.changes_seen,
Matt W207 revisions = outcome.revisions_added,
Matt W208 "indexed push"
Matt W209 );
Matt W210 Ok(())
Matt W211}
Matt W212
Matt W213fn init_tracing() {
Matt W214 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter};
Matt W215
Matt W216 let filter = EnvFilter::try_from_default_env()
Matt W217 .unwrap_or_else(|_| EnvFilter::new("info,df_worker=debug,df_index=debug"));
Matt W218 let json = !std::io::IsTerminal::is_terminal(&std::io::stdout());
Matt W219
Matt W220 let registry = tracing_subscriber::registry().with(filter);
Matt W221 if json {
Matt W222 registry.with(tracing_subscriber::fmt::layer().json()).init();
Matt W223 } else {
Matt W224 registry.with(tracing_subscriber::fmt::layer()).init();
Matt W225 }
Matt W226}
Matt W227
Matt W228async fn shutdown_signal() {
Matt W229 let ctrl_c = async {
Matt W230 tokio::signal::ctrl_c().await.expect("Ctrl+C handler");
Matt W231 };
Matt W232 #[cfg(unix)]
Matt W233 let terminate = async {
Matt W234 tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
Matt W235 .expect("SIGTERM handler")
Matt W236 .recv()
Matt W237 .await;
Matt W238 };
Matt W239 #[cfg(not(unix))]
Matt W240 let terminate = std::future::pending::<()>();
Matt W241
Matt W242 tokio::select! {
Matt W243 _ = ctrl_c => {}
Matt W244 _ = terminate => {}
Matt W245 }
Matt W246}

246 lines · Rust