1use super::allocation::{AllocationError, AllocationLedger, VisitedObjectSet};
4use super::attribute::attribute_page_decode_shape;
5use super::attribute_mutation::AttributeFormat;
6use super::codec::DecodedPageKind;
7use super::persistent_batch;
8use super::{
9 AttributeChild, AttributeEntry, AttributeName, AttributePage, CanonicalDecodeError,
10 DecodeLimits, decode_attribute_page,
11};
12use crate::async_storage::AsyncObjectStore;
13use crate::cancellation::CancellationToken;
14use crate::performance::{OperationFailure, WorkBudget, WorkCounters, WorkError};
15use crate::storage::{ObjectId, ObjectKind, ObjectReadRetention, ObjectStoreError};
16use std::mem::size_of;
17use thiserror::Error;
18
19#[derive(Clone, Debug, Eq, PartialEq)]
21pub struct AttributeLookup {
22 pub entry: Option<AttributeEntry>,
24 pub work: WorkCounters,
26}
27
28#[derive(Clone, Debug, Eq, PartialEq)]
30pub struct AttributeBatchLookup {
31 pub entries: Vec<Option<AttributeEntry>>,
33 pub work: WorkCounters,
35}
36
37pub fn lookup_attributes<S: crate::ImmediateObjectStore>(
44 store: &S,
45 root: ObjectId,
46 names: &[AttributeName],
47 maximum_queries: u32,
48 limits: DecodeLimits,
49 budget: WorkBudget,
50) -> Result<AttributeBatchLookup, AttributeLookupFailure> {
51 persistent_batch::lookup::<S, AttributeFormat>(
52 store,
53 root,
54 names,
55 maximum_queries,
56 limits,
57 budget,
58 )
59 .map(to_batch)
60 .map_err(map_batch_failure)
61}
62
63pub async fn lookup_attributes_async<S: AsyncObjectStore>(
69 store: &S,
70 root: ObjectId,
71 names: &[AttributeName],
72 maximum_queries: u32,
73 limits: DecodeLimits,
74 budget: WorkBudget,
75 cancellation: &CancellationToken,
76) -> Result<AttributeBatchLookup, AttributeLookupFailure> {
77 persistent_batch::lookup_async::<S, AttributeFormat>(
78 store,
79 root,
80 names,
81 maximum_queries,
82 limits,
83 budget,
84 cancellation,
85 )
86 .await
87 .map(to_batch)
88 .map_err(map_batch_failure)
89}
90
91fn to_batch(receipt: persistent_batch::Receipt<AttributeEntry>) -> AttributeBatchLookup {
92 AttributeBatchLookup {
93 entries: receipt.values,
94 work: receipt.work,
95 }
96}
97
98fn map_batch_failure(failure: persistent_batch::Failure) -> AttributeLookupFailure {
99 OperationFailure::new(map_batch_error(failure.error), *failure.work)
100}
101
102fn map_batch_error(error: persistent_batch::Error) -> AttributeLookupError {
103 match error {
104 persistent_batch::Error::Cancelled => AttributeLookupError::Cancelled,
105 persistent_batch::Error::Empty => AttributeLookupError::EmptyBatch,
106 persistent_batch::Error::TooManyQueries => AttributeLookupError::TooManyQueries,
107 persistent_batch::Error::WrongRootKind => AttributeLookupError::WrongRootKind,
108 persistent_batch::Error::InvalidLimits => AttributeLookupError::InvalidLimits,
109 persistent_batch::Error::HeightExceeded => AttributeLookupError::HeightExceeded,
110 persistent_batch::Error::CycleOrAlias => AttributeLookupError::CycleOrAlias,
111 persistent_batch::Error::ChildBoundsMismatch => AttributeLookupError::ChildBoundsMismatch,
112 persistent_batch::Error::InvalidRouting => AttributeLookupError::InvalidRouting,
113 persistent_batch::Error::AllocationFailed => AttributeLookupError::AllocationFailed,
114 persistent_batch::Error::Storage(ObjectStoreError::Cancelled) => {
115 AttributeLookupError::Cancelled
116 }
117 persistent_batch::Error::Storage(error) => error.into(),
118 persistent_batch::Error::Decode(error) => error.into(),
119 persistent_batch::Error::Work(error) => error.into(),
120 }
121}
122
123pub type AttributeLookupFailure = OperationFailure<AttributeLookupError>;
125
126pub fn lookup_attribute<S: crate::ImmediateObjectStore>(
133 store: &S,
134 root: ObjectId,
135 name: &AttributeName,
136 limits: DecodeLimits,
137 budget: WorkBudget,
138) -> Result<AttributeLookup, AttributeLookupFailure> {
139 crate::async_storage::poll_immediate(lookup_attribute_async(
140 store,
141 root,
142 name,
143 limits,
144 budget,
145 &CancellationToken::new(),
146 ))
147}
148
149pub async fn lookup_attribute_async<S: AsyncObjectStore>(
156 store: &S,
157 root: ObjectId,
158 name: &AttributeName,
159 limits: DecodeLimits,
160 budget: WorkBudget,
161 cancellation: &CancellationToken,
162) -> Result<AttributeLookup, AttributeLookupFailure> {
163 if cancellation.is_cancelled() {
164 return Err(OperationFailure::before_work(
165 AttributeLookupError::Cancelled,
166 ));
167 }
168 if root.kind != ObjectKind::AttributePage {
169 return Err(OperationFailure::before_work(
170 AttributeLookupError::WrongRootKind,
171 ));
172 }
173 if !limits.page_limits_valid(1) {
174 return Err(OperationFailure::before_work(
175 AttributeLookupError::InvalidLimits,
176 ));
177 }
178
179 let mut allocations = AllocationLedger::default();
180 let mut work = WorkCounters::default();
181 let mut visited = VisitedObjectSet::new(
182 usize::from(limits.maximum_page_height).min(
183 usize::try_from(limits.maximum_visited_pages)
184 .map_err(|_| OperationFailure::before_work(AttributeLookupError::InvalidLimits))?,
185 ),
186 &mut allocations,
187 &mut work,
188 budget,
189 )
190 .map_err(|error| failed(map_allocation(error), work))?;
191 let mut page = root;
192 let mut routing = RoutingState::default();
193
194 for _ in 0..limits.maximum_page_height {
195 let inserted = visited
196 .insert(page, &mut allocations, &mut work, budget)
197 .map_err(|error| failed(map_allocation(error), work))?;
198 if !inserted.inserted {
199 return Err(failed(AttributeLookupError::CycleOrAlias, work));
200 }
201
202 let decoded = read_page(
203 store,
204 page,
205 limits,
206 budget,
207 cancellation,
208 &mut allocations,
209 &mut work,
210 )
211 .await?;
212 let decoded_bytes = decoded.logical_bytes;
213
214 match decoded.page {
215 AttributePage::Leaf(mut entries) => {
216 validate_leaf(&entries, routing.lower.as_ref(), routing.upper.as_ref())
217 .map_err(|error| failed(error, work))?;
218 let (position, comparisons) = search_entries(&entries, name);
219 charge_items(&mut work, comparisons, budget)?;
220 let entry = position.ok().map(|index| entries.swap_remove(index));
221 work = work
222 .checked_add(WorkCounters {
223 items_returned: u64::from(entry.is_some()),
224 ..WorkCounters::default()
225 })
226 .map_err(|error| failed(error.into(), work))?;
227 work.verify(budget)
228 .map_err(|error| failed(error.into(), work))?;
229 allocations
230 .release(routing.logical_bytes)
231 .and_then(|()| allocations.release(decoded_bytes))
232 .map_err(|error| failed(map_allocation(error), work))?;
233 visited
234 .release(&mut allocations)
235 .map_err(|error| failed(map_allocation(error), work))?;
236 return Ok(AttributeLookup { entry, work });
237 }
238 AttributePage::Internal(children) => {
239 page = advance_internal(
240 children,
241 name,
242 &mut routing,
243 decoded_bytes,
244 &mut allocations,
245 &mut work,
246 budget,
247 )?;
248 }
249 }
250 }
251 Err(failed(AttributeLookupError::HeightExceeded, work))
252}
253
254#[derive(Default)]
255struct RoutingState {
256 lower: Option<AttributeName>,
257 upper: Option<AttributeName>,
258 logical_bytes: u64,
259}
260
261fn advance_internal(
262 children: Vec<AttributeChild>,
263 name: &AttributeName,
264 routing: &mut RoutingState,
265 decoded_bytes: u64,
266 allocations: &mut AllocationLedger,
267 work: &mut WorkCounters,
268 budget: WorkBudget,
269) -> Result<ObjectId, AttributeLookupFailure> {
270 validate_children(&children, routing.lower.as_ref(), routing.upper.as_ref())
271 .map_err(|error| failed(error, *work))?;
272 let (partition, comparisons) = upper_bound_children(&children, name);
273 charge_items(work, comparisons, budget)?;
274 let selected = partition.saturating_sub(1);
275 let mut selected_children = children.into_iter().skip(selected);
276 let child = selected_children
277 .next()
278 .ok_or_else(|| failed(AttributeLookupError::InvalidRouting, *work))?;
279 let next_upper = selected_children.next().map(|next| next.first_name);
280 let inherited_upper = routing.upper.take();
281 let old_lower_bytes = nested_name_bytes(routing.lower.as_ref());
282 let inherited_upper_bytes = nested_name_bytes(inherited_upper.as_ref());
283 let new_lower_bytes = nested_name_bytes(Some(&child.first_name));
284 let decoded_retained = nested_name_bytes(next_upper.as_ref())
285 .checked_add(new_lower_bytes)
286 .ok_or_else(|| failed(AttributeLookupError::AllocationFailed, *work))?;
287 allocations
288 .release(old_lower_bytes)
289 .and_then(|()| {
290 if next_upper.is_some() {
291 allocations.release(inherited_upper_bytes)
292 } else {
293 Ok(())
294 }
295 })
296 .and_then(|()| {
297 decoded_bytes
298 .checked_sub(decoded_retained)
299 .ok_or(AllocationError::ReleaseInvariant)
300 .and_then(|released| allocations.release(released))
301 })
302 .map_err(|error| failed(map_allocation(error), *work))?;
303 routing.lower = Some(child.first_name);
304 routing.upper = next_upper.or(inherited_upper);
305 routing.logical_bytes = new_lower_bytes
306 .checked_add(nested_name_bytes(routing.upper.as_ref()))
307 .ok_or_else(|| failed(AttributeLookupError::AllocationFailed, *work))?;
308 Ok(child.page)
309}
310
311fn nested_name_bytes(name: Option<&AttributeName>) -> u64 {
312 name.map_or(0, |bound| {
313 u64::try_from(bound.as_bytes().len()).unwrap_or(u64::MAX)
314 })
315}
316
317struct DecodedAttributePage {
318 page: AttributePage,
319 logical_bytes: u64,
320}
321
322async fn read_page<S: AsyncObjectStore>(
323 store: &S,
324 page: ObjectId,
325 limits: DecodeLimits,
326 budget: WorkBudget,
327 cancellation: &CancellationToken,
328 allocations: &mut AllocationLedger,
329 work: &mut WorkCounters,
330) -> Result<DecodedAttributePage, AttributeLookupFailure> {
331 let prospective = work
332 .checked_add(WorkCounters {
333 page_reads: 1,
334 ..WorkCounters::default()
335 })
336 .map_err(|error| failed(error.into(), *work))?;
337 let mut remaining = prospective
338 .remaining(budget)
339 .map_err(|error| failed(error.into(), *work))?;
340 remaining.peak_allocation_bytes = budget
341 .peak_allocation_bytes
342 .checked_sub(allocations.live_bytes())
343 .ok_or_else(|| failed(AttributeLookupError::Work(WorkError::Overflow), *work))?;
344 let receipt = AsyncObjectStore::read(
345 store,
346 page,
347 limits.maximum_page_object_bytes(),
348 remaining,
349 cancellation,
350 )
351 .await
352 .map_err(|failure| {
353 match merge_backend_work(prospective, *failure.work, allocations.live_bytes()) {
354 Ok(spent) => failed(map_storage(failure.error), spent),
355 Err(error) => failed(error.into(), prospective),
356 }
357 })?;
358 *work = merge_backend_work(prospective, receipt.work, allocations.live_bytes())
359 .map_err(|error| failed(error.into(), prospective))?;
360 work.verify(budget)
361 .map_err(|error| failed(error.into(), *work))?;
362
363 let retained_bytes = match receipt.value.retention {
364 ObjectReadRetention::Shared => 0,
365 ObjectReadRetention::Owned { logical_bytes } => logical_bytes,
366 };
367 allocations
368 .claim_bytes(retained_bytes, 0, work, budget)
369 .map_err(|error| failed(map_allocation(error), *work))?;
370 let shape = attribute_page_decode_shape(&receipt.value, limits)
371 .map_err(|error| failed(error.into(), *work))?;
372 charge_items(work, u64::try_from(shape.items).unwrap_or(u64::MAX), budget)?;
373 let decoded_work = work
374 .checked_add(WorkCounters {
375 bytes_copied: shape.nested_bytes,
376 ..WorkCounters::default()
377 })
378 .map_err(|error| failed(error.into(), *work))?;
379 decoded_work
380 .verify(budget)
381 .map_err(|error| failed(error.into(), *work))?;
382 *work = decoded_work;
383 let container_bytes = match shape.kind {
384 DecodedPageKind::Leaf => shape.items.checked_mul(size_of::<AttributeEntry>()),
385 DecodedPageKind::Internal => shape.items.checked_mul(size_of::<AttributeChild>()),
386 }
387 .map(crate::foundation::usize_to_u64)
388 .ok_or_else(|| failed(AttributeLookupError::AllocationFailed, *work))?;
389 let logical_bytes = container_bytes
390 .checked_add(shape.nested_bytes)
391 .ok_or_else(|| failed(AttributeLookupError::AllocationFailed, *work))?;
392 allocations
393 .claim_bytes(logical_bytes, u64::from(logical_bytes != 0), work, budget)
394 .map_err(|error| failed(map_allocation(error), *work))?;
395 let page = decode_attribute_page(&receipt.value, limits)
396 .map_err(|error| failed(error.into(), *work))?;
397 allocations
398 .release(retained_bytes)
399 .map_err(|error| failed(map_allocation(error), *work))?;
400 Ok(DecodedAttributePage {
401 page,
402 logical_bytes,
403 })
404}
405
406fn search_entries(entries: &[AttributeEntry], name: &AttributeName) -> (Result<usize, usize>, u64) {
407 let mut left = 0_usize;
408 let mut right = entries.len();
409 let mut comparisons = 0_u64;
410 while left < right {
411 comparisons = comparisons.saturating_add(1);
412 let middle = left + (right - left) / 2;
413 #[allow(
414 clippy::indexing_slicing,
415 reason = "binary-search invariant: the loop guard `left < right` with `right` \
416 initialized to entries.len() and only ever narrowed to `middle` keeps \
417 `middle = left + (right - left) / 2` strictly within `[left, right)`, so \
418 `middle < entries.len()` always holds here"
419 )]
420 let ordering = entries[middle].name.cmp(name);
421 match ordering {
422 std::cmp::Ordering::Less => left = middle + 1,
423 std::cmp::Ordering::Greater => right = middle,
424 std::cmp::Ordering::Equal => return (Ok(middle), comparisons),
425 }
426 }
427 (Err(left), comparisons)
428}
429
430fn upper_bound_children(children: &[AttributeChild], name: &AttributeName) -> (usize, u64) {
431 let mut left = 0_usize;
432 let mut right = children.len();
433 let mut comparisons = 0_u64;
434 while left < right {
435 comparisons = comparisons.saturating_add(1);
436 let middle = left + (right - left) / 2;
437 #[allow(
438 clippy::indexing_slicing,
439 reason = "binary-search invariant: the loop guard `left < right` with `right` \
440 initialized to children.len() and only ever narrowed to `middle` keeps \
441 `middle = left + (right - left) / 2` strictly within `[left, right)`, so \
442 `middle < children.len()` always holds here"
443 )]
444 let below_name = children[middle].first_name <= *name;
445 if below_name {
446 left = middle + 1;
447 } else {
448 right = middle;
449 }
450 }
451 (left, comparisons)
452}
453
454fn validate_leaf(
455 entries: &[AttributeEntry],
456 lower: Option<&AttributeName>,
457 upper: Option<&AttributeName>,
458) -> Result<(), AttributeLookupError> {
459 if lower.is_some() && entries.first().map(|entry| &entry.name) != lower {
460 return Err(AttributeLookupError::ChildBoundsMismatch);
461 }
462 if let Some(upper) = upper
463 && entries.last().is_some_and(|entry| entry.name >= *upper)
464 {
465 return Err(AttributeLookupError::ChildBoundsMismatch);
466 }
467 Ok(())
468}
469
470fn validate_children(
471 children: &[AttributeChild],
472 lower: Option<&AttributeName>,
473 upper: Option<&AttributeName>,
474) -> Result<(), AttributeLookupError> {
475 if lower.is_some() && children.first().map(|child| &child.first_name) != lower {
476 return Err(AttributeLookupError::ChildBoundsMismatch);
477 }
478 if let Some(upper) = upper
479 && children
480 .last()
481 .is_some_and(|child| child.first_name >= *upper)
482 {
483 return Err(AttributeLookupError::ChildBoundsMismatch);
484 }
485 Ok(())
486}
487
488fn charge_items(
489 work: &mut WorkCounters,
490 count: u64,
491 budget: WorkBudget,
492) -> Result<(), AttributeLookupFailure> {
493 let prospective = work
494 .checked_add(WorkCounters {
495 items_examined: count,
496 ..WorkCounters::default()
497 })
498 .map_err(|error| failed(error.into(), *work))?;
499 prospective
500 .verify(budget)
501 .map_err(|error| failed(error.into(), *work))?;
502 *work = prospective;
503 Ok(())
504}
505
506fn merge_backend_work(
507 prior: WorkCounters,
508 mut backend: WorkCounters,
509 live_bytes: u64,
510) -> Result<WorkCounters, WorkError> {
511 let simultaneous_peak = live_bytes
512 .checked_add(backend.peak_allocation_bytes)
513 .ok_or(WorkError::Overflow)?;
514 backend.peak_allocation_bytes = 0;
515 let mut merged = prior.checked_add(backend)?;
516 merged.peak_allocation_bytes = merged.peak_allocation_bytes.max(simultaneous_peak);
517 Ok(merged)
518}
519
520fn map_allocation(error: AllocationError) -> AttributeLookupError {
521 match error {
522 AllocationError::Work(error) => AttributeLookupError::Work(error),
523 AllocationError::Overflow
524 | AllocationError::ReleaseInvariant
525 | AllocationError::InvalidCapacity
526 | AllocationError::CapacityExceeded
527 | AllocationError::AllocationFailed => AttributeLookupError::AllocationFailed,
528 }
529}
530
531fn map_storage(error: ObjectStoreError) -> AttributeLookupError {
532 match error {
533 ObjectStoreError::Cancelled => AttributeLookupError::Cancelled,
534 error => AttributeLookupError::Storage(error),
535 }
536}
537
538fn failed(error: AttributeLookupError, work: WorkCounters) -> AttributeLookupFailure {
539 OperationFailure::new(error, work)
540}
541
542#[derive(Debug, Error)]
544pub enum AttributeLookupError {
545 #[error("attribute lookup was cancelled")]
547 Cancelled,
548 #[error("attribute lookup batch is empty")]
550 EmptyBatch,
551 #[error("attribute lookup batch exceeds its admitted query bound")]
553 TooManyQueries,
554 #[error("attribute lookup root is not an attribute page")]
556 WrongRootKind,
557 #[error("attribute lookup limits are invalid")]
559 InvalidLimits,
560 #[error("attribute tree exceeds its admitted height")]
562 HeightExceeded,
563 #[error("attribute tree contains a cycle or alias")]
565 CycleOrAlias,
566 #[error("attribute child bounds mismatch")]
568 ChildBoundsMismatch,
569 #[error("attribute tree routing is invalid")]
571 InvalidRouting,
572 #[error("attribute lookup scratch allocation failed")]
574 AllocationFailed,
575 #[error(transparent)]
577 Storage(#[from] ObjectStoreError),
578 #[error(transparent)]
580 Decode(#[from] CanonicalDecodeError),
581 #[error(transparent)]
583 Work(#[from] WorkError),
584}
585
586#[cfg(all(test, feature = "memory"))]
587#[path = "tests/attribute_access.rs"]
588mod tests;