From f97b8117eaa38d6a0c88e32fafe65a8e278bfa8d Mon Sep 17 00:00:00 2001 From: Robert Escriva Date: Wed, 2 Sep 2026 10:23:19 -0700 Subject: [PATCH 1/3] [ENH](wal3): pin fragment generations for conditional writes S3-backed conditional writes previously validated the committed log against a moving target: a concurrent publish could open and close a fragment generation between validation and append, forcing a retry. Introduce FragmentPin so a caller can hold the next fragment generation open while it reads the manifest and prepares its append, ensuring it validates against the generation its append will join. Add FragmentPin, a generation-scoped, must-use reservation: - acquire_fragment_pin waits for the active generation to finish, then reserves bytes against the fragment size limit and returns a pin. - Work submitted with a pin joins the pinned generation; dropping a pin releases its reservation so cancelled or invalidated requests make progress without an explicit cleanup call. - take_work holds selected work until every pin submits or drops, and bumps pin_generation when a batch is selected for write. Thread the pin through push_work, append_many_with_options, and LogWriterTrait, and add acquire_fragment_pin to the trait and its S3/replicated implementations. Supplying a pin disables transparent contention retries because a pin cannot transfer to a recovered writer generation. In log-service, S3-backed conditional pushes now acquire a fragment pin before validation and record the wait in a new conditional_write_fragment_pin_wait_us histogram; replicated logs stay serialized in Spanner and skip pinning. Factor the repeated writer shutdown blocks in handle_errors_and_contention into shutdown_epoch. Co-authored-by: AI --- rust/log-service/src/lib.rs | 213 +++++---- rust/wal3/src/interfaces/batch_manager.rs | 535 +++++++++++++++++++--- rust/wal3/src/interfaces/mod.rs | 30 +- rust/wal3/src/lib.rs | 26 +- rust/wal3/src/writer.rs | 109 +++-- 5 files changed, 712 insertions(+), 201 deletions(-) diff --git a/rust/log-service/src/lib.rs b/rust/log-service/src/lib.rs index cf5cd36c659..6e082e29966 100644 --- a/rust/log-service/src/lib.rs +++ b/rust/log-service/src/lib.rs @@ -481,6 +481,8 @@ pub struct Metrics { conditional_write_admission_rejections: opentelemetry::metrics::Counter, /// The number of conditional push retries after a stale required fragment start. conditional_write_required_start_retries: opentelemetry::metrics::Counter, + /// Time spent waiting to pin the next fragment generation. + conditional_write_fragment_pin_wait_us: opentelemetry::metrics::Histogram, /// The number of successful conditional writes with an inserted offset. conditional_write_success_with_offset: opentelemetry::metrics::Counter, /// The number of successful conditional writes with durable contention and no offset. @@ -528,6 +530,9 @@ impl Metrics { conditional_write_required_start_retries: meter .u64_counter("conditional_write_required_start_retries") .build(), + conditional_write_fragment_pin_wait_us: meter + .u64_histogram("conditional_write_fragment_pin_wait_us") + .build(), conditional_write_success_with_offset: meter .u64_counter("conditional_write_success_with_offset") .build(), @@ -2553,100 +2558,119 @@ impl LogServer { messages.push(buf); } let record_count = messages.len() as i32; - let first_inserted_record_offset = - if let Some(conditional_write) = conditional_write.as_ref() { - let admission_predicate = conditional_admission_predicate(conditional_write); - let mut append_result = None; - let mut validation_start = - LogPosition::from_offset(conditional_write.observed_log_offset); - let conditional_push_max_retries = self.config.conditional_push_retry_attempts(); - for attempt in 0..conditional_push_max_retries { - let validated_tail = self - .validate_committed_log_for_conditional_write( - topology_name.as_ref(), - collection_id, - conditional_write, - validation_start, - ) - .await?; - let append_options = AppendOptions::new(write_id_metadata.clone()) - .with_admission_predicate(Arc::clone(&admission_predicate)) - .with_required_fragment_start(validated_tail); - match log - .append_many_with_options(messages.clone(), Some(append_options)) - .await - { - Ok(offset) => { - self.metrics - .conditional_write_success_with_offset - .add(1, &[]); - tracing::debug!( - %collection_id, - first_inserted_record_offset = offset.offset(), - "conditional write append succeeded" - ); - append_result = Some(Some(offset)); - break; - } - Err(wal3::Error::LogContentionDurable) => { - self.metrics - .conditional_write_success_without_offset - .add(1, &[]); - tracing::debug!( - %collection_id, - "conditional write append durably contended" - ); - append_result = Some(None); - break; - } - Err(wal3::Error::AdmissionRejected) => { - self.metrics - .conditional_write_admission_rejections - .add(1, &[]); - self.metrics - .conditional_write_in_flight_conflicts - .add(1, &[]); - tracing::info!( - %collection_id, - "conditional write rejected by in-flight admission" - ); - return Err(Status::aborted(CONDITIONAL_WRITE_CONFLICT_MESSAGE)); - } - Err(wal3::Error::LogContentionRetry) - if attempt + 1 < conditional_push_max_retries => - { - validation_start = validated_tail; - self.metrics - .conditional_write_required_start_retries - .add(1, &[]); - tracing::info!( - %collection_id, - attempt = attempt + 1, - "conditional write missed required start; retrying validation" - ); - } - Err(err) => return Err(push_append_error_to_status(err)), - } - } - match append_result { - Some(append_result) => append_result, - None => { - return Err(Status::internal( - "conditional push retry loop exited without an append result", - )); - } - } - } else { - let append_options = AppendOptions::new(write_id_metadata); + let message_bytes = messages.iter().map(Vec::len).sum::(); + let first_inserted_record_offset = if let Some(conditional_write) = + conditional_write.as_ref() + { + let admission_predicate = conditional_admission_predicate(conditional_write); + let mut append_result = None; + let mut validation_start = + LogPosition::from_offset(conditional_write.observed_log_offset); + let conditional_push_max_retries = self.config.conditional_push_retry_attempts(); + for attempt in 0..conditional_push_max_retries { + // Replicated logs serialize conditional writes in Spanner. S3-backed logs pin + // their in-memory fragment before validation so concurrent transactions can join + // the same publish. + let fragment_pin = if topology_name.is_none() { + let pin_started = Instant::now(); + let pin_result = log + .acquire_fragment_pin(message_bytes) + .instrument(tracing::info_span!("acquire_conditional_fragment_pin")) + .await; + self.metrics.conditional_write_fragment_pin_wait_us.record( + u64::try_from(pin_started.elapsed().as_micros()).unwrap_or(u64::MAX), + &[], + ); + Some(pin_result.map_err(push_append_error_to_status)?) + } else { + None + }; + let validated_tail = self + .validate_committed_log_for_conditional_write( + topology_name.as_ref(), + collection_id, + conditional_write, + validation_start, + ) + .await?; + let append_options = AppendOptions::new(write_id_metadata.clone()) + .with_admission_predicate(Arc::clone(&admission_predicate)) + .with_required_fragment_start(validated_tail); match log - .append_many_with_options(messages, Some(append_options)) + .append_many_with_options(messages.clone(), Some(append_options), fragment_pin) .await { - Ok(offset) => Some(offset), - Err(wal3::Error::LogContentionDurable) => None, + Ok(offset) => { + self.metrics + .conditional_write_success_with_offset + .add(1, &[]); + tracing::debug!( + %collection_id, + first_inserted_record_offset = offset.offset(), + "conditional write append succeeded" + ); + append_result = Some(Some(offset)); + break; + } + Err(wal3::Error::LogContentionDurable) => { + self.metrics + .conditional_write_success_without_offset + .add(1, &[]); + tracing::debug!( + %collection_id, + "conditional write append durably contended" + ); + append_result = Some(None); + break; + } + Err(wal3::Error::AdmissionRejected) => { + self.metrics + .conditional_write_admission_rejections + .add(1, &[]); + self.metrics + .conditional_write_in_flight_conflicts + .add(1, &[]); + tracing::info!( + %collection_id, + "conditional write rejected by in-flight admission" + ); + return Err(Status::aborted(CONDITIONAL_WRITE_CONFLICT_MESSAGE)); + } + Err(wal3::Error::LogContentionRetry) + if attempt + 1 < conditional_push_max_retries => + { + validation_start = validated_tail; + self.metrics + .conditional_write_required_start_retries + .add(1, &[]); + tracing::info!( + %collection_id, + attempt = attempt + 1, + "conditional write missed required start; retrying validation" + ); + } Err(err) => return Err(push_append_error_to_status(err)), } - }; + } + match append_result { + Some(append_result) => append_result, + None => { + return Err(Status::internal( + "conditional push retry loop exited without an append result", + )); + } + } + } else { + let append_options = AppendOptions::new(write_id_metadata); + match log + .append_many_with_options(messages, Some(append_options), None) + .await + { + Ok(offset) => Some(offset), + Err(wal3::Error::LogContentionDurable) => None, + Err(err) => return Err(push_append_error_to_status(err)), + } + }; let first_inserted_record_offset = first_inserted_record_offset .map(|offset| { offset.offset().try_into().map_err(|_| { @@ -6753,18 +6777,27 @@ mod tests { #[async_trait::async_trait] impl LogWriterTrait for FakeLogWriter { + async fn acquire_fragment_pin( + &self, + _reserved_bytes: usize, + ) -> Result { + unreachable!("fake log writer does not pin fragments") + } + async fn append_with_options( &self, message: Vec, options: Option, ) -> Result { - self.append_many_with_options(vec![message], options).await + self.append_many_with_options(vec![message], options, None) + .await } async fn append_many_with_options( &self, _messages: Vec>, options: Option, + _pin: Option, ) -> Result { self.observed_options.lock().push(options); let mut append_results = self.append_results.lock(); diff --git a/rust/wal3/src/interfaces/batch_manager.rs b/rust/wal3/src/interfaces/batch_manager.rs index fcf09e0576c..716fbb82692 100644 --- a/rust/wal3/src/interfaces/batch_manager.rs +++ b/rust/wal3/src/interfaces/batch_manager.rs @@ -24,6 +24,9 @@ struct ManagerState { backoff: bool, next_write: Instant, writers_active: usize, + pin_generation: u64, + fragment_pins: usize, + pinned_bytes: usize, enqueued: Vec, admission_metadata: Vec>>, tearing_down: bool, @@ -67,6 +70,106 @@ impl ManagerState { fn finish_write(&mut self) { self.writers_active -= 1; } + + fn occupied_bytes(&self) -> usize { + self.enqueued + .iter() + .map(AppendWork::byte_count) + .sum::() + .saturating_add(self.pinned_bytes) + } + + fn can_reserve(&self, byte_count: usize, batch_size_bytes: usize) -> bool { + let occupied_bytes = self.occupied_bytes(); + occupied_bytes == 0 || occupied_bytes.saturating_add(byte_count) < batch_size_bytes + } +} + +//////////////////////////////////////// BatchCoordination //////////////////////////////////////// + +#[derive(Debug)] +struct BatchCoordination { + state: Mutex, + write_finished: tokio::sync::Notify, +} + +/////////////////////////////////////////// FragmentPin /////////////////////////////////////////// + +/// Keeps the next fragment open while an append is prepared. +/// +/// Pins are generation-scoped. Work submitted with a pin joins the fragment generation that was +/// open when the pin was acquired. Dropping a pin without submitting work releases it, allowing +/// validation failures and cancelled requests to make progress without an explicit cleanup call. +#[must_use = "dropping the pin allows its fragment generation to publish"] +pub struct FragmentPin { + coordination: Option>, + generation: u64, + reserved_bytes: usize, +} + +impl FragmentPin { + fn belongs_to(&self, coordination: &Arc) -> bool { + self.coordination + .as_ref() + .is_some_and(|pinned| Arc::ptr_eq(pinned, coordination)) + } + + fn reserves(&self, byte_count: usize) -> bool { + byte_count <= self.reserved_bytes + } + + fn consume( + mut self, + coordination: &Arc, + state: &mut ManagerState, + ) -> Result<(), Error> { + let Some(pinned) = self.coordination.take() else { + tracing::error!("consuming an already released wal3 fragment pin"); + return Err(Error::LogContentionRetry); + }; + if !Arc::ptr_eq(&pinned, coordination) || self.generation != state.pin_generation { + drop(pinned); + return Err(Error::LogContentionRetry); + } + release_pin_from_state(state, self.reserved_bytes); + Ok(()) + } +} + +impl std::fmt::Debug for FragmentPin { + fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + fmt.debug_struct("FragmentPin") + .field("generation", &self.generation) + .field("reserved_bytes", &self.reserved_bytes) + .field("released", &self.coordination.is_none()) + .finish() + } +} + +impl Drop for FragmentPin { + fn drop(&mut self) { + let Some(coordination) = self.coordination.take() else { + return; + }; + let mut state = coordination.state.lock().unwrap(); + if self.generation == state.pin_generation && state.fragment_pins > 0 { + release_pin_from_state(&mut state, self.reserved_bytes); + } else { + tracing::error!( + pin_generation = self.generation, + current_generation = state.pin_generation, + fragment_pins = state.fragment_pins, + "dropping stale wal3 fragment pin" + ); + } + drop(state); + coordination.write_finished.notify_one(); + } +} + +fn release_pin_from_state(state: &mut ManagerState, reserved_bytes: usize) { + state.fragment_pins -= 1; + state.pinned_bytes = state.pinned_bytes.saturating_sub(reserved_bytes); } fn count_compatible_required_fragment_starts( @@ -141,8 +244,7 @@ pub struct BatchManager> { options: LogWriterOptions, fragment_uploader: U, _fp_phantom: std::marker::PhantomData, - state: Mutex, - write_finished: tokio::sync::Notify, + coordination: Arc, } impl> BatchManager { @@ -152,29 +254,36 @@ impl> BatchManager { fragment_uploader, _fp_phantom: std::marker::PhantomData, options, - state: Mutex::new(ManagerState { - backoff: false, - next_write, - writers_active: 0, - enqueued: Vec::new(), - admission_metadata: Vec::new(), - tearing_down: false, + coordination: Arc::new(BatchCoordination { + state: Mutex::new(ManagerState { + backoff: false, + next_write, + writers_active: 0, + pin_generation: 0, + fragment_pins: 0, + pinned_bytes: 0, + enqueued: Vec::new(), + admission_metadata: Vec::new(), + tearing_down: false, + }), + write_finished: tokio::sync::Notify::new(), }), - write_finished: tokio::sync::Notify::new(), }) } pub fn count_waiters(&self) -> usize { - let state = self.state.lock().unwrap(); + let state = self.coordination.state.lock().unwrap(); state.enqueued.len() } pub fn debug_dump(&self) -> String { let mut output = "[batch manager]\n".to_string(); - let state = self.state.lock().unwrap(); + let state = self.coordination.state.lock().unwrap(); output += &format!("backoff: {:?}\n", state.backoff); output += &format!("next_write: {:?}\n", state.next_write); output += &format!("writers_active: {:?}\n", state.writers_active); + output += &format!("fragment_pins: {:?}\n", state.fragment_pins); + output += &format!("pinned_bytes: {:?}\n", state.pinned_bytes); output += &format!("enqueued: {}\n", state.enqueued.len()); output } @@ -192,16 +301,72 @@ impl> std::fmt::Debug for BatchMana impl> FragmentPublisher for BatchManager { type FragmentPointer = FP; + async fn acquire_fragment_pin(&self, reserved_bytes: usize) -> Result { + loop { + let notified = self.coordination.write_finished.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + { + // SAFETY(rescrv): Mutex poisoning. + let mut state = self.coordination.state.lock().unwrap(); + if state.tearing_down { + return Err(Error::LogContentionRetry); + } + if state.backoff { + return Err(Error::Backoff); + } + if state.writers_active == 0 { + if !state.can_reserve( + reserved_bytes, + self.options.throttle_fragment.batch_size_bytes, + ) { + return Err(Error::Backoff); + } + state.fragment_pins += 1; + state.pinned_bytes = state.pinned_bytes.saturating_add(reserved_bytes); + return Ok(FragmentPin { + coordination: Some(Arc::clone(&self.coordination)), + generation: state.pin_generation, + reserved_bytes, + }); + } + } + notified.await; + } + } + /// Enqueue work to be published. - async fn push_work(&self, work: AppendWork) { + async fn push_work(&self, work: AppendWork, pin: Option) { + let is_pinned = pin.is_some(); + if pin.as_ref().is_some_and(|pin| { + !pin.belongs_to(&self.coordination) || !pin.reserves(work.byte_count()) + }) { + let _ = work.tx.send(Err(Error::LogContentionRetry)); + return; + } + // SAFETY(rescrv): Mutex poisoning. - let mut state = self.state.lock().unwrap(); + let mut state = self.coordination.state.lock().unwrap(); + if let Some(pin) = pin { + if let Err(err) = pin.consume(&self.coordination, &mut state) { + let _ = work.tx.send(Err(err)); + self.coordination.write_finished.notify_one(); + return; + } + } if state.tearing_down { let _ = work.tx.send(Err(Error::LogContentionRetry)); - self.write_finished.notify_one(); - } else if state.backoff { + self.coordination.write_finished.notify_one(); + } else if state.backoff + || (!is_pinned + && state.fragment_pins > 0 + && !state.can_reserve( + work.byte_count(), + self.options.throttle_fragment.batch_size_bytes, + )) + { let _ = work.tx.send(Err(Error::Backoff)); - self.write_finished.notify_one(); + self.coordination.write_finished.notify_one(); } else { if let Some(admission_predicate) = work .options @@ -210,7 +375,7 @@ impl> FragmentPublisher for BatchMa { if !admission_predicate(&state.admission_metadata) { let _ = work.tx.send(Err(Error::AdmissionRejected)); - self.write_finished.notify_one(); + self.coordination.write_finished.notify_one(); return; } } @@ -225,13 +390,13 @@ impl> FragmentPublisher for BatchMa manifest_manager: &(dyn ManifestPublisher + Sync), ) -> Result, Vec)>, Error> { // SAFETY(rescrv): Mutex poisoning. - let mut state = self.state.lock().unwrap(); + let mut state = self.coordination.state.lock().unwrap(); // We're shutting down. Throw the work away. if state.tearing_down { state.enqueued.clear(); state.admission_metadata.clear(); - self.write_finished.notify_one(); + self.coordination.write_finished.notify_one(); return Ok(None); } @@ -241,6 +406,9 @@ impl> FragmentPublisher for BatchMa return Ok(None); } debug_assert_eq!(state.enqueued.len(), state.admission_metadata.len()); + if state.fragment_pins > 0 { + return Ok(None); + } let mut split_off = 0usize; let mut acc_count = 0usize; @@ -265,12 +433,12 @@ impl> FragmentPublisher for BatchMa if !did_split && state.next_write > Instant::now() { // This notify makes sure the background picks up the work and makes progress at end of // the batching interval. - self.write_finished.notify_one(); + self.coordination.write_finished.notify_one(); return Ok(None); } if split_off == 0 { // No work to do. - self.write_finished.notify_one(); + self.coordination.write_finished.notify_one(); return Ok(None); } let (compatible_split_off, required_fragment_start) = @@ -288,6 +456,7 @@ impl> FragmentPublisher for BatchMa // Cannot yet select for write. Notify will come from the timeout background is on. return Ok(None); }; + state.pin_generation = state.pin_generation.wrapping_add(1); if let (Some(required_fragment_start), Some(assigned_fragment_start)) = (required_fragment_start, pointer.fragment_start()) { @@ -301,7 +470,7 @@ impl> FragmentPublisher for BatchMa let work = take_selected_work(&mut state, split_off); state.finish_write(); if refresh_backoff_after_split(&mut state, &self.options) { - self.write_finished.notify_one(); + self.coordination.write_finished.notify_one(); } reject_selected_work(work, Error::LogContentionRetry); return Ok(None); @@ -309,26 +478,27 @@ impl> FragmentPublisher for BatchMa } let work = take_selected_work(&mut state, split_off); if refresh_backoff_after_split(&mut state, &self.options) { - self.write_finished.notify_one(); + self.coordination.write_finished.notify_one(); } Ok(Some((pointer, required_fragment_start, work))) } /// Finish the previous call to take_work. async fn finish_write(&self) { - self.state.lock().unwrap().finish_write(); - self.write_finished.notify_one(); + self.coordination.state.lock().unwrap().finish_write(); + self.coordination.write_finished.notify_waiters(); + self.coordination.write_finished.notify_one(); } /// Wait until take_work might have work. async fn wait_for_writable(&self) { - self.write_finished.notified().await; + self.coordination.write_finished.notified().await; } /// How long to sleep until take work might have work. fn until_next_time(&self) -> Duration { // SAFETY(rescrv): Mutex poisoning. - let state = self.state.lock().unwrap(); + let state = self.coordination.state.lock().unwrap(); let now = Instant::now(); if now < state.next_write { state.next_write - now @@ -378,7 +548,7 @@ impl> FragmentPublisher for BatchMa /// Start shutting down. The shutdown is split for historical and unprincipled reasons. fn shutdown_prepare(&self) { let enqueued = { - let mut state = self.state.lock().unwrap(); + let mut state = self.coordination.state.lock().unwrap(); state.tearing_down = true; state.admission_metadata.clear(); std::mem::take(&mut state.enqueued) @@ -390,7 +560,8 @@ impl> FragmentPublisher for BatchMa /// Finish shutting down. fn shutdown_finish(&self) { - self.write_finished.notify_one(); + self.coordination.write_finished.notify_waiters(); + self.coordination.write_finished.notify_one(); } async fn write_garbage( @@ -658,12 +829,10 @@ mod tests { ) -> tokio::sync::oneshot::Receiver> { let (tx, rx) = tokio::sync::oneshot::channel(); batch_manager - .push_work(AppendWork::new( - vec![message], - options, - tx, - tracing::Span::current(), - )) + .push_work( + AppendWork::new(vec![message], options, tx, tracing::Span::current()), + None, + ) .await; rx } @@ -675,12 +844,42 @@ mod tests { ) -> tokio::sync::oneshot::Receiver> { let (tx, rx) = tokio::sync::oneshot::channel(); batch_manager - .push_work(AppendWork::new( - vec![message], - options, - tx, - tracing::Span::current(), - )) + .push_work( + AppendWork::new(vec![message], options, tx, tracing::Span::current()), + None, + ) + .await; + rx + } + + async fn enqueue_s3_pinned( + batch_manager: &BatchManager<(FragmentSeqNo, LogPosition), NoopS3Uploader>, + message: &str, + required_fragment_start: u64, + pin: FragmentPin, + ) -> tokio::sync::oneshot::Receiver> { + let options = append_options(&[message]) + .with_required_fragment_start(LogPosition::from_offset(required_fragment_start)); + enqueue_s3_pinned_with_options(batch_manager, message, options, pin).await + } + + async fn enqueue_s3_pinned_with_options( + batch_manager: &BatchManager<(FragmentSeqNo, LogPosition), NoopS3Uploader>, + message: &str, + options: AppendOptions, + pin: FragmentPin, + ) -> tokio::sync::oneshot::Receiver> { + let (tx, rx) = tokio::sync::oneshot::channel(); + batch_manager + .push_work( + AppendWork::new( + vec![Vec::from(message)], + Some(options), + tx, + tracing::Span::current(), + ), + Some(pin), + ) .await; rx } @@ -974,6 +1173,227 @@ mod tests { assert_eq!(2, batch_manager.count_waiters()); } + #[tokio::test] + async fn fragment_pins_hold_compatible_work_until_every_pin_submits() { + let batch_manager = immediate_s3_batch_manager(); + let first_pin = batch_manager.acquire_fragment_pin(5).await.unwrap(); + let second_pin = batch_manager.acquire_fragment_pin(6).await.unwrap(); + let mut first_rx = enqueue_s3_pinned(&batch_manager, "first", 10, first_pin).await; + + let selected = batch_manager + .take_work(&FixedS3ManifestPublisher { + assigned_fragment_start: LogPosition::from_offset(10), + }) + .await + .expect("pinned work selection should not fail"); + assert!(selected.is_none()); + assert!(matches!( + first_rx.try_recv(), + Err(tokio::sync::oneshot::error::TryRecvError::Empty) + )); + + let second_rx = enqueue_s3_pinned(&batch_manager, "second", 10, second_pin).await; + let work = assert_s3_work( + take_s3_work(&batch_manager, 10).await, + 10, + Some(10), + &["first", "second"], + ); + commit_s3_work(&batch_manager, LogPosition::from_offset(10), work).await; + expect_committed(first_rx, 10).await; + expect_committed(second_rx, 11).await; + } + + #[tokio::test] + async fn dropping_fragment_pin_unblocks_submitted_work() { + let batch_manager = immediate_s3_batch_manager(); + let submitted_pin = batch_manager.acquire_fragment_pin(5).await.unwrap(); + let cancelled_pin = batch_manager.acquire_fragment_pin(9).await.unwrap(); + let submitted_rx = enqueue_s3_pinned(&batch_manager, "first", 10, submitted_pin).await; + + drop(cancelled_pin); + + let work = assert_s3_work( + take_s3_work(&batch_manager, 10).await, + 10, + Some(10), + &["first"], + ); + commit_s3_work(&batch_manager, LogPosition::from_offset(10), work).await; + expect_committed(submitted_rx, 10).await; + } + + #[tokio::test] + async fn pinned_admission_rejection_unblocks_compatible_work() { + let batch_manager = immediate_s3_batch_manager(); + let accepted_pin = batch_manager.acquire_fragment_pin(5).await.unwrap(); + let rejected_pin = batch_manager.acquire_fragment_pin(6).await.unwrap(); + let accepted_rx = enqueue_s3_pinned(&batch_manager, "first", 10, accepted_pin).await; + let predicate = Arc::new(|earlier_metadata: &[Vec>]| { + !earlier_metadata + .iter() + .flatten() + .any(|metadata| metadata.as_ref() == b"first") + }); + let rejected_options = append_options(&["second"]) + .with_admission_predicate(predicate) + .with_required_fragment_start(LogPosition::from_offset(10)); + let rejected_rx = enqueue_s3_pinned_with_options( + &batch_manager, + "second", + rejected_options, + rejected_pin, + ) + .await; + + assert!(matches!( + rejected_rx.await.expect("rejected append should complete"), + Err(Error::AdmissionRejected) + )); + let work = assert_s3_work( + take_s3_work(&batch_manager, 10).await, + 10, + Some(10), + &["first"], + ); + commit_s3_work(&batch_manager, LogPosition::from_offset(10), work).await; + expect_committed(accepted_rx, 10).await; + } + + #[tokio::test] + async fn fragment_pin_waits_for_publishing_generation() { + let batch_manager = immediate_s3_batch_manager(); + let active_rx = enqueue_s3_with_required_start(&batch_manager, "active", Some(10)).await; + let work = take_s3_work(&batch_manager, 10).await; + let work = assert_s3_work(work, 10, Some(10), &["active"]); + + assert!(tokio::time::timeout( + Duration::from_millis(1), + batch_manager.acquire_fragment_pin(4) + ) + .await + .is_err()); + + commit_s3_work(&batch_manager, LogPosition::from_offset(10), work).await; + expect_committed(active_rx, 10).await; + let next_pin = tokio::time::timeout( + Duration::from_secs(1), + batch_manager.acquire_fragment_pin(4), + ) + .await + .expect("next generation should become pinnable") + .expect("pin acquisition should succeed"); + drop(next_pin); + } + + #[tokio::test] + async fn publishing_generation_wakes_every_fragment_pin_waiter() { + let batch_manager = Arc::new(immediate_s3_batch_manager()); + let active_rx = enqueue_s3_with_required_start(&batch_manager, "active", Some(10)).await; + let work = assert_s3_work( + take_s3_work(&batch_manager, 10).await, + 10, + Some(10), + &["active"], + ); + let (first_ready_tx, first_ready_rx) = tokio::sync::oneshot::channel(); + let first_waiter = { + let batch_manager = Arc::clone(&batch_manager); + tokio::spawn(async move { + first_ready_tx.send(()).unwrap(); + batch_manager.acquire_fragment_pin(4).await + }) + }; + first_ready_rx.await.unwrap(); + let (second_ready_tx, second_ready_rx) = tokio::sync::oneshot::channel(); + let second_waiter = { + let batch_manager = Arc::clone(&batch_manager); + tokio::spawn(async move { + second_ready_tx.send(()).unwrap(); + batch_manager.acquire_fragment_pin(4).await + }) + }; + second_ready_rx.await.unwrap(); + + commit_s3_work(&batch_manager, LogPosition::from_offset(10), work).await; + expect_committed(active_rx, 10).await; + + let first_pin = tokio::time::timeout(Duration::from_secs(1), first_waiter) + .await + .expect("first waiter should wake") + .expect("first waiter should not panic") + .expect("first waiter should acquire a pin"); + let second_pin = tokio::time::timeout(Duration::from_secs(1), second_waiter) + .await + .expect("second waiter should wake") + .expect("second waiter should not panic") + .expect("second waiter should acquire a pin"); + drop(first_pin); + drop(second_pin); + } + + #[tokio::test] + async fn fragment_pin_reservations_honor_fragment_size_limit() { + let options = LogWriterOptions { + throttle_fragment: ThrottleOptions { + batch_size_bytes: 10, + throughput: 2_000_000, + batch_interval_us: 0, + ..ThrottleOptions::default() + }, + ..LogWriterOptions::default() + }; + let batch_manager = + BatchManager::<(FragmentSeqNo, LogPosition), _>::new(options, NoopS3Uploader).unwrap(); + let first_pin = batch_manager.acquire_fragment_pin(6).await.unwrap(); + + let second_result = batch_manager.acquire_fragment_pin(4).await; + + assert!(matches!(second_result, Err(Error::Backoff))); + drop(first_pin); + } + + #[tokio::test] + async fn fragment_pin_reservation_rejects_unpinned_work_that_would_fill_batch() { + let options = LogWriterOptions { + throttle_fragment: ThrottleOptions { + batch_size_bytes: 10, + throughput: 2_000_000, + batch_interval_us: 0, + ..ThrottleOptions::default() + }, + ..LogWriterOptions::default() + }; + let batch_manager = + BatchManager::<(FragmentSeqNo, LogPosition), _>::new(options, NoopS3Uploader).unwrap(); + let pin = batch_manager.acquire_fragment_pin(6).await.unwrap(); + + let unpinned_rx = enqueue_s3(&batch_manager, b"four".to_vec(), None).await; + + assert!(matches!( + unpinned_rx.await.expect("unpinned append should complete"), + Err(Error::Backoff) + )); + assert_eq!(0, batch_manager.count_waiters()); + drop(pin); + } + + #[tokio::test] + async fn fragment_pin_rejects_work_larger_than_its_reservation() { + let batch_manager = immediate_s3_batch_manager(); + let undersized_pin = batch_manager.acquire_fragment_pin(4).await.unwrap(); + + let rejected_rx = enqueue_s3_pinned(&batch_manager, "first", 10, undersized_pin).await; + + assert!(matches!( + rejected_rx.await.expect("rejected append should complete"), + Err(Error::LogContentionRetry) + )); + assert_eq!(0, batch_manager.count_waiters()); + let replacement_pin = batch_manager.acquire_fragment_pin(5).await.unwrap(); + drop(replacement_pin); + } + #[tokio::test] async fn mismatched_required_fragment_starts_split_selected_batch() { let batch_manager = immediate_s3_batch_manager(); @@ -1172,30 +1592,29 @@ mod tests { let (tx, _rx1) = tokio::sync::oneshot::channel(); let options1 = append_options(&["metadata-1"]); batch_manager - .push_work(AppendWork::new( - vec![vec![1]], - Some(options1.clone()), - tx, - tracing::Span::current(), - )) + .push_work( + AppendWork::new( + vec![vec![1]], + Some(options1.clone()), + tx, + tracing::Span::current(), + ), + None, + ) .await; let (tx, _rx2) = tokio::sync::oneshot::channel(); batch_manager - .push_work(AppendWork::new( - vec![vec![2, 3]], + .push_work( + AppendWork::new(vec![vec![2, 3]], None, tx, tracing::Span::current()), None, - tx, - tracing::Span::current(), - )) + ) .await; let (tx, _rx3) = tokio::sync::oneshot::channel(); batch_manager - .push_work(AppendWork::new( - vec![vec![4, 5, 6]], + .push_work( + AppendWork::new(vec![vec![4, 5, 6]], None, tx, tracing::Span::current()), None, - tx, - tracing::Span::current(), - )) + ) .await; let ((seq_no, log_position), required_fragment_start, work) = batch_manager .take_work(&manifest_manager) diff --git a/rust/wal3/src/interfaces/mod.rs b/rust/wal3/src/interfaces/mod.rs index c6612881432..e2e3477de4e 100644 --- a/rust/wal3/src/interfaces/mod.rs +++ b/rust/wal3/src/interfaces/mod.rs @@ -20,7 +20,7 @@ pub mod batch_manager; pub mod repl; pub mod s3; -pub use batch_manager::BatchManager; +pub use batch_manager::{BatchManager, FragmentPin}; ////////////////////////////////////////// FragmentPointer ///////////////////////////////////////// @@ -268,10 +268,17 @@ where { type FragmentPointer = FP; - async fn push_work(&self, work: AppendWork) { + async fn acquire_fragment_pin(&self, reserved_bytes: usize) -> Result { match self { - Self::Plain(publisher) => publisher.push_work(work).await, - Self::FaultInjecting(publisher) => publisher.push_work(work).await, + Self::Plain(publisher) => publisher.acquire_fragment_pin(reserved_bytes).await, + Self::FaultInjecting(publisher) => publisher.acquire_fragment_pin(reserved_bytes).await, + } + } + + async fn push_work(&self, work: AppendWork, pin: Option) { + match self { + Self::Plain(publisher) => publisher.push_work(work, pin).await, + Self::FaultInjecting(publisher) => publisher.push_work(work, pin).await, } } @@ -561,8 +568,19 @@ impl std::fmt::Debug for AppendWork { pub trait FragmentPublisher: Send + Sync + 'static { type FragmentPointer: FragmentPointer; - /// Enqueue work to be published. - async fn push_work(&self, work: AppendWork); + /// Keep the next fragment generation open while an append is prepared. + /// + /// `reserved_bytes` counts toward the fragment size limit until the pin is submitted or + /// dropped. + /// + /// # Errors + /// + /// Returns an error if the publisher is shutting down, applying backpressure, or cannot fit + /// the reservation in the open generation. + async fn acquire_fragment_pin(&self, reserved_bytes: usize) -> Result; + + /// Enqueue work to be published, consuming a fragment pin atomically when one is supplied. + async fn push_work(&self, work: AppendWork, pin: Option); /// Take enqueued work to be published. async fn take_work( &self, diff --git a/rust/wal3/src/lib.rs b/rust/wal3/src/lib.rs index 3c49e041eca..e1d356c1a40 100644 --- a/rust/wal3/src/lib.rs +++ b/rust/wal3/src/lib.rs @@ -36,7 +36,7 @@ pub use interfaces::s3::{ }; pub use interfaces::{ fragment_upload_replica_fault_label, AppendWork, BatchManager, - FaultInjectingFragmentManagerFactory, FragmentConsumer, FragmentManagerFactory, + FaultInjectingFragmentManagerFactory, FragmentConsumer, FragmentManagerFactory, FragmentPin, FragmentPointer, FragmentPublisher, FragmentUploadFault, FragmentUploadFaultInjector, FragmentUploader, ManifestConsumer, ManifestManagerFactory, ManifestPublisher, ManifestWitness, PositionWitness, FRAGMENT_UPLOAD_FAULT_LABEL, FRAGMENT_UPLOAD_REPLICA_FAULT_LABELS, @@ -1018,6 +1018,13 @@ pub fn fragment_path(prefix: &str, path: &str) -> String { /// Trait that provides a type-erased interface to LogWriter. #[async_trait::async_trait] pub trait LogWriterTrait: std::fmt::Debug + Send + Sync + 'static { + /// Keep the next fragment open while an append is prepared. + /// + /// # Errors + /// + /// Returns an error when the writer cannot reserve space in its next fragment generation. + async fn acquire_fragment_pin(&self, reserved_bytes: usize) -> Result; + /// Append a single message to the log. async fn append(&self, message: Vec) -> Result { self.append_with_options(message, None).await @@ -1032,14 +1039,20 @@ pub trait LogWriterTrait: std::fmt::Debug + Send + Sync + 'static { /// Append multiple messages to the log atomically. async fn append_many(&self, messages: Vec>) -> Result { - self.append_many_with_options(messages, None).await + self.append_many_with_options(messages, None, None).await } - /// Append multiple messages to the log atomically with options. + /// Append multiple messages atomically, optionally to a pinned fragment generation. + /// + /// # Errors + /// + /// Returns an error when the append is invalid, rejected, throttled, or cannot be made + /// durable. async fn append_many_with_options( &self, messages: Vec>, options: Option, + pin: Option, ) -> Result; /// Returns a possibly-stale copy of the manifest with a witness for verification. @@ -1106,8 +1119,13 @@ where &self, messages: Vec>, options: Option, + pin: Option, ) -> Result { - LogWriter::append_many_with_options(self, messages, options).await + LogWriter::append_many_with_options(self, messages, options, pin).await + } + + async fn acquire_fragment_pin(&self, reserved_bytes: usize) -> Result { + LogWriter::acquire_fragment_pin(self, reserved_bytes).await } async fn manifest_and_witness(&self) -> Result { diff --git a/rust/wal3/src/writer.rs b/rust/wal3/src/writer.rs index 1ac60572cdd..e9c7a4540b2 100644 --- a/rust/wal3/src/writer.rs +++ b/rust/wal3/src/writer.rs @@ -22,7 +22,7 @@ use crate::interfaces::{ }; use crate::{ parse_fragment_path, AppendOptions, AppendWork, BatchManager, CursorStore, CursorStoreOptions, - Error, ExponentialBackoff, Fragment, FragmentSeqNo, FragmentUuid, Garbage, + Error, ExponentialBackoff, Fragment, FragmentPin, FragmentSeqNo, FragmentUuid, Garbage, GarbageCollectionOptions, GarbageCollectionState, LogPosition, LogReader, LogReaderOptions, LogWriterOptions, Manifest, ManifestAndWitness, ManifestManager, }; @@ -264,35 +264,66 @@ impl< message: Vec, options: Option, ) -> Result { - self.append_many_with_options(vec![message], options).await + self.append_many_with_options(vec![message], options, None) + .await } #[tracing::instrument(skip(self, messages))] pub async fn append_many(&self, messages: Vec>) -> Result { - self.append_many_with_options(messages, None).await + self.append_many_with_options(messages, None, None).await } - #[tracing::instrument(skip(self, messages, options))] + /// Append messages atomically, optionally joining the generation held open by `pin`. + /// + /// A supplied pin is consumed by this call and disables transparent contention retries because + /// it cannot be transferred to a recovered writer generation. + /// + /// # Errors + /// + /// Returns an error when the batch is empty, rejected by admission, throttled, or cannot be + /// made durable. + #[tracing::instrument(skip(self, messages, options, pin))] pub async fn append_many_with_options( &self, messages: Vec>, options: Option, + pin: Option, ) -> Result { - let retry_contention_internally = match options.as_ref() { - Some(options) => options.required_fragment_start.is_none(), - None => true, - }; + let retry_contention_internally = pin.is_none() + && match options.as_ref() { + Some(options) => options.required_fragment_start.is_none(), + None => true, + }; + let mut pin = pin; let once_log_append_many = move |log: &Arc>| { let messages = messages.clone(); let options = options.clone(); + let pin = pin.take(); let log = Arc::clone(log); - async move { log.append(messages, options).await } + async move { log.append(messages, options, pin).await } }; self.handle_errors_and_contention(once_log_append_many, retry_contention_internally) .await } + /// Keep the next fragment open while an append is validated and prepared. + /// + /// Pins wait for an active fragment to finish, so a caller that acquires a pin before reading + /// the manifest validates against the generation its append will join. + /// + /// # Errors + /// + /// Returns an error when the writer cannot open, is shutting down, or cannot reserve + /// `reserved_bytes` in the next fragment. + pub async fn acquire_fragment_pin(&self, reserved_bytes: usize) -> Result { + let (writer, _) = self.ensure_open().await?; + writer + .batch_manager + .acquire_fragment_pin(reserved_bytes) + .await + } + pub async fn reader( &self, options: LogReaderOptions, @@ -384,7 +415,7 @@ impl< async fn handle_errors_and_contention>>( &self, - f: impl Fn(&Arc>) -> F, + mut f: impl FnMut(&Arc>) -> F, retry_contention_internally: bool, ) -> Result { for _ in 0..3 { @@ -394,15 +425,7 @@ impl< return Ok(out); } Err(Error::LogContentionDurable) => { - { - // SAFETY(rescrv): Mutex poisoning. - let mut inner = self.inner.lock().unwrap(); - if inner.epoch == epoch { - if let Some(writer) = inner.writer.take() { - writer.shutdown(); - } - } - } + self.shutdown_epoch(epoch); // Silence this error in favor of the one we got from f. if self.ensure_open().await.is_ok() { return Err(Error::LogContentionDurable); @@ -411,23 +434,11 @@ impl< } } Err(Error::LogContentionFailure) => { - // SAFETY(rescrv): Mutex poisoning. - let mut inner = self.inner.lock().unwrap(); - if inner.epoch == epoch { - if let Some(writer) = inner.writer.take() { - writer.shutdown(); - } - } + self.shutdown_epoch(epoch); return Err(Error::LogContentionFailure); } Err(Error::LogContentionRetry) => { - // SAFETY(rescrv): Mutex poisoning. - let mut inner = self.inner.lock().unwrap(); - if inner.epoch == epoch { - if let Some(writer) = inner.writer.take() { - writer.shutdown(); - } - } + self.shutdown_epoch(epoch); if !retry_contention_internally { return Err(Error::LogContentionRetry); } @@ -436,12 +447,7 @@ impl< return Err(Error::Backoff); } Err(err) => { - let mut inner = self.inner.lock().unwrap(); - if inner.epoch == epoch { - if let Some(writer) = inner.writer.take() { - writer.shutdown(); - } - } + self.shutdown_epoch(epoch); return Err(err); } } @@ -449,6 +455,16 @@ impl< Err(Error::LogContentionFailure) } + fn shutdown_epoch(&self, epoch: u64) { + // SAFETY(rescrv): Mutex poisoning. + let mut inner = self.inner.lock().unwrap(); + if inner.epoch == epoch { + if let Some(writer) = inner.writer.take() { + writer.shutdown(); + } + } + } + async fn ensure_open( &self, ) -> Result<(Arc>, u64), Error> { @@ -649,6 +665,7 @@ impl, MP: Manifes self: &Arc, messages: Vec>, options: Option, + pin: Option, ) -> Result { if messages.is_empty() { return Err(Error::EmptyBatch); @@ -657,9 +674,8 @@ impl, MP: Manifes let append_span_clone = append_span.clone(); async move { let (tx, rx) = tokio::sync::oneshot::channel(); - self.batch_manager - .push_work(AppendWork::new(messages, options, tx, append_span)) - .await; + let work = AppendWork::new(messages, options, tx, append_span); + self.batch_manager.push_work(work, pin).await; match self.batch_manager.take_work(&self.manifest_manager).await { Ok(Some(work)) => { let (pointer, required_fragment_start, work) = work; @@ -1691,7 +1707,14 @@ mod tests { impl crate::FragmentPublisher for TestFragmentPublisher { type FragmentPointer = (FragmentSeqNo, LogPosition); - async fn push_work(&self, _work: crate::AppendWork) { + async fn acquire_fragment_pin( + &self, + _reserved_bytes: usize, + ) -> Result { + unreachable!("acquire_fragment_pin is not used in this test") + } + + async fn push_work(&self, _work: crate::AppendWork, _pin: Option) { unreachable!("push_work is not used in this test") } From 45080029f62fa5e29c9f755be0fcb1ee53149bb7 Mon Sep 17 00:00:00 2001 From: Robert Escriva Date: Wed, 2 Sep 2026 11:02:27 -0700 Subject: [PATCH 2/3] contention --- rust/log-service/src/lib.rs | 70 +++++++++++++++++++++++++++++++++++-- 1 file changed, 68 insertions(+), 2 deletions(-) diff --git a/rust/log-service/src/lib.rs b/rust/log-service/src/lib.rs index 6e082e29966..bd5d4daef0b 100644 --- a/rust/log-service/src/lib.rs +++ b/rust/log-service/src/lib.rs @@ -2581,7 +2581,20 @@ impl LogServer { u64::try_from(pin_started.elapsed().as_micros()).unwrap_or(u64::MAX), &[], ); - Some(pin_result.map_err(push_append_error_to_status)?) + match pin_result { + Ok(pin) => Some(pin), + Err(wal3::Error::LogContentionRetry) + if attempt + 1 < conditional_push_max_retries => + { + tracing::info!( + %collection_id, + attempt = attempt + 1, + "conditional write fragment pin contended; retrying" + ); + continue; + } + Err(err) => return Err(push_append_error_to_status(err)), + } } else { None }; @@ -6763,6 +6776,8 @@ mod tests { #[derive(Debug)] struct FakeLogWriter { append_results: Mutex>>, + fragment_pin_errors: Mutex>, + fragment_pin_attempts: AtomicU64, observed_options: Mutex>>, } @@ -6770,6 +6785,17 @@ mod tests { fn new(append_results: Vec>) -> Arc { Arc::new(Self { append_results: Mutex::new(append_results), + fragment_pin_errors: Mutex::new(Vec::new()), + fragment_pin_attempts: AtomicU64::new(0), + observed_options: Mutex::new(Vec::new()), + }) + } + + fn with_fragment_pin_errors(fragment_pin_errors: Vec) -> Arc { + Arc::new(Self { + append_results: Mutex::new(Vec::new()), + fragment_pin_errors: Mutex::new(fragment_pin_errors), + fragment_pin_attempts: AtomicU64::new(0), observed_options: Mutex::new(Vec::new()), }) } @@ -6781,7 +6807,12 @@ mod tests { &self, _reserved_bytes: usize, ) -> Result { - unreachable!("fake log writer does not pin fragments") + self.fragment_pin_attempts.fetch_add(1, Ordering::Relaxed); + let mut fragment_pin_errors = self.fragment_pin_errors.lock(); + if fragment_pin_errors.is_empty() { + unreachable!("fake log writer does not pin fragments") + } + Err(fragment_pin_errors.remove(0)) } async fn append_with_options( @@ -7390,6 +7421,41 @@ mod tests { ); } + #[tokio::test] + async fn conditional_push_retries_fragment_pin_contention() { + let mut log_server = log_server_for_installed_writer_tests(); + log_server.config.conditional_push_max_retries = 2; + let collection_id = CollectionUuid::new(); + let fake_log = FakeLogWriter::with_fragment_pin_errors(vec![ + wal3::Error::LogContentionRetry, + wal3::Error::Backoff, + ]); + let fake_log_trait: Arc = fake_log.clone(); + let _handle = install_active_log(&log_server, collection_id, fake_log_trait).await; + + let err = log_server + .push_logs(Request::new(conditional_push_logs_request( + "dbname", + collection_id, + &[test_operation_record("doc-1")], + PushLogsCondition { + observed_log_offset: 0, + read_ids: vec![], + }, + ))) + .await + .expect_err("conditional push should return the second fragment pin error"); + + assert_eq!(2, fake_log.fragment_pin_attempts.load(Ordering::Relaxed)); + assert_eq!(Code::ResourceExhausted, err.code()); + assert_eq!( + Some("batching"), + err.metadata() + .get(BACKOFF_REASON_MD_KEY) + .and_then(|value| value.to_str().ok()) + ); + } + fn run_k8s_async_test(future: impl Future + Send + 'static) { let runtime = Runtime::new().unwrap(); std::thread::Builder::new() From 9be2a5381f011c3bc2fc4c24a6f7b1633b9904a8 Mon Sep 17 00:00:00 2001 From: Robert Escriva Date: Wed, 2 Sep 2026 11:33:07 -0700 Subject: [PATCH 3/3] Kimi's pass --- rust/log-service/src/lib.rs | 25 ++++++-- rust/wal3/src/interfaces/batch_manager.rs | 69 ++++++++++++++++------- rust/wal3/src/interfaces/mod.rs | 6 +- rust/wal3/src/lib.rs | 5 +- rust/wal3/src/writer.rs | 3 +- 5 files changed, 76 insertions(+), 32 deletions(-) diff --git a/rust/log-service/src/lib.rs b/rust/log-service/src/lib.rs index bd5d4daef0b..a3849d0495a 100644 --- a/rust/log-service/src/lib.rs +++ b/rust/log-service/src/lib.rs @@ -483,6 +483,8 @@ pub struct Metrics { conditional_write_required_start_retries: opentelemetry::metrics::Counter, /// Time spent waiting to pin the next fragment generation. conditional_write_fragment_pin_wait_us: opentelemetry::metrics::Histogram, + /// Time a fragment pin is held open across validation and the pinned append. + conditional_write_fragment_pin_hold_us: opentelemetry::metrics::Histogram, /// The number of successful conditional writes with an inserted offset. conditional_write_success_with_offset: opentelemetry::metrics::Counter, /// The number of successful conditional writes with durable contention and no offset. @@ -533,6 +535,9 @@ impl Metrics { conditional_write_fragment_pin_wait_us: meter .u64_histogram("conditional_write_fragment_pin_wait_us") .build(), + conditional_write_fragment_pin_hold_us: meter + .u64_histogram("conditional_write_fragment_pin_hold_us") + .build(), conditional_write_success_with_offset: meter .u64_counter("conditional_write_success_with_offset") .build(), @@ -2571,7 +2576,7 @@ impl LogServer { // Replicated logs serialize conditional writes in Spanner. S3-backed logs pin // their in-memory fragment before validation so concurrent transactions can join // the same publish. - let fragment_pin = if topology_name.is_none() { + let (fragment_pin, pin_acquired_at) = if topology_name.is_none() { let pin_started = Instant::now(); let pin_result = log .acquire_fragment_pin(message_bytes) @@ -2582,7 +2587,7 @@ impl LogServer { &[], ); match pin_result { - Ok(pin) => Some(pin), + Ok(pin) => (Some(pin), Some(Instant::now())), Err(wal3::Error::LogContentionRetry) if attempt + 1 < conditional_push_max_retries => { @@ -2596,7 +2601,7 @@ impl LogServer { Err(err) => return Err(push_append_error_to_status(err)), } } else { - None + (None, None) }; let validated_tail = self .validate_committed_log_for_conditional_write( @@ -2609,10 +2614,18 @@ impl LogServer { let append_options = AppendOptions::new(write_id_metadata.clone()) .with_admission_predicate(Arc::clone(&admission_predicate)) .with_required_fragment_start(validated_tail); - match log + let append_outcome = log .append_many_with_options(messages.clone(), Some(append_options), fragment_pin) - .await - { + .await; + // The pin is consumed by the append above; the hold spans validation plus the + // pinned append, which is how long publishing was held open for this log. + if let Some(acquired_at) = pin_acquired_at { + self.metrics.conditional_write_fragment_pin_hold_us.record( + u64::try_from(acquired_at.elapsed().as_micros()).unwrap_or(u64::MAX), + &[], + ); + } + match append_outcome { Ok(offset) => { self.metrics .conditional_write_success_with_offset diff --git a/rust/wal3/src/interfaces/batch_manager.rs b/rust/wal3/src/interfaces/batch_manager.rs index 716fbb82692..50d02c2f67c 100644 --- a/rust/wal3/src/interfaces/batch_manager.rs +++ b/rust/wal3/src/interfaces/batch_manager.rs @@ -127,8 +127,21 @@ impl FragmentPin { tracing::error!("consuming an already released wal3 fragment pin"); return Err(Error::LogContentionRetry); }; - if !Arc::ptr_eq(&pinned, coordination) || self.generation != state.pin_generation { - drop(pinned); + if !Arc::ptr_eq(&pinned, coordination) { + // Belongs to another publisher; reattach so Drop releases the reservation against + // the owning coordination rather than leaking it. + self.coordination = Some(pinned); + return Err(Error::LogContentionRetry); + } + if self.generation != state.pin_generation { + // Unreachable while take_work refuses to publish with pins outstanding, but release + // regardless: a leaked reservation wedges the publisher. + tracing::error!( + pin_generation = self.generation, + current_generation = state.pin_generation, + "consuming stale wal3 fragment pin" + ); + release_pin_from_state(state, self.reserved_bytes); return Err(Error::LogContentionRetry); } release_pin_from_state(state, self.reserved_bytes); @@ -152,9 +165,9 @@ impl Drop for FragmentPin { return; }; let mut state = coordination.state.lock().unwrap(); - if self.generation == state.pin_generation && state.fragment_pins > 0 { - release_pin_from_state(&mut state, self.reserved_bytes); - } else { + if self.generation != state.pin_generation || state.fragment_pins == 0 { + // Unreachable while take_work refuses to publish with pins outstanding, but release + // regardless: a leaked reservation wedges the publisher. tracing::error!( pin_generation = self.generation, current_generation = state.pin_generation, @@ -162,13 +175,16 @@ impl Drop for FragmentPin { "dropping stale wal3 fragment pin" ); } + release_pin_from_state(&mut state, self.reserved_bytes); drop(state); + coordination.write_finished.notify_waiters(); coordination.write_finished.notify_one(); } } fn release_pin_from_state(state: &mut ManagerState, reserved_bytes: usize) { - state.fragment_pins -= 1; + // Saturating so that a stale release on the defensive paths above cannot corrupt the counts. + state.fragment_pins = state.fragment_pins.saturating_sub(1); state.pinned_bytes = state.pinned_bytes.saturating_sub(reserved_bytes); } @@ -312,16 +328,16 @@ impl> FragmentPublisher for BatchMa if state.tearing_down { return Err(Error::LogContentionRetry); } - if state.backoff { - return Err(Error::Backoff); - } - if state.writers_active == 0 { - if !state.can_reserve( + // Wait for the active writer to finish and for the reservation to fit in the + // open generation rather than failing fast; capacity frees up as fragments + // publish and pins resolve. The backoff flag need not be checked: backoff + // implies occupied_bytes >= batch_size_bytes, which fails can_reserve anyway. + if state.writers_active == 0 + && state.can_reserve( reserved_bytes, self.options.throttle_fragment.batch_size_bytes, - ) { - return Err(Error::Backoff); - } + ) + { state.fragment_pins += 1; state.pinned_bytes = state.pinned_bytes.saturating_add(reserved_bytes); return Ok(FragmentPin { @@ -469,9 +485,11 @@ impl> FragmentPublisher for BatchMa ); let work = take_selected_work(&mut state, split_off); state.finish_write(); - if refresh_backoff_after_split(&mut state, &self.options) { - self.coordination.write_finished.notify_one(); - } + refresh_backoff_after_split(&mut state, &self.options); + // Rejecting the selection frees capacity and finishes the write inline; wake + // every fragment-pin waiter, since no finish_write notify will follow. + self.coordination.write_finished.notify_waiters(); + self.coordination.write_finished.notify_one(); reject_selected_work(work, Error::LogContentionRetry); return Ok(None); } @@ -1347,10 +1365,23 @@ mod tests { BatchManager::<(FragmentSeqNo, LogPosition), _>::new(options, NoopS3Uploader).unwrap(); let first_pin = batch_manager.acquire_fragment_pin(6).await.unwrap(); - let second_result = batch_manager.acquire_fragment_pin(4).await; + // A reservation that does not fit waits for capacity rather than failing. + assert!(tokio::time::timeout( + Duration::from_millis(1), + batch_manager.acquire_fragment_pin(4) + ) + .await + .is_err()); - assert!(matches!(second_result, Err(Error::Backoff))); drop(first_pin); + let second_pin = tokio::time::timeout( + Duration::from_secs(1), + batch_manager.acquire_fragment_pin(4), + ) + .await + .expect("dropping the first pin should free its reservation") + .expect("pin acquisition should succeed"); + drop(second_pin); } #[tokio::test] diff --git a/rust/wal3/src/interfaces/mod.rs b/rust/wal3/src/interfaces/mod.rs index e2e3477de4e..3a7eaab4157 100644 --- a/rust/wal3/src/interfaces/mod.rs +++ b/rust/wal3/src/interfaces/mod.rs @@ -571,12 +571,12 @@ pub trait FragmentPublisher: Send + Sync + 'static { /// Keep the next fragment generation open while an append is prepared. /// /// `reserved_bytes` counts toward the fragment size limit until the pin is submitted or - /// dropped. + /// dropped. Waits for the active fragment to publish and for the reservation to fit in the + /// open generation. /// /// # Errors /// - /// Returns an error if the publisher is shutting down, applying backpressure, or cannot fit - /// the reservation in the open generation. + /// Returns an error if the publisher is shutting down. async fn acquire_fragment_pin(&self, reserved_bytes: usize) -> Result; /// Enqueue work to be published, consuming a fragment pin atomically when one is supplied. diff --git a/rust/wal3/src/lib.rs b/rust/wal3/src/lib.rs index e1d356c1a40..1383211b1a7 100644 --- a/rust/wal3/src/lib.rs +++ b/rust/wal3/src/lib.rs @@ -1018,11 +1018,12 @@ pub fn fragment_path(prefix: &str, path: &str) -> String { /// Trait that provides a type-erased interface to LogWriter. #[async_trait::async_trait] pub trait LogWriterTrait: std::fmt::Debug + Send + Sync + 'static { - /// Keep the next fragment open while an append is prepared. + /// Keep the next fragment open while an append is prepared. Waits for the active fragment + /// to publish and for the reservation to fit in the open generation. /// /// # Errors /// - /// Returns an error when the writer cannot reserve space in its next fragment generation. + /// Returns an error when the writer cannot open or is shutting down. async fn acquire_fragment_pin(&self, reserved_bytes: usize) -> Result; /// Append a single message to the log. diff --git a/rust/wal3/src/writer.rs b/rust/wal3/src/writer.rs index e9c7a4540b2..133524d5561 100644 --- a/rust/wal3/src/writer.rs +++ b/rust/wal3/src/writer.rs @@ -314,8 +314,7 @@ impl< /// /// # Errors /// - /// Returns an error when the writer cannot open, is shutting down, or cannot reserve - /// `reserved_bytes` in the next fragment. + /// Returns an error when the writer cannot open or is shutting down. pub async fn acquire_fragment_pin(&self, reserved_bytes: usize) -> Result { let (writer, _) = self.ensure_open().await?; writer