optionally set caching headers for /latest/, invalidate caches after build

Author: syphar

Cargo.lock

@@ -73,8 +73,10 @@ tokio = { version = "1.0", features = ["rt-multi-thread"] }
 futures-util = "0.3.5"
 aws-config = "0.48.0"
 aws-sdk-s3 = "0.18.0"
+aws-sdk-cloudfront = "0.18.0"
 aws-smithy-types-convert = { version = "0.48.0", features = ["convert-chrono"] }
 http = "0.2.6"
+uuid = "1.1.2"
 
 # Data serialization and deserialization
 serde = { version = "1.0", features = ["derive"] }

Cargo.toml

@@ -16,6 +16,7 @@ use once_cell::sync::OnceCell;
 use sentry_log::SentryLogger;
 use structopt::StructOpt;
 use strum::VariantNames;
+use tokio::runtime::Runtime;
 
 fn main() {
     let _sentry_guard = if let Ok(sentry_dsn) = env::var("SENTRY_DSN") {
@@ -563,6 +564,7 @@ struct BinContext {
     metrics: OnceCell<Arc<Metrics>>,
     index: OnceCell<Arc<Index>>,
     repository_stats_updater: OnceCell<Arc<RepositoryStatsUpdater>>,
+    runtime: OnceCell<Arc<Runtime>>,
 }
 
 impl BinContext {
@@ -575,6 +577,7 @@ impl BinContext {
             metrics: OnceCell::new(),
             index: OnceCell::new(),
             repository_stats_updater: OnceCell::new(),
+            runtime: OnceCell::new(),
         }
     }
 
@@ -606,9 +609,11 @@ impl Context for BinContext {
             self.pool()?,
             self.metrics()?,
             self.config()?,
+            self.runtime()?,
         )?;
         fn config(self) -> Config = Config::from_env()?;
         fn metrics(self) -> Metrics = Metrics::new()?;
+        fn runtime(self) -> Runtime = Runtime::new()?;
         fn index(self) -> Index = {
             let config = self.config()?;
             let path = config.registry_index_path.clone();

src/bin/cratesfyi.rs

@@ -0,0 +1,53 @@
+use anyhow::Error;
+use aws_sdk_cloudfront::{
+    model::{InvalidationBatch, Paths},
+    Client, RetryConfig,
+};
+use tokio::runtime::Runtime;
+use uuid::Uuid;
+
+/// create a CloudFront invalidation request for a list of path patterns.
+/// patterns can be
+/// * `/filename.ext` (a specific path)
+/// * `/directory-path/file-name.*` (delete these files, all extensions)
+/// * `/directory-path/*` (invalidate all of the files in a directory, without subdirectories)
+/// * `/directory-path*` (recursive directory delete, including subdirectories)
+/// see https://docs.aws.amazon.com/AmazonCloudFront/latest/DeveloperGuide/Invalidation.html#invalidation-specifying-objects
+///
+/// Returns the caller reference that can be used to query the status of this
+/// invalidation request.
+pub(crate) fn create_cloudfront_invalidation(
+    runtime: &Runtime,
+    distribution_id: &str,
+    path_patterns: &[&str],
+) -> Result<Uuid, Error> {
+    let shared_config = runtime.block_on(aws_config::load_from_env());
+    let config_builder = aws_sdk_cloudfront::config::Builder::from(&shared_config)
+        .retry_config(RetryConfig::new().with_max_attempts(3));
+
+    let client = Client::from_conf(config_builder.build());
+    let path_patterns: Vec<_> = path_patterns.iter().cloned().map(String::from).collect();
+
+    let caller_reference = Uuid::new_v4();
+
+    runtime.block_on(async {
+        client
+            .create_invalidation()
+            .distribution_id(distribution_id)
+            .invalidation_batch(
+                InvalidationBatch::builder()
+                    .paths(
+                        Paths::builder()
+                            .quantity(path_patterns.len().try_into().unwrap())
+                            .set_items(Some(path_patterns))
+                            .build(),
+                    )
+                    .caller_reference(format!("{}", caller_reference))
+                    .build(),
+            )
+            .send()
+            .await
+    })?;
+
+    Ok(caller_reference)
+}

src/cache.rs

@@ -57,12 +57,26 @@ pub struct Config {
     // Content Security Policy
     pub(crate) csp_report_only: bool,
 
-    // Cache-Control header
+    // Cache-Control header, for versioned URLs.
     // If both are absent, don't generate the header. If only one is present,
     // generate just that directive. Values are in seconds.
     pub(crate) cache_control_stale_while_revalidate: Option<u32>,
     pub(crate) cache_control_max_age: Option<u32>,
 
+    // Cache-Control header, for /latest/ URLs.
+    // Same conditions as above apply.
+    pub(crate) cache_control_stale_while_revalidate_latest: Option<u32>,
+    pub(crate) cache_control_max_age_latest: Option<u32>,
+
+    // CloudFront distribution ID for the web server.
+    // Will be used for invalidation-requests.
+    pub(crate) cloudfront_distribution_id_web: Option<String>,
+
+    // when invalidating, either
+    // * invalidate the whole crate path (`/krate/*`)
+    // * invalidate only the `/latest/` url
+    pub(crate) cloudfront_only_invalidate_latest: bool,
+
     // Build params
     pub(crate) build_attempts: u16,
     pub(crate) rustwide_workspace: PathBuf,
@@ -141,6 +155,15 @@ impl Config {
             )?,
             cache_control_max_age: maybe_env("CACHE_CONTROL_MAX_AGE")?,
 
+            cache_control_stale_while_revalidate_latest: maybe_env(
+                "CACHE_CONTROL_STALE_WHILE_REVALIDATE_LATEST",
+            )?,
+            cache_control_max_age_latest: maybe_env("CACHE_CONTROL_MAX_AGE_LATEST")?,
+
+            cloudfront_distribution_id_web: maybe_env("CLOUDFRONT_DISTRIBUTION_ID_WEB")?,
+
+            cloudfront_only_invalidate_latest: env("CLOUDFRONT_ONLY_INVALIDATE_LATEST", false)?,
+
             local_archive_cache_path: env(
                 "DOCSRS_ARCHIVE_INDEX_CACHE_PATH",
                 prefix.join("archive_cache"),

src/config.rs

@@ -3,6 +3,7 @@ use crate::error::Result;
 use crate::repositories::RepositoryStatsUpdater;
 use crate::{BuildQueue, Config, Index, Metrics, Storage};
 use std::sync::Arc;
+use tokio::runtime::Runtime;
 
 pub trait Context {
     fn config(&self) -> Result<Arc<Config>>;
@@ -12,4 +13,5 @@ pub trait Context {
     fn metrics(&self) -> Result<Arc<Metrics>>;
     fn index(&self) -> Result<Arc<Index>>;
     fn repository_stats_updater(&self) -> Result<Arc<RepositoryStatsUpdater>>;
+    fn runtime(&self) -> Result<Arc<Runtime>>;
 }

src/context.rs

@@ -1,3 +1,4 @@
+use crate::cache;
 use crate::db::file::add_path_into_database;
 use crate::db::{
     add_build_into_database, add_doc_coverage, add_package_into_database,
@@ -25,6 +26,7 @@ use rustwide::{AlternativeRegistry, Build, Crate, Toolchain, Workspace, Workspac
 use std::collections::{HashMap, HashSet};
 use std::path::Path;
 use std::sync::Arc;
+use tokio::runtime::Runtime;
 
 const USER_AGENT: &str = "docs.rs builder (https://github.com/rust-lang/docs.rs)";
 const DUMMY_CRATE_NAME: &str = "empty-library";
@@ -43,6 +45,7 @@ pub struct RustwideBuilder {
     db: Pool,
     storage: Arc<Storage>,
     metrics: Arc<Metrics>,
+    runtime: Arc<Runtime>,
     index: Arc<Index>,
     rustc_version: String,
     repository_stats_updater: Arc<RepositoryStatsUpdater>,
@@ -80,6 +83,7 @@ impl RustwideBuilder {
             config,
             db: context.pool()?,
             storage: context.storage()?,
+            runtime: context.runtime()?,
             metrics: context.metrics()?,
             index: context.index()?,
             rustc_version: String::new(),
@@ -500,6 +504,20 @@ impl RustwideBuilder {
             .purge_from_cache(&self.workspace)
             .map_err(FailureError::compat)?;
         local_storage.close()?;
+        if let Some(distribution_id) = self.config.cloudfront_distribution_id_web.as_ref() {
+            if successful {
+                let invalidate_path = if self.config.cloudfront_only_invalidate_latest {
+                    format!("/{}/latest*", name)
+                } else {
+                    format!("/{}*", name)
+                };
+                cache::create_cloudfront_invalidation(
+                    &self.runtime,
+                    distribution_id,
+                    &[&invalidate_path],
+                )?;
+            }
+        }
         Ok(successful)
     }
 

src/docbuilder/rustwide_builder.rs

@@ -13,6 +13,7 @@ pub use self::storage::Storage;
 pub use self::web::Server;
 
 mod build_queue;
+pub mod cache;
 mod config;
 mod context;
 pub mod db;

src/lib.rs

@@ -21,6 +21,7 @@ use std::{
     path::{Path, PathBuf},
     sync::Arc,
 };
+use tokio::runtime::Runtime;
 
 const MAX_CONCURRENT_UPLOADS: usize = 1000;
 
@@ -113,14 +114,21 @@ pub struct Storage {
 }
 
 impl Storage {
-    pub fn new(pool: Pool, metrics: Arc<Metrics>, config: Arc<Config>) -> Result<Self> {
+    pub fn new(
+        pool: Pool,
+        metrics: Arc<Metrics>,
+        config: Arc<Config>,
+        runtime: Arc<Runtime>,
+    ) -> Result<Self> {
         Ok(Storage {
             config: config.clone(),
             backend: match config.storage_backend {
                 StorageKind::Database => {
                     StorageBackend::Database(DatabaseBackend::new(pool, metrics))
                 }
-                StorageKind::S3 => StorageBackend::S3(Box::new(S3Backend::new(metrics, &config)?)),
+                StorageKind::S3 => {
+                    StorageBackend::S3(Box::new(S3Backend::new(metrics, &config, runtime)?))
+                }
             },
         })
     }

src/storage/mod.rs

@@ -18,17 +18,19 @@ use tokio::runtime::Runtime;
 
 pub(super) struct S3Backend {
     client: Client,
-    runtime: Runtime,
+    runtime: Arc<Runtime>,
     bucket: String,
     metrics: Arc<Metrics>,
     #[cfg(test)]
     temporary: bool,
 }
 
 impl S3Backend {
-    pub(super) fn new(metrics: Arc<Metrics>, config: &Config) -> Result<Self, Error> {
-        let runtime = Runtime::new()?;
-
+    pub(super) fn new(
+        metrics: Arc<Metrics>,
+        config: &Config,
+        runtime: Arc<Runtime>,
+    ) -> Result<Self, Error> {
         let shared_config = runtime.block_on(aws_config::load_from_env());
         let mut config_builder = aws_sdk_s3::config::Builder::from(&shared_config)
             .retry_config(RetryConfig::new().with_max_attempts(3))

src/storage/s3.rs

@@ -17,6 +17,7 @@ use reqwest::{
     Method,
 };
 use std::{fs, net::SocketAddr, panic, sync::Arc, time::Duration};
+use tokio::runtime::Runtime;
 
 pub(crate) fn wrapper(f: impl FnOnce(&TestEnvironment) -> Result<()>) {
     let env = TestEnvironment::new();
@@ -114,6 +115,7 @@ pub(crate) struct TestEnvironment {
     db: OnceCell<TestDatabase>,
     storage: OnceCell<Arc<Storage>>,
     index: OnceCell<Arc<Index>>,
+    runtime: OnceCell<Arc<Runtime>>,
     metrics: OnceCell<Arc<Metrics>>,
     frontend: OnceCell<TestFrontend>,
     repository_stats_updater: OnceCell<Arc<RepositoryStatsUpdater>>,
@@ -140,6 +142,7 @@ impl TestEnvironment {
             index: OnceCell::new(),
             metrics: OnceCell::new(),
             frontend: OnceCell::new(),
+            runtime: OnceCell::new(),
             repository_stats_updater: OnceCell::new(),
         }
     }
@@ -216,8 +219,13 @@ impl TestEnvironment {
         self.storage
             .get_or_init(|| {
                 Arc::new(
-                    Storage::new(self.db().pool(), self.metrics(), self.config())
-                        .expect("failed to initialize the storage"),
+                    Storage::new(
+                        self.db().pool(),
+                        self.metrics(),
+                        self.config(),
+                        self.runtime(),
+                    )
+                    .expect("failed to initialize the storage"),
                 )
             })
             .clone()
@@ -228,6 +236,11 @@ impl TestEnvironment {
             .get_or_init(|| Arc::new(Metrics::new().expect("failed to initialize the metrics")))
             .clone()
     }
+    pub(crate) fn runtime(&self) -> Arc<Runtime> {
+        self.runtime
+            .get_or_init(|| Arc::new(Runtime::new().expect("failed to initialize runtime")))
+            .clone()
+    }
 
     pub(crate) fn index(&self) -> Arc<Index> {
         self.index
@@ -303,6 +316,10 @@ impl Context for TestEnvironment {
     fn repository_stats_updater(&self) -> Result<Arc<RepositoryStatsUpdater>> {
         Ok(self.repository_stats_updater())
     }
+
+    fn runtime(&self) -> Result<Arc<Runtime>> {
+        Ok(self.runtime())
+    }
 }
 
 pub(crate) struct TestDatabase {