1use std::collections::BTreeMap;
18use std::future::Future;
19use std::pin::Pin;
20use std::sync::Arc;
21use std::task::Poll;
22
23use futures_util::{StreamExt, stream::FuturesUnordered};
25
26use super::batching::{Batch, write_data};
28use super::manager::{Domain, Inner, Operation, Pending, Target};
29use super::snapshot::SnapshotAction;
30use super::staging::{FileAction, FileCompleted};
31use super::{RecoveryPosition, StorageManager};
32use crate::{BackendFailure, Error, Result};
33
34type Attempt =
41 Pin<Box<dyn Future<Output = (Action, usize, Option<u64>, u64, Option<crate::execution::Context>, Result<Completed>)> + Send>>;
42
43pub struct Worker(pub(crate) StorageManager<crate::storage::ErasedCatalog>);
46
47struct Running(Arc<Inner>);
48
49struct Failed {
50 action: Option<Action>,
51 source: Arc<Error>,
52 attempts: usize,
53 retry_at: Option<tokio::time::Instant>,
54}
55
56enum Action {
59 Write(Arc<Pending>, Option<Batch>),
60 Snapshot(SnapshotAction),
61 File(FileAction),
62 Delete(Arc<super::deletion::DeleteWork>),
63 Recovery(Target),
64 Prune(super::StorageCursor),
65}
66
67enum Completed {
68 Write { data: Option<Target>, saved: Option<super::StorageCursor>, recovery: Option<Target> },
69 Opened(u64, Result<super::BackendSnapshot>),
70 Released(u64),
71 File(FileCompleted),
72 Deleted,
73 Recovery(RecoveryPosition),
74 Pruned(bool),
75}
76
77impl Worker {
82 #[tracing::instrument(name = "storage.worker.cleanup", level = "debug", skip_all)]
83 pub async fn cleanup(&mut self) -> Result<()> {
84 if !self.0.inner.state.lock().unwrap().status.closed {
87 return Err(Error::NotReady { domain: "storage cleanup" });
88 }
89 Ok(())
90 }
91
92 #[tracing::instrument(name = "storage.worker.run", level = "debug", skip_all)]
93 pub async fn run(&mut self) -> Result<()> {
94 let inner = &self.0.inner;
95 let _running = Running(Arc::clone(inner));
96 let mut changed = inner.changed.subscribe();
97 let mut next_flush = inner.execution.now() + super::config::FLUSH_INTERVAL;
98 let mut retry_generation = 0;
99 let mut failures = BTreeMap::<u64, Failed>::new();
100 let mut next_failure = 0u64;
101 let mut active = FuturesUnordered::<Attempt>::new();
102 let mut files_running = 0usize;
103 let mut ordered_running = false;
104 let file_limit = inner.settings.request.max_active_file_operations.get();
105 let mut prefer_control = true;
106 let mut prefer_files = true;
107 let mut prefer_pruning = true;
108 let mut prefer_state = true;
109 loop {
110 changed.borrow_and_update();
111 let now = inner.execution.now();
112 let (action, failure_id, attempts, generation) = {
113 let mut state = inner.state.lock().unwrap();
114 if state.status.closed && active.is_empty() {
115 return Ok(());
116 }
117 if state.retry != retry_generation {
118 retry_generation = state.retry;
119 failures.clear();
120 state.status.failure = None;
121 }
122 if now >= next_flush {
123 state.request_flush();
124 next_flush = now + super::config::FLUSH_INTERVAL;
125 }
126 state.expire_files();
127 state.request_deletions();
128 state.collect();
129 if state.queue.bytes() >= inner.settings.target_batch_bytes() {
130 state.request_flush();
131 }
132 let retry = failures.iter().find_map(|(&id, failure)| {
133 let action = failure.action.as_ref()?;
134 let available = if action.independent() { files_running < file_limit } else { !ordered_running };
135 (available && failure.retry_at.is_some_and(|deadline| deadline <= now)).then_some(id)
136 });
137 let mut failure_id = None;
138 let mut attempts = 0;
139 let selected = if state.status.closed {
140 None
141 } else if let Some(id) = retry {
142 let failure = failures.get_mut(&id).expect("selected failed owner");
143 attempts = failure.attempts;
144 failure_id = Some(id);
145 let action = failure.action.take().expect("waiting failed action");
146 match action {
147 Action::File(file) => {
148 let id = file.id();
149 drop(file);
150 state.file_action(!ordered_running, Some(id)).map(Action::File)
151 }
152 action => Some(action),
153 }
154 } else {
155 let prune_through = if state.snapshot_opening() { 0 } else { state.status.saved.sequence };
157 let mut prune = (!ordered_running
158 && inner.has_objects
159 && state.ready
160 && !state.status.closing
161 && state.status.failure.is_none()
162 && prune_through > state.objects_pruned)
163 .then(|| Action::Prune(self.0.position(prune_through)));
164 let data = if prefer_pruning && prune.is_some() {
165 prune.take()
166 } else if ordered_running || !state.ready || state.status.failure.is_some() {
167 None
168 } else if let Some(deletion) = state.deletion_action() {
169 Some(Action::Delete(deletion))
170 } else if let Some(target) = state.recovery_request.filter(|target| target.index <= state.recovery_written.index) {
171 Some(Action::Recovery(target))
172 } else if let Some(boundary) = state.write_boundary() {
173 let mut pending = state.queue.iter().filter(|entry| entry.index <= boundary && state.needs_write(entry));
177 let first = pending.next();
178 let selected =
179 if state.data_request.is_some() && first.is_some_and(|entry| entry.operation.domain() == Domain::Recovery) {
180 pending.find(|entry| entry.operation.domain() == Domain::Data).or(first)
181 } else {
182 first
183 };
184 let ready_state = state
188 .recovery_request
189 .and_then(|target| {
190 state.queue.iter().find(|entry| {
191 entry.index <= target.index && entry.operation.domain() == Domain::Recovery && state.needs_write(entry)
192 })
193 })
194 .filter(|entry| {
195 entry.operation.required_data().is_some_and(|required| state.status.saved.covers(required).unwrap_or(false))
196 });
197 let selected = if prefer_state { ready_state.or(selected) } else { selected };
198 let selected = selected.map(|entry| entry.index);
199 drop(pending);
200 selected.map(|index| {
201 let entry = state.queue.iter().find(|entry| entry.index == index).expect("selected pending owner").clone();
202 Action::Write(entry, None)
203 })
204 } else {
205 None
206 };
207 let data = data.or(prune);
208 if prefer_control || data.is_none() {
211 let healthy = state.status.failure.is_none();
212 let file = if prefer_files && healthy && files_running < file_limit {
213 state.file_action(!ordered_running, None).map(Action::File)
214 } else {
215 None
216 };
217 let control = file
218 .or_else(|| if !ordered_running { state.snapshot_action(healthy).map(Action::Snapshot) } else { None })
219 .or_else(|| {
220 if healthy && files_running < file_limit {
221 state.file_action(!ordered_running, None).map(Action::File)
222 } else {
223 None
224 }
225 });
226 if let Some(action) = &control {
227 prefer_files = !matches!(action, Action::File(_));
228 }
229 control.or(data)
230 } else {
231 data
232 }
233 };
234 if selected.is_none()
235 && let Some(id) = failure_id.take()
236 {
237 failures.remove(&id);
238 state.status.failure = failures.values().next().map(|failure| Arc::clone(&failure.source));
239 }
240 let selected = match selected {
242 Some(Action::Write(entry, _)) => {
243 let index = entry.index;
244 drop(entry);
245 let batch = Batch::prepare(inner, &mut state, index);
246 let entry = state.queue.iter().find(|entry| entry.index == index).expect("selected pending owner").clone();
247 Some(Action::Write(entry, batch))
248 }
249 action => action,
250 };
251 if matches!(selected, Some(Action::Prune(_))) {
252 state.pruning_objects = true;
253 state.objects_prune_owed = false;
254 }
255 (selected, failure_id, attempts, retry_generation)
256 };
257 if let Some(action) = action {
258 if let Action::Write(entry, _) = &action {
259 prefer_state = entry.operation.domain() != Domain::Recovery;
260 }
261 prefer_pruning = !matches!(action, Action::Prune(_));
262 prefer_control = !matches!(action, Action::Snapshot(_) | Action::File(_));
263 if action.independent() {
264 files_running += 1;
265 } else {
266 ordered_running = true;
267 }
268 let (context, queued) = {
269 let mut state = inner.state.lock().unwrap();
270 match &action {
271 Action::Write(entry, _) => (Some(entry.context.clone()), None),
272 Action::Snapshot(SnapshotAction::Open(id, _, _) | SnapshotAction::Release(id, _, _)) => {
273 state.snapshots.get_mut(id).map(|held| (Some(held.context.clone()), held.queued.take())).unwrap_or((None, None))
274 }
275 Action::File(file) => state
276 .files
277 .get_mut(&file.id())
278 .map(|held| (Some(held.context.clone()), held.queued.take()))
279 .unwrap_or((None, None)),
280 Action::Delete(work) => (Some(work.context.clone()), work.queued.lock().unwrap().take()),
281 _ => (None, None),
282 }
283 };
284 let owner = Arc::clone(inner);
285 active.push(Box::pin(async move {
286 let mut action = action;
287 let future = async {
288 let span = tracing::debug_span!("storage.work", outcome = "cancelled");
289 if let Some(queued) = queued {
290 span.follows_from(queued.id());
291 }
292 let physical = async {
293 let result = perform(&owner, &action).await;
294 tracing::Span::current().record("outcome", if result.is_ok() { "completed" } else { "failed" });
295 result
296 };
297 tracing::Instrument::instrument(physical, span).await
298 };
299 let result = match &context {
300 Some(context) => context.scope(future).await,
301 None => future.await,
302 };
303 if let Action::Write(_, batch) = &mut action {
305 batch.take();
306 }
307 (action, attempts, failure_id, generation, context, result)
308 }));
309 continue;
310 }
311 let retry = failures.values().filter_map(|failure| failure.retry_at).filter(|deadline| *deadline > now).min();
312 let deadline = retry.map_or(next_flush, |retry| retry.min(next_flush));
313 let deadline = inner.state.lock().unwrap().file_deadline().map_or(deadline, |file| file.min(deadline));
314 let mut sleep = Box::pin(inner.execution.sleep(deadline.saturating_duration_since(inner.execution.now())));
315 let mut wake = Box::pin(changed.changed());
316 let outcome = std::future::poll_fn(|cx| {
317 if let Poll::Ready(Some(outcome)) = active.poll_next_unpin(cx) {
318 return Poll::Ready(Some(outcome));
319 }
320 if wake.as_mut().poll(cx).is_ready() || sleep.as_mut().poll(cx).is_ready() {
321 return Poll::Ready(None);
322 }
323 Poll::Pending
324 })
325 .await;
326 let Some((action, attempts, failure_id, generation, context, result)) = outcome else {
327 continue;
328 };
329 if action.independent() {
330 files_running -= 1;
331 } else {
332 ordered_running = false;
333 }
334 if matches!(action, Action::Prune(_)) {
335 inner.state.lock().unwrap().pruning_objects = false;
336 }
337 if let Ok(Completed::Opened(id, snapshot)) = result {
338 let mut state = inner.state.lock().unwrap();
339 state.snapshot_opened(id, snapshot);
340 if let Some(id) = failure_id {
341 failures.remove(&id);
342 }
343 state.status.failure = failures.values().next().map(|failure| Arc::clone(&failure.source));
344 drop(state);
345 inner.notify();
346 continue;
347 }
348 match result {
349 Ok(completed) => {
350 if let Some(id) = failure_id {
351 failures.remove(&id);
352 }
353 let mut retired_snapshot = None;
354 let mut state = inner.state.lock().unwrap();
355 state.status.failure = failures.values().next().map(|failure| Arc::clone(&failure.source));
357 match (action, completed) {
358 (Action::Delete(work), Completed::Deleted) => {
359 state.receipts.complete(&work.receipt().mutation, work.receipt().value.clone());
360 state.deletions.remove(&work.receipt().mutation);
361 }
362 (Action::Prune(through), Completed::Pruned(done)) => {
363 if done && !state.objects_prune_owed {
364 state.objects_pruned = through.sequence;
365 }
366 }
367 (Action::File(action), Completed::File(result)) => {
368 let publish = || {
369 let _publication = tracing::debug_span!("storage.file.publish").entered();
370 state.file_completed(action, result, super::config::FILE_SESSION_TIMEOUT, inner)
371 };
372 match &context {
373 Some(context) => context.sync(publish),
374 None => publish(),
375 }
376 }
377 (Action::Snapshot(SnapshotAction::Release(_, _, memory)), Completed::Released(id)) => {
378 state.snapshot_removed(id);
379 retired_snapshot = Some(memory);
380 }
381 (Action::Write(_, _), Completed::Write { data, saved, recovery }) => {
382 if let Some(target) = data {
383 let saved = saved.expect("data write returns durability");
384 let through = if saved.sequence == target.sequence {
385 target.index
386 } else {
387 state
388 .queue
389 .iter()
390 .filter(|entry| {
391 entry.operation.domain() == Domain::Data && entry.operation.sequence() <= saved.sequence
392 })
393 .map(|entry| entry.index)
394 .max()
395 .unwrap_or(state.data_flushed)
396 };
397 state.status.saved = saved;
398 state.data_flushed = through;
399 state.retire_data();
400 if state.data_request.is_some_and(|request| request.index <= through) {
401 state.data_request = None;
402 }
403 }
404 if let Some(target) = recovery {
405 state.recovery_written.index = target.index;
406 state.recovery_written.sequence = state.recovery_written.sequence.max(target.sequence);
407 }
408 }
409 (Action::Recovery(target), Completed::Recovery(saved)) => {
410 if let Some(progress) = &state.journal_saved
411 && let Some(frontier) =
412 state.queue.iter().filter(|entry| entry.index <= target.index).filter_map(|entry| entry.journal).max()
413 {
414 progress.send_replace(frontier);
415 }
416 state.status.recovery_saved = saved;
417 state.recovery_flushed = target.index;
418 state.queue.retire(Domain::Recovery, target.index);
419 state.recovery_request = None;
420 }
421 _ => unreachable!("completion matches its storage action"),
422 }
423 state.collect();
424 drop(state);
425 drop(retired_snapshot);
426 }
427 Err(error) => {
428 let attempts = if generation == retry_generation { attempts.saturating_add(1) } else { 1 };
429 let retries = match &action {
430 Action::Delete(_) => 0,
431 Action::Write(pending, _)
432 if pending.operation.domain() == Domain::Data && inner.state.lock().unwrap().deleting_data(pending.index) =>
433 {
434 0
435 }
436 Action::File(file) if file.caller_retries(&inner.state.lock().unwrap()) => 0,
437 _ => super::config::MAX_RETRIES,
438 };
439 let retry_at = if transient(&error) && attempts <= retries {
440 let multiplier = 1u32.checked_shl((attempts - 1).min(31) as u32).unwrap_or(u32::MAX);
441 Some(inner.execution.now() + super::config::RETRY_BASE.saturating_mul(multiplier).min(super::config::RETRY_CAP))
442 } else {
443 None
444 };
445 let source = Arc::new(error);
446 let mut state = inner.state.lock().unwrap();
447 if let Action::File(file) = &action {
448 state.file_failed(file, Arc::clone(&source), retry_at.is_some());
449 }
450 match &action {
451 Action::Delete(work) => state.deletion_failed(&work.receipt().mutation, Arc::clone(&source)),
452 Action::Write(pending, _) if pending.operation.domain() == Domain::Data => {
453 state.deletion_data_failed(pending.index, Arc::clone(&source));
454 }
455 _ => {}
456 }
457 let id = match failure_id.filter(|id| failures.contains_key(id)) {
458 Some(id) => id,
459 None => {
460 while failures.contains_key(&next_failure) {
461 next_failure = next_failure.wrapping_add(1);
462 }
463 let id = next_failure;
464 next_failure = next_failure.wrapping_add(1);
465 id
466 }
467 };
468 failures.insert(id, Failed { action: Some(action), source, attempts, retry_at });
469 state.status.failure = failures.values().next().map(|failure| Arc::clone(&failure.source));
470 state.collect();
471 }
472 }
473 inner.notify();
474 }
475 }
476}
477
478impl Action {
479 fn independent(&self) -> bool {
480 matches!(self, Self::File(file) if file.independent())
481 }
482}
483
484impl Drop for Running {
485 fn drop(&mut self) {
486 let mut state = self.0.state.lock().unwrap();
487 state.status.running = false;
488 state.pruning_objects = false;
489 if !state.status.closed {
490 state.status.failure = Some(Arc::new(Error::NotReady { domain: "storage worker" }));
491 }
492 self.0.notify();
493 }
494}
495
496impl Operation {
497 pub(super) fn encoded_batch_bytes(&self) -> Result<usize> {
498 fn sum<T: serde::Serialize>(values: &[T]) -> Result<usize> {
499 values.iter().try_fold(0usize, |bytes, value| {
500 bytes
501 .checked_add(super::manager::encoded_size(value)?)
502 .ok_or(Error::InvalidInput { field: "storage admission", reason: "encoded batch size overflow" })
503 })
504 }
505 match self {
506 Self::Schema(values) => values.as_slice().iter().try_fold(0usize, |bytes, value| {
507 bytes
508 .checked_add(value.encoded_bytes()?)
509 .ok_or(Error::InvalidInput { field: "storage admission", reason: "encoded schema size overflow" })
510 }),
511 Self::ObjectGroup { schemas, objects, .. } => {
512 for object in objects.as_slice() {
513 object.version.check_position(&object.position)?;
514 }
515 let schema_bytes = schemas.as_slice().iter().try_fold(0usize, |bytes, value| {
516 bytes
517 .checked_add(value.encoded_bytes()?)
518 .ok_or(Error::InvalidInput { field: "storage admission", reason: "encoded schema size overflow" })
519 })?;
520 schema_bytes
521 .checked_add(sum(objects.as_slice())?)
522 .ok_or(Error::InvalidInput { field: "storage admission", reason: "encoded batch size overflow" })
523 }
524 Self::Records(values) => sum(values.as_slice()),
525 Self::Stream(values) => sum(values.as_slice()),
526 Self::File(values) => sum(values.as_slice()),
527 Self::Recovery(_) | Self::Snapshot(_) => Ok(0),
528 }
529 }
530
531 pub(super) fn check_entry_bytes(&self, maximum: std::num::NonZeroUsize) -> Result<()> {
532 if self.encoded_batch_bytes()? > maximum.get() {
533 return Err(Error::InvalidInput {
534 field: "storage admission",
535 reason: "admitted batch exceeds max_batch_bytes; split before submission",
536 });
537 }
538 Ok(())
539 }
540
541 pub(super) fn required_data(&self) -> Option<&super::StorageCursor> {
542 match self {
543 Self::Recovery(entries) => {
544 entries.as_slice().iter().max_by_key(|entry| entry.required_data.sequence).map(|entry| &entry.required_data)
545 }
546 Self::Snapshot(snapshot) => Some(&snapshot.as_slice()[0].required_data),
547 _ => None,
548 }
549 }
550
551 pub(super) fn domain(&self) -> Domain {
552 match self {
553 Self::Recovery(_) | Self::Snapshot(_) => Domain::Recovery,
554 _ => Domain::Data,
555 }
556 }
557
558 pub(super) fn sequence(&self) -> u64 {
559 match self {
560 Self::Schema(schema) => schema.as_slice().last().unwrap().position.sequence,
561 Self::File(file) => file.as_slice().last().unwrap().position.sequence,
562 Self::ObjectGroup { objects, .. } => objects.as_slice().last().unwrap().position.sequence,
563 Self::Records(records) => records.as_slice().last().unwrap().position.sequence,
564 Self::Stream(rows) => rows.as_slice().last().unwrap().record.position.sequence,
565 Self::Recovery(entries) => entries.as_slice().last().unwrap().position.sequence,
566 Self::Snapshot(snapshot) => snapshot.as_slice()[0].through.sequence,
567 }
568 }
569}
570
571#[tracing::instrument(name = "storage.worker.perform", level = "debug", skip_all)]
576async fn perform(inner: &Arc<Inner>, action: &Action) -> Result<Completed> {
577 match action {
578 Action::Delete(work) => work.perform(inner).await.map(|()| Completed::Deleted),
579 Action::Prune(through) => inner
580 .bounded(
581 "primary.prune_objects",
582 inner.primary.prune_objects(through, super::config::OBJECT_PRUNE_ENTRIES, inner.object_retention),
583 )
584 .await
585 .map(Completed::Pruned),
586 Action::File(action) => super::staging::perform(inner, action).await.map(Completed::File),
587 Action::Snapshot(SnapshotAction::Open(id, at, memory)) => {
588 let opened = async {
589 let snapshot = inner.bounded("primary.open_snapshot", inner.primary.open_snapshot(at, memory.clone())).await?;
590 let bytes = snapshot
591 .token
592 .capacity()
593 .checked_mul(3)
594 .and_then(|bytes| bytes.checked_add(super::manager::ENTRY_OVERHEAD))
595 .ok_or(Error::ResourceExhausted { resource: "snapshot token memory" })?;
596 memory.grow_to(bytes).await.map_err(crate::capacity::Refusal::into_error)?;
597 Ok(snapshot)
598 }
599 .await;
600 Ok(Completed::Opened(*id, opened))
601 }
602 Action::Snapshot(SnapshotAction::Release(id, snapshot, _)) => {
603 if let Some(snapshot) = snapshot {
604 inner.bounded("primary.release_snapshot", inner.primary.release_snapshot(snapshot)).await?;
605 }
606 Ok(Completed::Released(*id))
607 }
608 Action::Write(entry, batch) => {
609 let future = async {
610 let operation = batch.as_ref().map_or(&entry.operation, |batch| &batch.operation);
611 let target = batch.as_ref().map_or(Target { index: entry.index, sequence: operation.sequence() }, |batch| batch.target);
612 match operation {
613 Operation::Recovery(entries) => {
614 let saved = inner.state.lock().unwrap().status.saved.clone();
615 for entry in entries.as_slice() {
616 if !saved.covers(&entry.required_data)? {
617 return Err(Error::InvalidInput {
618 field: "state evidence dependency",
619 reason: "journal cannot precede required primary persistence",
620 });
621 }
622 }
623 inner.bounded("recovery.write_entries", inner.recovery.write_entries(entries.clone())).await?;
624 Ok(Completed::Write { data: None, saved: None, recovery: Some(target) })
625 }
626 Operation::Snapshot(snapshot) => {
627 if snapshot.len() != 1 {
628 return Err(Error::InvalidInput {
629 field: "state snapshot",
630 reason: "snapshot write requires exactly one snapshot",
631 });
632 }
633 if !inner.state.lock().unwrap().status.saved.covers(&snapshot.as_slice()[0].required_data)? {
634 return Err(Error::InvalidInput {
635 field: "state snapshot dependency",
636 reason: "snapshot cannot precede required primary persistence",
637 });
638 }
639 inner.bounded("recovery.write_snapshot", inner.recovery.write_snapshot(snapshot.clone())).await?;
640 Ok(Completed::Write { data: None, saved: None, recovery: Some(target) })
641 }
642 _ => {
643 let saved = write_data(inner, operation).await?;
644 Ok(Completed::Write { data: Some(target), saved: Some(saved), recovery: None })
645 }
646 }
647 };
648 if let Some(batch) = batch {
649 batch.trace(inner, future).await
650 } else {
651 if let Some(queued) = entry.queued.lock().unwrap().take() {
652 tracing::Span::current().follows_from(queued.id());
653 }
654 future.await
655 }
656 }
657
658 Action::Recovery(target) => {
659 let through = RecoveryPosition { history: inner.recovery.history().clone(), sequence: target.sequence };
660 let saved = inner.bounded("recovery.flush", inner.recovery.flush(&through)).await?;
661 let written = inner.state.lock().unwrap().recovery_written.sequence;
662 if !saved.covers(&through)? || saved.sequence > written {
663 return Err(Error::InvalidInput {
664 field: "recovery flush receipt",
665 reason: "receipt does not cover the requested written frontier",
666 });
667 }
668 Ok(Completed::Recovery(saved))
669 }
670 }
671}
672
673fn transient(error: &Error) -> bool {
674 match error {
675 Error::Backend { kind: BackendFailure::Transient, .. } => true,
676 Error::Persistence { source } => transient(source),
677 _ => false,
678 }
679}
680
681