make Context.async_storage an async function

rust-lang · Nov 19, 2023 · 81a8398 · 81a8398
1 parent 62c620e
commit 81a8398
Show file tree

Hide file tree

Showing 6 changed files with 41 additions and 29 deletions.
diff --git a/src/bin/cratesfyi.rs b/src/bin/cratesfyi.rs
@@ -6,6 +6,7 @@ use std::str::FromStr;
 use std::sync::Arc;
 
 use anyhow::{anyhow, Context as _, Error, Result};
+use axum::async_trait;
 use clap::{Parser, Subcommand, ValueEnum};
 use docs_rs::cdn::CdnBackend;
 use docs_rs::db::{self, add_path_into_database, Overrides, Pool, PoolClient};
@@ -767,6 +768,7 @@ macro_rules! lazy {
     }
 }
 
+#[async_trait]
 impl Context for BinContext {
     lazy! {
         fn build_queue(self) -> BuildQueue = BuildQueue::new(
@@ -775,10 +777,13 @@ impl Context for BinContext {
             self.config()?,
             self.storage()?,
         );
-        fn storage(self) -> Storage = Storage::new(
-            self.async_storage()?,
-            self.runtime()?
-       );
+        fn storage(self) -> Storage = {
+            let runtime = self.runtime()?;
+            Storage::new(
+                runtime.block_on(self.async_storage())?,
+                runtime
+           )
+        };
         fn cdn(self) -> CdnBackend = CdnBackend::new(
             &self.config()?,
             &self.runtime()?,
@@ -824,11 +829,9 @@ impl Context for BinContext {
             .clone())
     }
 
-    fn async_storage(&self) -> Result<Arc<AsyncStorage>> {
-        Ok(Arc::new(self.runtime()?.block_on(AsyncStorage::new(
-            self.pool()?,
-            self.instance_metrics()?,
-            self.config()?,
-        ))?))
+    async fn async_storage(&self) -> Result<Arc<AsyncStorage>> {
+        Ok(Arc::new(
+            AsyncStorage::new(self.pool()?, self.instance_metrics()?, self.config()?).await?,
+        ))
     }
 }
diff --git a/src/context.rs b/src/context.rs
@@ -5,14 +5,16 @@ use crate::repositories::RepositoryStatsUpdater;
 use crate::{
     AsyncStorage, BuildQueue, Config, Index, InstanceMetrics, RegistryApi, ServiceMetrics, Storage,
 };
+use axum::async_trait;
 use std::sync::Arc;
 use tokio::runtime::Runtime;
 
+#[async_trait]
 pub trait Context {
     fn config(&self) -> Result<Arc<Config>>;
     fn build_queue(&self) -> Result<Arc<BuildQueue>>;
     fn storage(&self) -> Result<Arc<Storage>>;
-    fn async_storage(&self) -> Result<Arc<AsyncStorage>>;
+    async fn async_storage(&self) -> Result<Arc<AsyncStorage>>;
     fn cdn(&self) -> Result<Arc<CdnBackend>>;
     fn pool(&self) -> Result<Pool>;
     fn service_metrics(&self) -> Result<Arc<ServiceMetrics>>;

diff --git a/src/test/mod.rs b/src/test/mod.rs
@@ -9,9 +9,10 @@ use crate::storage::{AsyncStorage, Storage, StorageKind};
 use crate::web::{build_axum_app, cache, page::TemplateData};
 use crate::{BuildQueue, Config, Context, Index, InstanceMetrics, RegistryApi, ServiceMetrics};
 use anyhow::Context as _;
+use axum::async_trait;
 use fn_error_context::context;
 use futures_util::FutureExt;
-use once_cell::unsync::OnceCell;
+use once_cell::sync::OnceCell;
 use postgres::Client as Connection;
 use reqwest::{
     blocking::{Client, ClientBuilder, RequestBuilder, Response},
@@ -263,7 +264,7 @@ pub(crate) struct TestEnvironment {
     config: OnceCell<Arc<Config>>,
     db: tokio::sync::OnceCell<TestDatabase>,
     storage: OnceCell<Arc<Storage>>,
-    async_storage: OnceCell<Arc<AsyncStorage>>,
+    async_storage: tokio::sync::OnceCell<Arc<AsyncStorage>>,
     cdn: OnceCell<Arc<CdnBackend>>,
     index: OnceCell<Arc<Index>>,
     registry_api: OnceCell<Arc<RegistryApi>>,
@@ -298,7 +299,7 @@ impl TestEnvironment {
             config: OnceCell::new(),
             db: tokio::sync::OnceCell::new(),
             storage: OnceCell::new(),
-            async_storage: OnceCell::new(),
+            async_storage: tokio::sync::OnceCell::new(),
             cdn: OnceCell::new(),
             index: OnceCell::new(),
             registry_api: OnceCell::new(),
@@ -391,25 +392,29 @@ impl TestEnvironment {
             .clone()
     }
 
-    pub(crate) fn async_storage(&self) -> Arc<AsyncStorage> {
+    pub(crate) async fn async_storage(&self) -> Arc<AsyncStorage> {
         self.async_storage
-            .get_or_init(|| {
+            .get_or_init(|| async {
+                let db = self.async_db().await;
                 Arc::new(
-                    self.runtime()
-                        .block_on(AsyncStorage::new(
-                            self.db().pool(),
-                            self.instance_metrics(),
-                            self.config(),
-                        ))
+                    AsyncStorage::new(db.pool(), self.instance_metrics(), self.config())
+                        .await
                         .expect("failed to initialize the async storage"),
                 )
             })
+            .await
             .clone()
     }
 
     pub(crate) fn storage(&self) -> Arc<Storage> {
+        let runtime = self.runtime();
         self.storage
-            .get_or_init(|| Arc::new(Storage::new(self.async_storage(), self.runtime())))
+            .get_or_init(|| {
+                Arc::new(Storage::new(
+                    runtime.block_on(self.async_storage()),
+                    runtime,
+                ))
+            })
             .clone()
     }
 
@@ -515,6 +520,7 @@ impl TestEnvironment {
     }
 }
 
+#[async_trait]
 impl Context for TestEnvironment {
     fn config(&self) -> Result<Arc<Config>> {
         Ok(TestEnvironment::config(self))
@@ -528,8 +534,8 @@ impl Context for TestEnvironment {
         Ok(TestEnvironment::storage(self))
     }
 
-    fn async_storage(&self) -> Result<Arc<AsyncStorage>> {
-        Ok(TestEnvironment::async_storage(self))
+    async fn async_storage(&self) -> Result<Arc<AsyncStorage>> {
+        Ok(TestEnvironment::async_storage(self).await)
     }
 
     fn cdn(&self) -> Result<Arc<CdnBackend>> {

diff --git a/src/web/crate_details.rs b/src/web/crate_details.rs
@@ -1533,7 +1533,7 @@ mod tests {
                 .unwrap();
             assert!(matches!(
                 env.runtime()
-                    .block_on(details.fetch_readme(&env.async_storage())),
+                    .block_on(details.fetch_readme(&env.runtime().block_on(env.async_storage()))),
                 Ok(None)
             ));
             Ok(())

diff --git a/src/web/file.rs b/src/web/file.rs
@@ -78,7 +78,7 @@ mod tests {
             let mut file = env
                 .runtime()
                 .block_on(File::from_path(
-                    &env.async_storage(),
+                    &env.runtime().block_on(env.async_storage()),
                     "rustdoc/fake-package/1.0.0/fake-package/index.html",
                     &env.config(),
                 ))
@@ -125,7 +125,7 @@ mod tests {
 
             let file = |path| {
                 env.runtime().block_on(File::from_path(
-                    &env.async_storage(),
+                    &env.runtime().block_on(env.async_storage()),
                     &format!("rustdoc/dummy/0.1.0/{path}"),
                     &env.config(),
                 ))

diff --git a/src/web/mod.rs b/src/web/mod.rs
@@ -275,6 +275,7 @@ fn apply_middleware(
 ) -> Result<AxumRouter> {
     let config = context.config()?;
     let has_templates = template_data.is_some();
+    let async_storage = context.runtime()?.block_on(context.async_storage())?;
     Ok(router.layer(
         ServiceBuilder::new()
             .layer(TraceLayer::new_for_http())
@@ -293,7 +294,7 @@ fn apply_middleware(
             .layer(Extension(context.instance_metrics()?))
             .layer(Extension(context.config()?))
             .layer(Extension(context.storage()?))
-            .layer(Extension(context.async_storage()?))
+            .layer(Extension(async_storage))
             .layer(Extension(context.repository_stats_updater()?))
             .layer(option_layer(template_data.map(Extension)))
             .layer(middleware::from_fn(csp::csp_middleware))