chore: fix s3 and tests race condition, bump deps (#78)

Author: definfo

Cargo.lock

@@ -2,28 +2,29 @@
 authors = ["Alex Chi <iskyzh@gmail.com>"]
 edition = "2018"
 name = "mirror-intel"
+license = "Apache-2.0"
 version = "0.1.0"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-actix-http = "3.2"
-actix-web = { version = "4.1", features = ["macros"] }
-aws-config = { version = "1.8", features = ["behavior-version-latest"] }
-aws-sdk-s3 = "1.123"
+actix-http = "3.12"
+actix-web = { version = "4.13", features = ["macros"] }
+# NOTE: avoid massive deps used by aws-config
+# aws-config = { version = "1", features = ["behavior-version-latest"] }
+aws-sdk-s3 = { version = "1.122.0", features = ["behavior-version-latest"] }
 bytes = "1.11"
 figment = { version = "0.10", features = ["toml"] }
 futures = "0.3"
 futures-util = "0.3"
 lazy_static = "1.4"
-paste = "1"
 percent-encoding = "2.1"
 prometheus = "0.14"
 regex = "1"
-reqwest = { version = "0.11", features = ["stream"] }
+reqwest = { version = "0.13", default-features = false, features = ["stream"] }
 rstest = "0.26"
 serde = "1.0"
 thiserror = "2.0"
-tokio = { version = "1.0", features = ["full"] }
+tokio = { version = "1", features = ["full"] }
 tracing = "0.1"
 tracing-actix-web = "0.7"
 tracing-appender = "0.2"
@@ -33,6 +34,7 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 url = "2.2"
 
 [dev-dependencies]
-tempdir = "0.3"
+tempfile = "3"
 httpmock = "0.8"
 figment = { version = "0.10", features = ["test"] }
+serial_test = "3"

Cargo.toml

@@ -49,3 +49,4 @@ name = "jCloud S3"
 endpoint = "https://s3.jcloud.sjtu.edu.cn"
 website_endpoint = "https://s3.jcloud.sjtu.edu.cn"
 bucket = "899a892efef34b1b944a19981040f55b-oss01"
+sentinel_object_key = "sjtug-internal/mirror-intel/releases/download/v0.1.35/mirror-intel.tar.gz"

Rocket.toml

@@ -391,14 +391,13 @@ pub async fn download_artifacts(
 #[cfg(test)]
 mod tests {
     use httpmock::MockServer;
-    use tempdir::TempDir;
     use tokio::fs;
 
     use super::*;
 
     #[tokio::test]
     async fn must_download_payload_to_memory() {
-        let tmp_dir = TempDir::new("intel").unwrap();
+        let tmp_dir = tempfile::Builder::new().prefix("intel").tempdir().unwrap();
         let config = Config {
             buffer_path: tmp_dir.path().to_path_buf(),
             file_threshold_mb: 1,
@@ -430,7 +429,7 @@ mod tests {
 
     #[tokio::test]
     async fn must_download_payload_to_file() {
-        let tmp_dir = TempDir::new("intel").unwrap();
+        let tmp_dir = tempfile::Builder::new().prefix("intel").tempdir().unwrap();
         let config = Config {
             buffer_path: tmp_dir.path().to_path_buf(),
             file_threshold_mb: 0,
@@ -468,7 +467,7 @@ mod tests {
 
     #[tokio::test]
     async fn must_reject_large_payload() {
-        let tmp_dir = TempDir::new("intel").unwrap();
+        let tmp_dir = tempfile::Builder::new().prefix("intel").tempdir().unwrap();
         let config = Config {
             buffer_path: tmp_dir.path().to_path_buf(),
             file_threshold_mb: 0,

src/artifacts.rs

@@ -212,6 +212,8 @@ pub struct S3Config {
     pub website_endpoint: String,
     /// Bucket name.
     pub bucket: String,
+    /// [Test only] Object key used by get-object to check S3 availability.
+    pub sentinel_object_key: Option<String>,
 }
 
 /// Configuration for Github Release endpoint.
@@ -369,13 +371,16 @@ pub fn collect_config() -> Config {
 #[cfg(test)]
 mod tests {
     use figment::Jail;
+    use serial_test::serial;
 
     use crate::common::{
         collect_config, EndpointOverride, Endpoints, GithubReleaseConfig, S3Config,
     };
     use crate::Config;
 
+    #[allow(clippy::result_large_err)]
     #[test]
+    #[serial(cwd_env)]
     fn must_collect_config() {
         const MIRROR_INTEL_TOML: &str = include_str!("../tests/config/mirror-intel.toml");
         const ROCKET_TOML: &str = include_str!("../tests/config/Rocket.toml");
@@ -430,6 +435,10 @@ mod tests {
                     endpoint: "https://s3.jcloud.sjtu.edu.cn".into(),
                     website_endpoint: "https://s3.jcloud.sjtu.edu.cn".into(),
                     bucket: "899a892efef34b1b944a19981040f55b-oss01".into(),
+                    sentinel_object_key: Some(
+                        "sjtug-internal/mirror-intel/releases/download/v0.1.35/mirror-intel.tar.gz"
+                            .into(),
+                    ),
                 },
                 user_agent: "mirror-intel / 0.1 (siyuan.internal.sjtug.org)".into(),
                 file_threshold_mb: 4,

src/common.rs

@@ -7,6 +7,8 @@ use thiserror::Error;
 
 type PutObjectSdkError =
     aws_sdk_s3::error::SdkError<aws_sdk_s3::operation::put_object::PutObjectError>;
+type GetObjectSdkError =
+    aws_sdk_s3::error::SdkError<aws_sdk_s3::operation::get_object::GetObjectError>;
 type ListObjectsSdkError =
     aws_sdk_s3::error::SdkError<aws_sdk_s3::operation::list_objects::ListObjectsError>;
 
@@ -30,6 +32,8 @@ pub enum Error {
     InvalidRequest(()),
     #[error("Put Object Error {0}")]
     PutObjectError(Box<PutObjectSdkError>),
+    #[error("Get Object Error {0}")]
+    GetObjectsError(Box<GetObjectSdkError>),
     #[error("List Objects Error {0}")]
     ListObjectsError(Box<ListObjectsSdkError>),
     #[error("Timeout")]
@@ -44,6 +48,11 @@ impl From<PutObjectSdkError> for Error {
         Self::PutObjectError(Box::new(error))
     }
 }
+impl From<GetObjectSdkError> for Error {
+    fn from(error: GetObjectSdkError) -> Self {
+        Self::GetObjectsError(Box::new(error))
+    }
+}
 impl From<ListObjectsSdkError> for Error {
     fn from(error: ListObjectsSdkError) -> Self {
         Self::ListObjectsError(Box::new(error))

src/error.rs

@@ -450,10 +450,12 @@ mod tests {
     use actix_web::test::{call_service, init_service, TestRequest};
     use actix_web::App;
     use figment::providers::{Format, Toml};
+    use figment::util::map;
     use figment::Figment;
     use httpmock::MockServer;
     use reqwest::ClientBuilder;
     use rstest::rstest;
+    use serial_test::serial;
     use tokio::sync::mpsc::{channel, Receiver};
     use url::Url;
 
@@ -488,6 +490,14 @@ mod tests {
         let figment = Figment::new()
             .join(("address", "127.0.0.1"))
             .join(("port", 8000))
+            .join(("concurrent_download", 512))
+            .join(("max_pending_task", 16384))
+            .join((
+                "endpoints",
+                map!["sjtug_internal" => "https://github.com/sjtug"],
+            ))
+            .join(("s3.name", "Placeholder S3"))
+            .join(("s3.endpoint", server.base_url()))
             .join(("s3.website_endpoint", server.base_url()))
             .join(("s3.bucket", "bucket"))
             .join(("direct_stream_size_kb", 0))
@@ -591,6 +601,7 @@ mod tests {
     #[case(Method::HEAD, missing_object(), StatusCode::FOUND, | o: & Task, _c: & Config | o.upstream_url())]
     #[case(Method::GET, forbidden_object(), StatusCode::MOVED_PERMANENTLY, | o: & Task, _c: & Config | o.upstream_url())]
     #[case(Method::HEAD, forbidden_object(), StatusCode::MOVED_PERMANENTLY, | o: & Task, _c: & Config | o.upstream_url())]
+    #[serial(cwd_env)]
     #[tokio::test]
     async fn test_get_head(
         #[case] method: Method,
@@ -632,6 +643,7 @@ mod tests {
     #[case("/pytorch-wheels", is_no_route_for("pytorch-wheels"))]
     #[case("/pytorch-wheels/?mirror_intel_list", is_index_for("pytorch-wheels"))]
     #[case("/pytorch-wheels?mirror_intel_list", is_index_for("pytorch-wheels"))]
+    #[serial(cwd_env)]
     #[tokio::test]
     async fn test_index_list_page(#[case] url: &str, #[case] assert_f: impl FnOnce(&str)) {
         let (service, _config, _rx, _server) = make_service().await;
@@ -642,6 +654,7 @@ mod tests {
         assert_f(text);
     }
 
+    #[serial(cwd_env)]
     #[tokio::test]
     async fn test_url_segment() {
         // this case is to test if we could process escaped URL correctly
@@ -664,6 +677,7 @@ mod tests {
         );
     }
 
+    #[serial(cwd_env)]
     #[tokio::test]
     async fn test_url_segment_fail() {
         // this case is to test if we could process escaped URL correctly
@@ -682,6 +696,7 @@ mod tests {
         assert_eq!(resp.status(), StatusCode::NOT_FOUND);
     }
 
+    #[serial(cwd_env)]
     #[tokio::test]
     async fn test_url_segment_query() {
         // this case is to test if we could process escaped URL correctly
@@ -741,6 +756,7 @@ mod tests {
         assert_eq!(task.origin, "https://storage.googleapis.com/");
     }
 
+    #[serial(cwd_env)]
     #[tokio::test]
     async fn test_proxy_head() {
         // if an object doesn't exist in s3, we should temporarily redirect users to upstream

src/repos.rs

@@ -1,8 +1,7 @@
 //! S3 storage backend.
 use std::time::Duration;
 
-use aws_config::BehaviorVersion;
-use aws_sdk_s3::config::Region;
+use aws_sdk_s3::config::{BehaviorVersion, Region};
 use aws_sdk_s3::Client as S3Client;
 use tokio::time::timeout;
 
@@ -17,19 +16,11 @@ fn s3_region(s3_config: &S3Config) -> Region {
 ///
 /// The default credential provider is used.
 async fn get_s3_client(s3_config: &S3Config) -> S3Client {
-    let shared_config = aws_config::defaults(BehaviorVersion::latest())
-        .region(s3_region(s3_config))
-        .load()
-        .await;
-
-    let mut s3_builder = aws_sdk_s3::Config::builder()
+    let s3_builder = aws_sdk_s3::Config::builder()
         .region(s3_region(s3_config))
         .endpoint_url(s3_config.endpoint.clone())
-        .behavior_version(BehaviorVersion::latest())
         .force_path_style(true);
 
-    s3_builder.set_credentials_provider(shared_config.credentials_provider());
-
     S3Client::from_conf(s3_builder.build())
 }
 
@@ -41,7 +32,6 @@ pub fn get_anonymous_s3_client(s3_config: &S3Config) -> S3Client {
         aws_sdk_s3::Config::builder()
             .region(s3_region(s3_config))
             .endpoint_url(s3_config.endpoint.clone())
-            .behavior_version(BehaviorVersion::latest())
             .force_path_style(true)
             .allow_no_auth()
             .build(),
@@ -74,11 +64,28 @@ pub async fn check_s3(s3_config: &S3Config) -> Result<()> {
     timeout(Duration::from_secs(1), async move {
         let s3_client = get_s3_client(s3_config).await;
 
-        s3_client
-            .list_objects()
-            .bucket(s3_config.bucket.clone())
-            .send()
-            .await?;
+        // s3_client
+        //     .list_objects()
+        //     .bucket(s3_config.bucket.clone())
+        //     .send()
+        //     .await?;
+        if let Some(sentinel_object_key) = &s3_config.sentinel_object_key {
+            s3_client
+                .get_object()
+                .bucket(s3_config.bucket.clone())
+                .key(sentinel_object_key.clone())
+                .range("bytes=0-0")
+                .send()
+                .await?;
+        }
+        // NOTE: this can be too heavy for jCloud S3, thus we check only sentinel object instead of listing all objects.
+        else {
+            s3_client
+                .list_objects()
+                .bucket(s3_config.bucket.clone())
+                .send()
+                .await?;
+        }
 
         Ok::<(), Error>(())
     })

src/storage.rs

@@ -1,6 +1,6 @@
 use actix_web::body::{BodyStream, SizedStream};
 use actix_web::http::header::ContentType;
-use actix_web::http::{header, StatusCode, Uri};
+use actix_web::http::{StatusCode, Uri};
 use actix_web::{HttpResponse, Responder};
 use futures::stream::TryStreamExt;
 use tracing::debug;
@@ -115,15 +115,19 @@ impl IntelObject {
         if resp.status().is_success() {
             if let Some(content_length) = resp.content_length() {
                 if content_length <= below_size_kb * 1024 {
-                    let content_type = resp.headers().get(header::CONTENT_TYPE).cloned();
+                    let content_type = resp.headers().get("content-type").cloned();
                     let text = resp.text().await?;
                     let text = f(text);
 
                     return Ok(if let Some(content_type) = content_type {
-                        HttpResponse::Ok()
-                            .content_type(content_type)
-                            .body(text)
-                            .into()
+                        if let Ok(content_type) = content_type.to_str() {
+                            HttpResponse::Ok()
+                                .content_type(content_type)
+                                .body(text)
+                                .into()
+                        } else {
+                            HttpResponse::Ok().body(text).into()
+                        }
                     } else {
                         HttpResponse::Ok().body(text).into()
                     });
@@ -143,10 +147,12 @@ impl IntelObject {
             }
         };
 
-        let code = upstream_resp.status().normalize();
+        let code = upstream_resp.status().as_u16().normalize();
         let mut resp = HttpResponse::build(code);
-        if let Some(content_type) = upstream_resp.headers().get(header::CONTENT_TYPE) {
-            resp.content_type(content_type);
+        if let Some(content_type) = upstream_resp.headers().get("content-type") {
+            if let Ok(content_type) = content_type.to_str() {
+                resp.content_type(content_type);
+            }
         }
         let content_length = upstream_resp.content_length();
         let stream = upstream_resp
@@ -220,10 +226,16 @@ trait StatusCodeExt {
 
 impl StatusCodeExt for StatusCode {
     fn normalize(self) -> StatusCode {
-        if self.as_u16() == 499 {
-            Self::NOT_FOUND
+        self.as_u16().normalize()
+    }
+}
+
+impl StatusCodeExt for u16 {
+    fn normalize(self) -> StatusCode {
+        if self == 499 {
+            StatusCode::NOT_FOUND
         } else {
-            self
+            StatusCode::from_u16(self).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR)
         }
     }
 }
@@ -254,6 +266,7 @@ mod tests {
                 endpoint: "http://localhost:8081".to_string(),
                 website_endpoint: server.base_url(),
                 bucket: "bucket".to_string(),
+                sentinel_object_key: None,
             },
             ..Default::default()
         };

src/utils.rs

@@ -1,6 +1,7 @@
 [default]
 address = "0.0.0.0"
 concurrent_download = 256
+max_pending_task = 16384
 user_agent = "mirror-intel / 0.1 (siyuan.internal.sjtug.org)"
 file_threshold_mb = 4
 ignore_threshold_mb = 1024
@@ -48,3 +49,4 @@ name = "jCloud S3"
 endpoint = "https://s3.jcloud.sjtu.edu.cn"
 website_endpoint = "https://s3.jcloud.sjtu.edu.cn"
 bucket = "899a892efef34b1b944a19981040f55b-oss01"
+sentinel_object_key = "sjtug-internal/mirror-intel/releases/download/v0.1.35/mirror-intel.tar.gz"