move remaining state to db, allow multiple build servers

Author: syphar

src/bin/cratesfyi.rs

@@ -6,7 +6,9 @@ use std::sync::Arc;
 use anyhow::{anyhow, Context as _, Error, Result};
 use docs_rs::db::{self, add_path_into_database, Pool, PoolClient};
 use docs_rs::repositories::RepositoryStatsUpdater;
-use docs_rs::utils::{remove_crate_priority, set_crate_priority};
+use docs_rs::utils::{
+    get_config, queue_builder, remove_crate_priority, set_crate_priority, ConfigName,
+};
 use docs_rs::{
     BuildQueue, Config, Context, Index, Metrics, PackageKind, RustwideBuilder, Server, Storage,
 };
@@ -93,6 +95,18 @@ enum CommandLine {
         socket_addr: String,
     },
 
+    StartRegistryWatcher {
+        /// Enable or disable the repository stats updater
+        #[structopt(
+            long = "repository-stats-updater",
+            default_value = "disabled",
+            possible_values(Toggle::VARIANTS)
+        )]
+        repository_stats_updater: Toggle,
+    },
+
+    StartBuildServer,
+
     /// Starts the daemon
     Daemon {
         /// Enable or disable the registry watcher to automatically enqueue newly published crates
@@ -123,6 +137,20 @@ impl CommandLine {
 
         match self {
             Self::Build(build) => build.handle_args(ctx)?,
+            Self::StartRegistryWatcher {
+                repository_stats_updater,
+            } => {
+                if repository_stats_updater == Toggle::Enabled {
+                    docs_rs::utils::daemon::start_background_repository_stats_updater(&ctx)?;
+                }
+
+                docs_rs::utils::watch_registry(ctx.build_queue()?, ctx.config()?, ctx.index()?)?;
+            }
+            Self::StartBuildServer => {
+                let build_queue = ctx.build_queue()?;
+                let rustwide_builder = RustwideBuilder::init(&ctx)?;
+                queue_builder(rustwide_builder, build_queue)?;
+            }
             Self::StartWebServer { socket_addr } => {
                 // Blocks indefinitely
                 let _ = Server::start(Some(&socket_addr), &ctx)?;
@@ -336,10 +364,8 @@ impl BuildSubcommand {
                         .pool()?
                         .get()
                         .context("failed to get a database connection")?;
-                    let res =
-                        conn.query("SELECT * FROM config WHERE name = 'rustc_version';", &[])?;
 
-                    if !res.is_empty() {
+                    if get_config(&mut conn, ConfigName::RustcVersion)?.is_string() {
                         println!("update-toolchain was already called in the past, exiting");
                         return Ok(());
                     }

src/build_queue.rs

@@ -2,15 +2,14 @@ use crate::db::{delete_crate, Pool};
 use crate::docbuilder::PackageKind;
 use crate::error::Result;
 use crate::storage::Storage;
-use crate::utils::{get_crate_priority, report_error};
+use crate::utils::{get_config, get_crate_priority, report_error, set_config, ConfigName};
 use crate::{Config, Index, Metrics, RustwideBuilder};
 use anyhow::Context;
 
 use crates_index_diff::Change;
 use log::{debug, info};
 
-use std::fs;
-use std::path::PathBuf;
+use git2::Oid;
 use std::sync::Arc;
 
 #[derive(Debug, Clone, Eq, PartialEq, serde::Serialize)]
@@ -48,6 +47,31 @@ impl BuildQueue {
         }
     }
 
+    pub fn last_seen_reference(&self) -> Result<Option<Oid>> {
+        let mut conn = self.db.get()?;
+        if let Some(value) = get_config(&mut conn, ConfigName::LastSeenIndexReference)?.as_str() {
+            match Oid::from_str(value) {
+                Ok(oid) => return Ok(Some(oid)),
+                Err(err) => {
+                    log::error!("queue locked because of invalid last_seen_index_reference \"{}\" in database: {}", value, err);
+                    self.lock()?;
+                    return Ok(None);
+                }
+            }
+        }
+        Ok(None)
+    }
+
+    fn set_last_seen_reference(&self, oid: Oid) -> Result<()> {
+        let mut conn = self.db.get()?;
+        set_config(
+            &mut conn,
+            ConfigName::LastSeenIndexReference,
+            oid.to_string(),
+        )?;
+        Ok(())
+    }
+
     pub fn add_crate(
         &self,
         name: &str,
@@ -118,14 +142,36 @@ impl BuildQueue {
         f: impl FnOnce(&QueuedCrate) -> Result<()>,
     ) -> Result<()> {
         let mut conn = self.db.get()?;
-
-        let queued = self.queued_crates()?;
-        let to_process = match queued.get(0) {
+        let mut transaction = conn.transaction()?;
+
+        // fetch the next available crate from the queue table.
+        // We are using `SELECT FOR UPDATE` inside a transaction so
+        // the QueuedCrate is locked until we are finished with it.
+        // `SKIP LOCKED` here will enable another build-server to just
+        // skip over taken (=locked) rows and start building the first
+        // available one.
+        let to_process = match transaction
+            .query_opt(
+                "SELECT id, name, version, priority, registry
+                 FROM queue
+                 WHERE attempt < $1
+                 ORDER BY priority ASC, attempt ASC, id ASC
+                 LIMIT 1 
+                 FOR UPDATE SKIP LOCKED",
+                &[&self.max_attempts],
+            )?
+            .map(|row| QueuedCrate {
+                id: row.get("id"),
+                name: row.get("name"),
+                version: row.get("version"),
+                priority: row.get("priority"),
+                registry: row.get("registry"),
+            }) {
             Some(krate) => krate,
             None => return Ok(()),
         };
 
-        let res = f(to_process).with_context(|| {
+        let res = f(&to_process).with_context(|| {
             format!(
                 "Failed to build package {}-{} from queue",
                 to_process.name, to_process.version
@@ -134,15 +180,16 @@ impl BuildQueue {
         self.metrics.total_builds.inc();
         match res {
             Ok(()) => {
-                conn.execute("DELETE FROM queue WHERE id = $1;", &[&to_process.id])?;
+                transaction.execute("DELETE FROM queue WHERE id = $1;", &[&to_process.id])?;
             }
             Err(e) => {
                 // Increase attempt count
-                let rows = conn.query(
-                    "UPDATE queue SET attempt = attempt + 1 WHERE id = $1 RETURNING attempt;",
-                    &[&to_process.id],
-                )?;
-                let attempt: i32 = rows[0].get(0);
+                let attempt: i32 = transaction
+                    .query_one(
+                        "UPDATE queue SET attempt = attempt + 1 WHERE id = $1 RETURNING attempt;",
+                        &[&to_process.id],
+                    )?
+                    .get(0);
 
                 if attempt >= self.max_attempts {
                     self.metrics.failed_builds.inc();
@@ -152,39 +199,33 @@ impl BuildQueue {
             }
         }
 
+        transaction.commit()?;
+
         Ok(())
     }
 }
 
 /// Locking functions.
 impl BuildQueue {
-    pub(crate) fn lock_path(&self) -> PathBuf {
-        self.config.prefix.join("docsrs.lock")
-    }
+    /// Checks for the lock and returns whether it currently exists.
+    pub fn is_locked(&self) -> Result<bool> {
+        let mut conn = self.db.get()?;
 
-    /// Checks for the lock file and returns whether it currently exists.
-    pub fn is_locked(&self) -> bool {
-        self.lock_path().exists()
+        Ok(get_config(&mut conn, ConfigName::QueueLocked)?
+            .as_bool()
+            .unwrap_or(false))
     }
 
-    /// Creates a lock file. Daemon will check this lock file and stop operating if it exists.
+    /// lock the queue. Daemon will check this lock and stop operating if it exists.
     pub fn lock(&self) -> Result<()> {
-        let path = self.lock_path();
-        if !path.exists() {
-            fs::OpenOptions::new().write(true).create(true).open(path)?;
-        }
-
-        Ok(())
+        let mut conn = self.db.get()?;
+        set_config(&mut conn, ConfigName::QueueLocked, true)
     }
 
-    /// Removes lock file.
+    /// unlock the queue.
     pub fn unlock(&self) -> Result<()> {
-        let path = self.lock_path();
-        if path.exists() {
-            fs::remove_file(path)?;
-        }
-
-        Ok(())
+        let mut conn = self.db.get()?;
+        set_config(&mut conn, ConfigName::QueueLocked, false)
     }
 }
 
@@ -266,8 +307,15 @@ impl BuildQueue {
             }
         }
 
+        // store the last seen reference as git reference in
+        // the local crates.io index repo.
         diff.set_last_seen_reference(oid)?;
 
+        // additionally set the reference in the database
+        // so this survives recreating the registry watcher
+        // server.
+        self.set_last_seen_reference(oid)?;
+
         Ok(crates_added)
     }
 
@@ -559,4 +607,57 @@ mod tests {
             Ok(())
         });
     }
+
+    #[test]
+    fn test_last_seen_reference_in_db() {
+        crate::test::wrapper(|env| {
+            let queue = env.build_queue();
+            queue.unlock()?;
+            assert!(!queue.is_locked()?);
+            // initial db ref is empty
+            assert_eq!(queue.last_seen_reference()?, None);
+            assert!(!queue.is_locked()?);
+
+            let oid = git2::Oid::from_str("ffffffff")?;
+            queue.set_last_seen_reference(oid)?;
+
+            assert_eq!(queue.last_seen_reference()?, Some(oid));
+            assert!(!queue.is_locked()?);
+
+            Ok(())
+        });
+    }
+
+    #[test]
+    fn test_broken_db_reference_locks_queue() {
+        crate::test::wrapper(|env| {
+            let mut conn = env.db().conn();
+            set_config(&mut conn, ConfigName::LastSeenIndexReference, "invalid")?;
+
+            let queue = env.build_queue();
+            queue.unlock()?;
+            assert!(!queue.is_locked()?);
+            assert_eq!(queue.last_seen_reference()?, None);
+            assert!(queue.is_locked()?);
+
+            Ok(())
+        });
+    }
+
+    #[test]
+    fn test_queue_lock() {
+        crate::test::wrapper(|env| {
+            let queue = env.build_queue();
+            // unlocked without config
+            assert!(!queue.is_locked()?);
+
+            queue.lock()?;
+            assert!(queue.is_locked()?);
+
+            queue.unlock()?;
+            assert!(!queue.is_locked()?);
+
+            Ok(())
+        });
+    }
 }

src/docbuilder/rustwide_builder.rs

@@ -8,7 +8,9 @@ use crate::error::Result;
 use crate::index::api::ReleaseData;
 use crate::repositories::RepositoryStatsUpdater;
 use crate::storage::{rustdoc_archive_path, source_archive_path};
-use crate::utils::{copy_dir_all, parse_rustc_version, queue_builder, CargoMetadata};
+use crate::utils::{
+    copy_dir_all, parse_rustc_version, queue_builder, set_config, CargoMetadata, ConfigName,
+};
 use crate::{db::blacklist::is_blacklisted, utils::MetadataPackage};
 use crate::{Config, Context, Index, Metrics, Storage};
 use anyhow::{anyhow, bail, Error};
@@ -20,7 +22,6 @@ use rustwide::cmd::{Command, CommandError, SandboxBuilder, SandboxImage};
 use rustwide::logging::{self, LogStorage};
 use rustwide::toolchain::ToolchainError;
 use rustwide::{AlternativeRegistry, Build, Crate, Toolchain, Workspace, WorkspaceBuilder};
-use serde_json::Value;
 use std::collections::{HashMap, HashSet};
 use std::path::Path;
 use std::sync::Arc;
@@ -225,12 +226,12 @@ impl RustwideBuilder {
                         .tempdir()?;
                     copy_dir_all(source, &dest)?;
                     add_path_into_database(&self.storage, "", &dest)?;
-                    conn.query(
-                        "INSERT INTO config (name, value) VALUES ('rustc_version', $1) \
-                     ON CONFLICT (name) DO UPDATE SET value = $1;",
-                        &[&Value::String(self.rustc_version.clone())],
-                    )?;
 
+                    set_config(
+                        &mut conn,
+                        ConfigName::RustcVersion,
+                        self.rustc_version.clone(),
+                    )?;
                     Ok(())
                 })()
                 .map_err(|e| failure::Error::from_boxed_compat(e.into()))
@@ -806,6 +807,7 @@ pub(crate) struct BuildResult {
 mod tests {
     use super::*;
     use crate::test::{assert_redirect, assert_success, wrapper};
+    use serde_json::Value;
 
     #[test]
     #[ignore]

src/metrics/mod.rs

@@ -158,7 +158,7 @@ impl Metrics {
         self.idle_db_connections.set(pool.idle_connections() as i64);
         self.used_db_connections.set(pool.used_connections() as i64);
         self.max_db_connections.set(pool.max_size() as i64);
-        self.queue_is_locked.set(queue.is_locked() as i64);
+        self.queue_is_locked.set(queue.is_locked()? as i64);
 
         self.queued_crates_count.set(queue.pending_count()? as i64);
         self.prioritized_crates_count