From 21ba72bc648b10c323b18a835d6d5c4c193ae172 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 29 Aug 2025 00:09:22 +0000 Subject: [PATCH] Simplify RocksDB retry logic and remove unnecessary async complexity Co-authored-by: miles.frankel --- src/storage.rs | 310 ++++++++++++------------------------------------- 1 file changed, 74 insertions(+), 236 deletions(-) diff --git a/src/storage.rs b/src/storage.rs index 61cb6af..c581bed 100644 --- a/src/storage.rs +++ b/src/storage.rs @@ -771,170 +771,62 @@ impl RetriedStorage { visibility_timeout_secs: u64, ) -> Result<()> { let inner = std::sync::Arc::clone(&self.inner); - - type OwnedBuffers = (std::sync::Arc<[u8]>, std::sync::Arc<[u8]>); - struct State { - attempt: u32, - owned: Option, - } - let state = std::sync::Arc::new(std::sync::Mutex::new(State { - attempt: 0, - owned: None, - })); - - with_rocksdb_busy_retry_async("add_available_item_from_parts", || { - let inner_ref = std::sync::Arc::clone(&inner); - let state_ref = std::sync::Arc::clone(&state); - async move { - { - let mut guard = state_ref.lock().unwrap(); - if guard.attempt > 0 && guard.owned.is_none() { - guard.owned = Some(( - std::sync::Arc::<[u8]>::from(id.to_vec()), - std::sync::Arc::<[u8]>::from(contents.to_vec()), - )); - } - } - if let Some((id_arc, contents_arc)) = { - let guard = state_ref.lock().unwrap(); - guard.owned.as_ref().cloned() - } { - tokio::task::spawn_blocking(move || { - inner_ref.add_available_item_from_parts( - &id_arc, - &contents_arc, - visibility_timeout_secs, - ) - }) - .await - .map_err(Into::::into)??; - let mut guard = state_ref.lock().unwrap(); - guard.attempt = guard.attempt.saturating_add(1); - Ok(()) - } else { - // First attempt: zero-copy direct or block_in_place depending on runtime - let res = match tokio::runtime::Handle::try_current() { - Ok(handle) - if handle.runtime_flavor() - == tokio::runtime::RuntimeFlavor::CurrentThread => - { - inner_ref.add_available_item_from_parts( - id, - contents, - visibility_timeout_secs, - ) - } - _ => tokio::task::block_in_place(move || { - inner_ref.add_available_item_from_parts( - id, - contents, - visibility_timeout_secs, - ) - }), - }; - // Increment attempt after finishing - let mut guard = state_ref.lock().unwrap(); - guard.attempt = guard.attempt.saturating_add(1); - res.map(|_| ()) - } - } + // Own inputs once to avoid per-retry allocations and allow a single blocking task + let id_owned: std::sync::Arc<[u8]> = std::sync::Arc::from(id.to_vec()); + let contents_owned: std::sync::Arc<[u8]> = std::sync::Arc::from(contents.to_vec()); + + tokio::task::spawn_blocking(move || { + with_rocksdb_busy_retry_blocking("add_available_item_from_parts", || { + inner.add_available_item_from_parts( + id_owned.as_ref(), + contents_owned.as_ref(), + visibility_timeout_secs, + ) + }) }) .await + .map_err(Into::::into)??; + Ok(()) } pub async fn add_available_items_from_parts<'a, I>(&self, items: I) -> Result<()> where I: IntoIterator, { - type BorrowedItem<'b> = (&'b [u8], (&'b [u8], u64)); type OwnedItem = (std::sync::Arc<[u8]>, (std::sync::Arc<[u8]>, u64)); - - let borrowed: Vec> = items.into_iter().collect(); + let owned: Vec = items + .into_iter() + .map(|(id, (c, v))| { + ( + std::sync::Arc::<[u8]>::from(id), + (std::sync::Arc::<[u8]>::from(c), v), + ) + }) + .collect(); let inner = std::sync::Arc::clone(&self.inner); - let borrowed_arc = std::sync::Arc::new(borrowed); - - struct State { - attempt: u32, - owned: Option>>, - } - let state = std::sync::Arc::new(std::sync::Mutex::new(State { - attempt: 0, - owned: None, - })); - - with_rocksdb_busy_retry_async("add_available_items_from_parts", || { - let inner_ref = std::sync::Arc::clone(&inner); - let state_ref = std::sync::Arc::clone(&state); - let borrowed_ref = std::sync::Arc::clone(&borrowed_arc); - async move { - { - let mut guard = state_ref.lock().unwrap(); - if guard.attempt > 0 && guard.owned.is_none() { - let owned_vec: Vec = borrowed_ref - .iter() - .map(|(id, (c, v))| { - ( - std::sync::Arc::<[u8]>::from(*id), - (std::sync::Arc::<[u8]>::from(*c), *v), - ) - }) - .collect(); - guard.owned = Some(std::sync::Arc::new(owned_vec)); - } - } - - if let Some(owned_ref) = { state_ref.lock().unwrap().owned.as_ref().cloned() } { - tokio::task::spawn_blocking(move || { - let iter = owned_ref - .iter() - .map(|(id, (c, v))| (id.as_ref(), (c.as_ref(), *v))); - inner_ref.add_available_items_from_parts(iter) - }) - .await - .map_err(Into::::into)??; - let mut guard = state_ref.lock().unwrap(); - guard.attempt = guard.attempt.saturating_add(1); - Ok(()) - } else { - let res = match tokio::runtime::Handle::try_current() { - Ok(handle) - if handle.runtime_flavor() - == tokio::runtime::RuntimeFlavor::CurrentThread => - { - let iter = borrowed_ref.iter().map(|(id, (c, v))| ((*id), ((*c), *v))); - inner_ref.add_available_items_from_parts(iter) - } - _ => { - let borrowed_ref = std::sync::Arc::clone(&borrowed_ref); - tokio::task::block_in_place(move || { - let iter = - borrowed_ref.iter().map(|(id, (c, v))| ((*id), ((*c), *v))); - inner_ref.add_available_items_from_parts(iter) - }) - } - }; - let mut guard = state_ref.lock().unwrap(); - guard.attempt = guard.attempt.saturating_add(1); - res.map(|_| ()) - } - } + tokio::task::spawn_blocking(move || { + with_rocksdb_busy_retry_blocking("add_available_items_from_parts", || { + let iter = owned + .iter() + .map(|(id, (c, v))| (id.as_ref(), (c.as_ref(), *v))); + inner.add_available_items_from_parts(iter) + }) }) .await + .map_err(Into::::into)??; + Ok(()) } pub async fn peek_next_visibility_ts_secs(&self) -> Result> { let inner = std::sync::Arc::clone(&self.inner); - with_rocksdb_busy_retry_async("peek_next_visibility_ts_secs", || { - let inner_ref = std::sync::Arc::clone(&inner); - async move { - let out = - tokio::task::spawn_blocking(move || inner_ref.peek_next_visibility_ts_secs()) - .await - .map_err(Into::::into)??; - Ok(out) - } + let out = tokio::task::spawn_blocking(move || { + with_rocksdb_busy_retry_blocking("peek_next_visibility_ts_secs", || { + inner.peek_next_visibility_ts_secs() + }) }) .await + .map_err(Into::::into)??; + Ok(out) } pub async fn get_next_available_entries( @@ -942,17 +834,14 @@ impl RetriedStorage { n: usize, ) -> Result<(Lease, Vec)> { let inner = std::sync::Arc::clone(&self.inner); - with_rocksdb_busy_retry_async("get_next_available_entries", || { - let inner_ref = std::sync::Arc::clone(&inner); - async move { - let out = - tokio::task::spawn_blocking(move || inner_ref.get_next_available_entries(n)) - .await - .map_err(Into::::into)??; - Ok(out) - } + let out = tokio::task::spawn_blocking(move || { + with_rocksdb_busy_retry_blocking("get_next_available_entries", || { + inner.get_next_available_entries(n) + }) }) .await + .map_err(Into::::into)??; + Ok(out) } pub async fn get_next_available_entries_with_lease( @@ -961,110 +850,59 @@ impl RetriedStorage { lease_validity_secs: u64, ) -> Result<(Lease, Vec)> { let inner = std::sync::Arc::clone(&self.inner); - with_rocksdb_busy_retry_async("get_next_available_entries_with_lease", || { - let inner_ref = std::sync::Arc::clone(&inner); - async move { - let out = tokio::task::spawn_blocking(move || { - inner_ref.get_next_available_entries_with_lease(n, lease_validity_secs) - }) - .await - .map_err(Into::::into)??; - Ok(out) - } + let out = tokio::task::spawn_blocking(move || { + with_rocksdb_busy_retry_blocking("get_next_available_entries_with_lease", || { + inner.get_next_available_entries_with_lease(n, lease_validity_secs) + }) }) .await + .map_err(Into::::into)??; + Ok(out) } pub async fn remove_in_progress_item(&self, id: &[u8], lease: &Lease) -> Result { let lease_copy = *lease; let inner = std::sync::Arc::clone(&self.inner); - - struct State { - attempt: u32, - owned_id: Option>, - } - let state = std::sync::Arc::new(std::sync::Mutex::new(State { - attempt: 0, - owned_id: None, - })); - - with_rocksdb_busy_retry_async("remove_in_progress_item", || { - let inner_ref = std::sync::Arc::clone(&inner); - let state_ref = std::sync::Arc::clone(&state); - async move { - { - let mut guard = state_ref.lock().unwrap(); - if guard.attempt > 0 && guard.owned_id.is_none() { - guard.owned_id = Some(std::sync::Arc::<[u8]>::from(id.to_vec())); - } - } - - if let Some(id_arc) = { state_ref.lock().unwrap().owned_id.as_ref().cloned() } { - let out = tokio::task::spawn_blocking(move || { - inner_ref.remove_in_progress_item(&id_arc, &lease_copy) - }) - .await - .map_err(Into::::into)??; - let mut guard = state_ref.lock().unwrap(); - guard.attempt = guard.attempt.saturating_add(1); - Ok(out) - } else { - let out = match tokio::runtime::Handle::try_current() { - Ok(handle) - if handle.runtime_flavor() - == tokio::runtime::RuntimeFlavor::CurrentThread => - { - inner_ref.remove_in_progress_item(id, &lease_copy) - } - _ => tokio::task::block_in_place(move || { - inner_ref.remove_in_progress_item(id, &lease_copy) - }), - }?; - let mut guard = state_ref.lock().unwrap(); - guard.attempt = guard.attempt.saturating_add(1); - Ok(out) - } - } + let id_owned: std::sync::Arc<[u8]> = std::sync::Arc::from(id.to_vec()); + let out = tokio::task::spawn_blocking(move || { + with_rocksdb_busy_retry_blocking("remove_in_progress_item", || { + inner.remove_in_progress_item(id_owned.as_ref(), &lease_copy) + }) }) .await + .map_err(Into::::into)??; + Ok(out) } pub async fn expire_due_leases(&self) -> Result { let inner = std::sync::Arc::clone(&self.inner); - with_rocksdb_busy_retry_async("expire_due_leases", || { - let inner_ref = std::sync::Arc::clone(&inner); - async move { - let out = tokio::task::spawn_blocking(move || inner_ref.expire_due_leases()) - .await - .map_err(Into::::into)??; - Ok(out) - } + let out = tokio::task::spawn_blocking(move || { + with_rocksdb_busy_retry_blocking("expire_due_leases", || inner.expire_due_leases()) }) .await + .map_err(Into::::into)??; + Ok(out) } pub async fn extend_lease(&self, lease: &Lease, lease_validity_secs: u64) -> Result { let lease = *lease; let inner = std::sync::Arc::clone(&self.inner); - with_rocksdb_busy_retry_async("extend_lease", || { - let inner_ref = std::sync::Arc::clone(&inner); - async move { - let out = tokio::task::spawn_blocking(move || { - inner_ref.extend_lease(&lease, lease_validity_secs) - }) - .await - .map_err(Into::::into)??; - Ok(out) - } + let out = tokio::task::spawn_blocking(move || { + with_rocksdb_busy_retry_blocking("extend_lease", || { + inner.extend_lease(&lease, lease_validity_secs) + }) }) .await + .map_err(Into::::into)??; + Ok(out) } } -async fn with_rocksdb_busy_retry_async(name: &str, mut f: F) -> Result +// removed unused async retry helper (consolidated to blocking variant) + +fn with_rocksdb_busy_retry_blocking(name: &str, mut f: F) -> Result where - Fut: std::future::Future>, - F: FnMut() -> Fut, + F: FnMut() -> Result, { const MAX_RETRIES: u32 = 15; const BASE_DELAY_MS: u64 = 10; @@ -1072,7 +910,7 @@ where let mut attempt: u32 = 0; loop { - match f().await { + match f() { Ok(val) => return Ok(val), Err(e) => { let is_busy = matches!( @@ -1093,8 +931,8 @@ where } else { rand::thread_rng().gen_range(0..=ceiling) }; - tracing::debug!(operation = %name, attempt = attempt + 1, delay_ms = jitter_ms, "RocksDB Busy: backing off with jitter"); - tokio::time::sleep(std::time::Duration::from_millis(jitter_ms)).await; + tracing::debug!(operation = %name, attempt = attempt + 1, delay_ms = jitter_ms, "RocksDB Busy: backing off with jitter (blocking)"); + std::thread::sleep(std::time::Duration::from_millis(jitter_ms)); attempt = attempt.saturating_add(1); } }