A gentle deadlock monitor

.github/workflows/publish-server.yml

@@ -4,12 +4,13 @@ on:
   push:
     branches:
       - "main"
+      - "as/add-timeout"
     tags:
       - libsql-server-v*.*.*
 
 env:
   REGISTRY: ghcr.io
-  IMAGE_NAME: tursodatabase/libsql-server
+  IMAGE_NAME: spice-labs-inc/libsql-server
 
 jobs:
   build-amd64:

Cargo.toml

@@ -28,6 +28,7 @@ exclude = [
 version = "0.10.0-pre.4"
 authors = ["the libSQL authors"]
 edition = "2021"
+rust-version = "1.85"
 license = "MIT"
 repository = "https://github.com/tursodatabase/libsql"
 

bottomless-cli/src/main.rs

@@ -133,8 +133,23 @@ async fn detect_database(options: &Cli, namespace: &str) -> Result<(String, Stri
     let database = match options.database.clone() {
         Some(db) => db,
         None => {
+            let timeout_config = aws_smithy_types::timeout::TimeoutConfig::builder()
+                .connect_timeout(std::time::Duration::from_secs(
+                    std::env::var("LIBSQL_BOTTOMLESS_S3_CONNECT_TIMEOUT_SECS")
+                        .ok()
+                        .and_then(|s| s.parse().ok())
+                        .unwrap_or(5),
+                ))
+                .read_timeout(std::time::Duration::from_secs(
+                    std::env::var("LIBSQL_BOTTOMLESS_S3_READ_TIMEOUT_SECS")
+                        .ok()
+                        .and_then(|s| s.parse().ok())
+                        .unwrap_or(5),
+                ))
+                .build();
             let client = Client::from_conf({
-                let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
+                let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest())
+                    .timeout_config(timeout_config);
                 if let Some(endpoint) = options.endpoint.clone() {
                     loader = loader.endpoint_url(endpoint);
                 }

bottomless/Cargo.toml

@@ -13,6 +13,7 @@ anyhow = "1.0.66"
 async-compression = { version = "0.4.4", features = ["tokio", "gzip", "zstd"] }
 aws-config = { version = "1" }
 aws-sdk-s3 = { version = "1" }
+aws-smithy-types = "1"
 bytes = "1"
 libsql-sys = { path = "../libsql-sys" }
 libsql_replication = { path = "../libsql-replication" }

bottomless/src/replicator.rs

@@ -121,6 +121,10 @@ pub struct Options {
     pub s3_max_parallelism: usize,
     /// Max number of retries for S3 operations
     pub s3_max_retries: u32,
+    /// Max seconds to wait for first byte of S3 response
+    pub s3_read_timeout_secs: u64,
+    /// Max seconds to wait for S3 connection establishment
+    pub s3_connect_timeout_secs: u64,
     /// Skip snapshot upload per checkpoint.
     pub skip_snapshot: bool,
     /// Skip uploading snapshots on shutdown
@@ -145,8 +149,15 @@ impl Options {
             "LIBSQL_BOTTOMLESS_AWS_SECRET_ACCESS_KEY was not set"
         ))?;
         let session_token: Option<String> = self.session_token.clone();
+
+        let timeout_config = aws_smithy_types::timeout::TimeoutConfig::builder()
+            .connect_timeout(Duration::from_secs(self.s3_connect_timeout_secs))
+            .read_timeout(Duration::from_secs(self.s3_read_timeout_secs))
+            .build();
+
         let conf = loader
             .behavior_version(BehaviorVersion::latest())
+            .timeout_config(timeout_config)
             .region(Region::new(region))
             .credentials_provider(SharedCredentialsProvider::new(Credentials::new(
                 access_key_id,
@@ -233,6 +244,10 @@ impl Options {
             ),
         };
         let s3_max_retries = env_var_or("LIBSQL_BOTTOMLESS_S3_MAX_RETRIES", 10).parse::<u32>()?;
+        let s3_read_timeout_secs =
+            env_var_or("LIBSQL_BOTTOMLESS_S3_READ_TIMEOUT_SECS", 5).parse::<u64>()?;
+        let s3_connect_timeout_secs =
+            env_var_or("LIBSQL_BOTTOMLESS_S3_CONNECT_TIMEOUT_SECS", 5).parse::<u64>()?;
         let cipher = match encryption_cipher {
             Some(cipher) => Cipher::from_str(&cipher)?,
             None => Cipher::default(),
@@ -261,6 +276,8 @@ impl Options {
             region,
             bucket_name,
             s3_max_retries,
+            s3_read_timeout_secs,
+            s3_connect_timeout_secs,
             skip_snapshot,
             skip_shutdown_upload,
         })
@@ -1090,37 +1107,29 @@ impl Replicator {
         }
         let generation = self.generation()?;
         let start_ts = Instant::now();
-        let client = self.client.clone();
         let change_counter = self.read_change_counter()?;
-        let snapshot_req = client.put_object().bucket(self.bucket.clone()).key(format!(
-            "{}-{}/db.{}",
-            self.db_name, generation, self.use_compression
-        ));
 
         /* FIXME: we can't rely on the change counter in WAL mode:
          ** "In WAL mode, changes to the database are detected using the wal-index and
          ** so the change counter is not needed. Hence, the change counter might not be
          ** incremented on each transaction in WAL mode."
          ** Instead, we need to consult WAL checksums.
          */
-        let change_counter_key = format!("{}-{}/.changecounter", self.db_name, generation);
-        let change_counter_req = self
-            .client
-            .put_object()
-            .bucket(&self.bucket)
-            .key(change_counter_key)
-            .body(ByteStream::from(Bytes::copy_from_slice(
-                change_counter.as_ref(),
-            )));
         let snapshot_notifier = self.snapshot_notifier.clone();
         let compression = self.use_compression;
         let db_path = PathBuf::from(self.db_path.clone());
         let db_name = self.db_name.clone();
+        let bucket = self.bucket.clone();
+        let client = self.client.clone();
         let handle = tokio::spawn(async move {
             tracing::trace!("Start snapshotting generation {}", generation);
             let start = Instant::now();
-            let body = match Self::maybe_compress_main_db_file(&db_path, compression).await {
-                Ok(file) => file,
+            let body_path = match Self::maybe_compress_main_db_file(&db_path, compression).await {
+                Ok(_) => match compression {
+                    CompressionKind::None => db_path.clone(),
+                    CompressionKind::Gzip => Self::db_compressed_path(&db_path, "gz"),
+                    CompressionKind::Zstd => Self::db_compressed_path(&db_path, "zstd"),
+                },
                 Err(e) => {
                     tracing::error!(
                         "Failed to compress db file (generation `{}`, path: `{}`): {:?}",
@@ -1132,25 +1141,58 @@ impl Replicator {
                     return;
                 }
             };
-            let mut result = snapshot_req.body(body).send().await;
-            if let Err(e) = result {
+            let snapshot_key = format!("{}-{}/db.{}", db_name, generation, compression);
+            let mut result = client
+                .put_object()
+                .bucket(&bucket)
+                .key(&snapshot_key)
+                .body(ByteStream::from_path(&body_path).await.unwrap())
+                .send()
+                .await;
+            let mut retries = 0;
+            while let Err(e) = result {
+                retries += 1;
                 tracing::error!(
-                    "Failed to upload snapshot for generation {}: {:?}",
+                    "Failed to upload snapshot for generation {}: {:?}, will retry after 1 second (attempt {})",
                     generation,
-                    e
+                    e,
+                    retries
                 );
-                let _ = snapshot_notifier.send(Err(e.into()));
-                return;
+                tokio::time::sleep(Duration::from_millis(1000)).await;
+                result = client
+                    .put_object()
+                    .bucket(&bucket)
+                    .key(&snapshot_key)
+                    .body(ByteStream::from_path(&body_path).await.unwrap())
+                    .send()
+                    .await;
             }
-            result = change_counter_req.send().await;
-            if let Err(e) = result {
+            let change_counter_key = format!("{}-{}/.changecounter", db_name, generation);
+            result = client
+                .put_object()
+                .bucket(&bucket)
+                .key(change_counter_key.clone())
+                .body(ByteStream::from(Bytes::copy_from_slice(
+                    change_counter.as_ref(),
+                )))
+                .send()
+                .await;
+            while let Err(e) = result {
                 tracing::error!(
-                    "Failed to upload change counter for generation {}: {:?}",
+                    "Failed to upload change counter for generation {}: {:?}, will retry after 1 second",
                     generation,
                     e
                 );
-                let _ = snapshot_notifier.send(Err(e.into()));
-                return;
+                tokio::time::sleep(Duration::from_millis(1000)).await;
+                result = client
+                    .put_object()
+                    .bucket(&bucket)
+                    .key(change_counter_key.clone())
+                    .body(ByteStream::from(Bytes::copy_from_slice(
+                        change_counter.as_ref(),
+                    )))
+                    .send()
+                    .await;
             }
             let _ = snapshot_notifier.send(Ok(Some(generation)));
             let elapsed = Instant::now() - start;
@@ -2184,3 +2226,50 @@ impl std::fmt::Display for CompressionKind {
         }
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use std::time::Duration;
+
+    #[tokio::test]
+    async fn client_config_applies_timeouts() {
+        let options = Options {
+            create_bucket_if_not_exists: true,
+            verify_crc: false,
+            use_compression: CompressionKind::None,
+            encryption_config: None,
+            aws_endpoint: Some("http://localhost:9000".to_string()),
+            access_key_id: Some("test".to_string()),
+            secret_access_key: Some("test".to_string()),
+            session_token: None,
+            region: Some("us-east-1".to_string()),
+            db_id: None,
+            bucket_name: "test".to_string(),
+            max_frames_per_batch: 1000,
+            max_batch_interval: Duration::from_secs(1),
+            s3_max_parallelism: 1,
+            s3_max_retries: 1,
+            s3_read_timeout_secs: 7,
+            s3_connect_timeout_secs: 3,
+            skip_snapshot: false,
+            skip_shutdown_upload: false,
+        };
+
+        let config = options.client_config().await.unwrap();
+        let timeout_config = config
+            .timeout_config()
+            .expect("timeout_config should be set");
+
+        assert_eq!(
+            timeout_config.connect_timeout(),
+            Some(Duration::from_secs(3)),
+            "connect_timeout should be 3s"
+        );
+        assert_eq!(
+            timeout_config.read_timeout(),
+            Some(Duration::from_secs(7)),
+            "read_timeout should be 7s"
+        );
+    }
+}