Repository: chroma-core/chroma
Pull request #4476: [BUG]  Catch a bug that causes wal3 to hang under contention.
Author: login: rescrv
Created: 2025-05-07T15:05:33Z | Merged: 2025-05-07T17:46:40Z

## Description of changes

If two wal3 instances try writing the same log, only one will succeed,
and the other will fail.  This can, under rare situtation, lead to a log
failing recovery due to contention, and thus failing to open.
Contention should be a transient thing and the log should handle it.

Included is a claude-written test that plays ping-pong between two log
writers in a coordinated fashion.

## Test plan

- [X] Tests pass locally with `pytest` for python, `yarn test` for js, `cargo test` for rust

## Documentation Changes

N/A



Diff:
diff --git a/rust/wal3/src/manifest_manager.rs b/rust/wal3/src/manifest_manager.rs
index aff74239b72..0091a783985 100644
--- a/rust/wal3/src/manifest_manager.rs
+++ b/rust/wal3/src/manifest_manager.rs
@@ -171,13 +171,13 @@ impl ManifestManager {
     /// Recover from a fault in writing.  It is possible that fragments have been written that are
     /// not referenced by the manifest.  Scout ahead until an empty slot is observed.  Then write
     /// the manifest that includes the new fragments.
-    pub async fn recover(&mut self, mark_dirty: &dyn MarkDirty) -> Result<(), Error> {
-        let mut next_seq_no_to_apply = {
-            // SAFETY(rescrv):  Mutex poisoning.
-            let staging = self.staging.lock().unwrap();
-            staging.next_seq_no_to_apply
-        };
+    pub async fn recover(&self, mark_dirty: &dyn MarkDirty) -> Result<(), Error> {
         loop {
+            let next_seq_no_to_apply = {
+                // SAFETY(rescrv):  Mutex poisoning.
+                let staging = self.staging.lock().unwrap();
+                staging.next_seq_no_to_apply
+            };
             let next_fragment = read_fragment(
                 &self.storage,
                 &self.prefix,
@@ -189,16 +189,10 @@ impl ManifestManager {
                     .mark_dirty(fragment.start, (fragment.limit - fragment.start) as usize)
                     .await?;
                 self.publish_fragment(fragment).await?;
-                next_seq_no_to_apply += 1;
             } else {
                 break;
             }
         }
-        {
-            // SAFETY(rescrv):  Mutex poisoning.
-            let mut staging = self.staging.lock().unwrap();
-            staging.next_seq_no_to_apply = next_seq_no_to_apply;
-        }
         Ok(())
     }
 
diff --git a/rust/wal3/src/writer.rs b/rust/wal3/src/writer.rs
index 39095217350..36a7046eccf 100644
--- a/rust/wal3/src/writer.rs
+++ b/rust/wal3/src/writer.rs
@@ -152,29 +152,36 @@ impl LogWriter {
     pub async fn append_many(&self, messages: Vec<Vec<u8>>) -> Result<LogPosition, Error> {
         // SAFETY(rescrv):  Mutex poisoning.
         let inner = { self.inner.lock().unwrap().clone() };
-        if let Some(ref epoch_writer) = inner {
+        if let Some(epoch_writer) = inner {
             let res = epoch_writer.writer.append(messages.clone()).await;
             if matches!(res, Err(Error::LogContention)) {
-                let writer = OnceLogWriter::open(
-                    epoch_writer.writer.options.clone(),
-                    epoch_writer.writer.storage.clone(),
-                    epoch_writer.writer.prefix.clone(),
-                    self.writer.clone(),
-                    Arc::clone(&self.mark_dirty),
-                )
-                .await?;
-                // SAFETY(rescrv):  Mutex poisoning.
-                let mut inner = self.inner.lock().unwrap();
-                if let Some(second) = inner.as_mut() {
-                    if second.epoch == epoch_writer.epoch {
-                        second.epoch += 1;
-                        second.writer.shutdown();
-                        second.writer = writer;
+                loop {
+                    let writer = OnceLogWriter::open(
+                        epoch_writer.writer.options.clone(),
+                        epoch_writer.writer.storage.clone(),
+                        epoch_writer.writer.prefix.clone(),
+                        self.writer.clone(),
+                        Arc::clone(&self.mark_dirty),
+                    )
+                    .await;
+                    let writer = match writer {
+                        Ok(writer) => writer,
+                        Err(Error::LogContention) => continue,
+                        Err(err) => return Err(err),
+                    };
+                    // SAFETY(rescrv):  Mutex poisoning.
+                    let mut inner = self.inner.lock().unwrap();
+                    if let Some(second) = inner.as_mut() {
+                        if second.epoch == epoch_writer.epoch {
+                            second.epoch += 1;
+                            second.writer.shutdown();
+                            second.writer = writer;
+                        }
+                        return Err(Error::LogContention);
+                    } else {
+                        // This should never happen, so just be polite with an error.
+                        return Err(Error::LogClosed);
                     }
-                    Err(Error::LogContention)
-                } else {
-                    // This should never happen, so just be polite with an error.
-                    Err(Error::LogClosed)
                 }
             } else {
                 res
@@ -257,7 +264,7 @@ impl OnceLogWriter {
     ) -> Result<Arc<Self>, Error> {
         let done = AtomicBool::new(false);
         let batch_manager = BatchManager::new(options.throttle_fragment).ok_or(Error::Internal)?;
-        let mut manifest_manager = ManifestManager::new(
+        let manifest_manager = ManifestManager::new(
             options.throttle_manifest,
             options.snapshot_manifest,
             Arc::clone(&storage),
diff --git a/rust/wal3/tests/test_k8s_integration_99_ping_pong_contention.rs b/rust/wal3/tests/test_k8s_integration_99_ping_pong_contention.rs
new file mode 100644
index 00000000000..49f1638af25
--- /dev/null
+++ b/rust/wal3/tests/test_k8s_integration_99_ping_pong_contention.rs
@@ -0,0 +1,178 @@
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
+use std::time::Duration;
+
+use tokio::sync::Mutex;
+
+use chroma_storage::s3_client_for_test_with_new_bucket;
+
+use wal3::{Error, LogWriter, LogWriterOptions, Manifest};
+
+pub mod common;
+
+async fn writer_thread(
+    writer: Arc<LogWriter>,
+    mutex: Arc<Mutex<()>>,
+    running: Arc<AtomicUsize>,
+    thread_id: usize,
+    iterations: usize,
+) -> (usize, usize) {
+    let mut successful_writes = 0;
+    let mut contention_errors = 0;
+    println!(
+        "writer {thread_id} also known as {:?}",
+        &*writer as *const LogWriter
+    );
+
+    for i in 0..iterations {
+        let message = format!("Message from writer{}: {}", thread_id, i).into_bytes();
+
+        // Acquire the mutex, do a write, release the mutex
+        loop {
+            // Using tokio::sync::Mutex which is safe to use with .await
+            let _guard = mutex.lock().await;
+            println!("writer {thread_id} grabs lock in iteration {i}");
+            // We have the lock, do a write
+            match writer.append(message.clone()).await {
+                Ok(_) => {
+                    println!("writer {thread_id} succeeds in iteration {i}");
+                    successful_writes += 1;
+                    // Release mutex (implicit) and break
+                    break;
+                }
+                Err(Error::LogContention) => {
+                    println!("writer {thread_id} sees contention preventing {i}");
+                    contention_errors += 1;
+                }
+                Err(e) => panic!("Unexpected error: {:?}", e),
+            }
+        }
+
+        println!("writer {thread_id} goes for contention {i}");
+
+        // Now write until we hit LogContention
+        while i + 1 < iterations && running.load(Ordering::Relaxed) > 1 {
+            // NOTE(rescrv):
+            // Default batching interval is 100ms.
+            // This must be at least that to prevent waste.
+            tokio::time::sleep(Duration::from_millis(150)).await;
+            match writer.append(message.clone()).await {
+                Ok(_) => {
+                    successful_writes += 1;
+                }
+                Err(Error::LogContention) => {
+                    println!("writer {thread_id} contends without lock in iteration {i}");
+                    contention_errors += 1;
+                    // Got contention, now we go back to the top and wait for mutex
+                    break;
+                }
+                Err(e) => panic!("Unexpected error: {:?}", e),
+            }
+        }
+    }
+
+    running.fetch_sub(1, Ordering::Relaxed);
+    (successful_writes, contention_errors)
+}
+
+#[tokio::test]
+async fn test_k8s_integration_ping_pong_contention() {
+    // Create a shared storage for both threads to use
+    let storage = Arc::new(s3_client_for_test_with_new_bucket().await);
+    let prefix = "test_k8s_integration_ping_pong_contention";
+    let writer_name = "init";
+
+    // Initialize the log
+    Manifest::initialize(&LogWriterOptions::default(), &storage, prefix, writer_name)
+        .await
+        .unwrap();
+
+    // Create a shared mutex that our two threads will use to coordinate access
+    let mutex = Arc::new(Mutex::new(()));
+
+    // Create two writers that will contend with each other
+    let writer1 = Arc::new(
+        LogWriter::open(
+            LogWriterOptions::default(),
+            Arc::clone(&storage),
+            prefix,
+            "writer1",
+            (),
+        )
+        .await
+        .unwrap(),
+    );
+
+    let writer2 = Arc::new(
+        LogWriter::open(
+            LogWriterOptions::default(),
+            Arc::clone(&storage),
+            prefix,
+            "writer2",
+            (),
+        )
+        .await
+        .unwrap(),
+    );
+
+    // Set a timer to make sure the test only runs for 2 minutes.
+    let fail = tokio::spawn(async move {
+        tokio::time::sleep(Duration::from_secs(120)).await;
+        eprintln!("Taking down the test");
+        std::process::exit(13);
+    });
+
+    let running = Arc::new(AtomicUsize::new(2));
+
+    // Launch both threads using the same writer_thread function
+    let handle1 = tokio::spawn(writer_thread(
+        Arc::clone(&writer1),
+        Arc::clone(&mutex),
+        Arc::clone(&running),
+        1,
+        30,
+    ));
+
+    let handle2 = tokio::spawn(writer_thread(
+        Arc::clone(&writer2),
+        Arc::clone(&mutex),
+        Arc::clone(&running),
+        2,
+        30,
+    ));
+
+    // Wait for both threads to complete
+    let (writer1_results, writer2_results) = tokio::join!(handle1, handle2);
+    fail.abort();
+
+    // Examine results
+    let (writer1_successes, writer1_contentions) = writer1_results.unwrap();
+    let (writer2_successes, writer2_contentions) = writer2_results.unwrap();
+
+    println!(
+        "Writer 1: {} successful writes, {} contentions",
+        writer1_successes, writer1_contentions
+    );
+    println!(
+        "Writer 2: {} successful writes, {} contentions",
+        writer2_successes, writer2_contentions
+    );
+
+    // Assert some things about the test
+    assert!(
+        writer1_successes > 0,
+        "Writer 1 should have some successful writes"
+    );
+    assert!(
+        writer2_successes > 0,
+        "Writer 2 should have some successful writes"
+    );
+    assert!(
+        writer1_contentions > 0,
+        "Writer 1 should encounter contention"
+    );
+    assert!(
+        writer2_contentions > 0,
+        "Writer 2 should encounter contention"
+    );
+}


Reviews:
user: login: propel-code-bot[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T15:05:58Z
user: login: rescrv
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T15:07:14Z
user: login: sanketkedia
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T17:44:52Z

Review comments:
diff_hunk: @@ -152,29 +152,36 @@ impl LogWriter {
     pub async fn append_many(&self, messages: Vec<Vec<u8>>) -> Result<LogPosition, Error> {
         // SAFETY(rescrv):  Mutex poisoning.
         let inner = { self.inner.lock().unwrap().clone() };
-        if let Some(ref epoch_writer) = inner {
+        if let Some(epoch_writer) = inner {
             let res = epoch_writer.writer.append(messages.clone()).await;
             if matches!(res, Err(Error::LogContention)) {
-                let writer = OnceLogWriter::open(
-                    epoch_writer.writer.options.clone(),
-                    epoch_writer.writer.storage.clone(),
-                    epoch_writer.writer.prefix.clone(),
-                    self.writer.clone(),
-                    Arc::clone(&self.mark_dirty),
-                )
-                .await?;
-                // SAFETY(rescrv):  Mutex poisoning.
-                let mut inner = self.inner.lock().unwrap();
-                if let Some(second) = inner.as_mut() {
-                    if second.epoch == epoch_writer.epoch {
-                        second.epoch += 1;
-                        second.writer.shutdown();
-                        second.writer = writer;
+                loop {
+                    let writer = OnceLogWriter::open(
+                        epoch_writer.writer.options.clone(),
+                        epoch_writer.writer.storage.clone(),
+                        epoch_writer.writer.prefix.clone(),
+                        self.writer.clone(),
+                        Arc::clone(&self.mark_dirty),
+                    )
+                    .await;
+                    let writer = match writer {
+                        Ok(writer) => writer,
+                        Err(Error::LogContention) => continue,
+                        Err(err) => return Err(err),
+                    };
path: rust/wal3/src/writer.rs
user: login: propel-code-bot[bot]
body: [**BestPractice**]

The writer's contention handling has been improved with a loop that retries opening a new writer, but it's missing a backoff mechanism. Without a backoff, this could lead to tight CPU-bound loops when contention is persistent. Consider adding a small backoff delay that increases with consecutive failures.
created_at: 2025-05-07T15:05:57Z
author_association: CONTRIBUTOR
diff_hunk: @@ -152,29 +152,36 @@ impl LogWriter {
     pub async fn append_many(&self, messages: Vec<Vec<u8>>) -> Result<LogPosition, Error> {
         // SAFETY(rescrv):  Mutex poisoning.
         let inner = { self.inner.lock().unwrap().clone() };
-        if let Some(ref epoch_writer) = inner {
+        if let Some(epoch_writer) = inner {
             let res = epoch_writer.writer.append(messages.clone()).await;
             if matches!(res, Err(Error::LogContention)) {
-                let writer = OnceLogWriter::open(
-                    epoch_writer.writer.options.clone(),
-                    epoch_writer.writer.storage.clone(),
-                    epoch_writer.writer.prefix.clone(),
-                    self.writer.clone(),
-                    Arc::clone(&self.mark_dirty),
-                )
-                .await?;
-                // SAFETY(rescrv):  Mutex poisoning.
-                let mut inner = self.inner.lock().unwrap();
-                if let Some(second) = inner.as_mut() {
-                    if second.epoch == epoch_writer.epoch {
-                        second.epoch += 1;
-                        second.writer.shutdown();
-                        second.writer = writer;
+                loop {
+                    let writer = OnceLogWriter::open(
+                        epoch_writer.writer.options.clone(),
+                        epoch_writer.writer.storage.clone(),
+                        epoch_writer.writer.prefix.clone(),
+                        self.writer.clone(),
+                        Arc::clone(&self.mark_dirty),
+                    )
+                    .await;
+                    let writer = match writer {
+                        Ok(writer) => writer,
+                        Err(Error::LogContention) => continue,
+                        Err(err) => return Err(err),
+                    };
path: rust/wal3/src/writer.rs
user: login: rescrv
body: The only way this happens is if the application is repeatedly contending two writers.  Our arch won't allow this.
created_at: 2025-05-07T15:07:14Z
author_association: CONTRIBUTOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-05-07T15:05:55Z
body: # Reviewer Checklist
  Please leverage this checklist to ensure your code review is thorough before approving
  ## Testing, Bugs, Errors, Logs, Documentation
  - [ ] Can you think of any use case in which the code does not behave as intended? Have they been tested?
  - [ ] Can you think of any inputs or external events that could break the code? Is user input validated and safe? Have they been tested?
  - [ ] If appropriate, are there adequate property based tests?
  - [ ] If appropriate, are there adequate unit tests?
  - [ ] Should any logging, debugging, tracing information be added or removed?
  - [ ] Are error messages user-friendly?
  - [ ] Have all documentation changes needed been made?
  - [ ] Have all non-obvious changes been commented?
  ## System Compatibility
  - [ ] Are there any potential impacts on other parts of the system or backward compatibility?
  - [ ] Does this change intersect with any items on our roadmap, and if so, is there a plan for fitting them together?
  ## Quality
  - [ ] Is this code of a unexpectedly high quality (Readability, Modularity, Intuitiveness)
author_association: NONE
user: login: propel-code-bot[bot]
created_at: 2025-05-07T15:06:05Z
body: <!-- Summary by @propel-code-bot -->

**Fix for WAL3 Hang Issues Under Log Writer Contention**

This PR addresses a bug in WAL3 where log writing could hang due to contention between multiple writers. When two WAL3 instances attempt to write to the same log, one succeeds while the other fails, but under rare circumstances this could lead to recovery failure. The fix primarily involves improving how contention is handled by implementing retry logic in the writer and simplifying the manifest recovery process.

**Key Changes:**
• Added retry loop in LogWriter.append_many() to handle LogContention errors
• Simplified the ManifestManager.recover() method to avoid state tracking issues
• Added a new test that creates contention between two writers


**Affected Areas:**
• rust/wal3/src/writer.rs
• rust/wal3/src/manifest_manager.rs


*This summary was automatically generated by @propel-code-bot*
author_association: CONTRIBUTOR