harmos_storage/state/
reader.rs1use std::num::NonZeroUsize;
18
19use tokio::sync::{RwLockReadGuard, watch};
21
22use super::{RecoveryPage, RecoverySnapshot};
24use crate::storage::capacity::{Pool, Reservation};
25use crate::storage::manager::{ENTRY_OVERHEAD, Lifecycle, encoded_size, recovery_preparation_bytes};
26use crate::storage::{ReadLimits, RecoveryPosition, StorageManager};
27use crate::{Error, Result};
28use harmos_domain::state::Position;
29
30pub struct RecoveryReader<'a, C> {
39 context: crate::execution::Context,
40 storage: &'a StorageManager<C>,
41 _lifecycle: RwLockReadGuard<'a, Lifecycle>,
42 through: RecoveryPosition,
43 limits: ReadLimits,
44 page_bytes: usize,
45 snapshot_bound: Option<NonZeroUsize>,
46}
47
48impl<C> StorageManager<C> {
55 #[tracing::instrument(name = "storage.state.reader.recovery_reader", level = "debug", skip_all)]
56 pub async fn recovery_reader(&self, limits: ReadLimits, snapshot_bytes: Option<NonZeroUsize>) -> Result<RecoveryReader<'_, C>> {
57 let lifecycle = tracing::Instrument::instrument(self.inner.lifecycle.read(), tracing::trace_span!("storage.lock.wait")).await;
58 let invalid = |reason| Error::InvalidInput { field: "state recovery", reason };
59 let snapshot_bound = snapshot_bytes
60 .map(|bytes| {
61 recovery_preparation_bytes(bytes.get(), self.inner.recovery.history(), self.inner.primary.history())
62 .map(|charged| NonZeroUsize::new(charged).expect("positive snapshot budget"))
63 })
64 .transpose()?;
65 let page_bytes = limits
66 .bytes
67 .get()
68 .checked_mul(2)
69 .and_then(|n| n.checked_add(limits.rows.get().checked_mul(ENTRY_OVERHEAD)?))
70 .ok_or_else(|| invalid("recovery read size overflow"))?;
71 let mut bytes = page_bytes;
72 if let Some(bound) = snapshot_bound {
73 bytes = bytes.max(bound.get().checked_mul(2).ok_or_else(|| invalid("snapshot read size overflow"))?);
74 }
75 let through = {
76 let mut current = self.inner.state.lock().unwrap();
77 current.collect();
78 if current.status.running || current.status.closing || current.journal_saved.is_some() {
79 return Err(invalid("managed recovery requires a dormant acquired manager"));
80 }
81 if bytes > self.inner.settings.request.max_pending_write_bytes.get().saturating_sub(current.status.pending_bytes) {
82 return Err(invalid("recovery page exceeds shared memory budget"));
83 }
84 current.status.recovery_saved.clone()
85 };
86 Ok(RecoveryReader {
87 context: crate::execution::Context::current(),
88 storage: self,
89 _lifecycle: lifecycle,
90 through,
91 limits,
92 page_bytes,
93 snapshot_bound,
94 })
95 }
96}
97
98impl<C> RecoveryReader<'_, C> {
99 pub fn through(&self) -> &RecoveryPosition {
100 &self.through
101 }
102
103 pub async fn snapshot(&self) -> Result<Option<crate::ReadBatch<RecoverySnapshot>>> {
104 let context = self.context.clone();
105 let span = context.sync(|| tracing::trace_span!("storage.reader.page"));
106 let future = async {
107 let Some(bound) = self.snapshot_bound else {
108 return Ok(None);
109 };
110 let reservation = self.reserve(bound.get().saturating_mul(2))?.retain();
111 let snapshot = self
112 .storage
113 .inner
114 .bounded("recovery.read_snapshot", self.storage.inner.recovery.read_snapshot(bound, reservation.clone()))
115 .await?;
116 if let Some(snapshot) = &snapshot
117 && (!self.through.covers(&snapshot.through)? || encoded_size(snapshot)? > bound.get())
118 {
119 return Err(Error::InvalidInput { field: "state snapshot", reason: "invalid managed snapshot boundary or size" });
120 }
121 if let Some(snapshot) = &snapshot {
122 self.storage.check_state_dependency(&snapshot.required_data)?;
123 }
124 Ok(snapshot.map(|snapshot| crate::ReadBatch::retained(snapshot, reservation)))
125 };
126 context.scope(tracing::Instrument::instrument(future, span)).await
127 }
128
129 pub async fn page(&self, after: &RecoveryPosition) -> Result<crate::ReadBatch<RecoveryPage>> {
130 let context = self.context.clone();
131 let span = context.sync(|| tracing::trace_span!("storage.reader.page"));
132 let future = async {
133 let reservation = self.reserve(self.page_bytes)?.retain();
134 let page = self
135 .storage
136 .inner
137 .bounded(
138 "recovery.read_entries",
139 self.storage.inner.recovery.read_entries(after, &self.through, self.limits, reservation.clone()),
140 )
141 .await?;
142 page.check(after, &self.through, self.limits, |entry| encoded_size(entry).unwrap_or(usize::MAX))?;
143 for entry in &page.entries {
144 self.storage.check_state_dependency(&entry.required_data)?;
145 }
146 Ok(crate::ReadBatch::retained(page, reservation))
147 };
148 context.scope(tracing::Instrument::instrument(future, span)).await
149 }
150
151 fn reserve(&self, bytes: usize) -> Result<Reservation> {
152 let mut state = self.storage.inner.state.lock().unwrap();
153 state.collect();
154 if bytes > self.storage.inner.settings.request.max_pending_write_bytes.get().saturating_sub(state.status.pending_bytes) {
155 return Err(Error::ResourceExhausted { resource: "retained recovery read memory" });
156 }
157 state.scratch_bytes += bytes;
158 state.collect();
159 Ok(Reservation::charged_lifecycle(&self.storage.inner, Pool::Write, bytes))
160 }
161
162 pub fn finish(self, position: Position) -> watch::Receiver<Position> {
163 let (saved, durable) = watch::channel(position);
164 let mut current = self.storage.inner.state.lock().unwrap();
165 current.journal_required_data = current.status.saved.clone();
166 current.journal_saved = Some(saved);
167 current.journal_limit = self.limits.bytes.get();
168 durable
169 }
170}
171
172