storage: use the shared instance for uploading

Author: pietroalbini

src/bin/cratesfyi.rs

@@ -119,14 +119,13 @@ impl CommandLine {
                 socket_addr,
                 reload_templates,
             } => {
-                let storage = Storage::new(ctx.pool()?);
                 Server::start(
                     Some(&socket_addr),
                     reload_templates,
                     ctx.pool()?,
                     ctx.config()?,
                     ctx.build_queue()?,
-                    Arc::new(storage),
+                    ctx.storage()?,
                 )?;
             }
             Self::Daemon {
@@ -141,6 +140,7 @@ impl CommandLine {
                     ctx.config()?,
                     ctx.pool()?,
                     ctx.build_queue()?,
+                    ctx.storage()?,
                     registry_watcher == Toggle::Enabled,
                 )?;
             }
@@ -350,7 +350,7 @@ impl BuildSubcommand {
             Self::World => {
                 docbuilder.load_cache().context("Failed to load cache")?;
 
-                let mut builder = RustwideBuilder::init(ctx.pool()?)?;
+                let mut builder = RustwideBuilder::init(ctx.pool()?, ctx.storage()?)?;
                 builder
                     .build_world(&mut docbuilder)
                     .context("Failed to build world")?;
@@ -364,8 +364,8 @@ impl BuildSubcommand {
                 local,
             } => {
                 docbuilder.load_cache().context("Failed to load cache")?;
-                let mut builder =
-                    RustwideBuilder::init(ctx.pool()?).context("failed to initialize rustwide")?;
+                let mut builder = RustwideBuilder::init(ctx.pool()?, ctx.storage()?)
+                    .context("failed to initialize rustwide")?;
 
                 if let Some(path) = local {
                     builder
@@ -401,14 +401,14 @@ impl BuildSubcommand {
                     }
                 }
 
-                let mut builder = RustwideBuilder::init(ctx.pool()?)?;
+                let mut builder = RustwideBuilder::init(ctx.pool()?, ctx.storage()?)?;
                 builder
                     .update_toolchain()
                     .context("failed to update toolchain")?;
             }
 
             Self::AddEssentialFiles => {
-                let mut builder = RustwideBuilder::init(ctx.pool()?)?;
+                let mut builder = RustwideBuilder::init(ctx.pool()?, ctx.storage()?)?;
                 builder
                     .add_essential_files()
                     .context("failed to add essential files")?;
@@ -473,8 +473,8 @@ impl DatabaseSubcommand {
             }
 
             Self::AddDirectory { directory, prefix } => {
-                add_path_into_database(ctx.pool()?, &prefix, directory)
-                    .context("Failed to add directory into database")?;
+                add_path_into_database(&*ctx.storage()?, &prefix, directory)
+                    .expect("Failed to add directory into database");
             }
 
             // FIXME: This is actually util command not database
@@ -557,6 +557,7 @@ enum DeleteSubcommand {
 
 struct Context {
     build_queue: OnceCell<Arc<BuildQueue>>,
+    storage: OnceCell<Arc<Storage>>,
     config: OnceCell<Arc<Config>>,
     pool: OnceCell<Pool>,
 }
@@ -565,6 +566,7 @@ impl Context {
     fn new() -> Self {
         Self {
             build_queue: OnceCell::new(),
+            storage: OnceCell::new(),
             config: OnceCell::new(),
             pool: OnceCell::new(),
         }
@@ -579,6 +581,13 @@ impl Context {
             .clone())
     }
 
+    fn storage(&self) -> Result<Arc<Storage>, Error> {
+        Ok(self
+            .storage
+            .get_or_try_init::<_, Error>(|| Ok(Arc::new(Storage::new(self.pool()?))))?
+            .clone())
+    }
+
     fn config(&self) -> Result<Arc<Config>, Error> {
         Ok(self
             .config

src/db/file.rs

@@ -4,7 +4,6 @@
 //! They are using so many inodes and it is better to store them in database instead of
 //! filesystem. This module is adding files into database and retrieving them.
 
-use crate::db::Pool;
 use crate::error::Result;
 use crate::storage::{CompressionAlgorithms, Storage};
 
@@ -13,20 +12,16 @@ use std::path::{Path, PathBuf};
 
 /// Store all files in a directory and return [[mimetype, filename]] as Json
 ///
-/// If there is an S3 Client configured, store files into an S3 bucket;
-/// otherwise, stores files into the 'files' table of the local database.
-///
 /// The mimetype is detected using `magic`.
 ///
 /// Note that this function is used for uploading both sources
 /// and files generated by rustdoc.
 pub fn add_path_into_database<P: AsRef<Path>>(
-    pool: Pool,
+    storage: &Storage,
     prefix: &str,
     path: P,
 ) -> Result<(Value, CompressionAlgorithms)> {
-    let mut backend = Storage::new(pool);
-    let (file_list, algorithms) = backend.store_all(prefix, path.as_ref())?;
+    let (file_list, algorithms) = storage.store_all(prefix, path.as_ref())?;
     Ok((
         file_list_to_json(file_list.into_iter().collect())?,
         algorithms,

src/docbuilder/rustwide_builder.rs

@@ -6,6 +6,7 @@ use crate::db::{add_build_into_database, add_package_into_database, Pool};
 use crate::docbuilder::{crates::crates_from_path, Limits};
 use crate::error::Result;
 use crate::storage::CompressionAlgorithms;
+use crate::storage::Storage;
 use crate::utils::{copy_doc_dir, parse_rustc_version, CargoMetadata};
 use failure::ResultExt;
 use log::{debug, info, warn, LevelFilter};
@@ -17,6 +18,7 @@ use serde_json::Value;
 use std::borrow::Cow;
 use std::collections::HashSet;
 use std::path::Path;
+use std::sync::Arc;
 
 const USER_AGENT: &str = "docs.rs builder (https://github.com/rust-lang/docs.rs)";
 const DEFAULT_RUSTWIDE_WORKSPACE: &str = ".rustwide";
@@ -67,12 +69,13 @@ pub struct RustwideBuilder {
     workspace: Workspace,
     toolchain: Toolchain,
     db: Pool,
+    storage: Arc<Storage>,
     rustc_version: String,
     cpu_limit: Option<u32>,
 }
 
 impl RustwideBuilder {
-    pub fn init(db: Pool) -> Result<Self> {
+    pub fn init(db: Pool, storage: Arc<Storage>) -> Result<Self> {
         use rustwide::cmd::SandboxImage;
         let env_workspace_path = ::std::env::var("CRATESFYI_RUSTWIDE_WORKSPACE");
         let workspace_path = env_workspace_path
@@ -107,6 +110,7 @@ impl RustwideBuilder {
             workspace,
             toolchain,
             db,
+            storage,
             rustc_version: String::new(),
             cpu_limit,
         })
@@ -244,7 +248,7 @@ impl RustwideBuilder {
                     })?;
                 }
 
-                add_path_into_database(self.db.clone(), "", &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;",
@@ -349,7 +353,7 @@ impl RustwideBuilder {
                     debug!("adding sources into database");
                     let prefix = format!("sources/{}/{}", name, version);
                     let (files, new_algs) =
-                        add_path_into_database(self.db.clone(), &prefix, build.host_source_dir())?;
+                        add_path_into_database(&self.storage, &prefix, build.host_source_dir())?;
                     files_list = Some(files);
                     algs.extend(new_algs);
 
@@ -581,7 +585,7 @@ impl RustwideBuilder {
     ) -> Result<CompressionAlgorithms> {
         debug!("Adding documentation into database");
         add_path_into_database(
-            self.db.clone(),
+            &self.storage,
             &format!("rustdoc/{}/{}", name, version),
             local_storage,
         )

src/storage/mod.rs

@@ -155,7 +155,7 @@ impl Storage {
     //
     // This returns (map<filename, mime type>, set<compression algorithms>).
     pub(crate) fn store_all(
-        &mut self,
+        &self,
         prefix: &str,
         root_dir: &Path,
     ) -> Result<(HashMap<PathBuf, String>, HashSet<CompressionAlgorithm>), Error> {
@@ -301,7 +301,7 @@ mod test {
         }
         wrapper(|env| {
             let db = env.db();
-            let mut backend = Storage {
+            let backend = Storage {
                 backend: StorageBackend::Database(DatabaseBackend::new(db.pool())),
             };
             let (stored_files, _algs) = backend.store_all("", dir.path()).unwrap();
@@ -335,7 +335,7 @@ mod test {
         }
         wrapper(|env| {
             let db = env.db();
-            let mut backend = Storage {
+            let backend = Storage {
                 backend: StorageBackend::Database(DatabaseBackend::new(db.pool())),
             };
             let (stored_files, _algs) = backend.store_all("rustdoc", dir.path()).unwrap();

src/test/fakes.rs

@@ -1,13 +1,16 @@
 use super::TestDatabase;
 use crate::docbuilder::BuildResult;
 use crate::index::api::RegistryCrateData;
+use crate::storage::Storage;
 use crate::utils::{Dependency, MetadataPackage, Target};
 use chrono::{DateTime, Utc};
 use failure::Error;
+use std::sync::Arc;
 
 #[must_use = "FakeRelease does nothing until you call .create()"]
 pub(crate) struct FakeRelease<'a> {
     db: &'a TestDatabase,
+    storage: Arc<Storage>,
     package: MetadataPackage,
     build_result: BuildResult,
     /// name, content
@@ -24,9 +27,10 @@ pub(crate) struct FakeRelease<'a> {
 }
 
 impl<'a> FakeRelease<'a> {
-    pub(super) fn new(db: &'a TestDatabase) -> Self {
+    pub(super) fn new(db: &'a TestDatabase, storage: Arc<Storage>) -> Self {
         FakeRelease {
             db,
+            storage,
             package: MetadataPackage {
                 id: "fake-package-id".into(),
                 name: "fake-package".into(),
@@ -175,6 +179,7 @@ impl<'a> FakeRelease<'a> {
         let mut source_meta = None;
         let mut algs = HashSet::new();
         if self.build_result.successful {
+            let storage = self.storage.clone();
             let upload_files = |prefix: &str, files: &[(&str, &[u8])], target: Option<&str>| {
                 let mut path_prefix = tempdir.path().join(prefix);
                 if let Some(target) = target {
@@ -200,7 +205,7 @@ impl<'a> FakeRelease<'a> {
                     target.unwrap_or("")
                 );
                 log::debug!("adding directory {} from {}", prefix, path_prefix.display());
-                crate::db::add_path_into_database(db.pool(), &prefix, path_prefix)
+                crate::db::add_path_into_database(&storage, &prefix, path_prefix)
             };
 
             let index = [&package.name, "index.html"].join("/");

src/test/mod.rs

@@ -179,7 +179,7 @@ impl TestEnvironment {
     }
 
     pub(crate) fn fake_release(&self) -> fakes::FakeRelease {
-        fakes::FakeRelease::new(self.db())
+        fakes::FakeRelease::new(self.db(), self.storage())
     }
 }
 

src/utils/daemon.rs

@@ -47,6 +47,7 @@ pub fn start_daemon(
     config: Arc<Config>,
     db: Pool,
     build_queue: Arc<BuildQueue>,
+    storage: Arc<Storage>,
     enable_registry_watcher: bool,
 ) -> Result<(), Error> {
     const CRATE_VARIABLES: &[&str] = &["CRATESFYI_PREFIX"];
@@ -71,12 +72,13 @@ pub fn start_daemon(
     // build new crates every minute
     let cloned_db = db.clone();
     let cloned_build_queue = build_queue.clone();
+    let cloned_storage = storage.clone();
     thread::Builder::new()
         .name("build queue reader".to_string())
         .spawn(move || {
             let doc_builder =
                 DocBuilder::new(opts(), cloned_db.clone(), cloned_build_queue.clone());
-            queue_builder(doc_builder, cloned_db, cloned_build_queue).unwrap();
+            queue_builder(doc_builder, cloned_db, cloned_build_queue, cloned_storage).unwrap();
         })
         .unwrap();
 
@@ -111,8 +113,6 @@ pub fn start_daemon(
     // at least start web server
     info!("Starting web server");
 
-    let storage = Arc::new(Storage::new(db.clone()));
-
     crate::Server::start(None, false, db, config, build_queue, storage)?;
     Ok(())
 }

src/utils/queue_builder.rs

@@ -1,4 +1,6 @@
-use crate::{db::Pool, docbuilder::RustwideBuilder, utils::pubsubhubbub, BuildQueue, DocBuilder};
+use crate::{
+    db::Pool, docbuilder::RustwideBuilder, utils::pubsubhubbub, BuildQueue, DocBuilder, Storage,
+};
 use failure::Error;
 use log::{debug, error, info, warn};
 use std::panic::{catch_unwind, AssertUnwindSafe};
@@ -11,6 +13,7 @@ pub fn queue_builder(
     mut doc_builder: DocBuilder,
     db: Pool,
     build_queue: Arc<BuildQueue>,
+    storage: Arc<Storage>,
 ) -> Result<(), Error> {
     /// Represents the current state of the builder thread.
     enum BuilderState {
@@ -25,7 +28,7 @@ pub fn queue_builder(
         QueueInProgress(usize),
     }
 
-    let mut builder = RustwideBuilder::init(db)?;
+    let mut builder = RustwideBuilder::init(db, storage)?;
 
     let mut status = BuilderState::Fresh;